From 012c8b95f97ba91ffa513cbf1a50cbc47a067e82 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 03:06:20 +0000 Subject: [PATCH] feat(auth): move auth policy and effect reads into Rust --- .github/workflows/test-unit.yml | 1 + litellm-rust/Cargo.lock | 20 + litellm-rust/Cargo.toml | 1 + ...de6d938c46a3108c2385c52d80e86c4473b96.json | 23 + ...bc5f270be24cd01b8c39bceea2de70d44608b.json | 24 + ...a2a45bd52705356fbf94128776db52032fd05.json | 23 + ...7398561f205596a12e7f6eb71dd5ecbf7d49a.json | 23 + ...c78c7c6a44ada9e601f6c74af1e9debd06f21.json | 23 + ...2c8ef734194248e2a078af73776398a256090.json | 23 + ...1b53e073a2b8f1e5b9007793b1c89df1ae24f.json | 23 + ...4f571c6fae8cd6d09a7f67fd3f41025e260be.json | 23 + ...837fa0adb227764bcaff1c2267f9a4624f99c.json | 24 + ...723f7346cf3192e7e5eb6d54f0e4b444b20e7.json | 23 + ...20aac76ac927183a15738d776ff9e104375af.json | 23 + ...26d0ffba418b154ae4e511287256db5168177.json | 23 + ...cf1ff4c7fb7a158a0af6694c4cc3da0d66dd1.json | 23 + ...5926a45b69c60677f527f227a2f8a58020b5b.json | 25 + ...27eae2c4d796d2891f6f3816289f8f71d7ea1.json | 23 + ...94c7e1bafada2968e53f4ecd16cd49d82cd5b.json | 23 + ...ee9dd8dcb36772073095414df0d4493a3e363.json | 23 + ...223c5ce05000eaf8cb6888a66033a94fe7a09.json | 23 + ...8e012ff43436560342357afabbc0d5b961c44.json | 24 + ...e14a6dcec24dbe6c5063ef9ef75cb196d886c.json | 23 + ...fee4c56e187d9f9e3d26d625fc3432e9d5652.json | 25 + ...93f983adaeaf5ab4166a90ab4a89ae39e2b0b.json | 25 + ...ec5df311bb3953f9756a5181821af9339cdba.json | 23 + ...9860f7a2da7d3d6d40c209322e0893b99ac69.json | 23 + ...7b51a9910fb27160356d9cc5dd03eb16985b2.json | 23 + ...d5b7f5ea9ad7a158c5824ab8fc42e1e67dbdb.json | 23 + ...1a0cb3af72f78c4c6f90890131faa1a15d3db.json | 23 + ...2b5719d6fcbc961b5234cd5ad3100674b293b.json | 23 + ...0b4bd94727d6f882992ff4b8419be945175ef.json | 23 + ...9ba9059a84dbd20d9e148c72f23a62f52a981.json | 23 + ...c97afbec425ba252db103d956f15c43c23c28.json | 23 + ...494d92ea3bf7d09bfeb6f6aa30a7bff7da7bf.json | 23 + ...e562f532616ee7d81288ffdcc9ba34d4762c0.json | 24 + ...e8b674266007e4b11e388f722c9418f8ccea0.json | 23 + ...2f48ade878aae01105b6d7181f95cd5a16864.json | 25 + ...12422927cedd9981b6b3617aa61b1086ab8f9.json | 23 + ...33c8443365b9262f2827e1c0b09f715b0227b.json | 27 + ...9e5d79069036df1c90481aa17b5ee7c6e4293.json | 23 + ...02f70cecce44f8e5d6673d2521a57de3513d8.json | 23 + ...d1dcafc1f0cf69e33cb58ebe8642b75ee8d36.json | 23 + litellm-rust/crates/db/Cargo.toml | 5 +- .../crates/db/queries/auth/access_group.sql | 9 + litellm-rust/crates/db/queries/auth/agent.sql | 16 + .../db/queries/auth/agent_access_group.sql | 4 + .../crates/db/queries/auth/agent_identity.sql | 2 + .../crates/db/queries/auth/agent_name.sql | 16 + .../crates/db/queries/auth/budget.sql | 4 + .../crates/db/queries/auth/budget_window.sql | 1 + .../crates/db/queries/auth/end_user.sql | 13 + .../auth/end_user_restricted_registry.sql | 3 + .../crates/db/queries/auth/fixture.sql | 58 + litellm-rust/crates/db/queries/auth/key.sql | 50 + .../db/queries/auth/managed_vector_store.sql | 4 + .../crates/db/queries/auth/member.sql | 2 + .../db/queries/auth/model_access_group.sql | 4 + .../auth/model_access_group_registry.sql | 1 + .../db/queries/auth/object_permission.sql | 13 + litellm-rust/crates/db/queries/auth/org.sql | 13 + .../crates/db/queries/auth/org_alias.sql | 10 + .../crates/db/queries/auth/prefetch.sql | 42 + .../crates/db/queries/auth/project.sql | 13 + .../crates/db/queries/auth/retired_agent.sql | 4 + .../queries/auth/retired_agent_identity.sql | 1 + .../crates/db/queries/auth/spend_end_user.sql | 1 + .../crates/db/queries/auth/spend_key.sql | 1 + .../crates/db/queries/auth/spend_member.sql | 1 + .../queries/auth/spend_model_access_group.sql | 1 + .../crates/db/queries/auth/spend_org.sql | 1 + .../crates/db/queries/auth/spend_project.sql | 1 + .../crates/db/queries/auth/spend_tag.sql | 1 + .../crates/db/queries/auth/spend_team.sql | 1 + .../crates/db/queries/auth/spend_user.sql | 1 + litellm-rust/crates/db/queries/auth/tag.sql | 4 + .../crates/db/queries/auth/tag_registry.sql | 1 + litellm-rust/crates/db/queries/auth/team.sql | 19 + .../crates/db/queries/auth/team_alias.sql | 16 + litellm-rust/crates/db/queries/auth/user.sql | 7 + .../crates/db/queries/auth/user_email.sql | 4 + .../crates/db/queries/auth/user_sso.sql | 4 + .../db/queries/auth/verified_subject.sql | 1 + .../db/queries/auth/window_key_spend.sql | 1 + .../db/queries/auth/window_team_spend.sql | 1 + litellm-rust/crates/db/src/auth.rs | 814 ++++ litellm-rust/crates/db/src/error.rs | 49 +- litellm-rust/crates/db/src/executor.rs | 84 + litellm-rust/crates/db/src/lib.rs | 5 + litellm-rust/crates/db/src/pool.rs | 495 ++ litellm-rust/crates/python-bridge/Cargo.toml | 2 + litellm-rust/crates/python-bridge/src/db.rs | 171 +- litellm-rust/crates/python-bridge/src/lib.rs | 135 +- .../src/redis/store/auth_cache.rs | 92 + .../crates/python-bridge/src/request_auth.rs | 1086 +++++ .../crates/redis-client/src/client.rs | 4 + litellm-rust/crates/redis-client/src/lib.rs | 1 + .../crates/redis-client/src/pubsub.rs | 89 +- .../crates/redis-client/tests/pubsub.rs | 81 +- litellm-rust/crates/redis-store/Cargo.toml | 1 + .../redis-store/src/auth_cache/local.rs | 108 + .../crates/redis-store/src/auth_cache/mod.rs | 293 +- .../redis-store/src/auth_cache/types.rs | 11 +- .../crates/redis-store/src/dual_cache/mod.rs | 1 + .../redis-store/src/dual_cache/write.rs | 81 +- litellm-rust/crates/redis-store/src/memory.rs | 5 +- .../crates/request-auth-core/Cargo.toml | 17 + .../crates/request-auth-core/src/access.rs | 714 +++ .../crates/request-auth-core/src/admission.rs | 183 + .../request-auth-core/src/agent_access.rs | 269 ++ .../request-auth-core/src/agent_ceilings.rs | 151 + .../request-auth-core/src/agent_grants.rs | 408 ++ .../request-auth-core/src/auth_context.rs | 169 + .../request-auth-core/src/budget_alerts.rs | 233 + .../request-auth-core/src/budget_context.rs | 484 ++ .../request-auth-core/src/budget_lookup.rs | 284 ++ .../request-auth-core/src/budget_pricing.rs | 261 ++ .../src/claude_membership.rs | 100 + .../src/claude_normalization.rs | 176 + .../request-auth-core/src/codex_models.rs | 427 ++ .../src/codex_normalization.rs | 201 + .../request-auth-core/src/codex_turn.rs | 230 + .../request-auth-core/src/common_budgets.rs | 230 + .../request-auth-core/src/common_flow.rs | 358 ++ .../request-auth-core/src/credential.rs | 225 + .../request-auth-core/src/end_user_context.rs | 428 ++ .../crates/request-auth-core/src/error.rs | 9 + .../request-auth-core/src/fallback_access.rs | 145 + .../src/fallback_candidate.rs | 322 ++ .../crates/request-auth-core/src/key.rs | 442 ++ .../request-auth-core/src/key_model_access.rs | 88 + .../crates/request-auth-core/src/lib.rs | 62 + .../crates/request-auth-core/src/managed.rs | 348 ++ .../request-auth-core/src/managed_models.rs | 253 + .../request-auth-core/src/member_models.rs | 77 + .../request-auth-core/src/model_access.rs | 726 +++ .../src/model_attribution.rs | 147 + .../src/model_budget_aliases.rs | 250 + .../src/model_budget_data.rs | 331 ++ .../request-auth-core/src/model_budgets.rs | 312 ++ .../request-auth-core/src/model_groups.rs | 73 + .../request-auth-core/src/model_helpers.rs | 148 + .../request-auth-core/src/model_projection.rs | 299 ++ .../request-auth-core/src/org_routes.rs | 304 ++ .../request-auth-core/src/owner_teams.rs | 173 + .../crates/request-auth-core/src/patterns.rs | 363 ++ .../request-auth-core/src/personal_budget.rs | 186 + .../src/personal_model_listing.rs | 81 + .../src/personal_model_request.rs | 86 + .../request-auth-core/src/personal_models.rs | 120 + .../request-auth-core/src/personal_owner.rs | 87 + .../crates/request-auth-core/src/protocol.rs | 463 ++ .../request-auth-core/src/reservation.rs | 312 ++ .../request-auth-core/src/resolved_models.rs | 219 + .../request-auth-core/src/route_access.rs | 567 +++ .../request-auth-core/src/route_enable.rs | 180 + .../request-auth-core/src/route_roles.rs | 325 ++ .../request-auth-core/src/route_scope.rs | 97 + .../src/spend_counter_policy.rs | 299 ++ .../request-auth-core/src/team_aliases.rs | 77 + .../request-auth-core/src/team_override.rs | 70 + .../request-auth-core/src/unpriced_models.rs | 144 + litellm/_auth/AGENTS.md | 18 + litellm/_auth/__init__.py | 0 litellm/_auth/access.py | 248 + litellm/_auth/admission.py | 77 + litellm/_auth/agent_access.py | 68 + litellm/_auth/agent_ceilings.py | 60 + litellm/_auth/agent_grants.py | 128 + litellm/_auth/auth_context.py | 74 + litellm/_auth/budget_alerts.py | 69 + litellm/_auth/budget_context.py | 174 + litellm/_auth/budget_lookup.py | 90 + litellm/_auth/budget_pricing.py | 90 + litellm/_auth/claude_normalization.py | 74 + litellm/_auth/codex_models.py | 74 + litellm/_auth/codex_normalization.py | 64 + litellm/_auth/codex_turn.py | 100 + litellm/_auth/common_budgets.py | 115 + litellm/_auth/common_flow.py | 148 + litellm/_auth/credential.py | 141 + litellm/_auth/decision.py | 460 ++ litellm/_auth/driver.py | 274 ++ litellm/_auth/end_user_context.py | 171 + litellm/_auth/failure.py | 39 + litellm/_auth/fallback_access.py | 58 + litellm/_auth/fallback_candidate.py | 88 + litellm/_auth/key.py | 12 + litellm/_auth/key_model_access.py | 45 + litellm/_auth/managed.py | 129 + litellm/_auth/managed_models.py | 95 + litellm/_auth/member_models.py | 32 + litellm/_auth/model_access.py | 261 ++ litellm/_auth/model_attribution.py | 46 + litellm/_auth/model_budget_aliases.py | 59 + litellm/_auth/model_budget_data.py | 94 + litellm/_auth/model_groups.py | 48 + litellm/_auth/model_helpers.py | 37 + litellm/_auth/model_projection.py | 27 + litellm/_auth/org_routes.py | 95 + litellm/_auth/owner_teams.py | 53 + litellm/_auth/personal_budget.py | 66 + litellm/_auth/personal_model_listing.py | 32 + litellm/_auth/personal_model_request.py | 33 + litellm/_auth/personal_models.py | 45 + litellm/_auth/personal_owner.py | 36 + litellm/_auth/reservation.py | 131 + litellm/_auth/resolved_models.py | 78 + litellm/_auth/route_enable.py | 52 + litellm/_auth/route_roles.py | 120 + litellm/_auth/route_scope.py | 40 + litellm/_auth/routes.py | 146 + litellm/_auth/runtime.py | 23 + litellm/_auth/rust_core.py | 270 ++ litellm/_auth/spend_counter_policy.py | 129 + litellm/_auth/team_aliases.py | 34 + litellm/_auth/team_override.py | 30 + litellm/_auth/unpriced_models.py | 53 + .../auth/agent_access_groups.py | 62 +- .../auth/agent_permission_handler.py | 664 ++- .../auth/managed_authorization.py | 290 +- .../proxy/agent_endpoints/identity_store.py | 11 + litellm/proxy/auth/agent_ceiling_adapter.py | 156 + litellm/proxy/auth/auth_checks.py | 4091 +++++------------ .../proxy/auth/auth_checks_organization.py | 130 +- litellm/proxy/auth/auth_denial.py | 152 + litellm/proxy/auth/auth_object_prefetch.py | 52 +- litellm/proxy/auth/budget_alert_adapter.py | 178 + litellm/proxy/auth/budget_context_adapter.py | 385 ++ litellm/proxy/auth/budget_lookup_adapter.py | 144 + litellm/proxy/auth/budget_pricing_adapter.py | 180 + litellm/proxy/auth/budget_throttle.py | 31 +- litellm/proxy/auth/budget_windows_adapter.py | 101 + .../auth/claude_normalization_adapter.py | 174 + litellm/proxy/auth/codex_model_adapter.py | 168 + .../proxy/auth/codex_normalization_adapter.py | 169 + litellm/proxy/auth/codex_turn_adapter.py | 226 + litellm/proxy/auth/common_budget_adapter.py | 432 ++ litellm/proxy/auth/common_budget_lookups.py | 240 + .../proxy/auth/end_user_context_adapter.py | 406 ++ litellm/proxy/auth/end_user_context_io.py | 64 + litellm/proxy/auth/fallback_access_adapter.py | 154 + .../proxy/auth/fallback_candidate_adapter.py | 320 ++ .../proxy/auth/key_budget_context_adapter.py | 95 + .../proxy/auth/key_model_access_adapter.py | 161 + litellm/proxy/auth/managed_model_adapter.py | 366 ++ litellm/proxy/auth/member_model_adapter.py | 132 + litellm/proxy/auth/model_access_adapter.py | 310 ++ .../proxy/auth/model_attribution_adapter.py | 213 + litellm/proxy/auth/model_budget_adapter.py | 306 ++ litellm/proxy/auth/model_checks.py | 17 +- litellm/proxy/auth/model_group_adapter.py | 94 + litellm/proxy/auth/model_helper_adapter.py | 55 + litellm/proxy/auth/model_policy_read.py | 11 + .../proxy/auth/model_projection_adapter.py | 121 + litellm/proxy/auth/org_route_adapter.py | 130 + litellm/proxy/auth/owner_team_adapter.py | 104 + litellm/proxy/auth/personal_budget_adapter.py | 280 ++ litellm/proxy/auth/personal_model_adapter.py | 192 + .../auth/personal_model_listing_adapter.py | 111 + .../auth/personal_model_request_adapter.py | 93 + litellm/proxy/auth/personal_owner_adapter.py | 129 + litellm/proxy/auth/resolved_model_adapter.py | 286 ++ litellm/proxy/auth/resolvers/store.py | 35 +- litellm/proxy/auth/route_checks.py | 721 +-- litellm/proxy/auth/route_enable_adapter.py | 130 + litellm/proxy/auth/route_policy_adapter.py | 256 ++ litellm/proxy/auth/route_role_adapter.py | 163 + litellm/proxy/auth/route_scope_adapter.py | 73 + litellm/proxy/auth/spend_counter_adapter.py | 34 + litellm/proxy/auth/team_alias_adapter.py | 33 + litellm/proxy/auth/team_override_adapter.py | 102 + litellm/proxy/auth/unpriced_model_adapter.py | 128 + litellm/proxy/auth/user_api_key_auth.py | 3538 +++++++------- .../auth_cache_invalidation_pubsub.py | 34 +- .../common_utils/registry_read_through.py | 14 +- .../proxy/common_utils/user_api_key_cache.py | 45 +- litellm/proxy/db/spend_counter_reseed.py | 63 + .../proxy/hooks/model_max_budget_limiter.py | 123 +- .../key_management_endpoints.py | 6 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 8 +- .../codex_passthrough_endpoints.py | 149 +- .../llm_passthrough_endpoints.py | 15 +- litellm/proxy/proxy_server.py | 70 +- .../spend_tracking/budget_reservation.py | 235 +- .../spend_tracking/spend_tracking_utils.py | 12 +- litellm/proxy/utils.py | 41 +- litellm/redis/stores/auth_cache/__init__.py | 18 +- litellm/redis/stores/auth_cache/handle.py | 6 +- litellm/repositories/database.py | 76 +- litellm/repositories/database_settings.py | 64 + .../repositories/native_auth_repository.py | 194 + litellm/rust_bridge/_native.pyi | 147 +- litellm/types/proxy/auth/user_api_key_auth.py | 10 +- tests/integration/_support/database_relay.py | 6 +- tests/integration/_support/process.py | 22 +- .../authorization/test_native_auth_io.py | 149 + .../caching/test_auth_cache_native.py | 237 + tests/unit/_auth/__init__.py | 0 tests/unit/_auth/budget_reference.py | 117 + tests/unit/_auth/conftest.py | 28 + tests/unit/_auth/model_alias_reference.py | 38 + tests/unit/_auth/raw_answers.py | 40 + tests/unit/_auth/raw_model_reference.py | 52 + tests/unit/_auth/reference_core.py | 266 ++ tests/unit/_auth/reference_model_budgets.py | 175 + tests/unit/_auth/reference_policy.py | 186 + tests/unit/_auth/shadow.py | 115 + tests/unit/_auth/support/claude_reference.py | 103 + tests/unit/_auth/support/codex_reference.py | 68 + .../unit/_auth/support/fallback_reference.py | 68 + tests/unit/_auth/support/model_reference.py | 91 + .../_auth/support/model_registry_reference.py | 189 + .../_auth/support/route_enable_reference.py | 39 + tests/unit/_auth/support/route_reference.py | 1233 +++++ tests/unit/_auth/test_access.py | 110 + tests/unit/_auth/test_admission.py | 21 + tests/unit/_auth/test_agent_ceilings.py | 91 + tests/unit/_auth/test_agent_grants.py | 193 + tests/unit/_auth/test_auth_context.py | 64 + tests/unit/_auth/test_budget_pricing.py | 78 + tests/unit/_auth/test_claude_normalization.py | 264 ++ tests/unit/_auth/test_codex_models.py | 140 + tests/unit/_auth/test_codex_normalization.py | 168 + tests/unit/_auth/test_codex_turn.py | 226 + tests/unit/_auth/test_common_budgets.py | 134 + tests/unit/_auth/test_common_flow.py | 147 + tests/unit/_auth/test_credential.py | 161 + tests/unit/_auth/test_failure.py | 36 + tests/unit/_auth/test_fallback_access.py | 106 + tests/unit/_auth/test_fallback_candidate.py | 211 + tests/unit/_auth/test_key.py | 35 + tests/unit/_auth/test_key_model_access.py | 80 + tests/unit/_auth/test_managed.py | 84 + tests/unit/_auth/test_managed_models.py | 116 + tests/unit/_auth/test_member_models.py | 56 + tests/unit/_auth/test_model_access.py | 285 ++ tests/unit/_auth/test_model_attribution.py | 87 + tests/unit/_auth/test_model_budget_aliases.py | 66 + tests/unit/_auth/test_model_budget_data.py | 111 + tests/unit/_auth/test_model_budgets.py | 332 ++ tests/unit/_auth/test_model_groups.py | 51 + tests/unit/_auth/test_model_helpers.py | 105 + tests/unit/_auth/test_owner_teams.py | 91 + tests/unit/_auth/test_personal_budget.py | 149 + .../unit/_auth/test_personal_model_listing.py | 66 + .../unit/_auth/test_personal_model_request.py | 57 + tests/unit/_auth/test_personal_models.py | 110 + tests/unit/_auth/test_personal_owner.py | 61 + tests/unit/_auth/test_python_core.py | 637 +++ tests/unit/_auth/test_reservation.py | 123 + tests/unit/_auth/test_resolved_models.py | 156 + tests/unit/_auth/test_route_enable.py | 75 + tests/unit/_auth/test_route_scope.py | 61 + tests/unit/_auth/test_routes.py | 191 + tests/unit/_auth/test_runtime.py | 22 + tests/unit/_auth/test_rust_core.py | 127 + tests/unit/_auth/test_spend_counter_policy.py | 81 + tests/unit/_auth/test_team_aliases.py | 49 + tests/unit/_auth/test_team_override.py | 74 + tests/unit/_auth/test_unpriced_models.py | 106 + .../auth/test_agent_permission_handler.py | 46 +- .../agent_endpoints/test_identity_store.py | 36 +- .../unit/proxy/auth/budget_alert_reference.py | 42 + tests/unit/proxy/auth/budget_reference.py | 44 + .../auth/phase1_budget_auth_reference.py | 915 ++++ .../auth/phase1_budget_context_reference.py | 352 ++ .../auth/phase1_budget_lookup_reference.py | 341 ++ .../auth/phase1_budget_route_reference.py | 151 + .../auth/phase1_budget_throttle_reference.py | 45 + tests/unit/proxy/auth/route_role_reference.py | 81 + tests/unit/proxy/auth/test_auth_checks.py | 5 +- ...st_auth_checks_object_access_and_lookup.py | 266 +- .../proxy/auth/test_auth_object_prefetch.py | 79 + .../proxy/auth/test_budget_alert_adapter.py | 125 + .../proxy/auth/test_budget_context_adapter.py | 335 ++ .../proxy/auth/test_budget_lookup_adapter.py | 276 ++ .../proxy/auth/test_budget_pricing_adapter.py | 121 + .../proxy/auth/test_common_budget_adapter.py | 297 ++ .../proxy/auth/test_common_budget_lookups.py | 44 + .../auth/test_custom_auth_end_user_budget.py | 2 +- .../auth/test_end_user_context_adapter.py | 437 ++ .../auth/test_key_budget_context_adapter.py | 108 + .../proxy/auth/test_managed_model_adapter.py | 84 + .../proxy/auth/test_model_budget_adapter.py | 307 ++ .../unit/proxy/auth/test_model_policy_read.py | 18 + .../auth/test_model_projection_adapter.py | 203 + .../unit/proxy/auth/test_org_route_adapter.py | 101 + .../test_personal_key_team_model_access.py | 10 +- tests/unit/proxy/auth/test_resolvers_store.py | 41 +- .../proxy/auth/test_route_enable_adapter.py | 140 + .../proxy/auth/test_route_role_adapter.py | 126 + .../test_router_override_fallback_auth.py | 186 +- .../proxy/auth/test_spend_counter_adapter.py | 49 + .../unit/proxy/auth/test_user_api_key_auth.py | 201 +- .../test_user_api_key_auth_request_flow.py | 431 +- .../test_auth_cache_invalidation_pubsub.py | 12 +- .../test_registry_read_through.py | 80 +- .../common_utils/test_user_api_key_cache.py | 105 +- .../proxy/db/test_spend_counter_reseed.py | 96 + ...test_unit_test_max_model_budget_limiter.py | 141 +- .../test_spend_tracking_utils.py | 32 - .../proxy/test_tools_allowlist_enforcement.py | 63 +- tests/unit/redis/stores/auth_cache_double.py | 12 + .../repositories/test_database_settings.py | 62 + .../test_native_auth_repository.py | 149 + 406 files changed, 50782 insertions(+), 7667 deletions(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-00ddffb253f55705f0b742827f9de6d938c46a3108c2385c52d80e86c4473b96.json create mode 100644 litellm-rust/crates/db/.sqlx/query-0136e6b22d40ef0181f3f122c5abc5f270be24cd01b8c39bceea2de70d44608b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-16fec3e3896126e8bd14726b346a2a45bd52705356fbf94128776db52032fd05.json create mode 100644 litellm-rust/crates/db/.sqlx/query-1d689acb624e95e13cdadde37817398561f205596a12e7f6eb71dd5ecbf7d49a.json create mode 100644 litellm-rust/crates/db/.sqlx/query-1f9b9a6524e417b18865fda7b49c78c7c6a44ada9e601f6c74af1e9debd06f21.json create mode 100644 litellm-rust/crates/db/.sqlx/query-21e8660e197f5f9132584f22da32c8ef734194248e2a078af73776398a256090.json create mode 100644 litellm-rust/crates/db/.sqlx/query-26a5fa1687c580afc0e1c1f7ca71b53e073a2b8f1e5b9007793b1c89df1ae24f.json create mode 100644 litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json create mode 100644 litellm-rust/crates/db/.sqlx/query-2c4e1f73d4170dd3c6a2b077fbd837fa0adb227764bcaff1c2267f9a4624f99c.json create mode 100644 litellm-rust/crates/db/.sqlx/query-31061f5f7b219176565eb4b3ebe723f7346cf3192e7e5eb6d54f0e4b444b20e7.json create mode 100644 litellm-rust/crates/db/.sqlx/query-3210aa55f379884ca4232cf2a0b20aac76ac927183a15738d776ff9e104375af.json create mode 100644 litellm-rust/crates/db/.sqlx/query-34107587736bee7d817b940717c26d0ffba418b154ae4e511287256db5168177.json create mode 100644 litellm-rust/crates/db/.sqlx/query-36c30969226b767cf7d4eec854ccf1ff4c7fb7a158a0af6694c4cc3da0d66dd1.json create mode 100644 litellm-rust/crates/db/.sqlx/query-3b2157705b66ebcf6ee0c9805f15926a45b69c60677f527f227a2f8a58020b5b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-498268b2d6edf244dfbe6d2c33327eae2c4d796d2891f6f3816289f8f71d7ea1.json create mode 100644 litellm-rust/crates/db/.sqlx/query-4ddd384812d4cc34d601873e73294c7e1bafada2968e53f4ecd16cd49d82cd5b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-56a5d90aee114ff352a31fbec84ee9dd8dcb36772073095414df0d4493a3e363.json create mode 100644 litellm-rust/crates/db/.sqlx/query-57d69d9d2ad5803d123717bd788223c5ce05000eaf8cb6888a66033a94fe7a09.json create mode 100644 litellm-rust/crates/db/.sqlx/query-65e9b4dc09101e63092902802978e012ff43436560342357afabbc0d5b961c44.json create mode 100644 litellm-rust/crates/db/.sqlx/query-667556bca77708a5c027641bb3be14a6dcec24dbe6c5063ef9ef75cb196d886c.json create mode 100644 litellm-rust/crates/db/.sqlx/query-683dd0c1cdf39792c5b0c7ca482fee4c56e187d9f9e3d26d625fc3432e9d5652.json create mode 100644 litellm-rust/crates/db/.sqlx/query-6af954d0490c3cea037cd7bd00993f983adaeaf5ab4166a90ab4a89ae39e2b0b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-786700d34873fa39e57ca9b0c1bec5df311bb3953f9756a5181821af9339cdba.json create mode 100644 litellm-rust/crates/db/.sqlx/query-847628007ea76e788caf20d35219860f7a2da7d3d6d40c209322e0893b99ac69.json create mode 100644 litellm-rust/crates/db/.sqlx/query-8637bd849446309e631a4eaa47f7b51a9910fb27160356d9cc5dd03eb16985b2.json create mode 100644 litellm-rust/crates/db/.sqlx/query-8d8ac844d36122a8e652c494519d5b7f5ea9ad7a158c5824ab8fc42e1e67dbdb.json create mode 100644 litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json create mode 100644 litellm-rust/crates/db/.sqlx/query-ab4b895f3a1ee6ab9e86f13b4f02b5719d6fcbc961b5234cd5ad3100674b293b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-b0d27eac417a4044836ce7ba4450b4bd94727d6f882992ff4b8419be945175ef.json create mode 100644 litellm-rust/crates/db/.sqlx/query-b23d9fe187171500ebd8593571e9ba9059a84dbd20d9e148c72f23a62f52a981.json create mode 100644 litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json create mode 100644 litellm-rust/crates/db/.sqlx/query-b840de2b7f4c4dfc7c2909062a4494d92ea3bf7d09bfeb6f6aa30a7bff7da7bf.json create mode 100644 litellm-rust/crates/db/.sqlx/query-c265904d0f03ab912de3cfd5adae562f532616ee7d81288ffdcc9ba34d4762c0.json create mode 100644 litellm-rust/crates/db/.sqlx/query-c720cb3a555cd750cd929003591e8b674266007e4b11e388f722c9418f8ccea0.json create mode 100644 litellm-rust/crates/db/.sqlx/query-c879e82114e82685c3613fa4bc82f48ade878aae01105b6d7181f95cd5a16864.json create mode 100644 litellm-rust/crates/db/.sqlx/query-cc5bac434a523f8deb0416b806412422927cedd9981b6b3617aa61b1086ab8f9.json create mode 100644 litellm-rust/crates/db/.sqlx/query-d804605eae54eaeaa01c7f511cb33c8443365b9262f2827e1c0b09f715b0227b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-dbe6bf2326dcb74823f6bfb90149e5d79069036df1c90481aa17b5ee7c6e4293.json create mode 100644 litellm-rust/crates/db/.sqlx/query-e6b5061b738fd353df83eaef9dc02f70cecce44f8e5d6673d2521a57de3513d8.json create mode 100644 litellm-rust/crates/db/.sqlx/query-f6ed52b7b550c41a8d085f1a66bd1dcafc1f0cf69e33cb58ebe8642b75ee8d36.json create mode 100644 litellm-rust/crates/db/queries/auth/access_group.sql create mode 100644 litellm-rust/crates/db/queries/auth/agent.sql create mode 100644 litellm-rust/crates/db/queries/auth/agent_access_group.sql create mode 100644 litellm-rust/crates/db/queries/auth/agent_identity.sql create mode 100644 litellm-rust/crates/db/queries/auth/agent_name.sql create mode 100644 litellm-rust/crates/db/queries/auth/budget.sql create mode 100644 litellm-rust/crates/db/queries/auth/budget_window.sql create mode 100644 litellm-rust/crates/db/queries/auth/end_user.sql create mode 100644 litellm-rust/crates/db/queries/auth/end_user_restricted_registry.sql create mode 100644 litellm-rust/crates/db/queries/auth/fixture.sql create mode 100644 litellm-rust/crates/db/queries/auth/key.sql create mode 100644 litellm-rust/crates/db/queries/auth/managed_vector_store.sql create mode 100644 litellm-rust/crates/db/queries/auth/member.sql create mode 100644 litellm-rust/crates/db/queries/auth/model_access_group.sql create mode 100644 litellm-rust/crates/db/queries/auth/model_access_group_registry.sql create mode 100644 litellm-rust/crates/db/queries/auth/object_permission.sql create mode 100644 litellm-rust/crates/db/queries/auth/org.sql create mode 100644 litellm-rust/crates/db/queries/auth/org_alias.sql create mode 100644 litellm-rust/crates/db/queries/auth/prefetch.sql create mode 100644 litellm-rust/crates/db/queries/auth/project.sql create mode 100644 litellm-rust/crates/db/queries/auth/retired_agent.sql create mode 100644 litellm-rust/crates/db/queries/auth/retired_agent_identity.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_end_user.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_key.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_member.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_model_access_group.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_org.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_project.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_tag.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_team.sql create mode 100644 litellm-rust/crates/db/queries/auth/spend_user.sql create mode 100644 litellm-rust/crates/db/queries/auth/tag.sql create mode 100644 litellm-rust/crates/db/queries/auth/tag_registry.sql create mode 100644 litellm-rust/crates/db/queries/auth/team.sql create mode 100644 litellm-rust/crates/db/queries/auth/team_alias.sql create mode 100644 litellm-rust/crates/db/queries/auth/user.sql create mode 100644 litellm-rust/crates/db/queries/auth/user_email.sql create mode 100644 litellm-rust/crates/db/queries/auth/user_sso.sql create mode 100644 litellm-rust/crates/db/queries/auth/verified_subject.sql create mode 100644 litellm-rust/crates/db/queries/auth/window_key_spend.sql create mode 100644 litellm-rust/crates/db/queries/auth/window_team_spend.sql create mode 100644 litellm-rust/crates/db/src/auth.rs create mode 100644 litellm-rust/crates/db/src/executor.rs create mode 100644 litellm-rust/crates/db/src/pool.rs create mode 100644 litellm-rust/crates/python-bridge/src/request_auth.rs create mode 100644 litellm-rust/crates/redis-store/src/auth_cache/local.rs create mode 100644 litellm-rust/crates/request-auth-core/Cargo.toml create mode 100644 litellm-rust/crates/request-auth-core/src/access.rs create mode 100644 litellm-rust/crates/request-auth-core/src/admission.rs create mode 100644 litellm-rust/crates/request-auth-core/src/agent_access.rs create mode 100644 litellm-rust/crates/request-auth-core/src/agent_ceilings.rs create mode 100644 litellm-rust/crates/request-auth-core/src/agent_grants.rs create mode 100644 litellm-rust/crates/request-auth-core/src/auth_context.rs create mode 100644 litellm-rust/crates/request-auth-core/src/budget_alerts.rs create mode 100644 litellm-rust/crates/request-auth-core/src/budget_context.rs create mode 100644 litellm-rust/crates/request-auth-core/src/budget_lookup.rs create mode 100644 litellm-rust/crates/request-auth-core/src/budget_pricing.rs create mode 100644 litellm-rust/crates/request-auth-core/src/claude_membership.rs create mode 100644 litellm-rust/crates/request-auth-core/src/claude_normalization.rs create mode 100644 litellm-rust/crates/request-auth-core/src/codex_models.rs create mode 100644 litellm-rust/crates/request-auth-core/src/codex_normalization.rs create mode 100644 litellm-rust/crates/request-auth-core/src/codex_turn.rs create mode 100644 litellm-rust/crates/request-auth-core/src/common_budgets.rs create mode 100644 litellm-rust/crates/request-auth-core/src/common_flow.rs create mode 100644 litellm-rust/crates/request-auth-core/src/credential.rs create mode 100644 litellm-rust/crates/request-auth-core/src/end_user_context.rs create mode 100644 litellm-rust/crates/request-auth-core/src/error.rs create mode 100644 litellm-rust/crates/request-auth-core/src/fallback_access.rs create mode 100644 litellm-rust/crates/request-auth-core/src/fallback_candidate.rs create mode 100644 litellm-rust/crates/request-auth-core/src/key.rs create mode 100644 litellm-rust/crates/request-auth-core/src/key_model_access.rs create mode 100644 litellm-rust/crates/request-auth-core/src/lib.rs create mode 100644 litellm-rust/crates/request-auth-core/src/managed.rs create mode 100644 litellm-rust/crates/request-auth-core/src/managed_models.rs create mode 100644 litellm-rust/crates/request-auth-core/src/member_models.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_access.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_attribution.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_budget_aliases.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_budget_data.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_budgets.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_groups.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_helpers.rs create mode 100644 litellm-rust/crates/request-auth-core/src/model_projection.rs create mode 100644 litellm-rust/crates/request-auth-core/src/org_routes.rs create mode 100644 litellm-rust/crates/request-auth-core/src/owner_teams.rs create mode 100644 litellm-rust/crates/request-auth-core/src/patterns.rs create mode 100644 litellm-rust/crates/request-auth-core/src/personal_budget.rs create mode 100644 litellm-rust/crates/request-auth-core/src/personal_model_listing.rs create mode 100644 litellm-rust/crates/request-auth-core/src/personal_model_request.rs create mode 100644 litellm-rust/crates/request-auth-core/src/personal_models.rs create mode 100644 litellm-rust/crates/request-auth-core/src/personal_owner.rs create mode 100644 litellm-rust/crates/request-auth-core/src/protocol.rs create mode 100644 litellm-rust/crates/request-auth-core/src/reservation.rs create mode 100644 litellm-rust/crates/request-auth-core/src/resolved_models.rs create mode 100644 litellm-rust/crates/request-auth-core/src/route_access.rs create mode 100644 litellm-rust/crates/request-auth-core/src/route_enable.rs create mode 100644 litellm-rust/crates/request-auth-core/src/route_roles.rs create mode 100644 litellm-rust/crates/request-auth-core/src/route_scope.rs create mode 100644 litellm-rust/crates/request-auth-core/src/spend_counter_policy.rs create mode 100644 litellm-rust/crates/request-auth-core/src/team_aliases.rs create mode 100644 litellm-rust/crates/request-auth-core/src/team_override.rs create mode 100644 litellm-rust/crates/request-auth-core/src/unpriced_models.rs create mode 100644 litellm/_auth/AGENTS.md create mode 100644 litellm/_auth/__init__.py create mode 100644 litellm/_auth/access.py create mode 100644 litellm/_auth/admission.py create mode 100644 litellm/_auth/agent_access.py create mode 100644 litellm/_auth/agent_ceilings.py create mode 100644 litellm/_auth/agent_grants.py create mode 100644 litellm/_auth/auth_context.py create mode 100644 litellm/_auth/budget_alerts.py create mode 100644 litellm/_auth/budget_context.py create mode 100644 litellm/_auth/budget_lookup.py create mode 100644 litellm/_auth/budget_pricing.py create mode 100644 litellm/_auth/claude_normalization.py create mode 100644 litellm/_auth/codex_models.py create mode 100644 litellm/_auth/codex_normalization.py create mode 100644 litellm/_auth/codex_turn.py create mode 100644 litellm/_auth/common_budgets.py create mode 100644 litellm/_auth/common_flow.py create mode 100644 litellm/_auth/credential.py create mode 100644 litellm/_auth/decision.py create mode 100644 litellm/_auth/driver.py create mode 100644 litellm/_auth/end_user_context.py create mode 100644 litellm/_auth/failure.py create mode 100644 litellm/_auth/fallback_access.py create mode 100644 litellm/_auth/fallback_candidate.py create mode 100644 litellm/_auth/key.py create mode 100644 litellm/_auth/key_model_access.py create mode 100644 litellm/_auth/managed.py create mode 100644 litellm/_auth/managed_models.py create mode 100644 litellm/_auth/member_models.py create mode 100644 litellm/_auth/model_access.py create mode 100644 litellm/_auth/model_attribution.py create mode 100644 litellm/_auth/model_budget_aliases.py create mode 100644 litellm/_auth/model_budget_data.py create mode 100644 litellm/_auth/model_groups.py create mode 100644 litellm/_auth/model_helpers.py create mode 100644 litellm/_auth/model_projection.py create mode 100644 litellm/_auth/org_routes.py create mode 100644 litellm/_auth/owner_teams.py create mode 100644 litellm/_auth/personal_budget.py create mode 100644 litellm/_auth/personal_model_listing.py create mode 100644 litellm/_auth/personal_model_request.py create mode 100644 litellm/_auth/personal_models.py create mode 100644 litellm/_auth/personal_owner.py create mode 100644 litellm/_auth/reservation.py create mode 100644 litellm/_auth/resolved_models.py create mode 100644 litellm/_auth/route_enable.py create mode 100644 litellm/_auth/route_roles.py create mode 100644 litellm/_auth/route_scope.py create mode 100644 litellm/_auth/routes.py create mode 100644 litellm/_auth/runtime.py create mode 100644 litellm/_auth/rust_core.py create mode 100644 litellm/_auth/spend_counter_policy.py create mode 100644 litellm/_auth/team_aliases.py create mode 100644 litellm/_auth/team_override.py create mode 100644 litellm/_auth/unpriced_models.py create mode 100644 litellm/proxy/auth/agent_ceiling_adapter.py create mode 100644 litellm/proxy/auth/auth_denial.py create mode 100644 litellm/proxy/auth/budget_alert_adapter.py create mode 100644 litellm/proxy/auth/budget_context_adapter.py create mode 100644 litellm/proxy/auth/budget_lookup_adapter.py create mode 100644 litellm/proxy/auth/budget_pricing_adapter.py create mode 100644 litellm/proxy/auth/budget_windows_adapter.py create mode 100644 litellm/proxy/auth/claude_normalization_adapter.py create mode 100644 litellm/proxy/auth/codex_model_adapter.py create mode 100644 litellm/proxy/auth/codex_normalization_adapter.py create mode 100644 litellm/proxy/auth/codex_turn_adapter.py create mode 100644 litellm/proxy/auth/common_budget_adapter.py create mode 100644 litellm/proxy/auth/common_budget_lookups.py create mode 100644 litellm/proxy/auth/end_user_context_adapter.py create mode 100644 litellm/proxy/auth/end_user_context_io.py create mode 100644 litellm/proxy/auth/fallback_access_adapter.py create mode 100644 litellm/proxy/auth/fallback_candidate_adapter.py create mode 100644 litellm/proxy/auth/key_budget_context_adapter.py create mode 100644 litellm/proxy/auth/key_model_access_adapter.py create mode 100644 litellm/proxy/auth/managed_model_adapter.py create mode 100644 litellm/proxy/auth/member_model_adapter.py create mode 100644 litellm/proxy/auth/model_access_adapter.py create mode 100644 litellm/proxy/auth/model_attribution_adapter.py create mode 100644 litellm/proxy/auth/model_budget_adapter.py create mode 100644 litellm/proxy/auth/model_group_adapter.py create mode 100644 litellm/proxy/auth/model_helper_adapter.py create mode 100644 litellm/proxy/auth/model_policy_read.py create mode 100644 litellm/proxy/auth/model_projection_adapter.py create mode 100644 litellm/proxy/auth/org_route_adapter.py create mode 100644 litellm/proxy/auth/owner_team_adapter.py create mode 100644 litellm/proxy/auth/personal_budget_adapter.py create mode 100644 litellm/proxy/auth/personal_model_adapter.py create mode 100644 litellm/proxy/auth/personal_model_listing_adapter.py create mode 100644 litellm/proxy/auth/personal_model_request_adapter.py create mode 100644 litellm/proxy/auth/personal_owner_adapter.py create mode 100644 litellm/proxy/auth/resolved_model_adapter.py create mode 100644 litellm/proxy/auth/route_enable_adapter.py create mode 100644 litellm/proxy/auth/route_policy_adapter.py create mode 100644 litellm/proxy/auth/route_role_adapter.py create mode 100644 litellm/proxy/auth/route_scope_adapter.py create mode 100644 litellm/proxy/auth/spend_counter_adapter.py create mode 100644 litellm/proxy/auth/team_alias_adapter.py create mode 100644 litellm/proxy/auth/team_override_adapter.py create mode 100644 litellm/proxy/auth/unpriced_model_adapter.py create mode 100644 litellm/repositories/database_settings.py create mode 100644 litellm/repositories/native_auth_repository.py create mode 100644 tests/integration/authorization/test_native_auth_io.py create mode 100644 tests/integration/caching/test_auth_cache_native.py create mode 100644 tests/unit/_auth/__init__.py create mode 100644 tests/unit/_auth/budget_reference.py create mode 100644 tests/unit/_auth/conftest.py create mode 100644 tests/unit/_auth/model_alias_reference.py create mode 100644 tests/unit/_auth/raw_answers.py create mode 100644 tests/unit/_auth/raw_model_reference.py create mode 100644 tests/unit/_auth/reference_core.py create mode 100644 tests/unit/_auth/reference_model_budgets.py create mode 100644 tests/unit/_auth/reference_policy.py create mode 100644 tests/unit/_auth/shadow.py create mode 100644 tests/unit/_auth/support/claude_reference.py create mode 100644 tests/unit/_auth/support/codex_reference.py create mode 100644 tests/unit/_auth/support/fallback_reference.py create mode 100644 tests/unit/_auth/support/model_reference.py create mode 100644 tests/unit/_auth/support/model_registry_reference.py create mode 100644 tests/unit/_auth/support/route_enable_reference.py create mode 100644 tests/unit/_auth/support/route_reference.py create mode 100644 tests/unit/_auth/test_access.py create mode 100644 tests/unit/_auth/test_admission.py create mode 100644 tests/unit/_auth/test_agent_ceilings.py create mode 100644 tests/unit/_auth/test_agent_grants.py create mode 100644 tests/unit/_auth/test_auth_context.py create mode 100644 tests/unit/_auth/test_budget_pricing.py create mode 100644 tests/unit/_auth/test_claude_normalization.py create mode 100644 tests/unit/_auth/test_codex_models.py create mode 100644 tests/unit/_auth/test_codex_normalization.py create mode 100644 tests/unit/_auth/test_codex_turn.py create mode 100644 tests/unit/_auth/test_common_budgets.py create mode 100644 tests/unit/_auth/test_common_flow.py create mode 100644 tests/unit/_auth/test_credential.py create mode 100644 tests/unit/_auth/test_failure.py create mode 100644 tests/unit/_auth/test_fallback_access.py create mode 100644 tests/unit/_auth/test_fallback_candidate.py create mode 100644 tests/unit/_auth/test_key.py create mode 100644 tests/unit/_auth/test_key_model_access.py create mode 100644 tests/unit/_auth/test_managed.py create mode 100644 tests/unit/_auth/test_managed_models.py create mode 100644 tests/unit/_auth/test_member_models.py create mode 100644 tests/unit/_auth/test_model_access.py create mode 100644 tests/unit/_auth/test_model_attribution.py create mode 100644 tests/unit/_auth/test_model_budget_aliases.py create mode 100644 tests/unit/_auth/test_model_budget_data.py create mode 100644 tests/unit/_auth/test_model_budgets.py create mode 100644 tests/unit/_auth/test_model_groups.py create mode 100644 tests/unit/_auth/test_model_helpers.py create mode 100644 tests/unit/_auth/test_owner_teams.py create mode 100644 tests/unit/_auth/test_personal_budget.py create mode 100644 tests/unit/_auth/test_personal_model_listing.py create mode 100644 tests/unit/_auth/test_personal_model_request.py create mode 100644 tests/unit/_auth/test_personal_models.py create mode 100644 tests/unit/_auth/test_personal_owner.py create mode 100644 tests/unit/_auth/test_python_core.py create mode 100644 tests/unit/_auth/test_reservation.py create mode 100644 tests/unit/_auth/test_resolved_models.py create mode 100644 tests/unit/_auth/test_route_enable.py create mode 100644 tests/unit/_auth/test_route_scope.py create mode 100644 tests/unit/_auth/test_routes.py create mode 100644 tests/unit/_auth/test_runtime.py create mode 100644 tests/unit/_auth/test_rust_core.py create mode 100644 tests/unit/_auth/test_spend_counter_policy.py create mode 100644 tests/unit/_auth/test_team_aliases.py create mode 100644 tests/unit/_auth/test_team_override.py create mode 100644 tests/unit/_auth/test_unpriced_models.py create mode 100644 tests/unit/proxy/auth/budget_alert_reference.py create mode 100644 tests/unit/proxy/auth/budget_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_auth_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_context_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_lookup_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_route_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_throttle_reference.py create mode 100644 tests/unit/proxy/auth/route_role_reference.py create mode 100644 tests/unit/proxy/auth/test_budget_alert_adapter.py create mode 100644 tests/unit/proxy/auth/test_budget_context_adapter.py create mode 100644 tests/unit/proxy/auth/test_budget_lookup_adapter.py create mode 100644 tests/unit/proxy/auth/test_budget_pricing_adapter.py create mode 100644 tests/unit/proxy/auth/test_common_budget_adapter.py create mode 100644 tests/unit/proxy/auth/test_common_budget_lookups.py create mode 100644 tests/unit/proxy/auth/test_end_user_context_adapter.py create mode 100644 tests/unit/proxy/auth/test_key_budget_context_adapter.py create mode 100644 tests/unit/proxy/auth/test_managed_model_adapter.py create mode 100644 tests/unit/proxy/auth/test_model_budget_adapter.py create mode 100644 tests/unit/proxy/auth/test_model_policy_read.py create mode 100644 tests/unit/proxy/auth/test_model_projection_adapter.py create mode 100644 tests/unit/proxy/auth/test_org_route_adapter.py create mode 100644 tests/unit/proxy/auth/test_route_enable_adapter.py create mode 100644 tests/unit/proxy/auth/test_route_role_adapter.py create mode 100644 tests/unit/proxy/auth/test_spend_counter_adapter.py create mode 100644 tests/unit/repositories/test_database_settings.py create mode 100644 tests/unit/repositories/test_native_auth_repository.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index fb29a92e306..1d56f8b4a29 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -475,6 +475,7 @@ jobs: - shard: rust-bridge-harness test-path: >- + tests/unit/_auth tests/test_litellm_rust/test_dispatch.py tests/test_litellm_rust/test_public_entries.py tests/test_litellm_rust/chat_completions/test_dispatch.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 48b6730b0c6..7565f7911fa 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3768,6 +3768,8 @@ name = "litellm-db" version = "0.1.0" dependencies = [ "chrono", + "futures-util", + "litellm-secret-value", "macro_rules_attribute", "rstest", "serde", @@ -3776,6 +3778,7 @@ dependencies = [ "strum", "thiserror 2.0.19", "tokio", + "veil", ] [[package]] @@ -4389,6 +4392,7 @@ dependencies = [ "litellm-callbacks-legacy-python", "litellm-circuit-breaker", "litellm-core-utils", + "litellm-db", "litellm-error", "litellm-host", "litellm-host-python", @@ -4408,6 +4412,7 @@ dependencies = [ "litellm-redis-resolve", "litellm-redis-store", "litellm-redis-types", + "litellm-request-auth-core", "litellm-router-core", "litellm-router-types", "litellm-secret-value", @@ -4534,6 +4539,7 @@ dependencies = [ "litellm-redis-testing", "litellm-secret-value", "litellm-tracing", + "moka", "rstest", "strum", "thiserror 2.0.19", @@ -4562,6 +4568,20 @@ dependencies = [ "url", ] +[[package]] +name = "litellm-request-auth-core" +version = "0.1.0" +dependencies = [ + "fancy-regex 0.19.2", + "litellm-secret-value", + "rstest", + "serde", + "sha2 0.10.9", + "strum", + "thiserror 2.0.19", + "veil", +] + [[package]] name = "litellm-router" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 72d6694ba79..e8a7ade51e8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -9,6 +9,7 @@ license = "MIT" repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] +litellm-request-auth-core = { path = "crates/request-auth-core" } litellm-config = { path = "crates/config" } litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } diff --git a/litellm-rust/crates/db/.sqlx/query-00ddffb253f55705f0b742827f9de6d938c46a3108c2385c52d80e86c4473b96.json b/litellm-rust/crates/db/.sqlx/query-00ddffb253f55705f0b742827f9de6d938c46a3108c2385c52d80e86c4473b96.json new file mode 100644 index 00000000000..eeacb0381a9 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-00ddffb253f55705f0b742827f9de6d938c46a3108c2385c52d80e86c4473b96.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_TeamTable\" WHERE team_id=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "00ddffb253f55705f0b742827f9de6d938c46a3108c2385c52d80e86c4473b96" +} diff --git a/litellm-rust/crates/db/.sqlx/query-0136e6b22d40ef0181f3f122c5abc5f270be24cd01b8c39bceea2de70d44608b.json b/litellm-rust/crates/db/.sqlx/query-0136e6b22d40ef0181f3f122c5abc5f270be24cd01b8c39bceea2de70d44608b.json new file mode 100644 index 00000000000..81c61f458ba --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-0136e6b22d40ef0181f3f122c5abc5f270be24cd01b8c39bceea2de70d44608b.json @@ -0,0 +1,24 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS \"data!: AuthData\" FROM \"LiteLLM_SpendLogs\" WHERE \"team_id\"=$1 AND \"startTime\" >= $2\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Timestamp" + ] + }, + "nullable": [ + null + ] + }, + "hash": "0136e6b22d40ef0181f3f122c5abc5f270be24cd01b8c39bceea2de70d44608b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-16fec3e3896126e8bd14726b346a2a45bd52705356fbf94128776db52032fd05.json b/litellm-rust/crates/db/.sqlx/query-16fec3e3896126e8bd14726b346a2a45bd52705356fbf94128776db52032fd05.json new file mode 100644 index 00000000000..cf809ef03de --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-16fec3e3896126e8bd14726b346a2a45bd52705356fbf94128776db52032fd05.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(a) AS \"data!: AuthData\"\nFROM \"LiteLLM_RetiredAgent\" a\nWHERE a.original_agent_id = ANY($1::text[])\nORDER BY a.original_agent_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "16fec3e3896126e8bd14726b346a2a45bd52705356fbf94128776db52032fd05" +} diff --git a/litellm-rust/crates/db/.sqlx/query-1d689acb624e95e13cdadde37817398561f205596a12e7f6eb71dd5ecbf7d49a.json b/litellm-rust/crates/db/.sqlx/query-1d689acb624e95e13cdadde37817398561f205596a12e7f6eb71dd5ecbf7d49a.json new file mode 100644 index 00000000000..005f05f2f79 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-1d689acb624e95e13cdadde37817398561f205596a12e7f6eb71dd5ecbf7d49a.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_OrganizationTable\" WHERE organization_id=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "1d689acb624e95e13cdadde37817398561f205596a12e7f6eb71dd5ecbf7d49a" +} diff --git a/litellm-rust/crates/db/.sqlx/query-1f9b9a6524e417b18865fda7b49c78c7c6a44ada9e601f6c74af1e9debd06f21.json b/litellm-rust/crates/db/.sqlx/query-1f9b9a6524e417b18865fda7b49c78c7c6a44ada9e601f6c74af1e9debd06f21.json new file mode 100644 index 00000000000..6111597cc95 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-1f9b9a6524e417b18865fda7b49c78c7c6a44ada9e601f6c74af1e9debd06f21.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_ModelAccessGroupBudgetTable\" WHERE access_group_name=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "1f9b9a6524e417b18865fda7b49c78c7c6a44ada9e601f6c74af1e9debd06f21" +} diff --git a/litellm-rust/crates/db/.sqlx/query-21e8660e197f5f9132584f22da32c8ef734194248e2a078af73776398a256090.json b/litellm-rust/crates/db/.sqlx/query-21e8660e197f5f9132584f22da32c8ef734194248e2a078af73776398a256090.json new file mode 100644 index 00000000000..e8ae4311441 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-21e8660e197f5f9132584f22da32c8ef734194248e2a078af73776398a256090.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(g) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = g.budget_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_ModelAccessGroupBudgetTable\" g\nWHERE g.access_group_name = ANY($1::text[])\nORDER BY g.access_group_name\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "21e8660e197f5f9132584f22da32c8ef734194248e2a078af73776398a256090" +} diff --git a/litellm-rust/crates/db/.sqlx/query-26a5fa1687c580afc0e1c1f7ca71b53e073a2b8f1e5b9007793b1c89df1ae24f.json b/litellm-rust/crates/db/.sqlx/query-26a5fa1687c580afc0e1c1f7ca71b53e073a2b8f1e5b9007793b1c89df1ae24f.json new file mode 100644 index 00000000000..f4ae0464445 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-26a5fa1687c580afc0e1c1f7ca71b53e073a2b8f1e5b9007793b1c89df1ae24f.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),\n 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM \"LiteLLM_OrganizationMembership\" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS \"data!: AuthData\"\nFROM \"LiteLLM_UserTable\" u\nWHERE u.user_id = ANY($1::text[])\nORDER BY u.user_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "26a5fa1687c580afc0e1c1f7ca71b53e073a2b8f1e5b9007793b1c89df1ae24f" +} diff --git a/litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json b/litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json new file mode 100644 index 00000000000..3c8e9b5c55f --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(v) AS \"data!: AuthData\"\nFROM \"LiteLLM_ManagedVectorStoresTable\" v\nWHERE v.vector_store_id = ANY($1::text[])\nORDER BY v.vector_store_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be" +} diff --git a/litellm-rust/crates/db/.sqlx/query-2c4e1f73d4170dd3c6a2b077fbd837fa0adb227764bcaff1c2267f9a4624f99c.json b/litellm-rust/crates/db/.sqlx/query-2c4e1f73d4170dd3c6a2b077fbd837fa0adb227764bcaff1c2267f9a4624f99c.json new file mode 100644 index 00000000000..f9dac05e937 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-2c4e1f73d4170dd3c6a2b077fbd837fa0adb227764bcaff1c2267f9a4624f99c.json @@ -0,0 +1,24 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_TeamMembership\" WHERE user_id=$1 AND team_id=$2\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "2c4e1f73d4170dd3c6a2b077fbd837fa0adb227764bcaff1c2267f9a4624f99c" +} diff --git a/litellm-rust/crates/db/.sqlx/query-31061f5f7b219176565eb4b3ebe723f7346cf3192e7e5eb6d54f0e4b444b20e7.json b/litellm-rust/crates/db/.sqlx/query-31061f5f7b219176565eb4b3ebe723f7346cf3192e7e5eb6d54f0e4b444b20e7.json new file mode 100644 index 00000000000..1de9e164903 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-31061f5f7b219176565eb4b3ebe723f7346cf3192e7e5eb6d54f0e4b444b20e7.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = o.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = o.object_permission_id)) AS \"data!: AuthData\" FROM \"LiteLLM_OrganizationTable\" o WHERE o.organization_alias = $1 ORDER BY o.organization_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "31061f5f7b219176565eb4b3ebe723f7346cf3192e7e5eb6d54f0e4b444b20e7" +} diff --git a/litellm-rust/crates/db/.sqlx/query-3210aa55f379884ca4232cf2a0b20aac76ac927183a15738d776ff9e104375af.json b/litellm-rust/crates/db/.sqlx/query-3210aa55f379884ca4232cf2a0b20aac76ac927183a15738d776ff9e104375af.json new file mode 100644 index 00000000000..ed1160e09dd --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-3210aa55f379884ca4232cf2a0b20aac76ac927183a15738d776ff9e104375af.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = o.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = o.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_OrganizationTable\" o\nWHERE o.organization_id = ANY($1::text[])\nORDER BY o.organization_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "3210aa55f379884ca4232cf2a0b20aac76ac927183a15738d776ff9e104375af" +} diff --git a/litellm-rust/crates/db/.sqlx/query-34107587736bee7d817b940717c26d0ffba418b154ae4e511287256db5168177.json b/litellm-rust/crates/db/.sqlx/query-34107587736bee7d817b940717c26d0ffba418b154ae4e511287256db5168177.json new file mode 100644 index 00000000000..0531ad4f9f1 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-34107587736bee7d817b940717c26d0ffba418b154ae4e511287256db5168177.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(e) || jsonb_build_object('models', COALESCE(to_jsonb(e.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = e.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = e.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_EndUserTable\" e\nWHERE e.user_id = ANY($1::text[])\nORDER BY e.user_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "34107587736bee7d817b940717c26d0ffba418b154ae4e511287256db5168177" +} diff --git a/litellm-rust/crates/db/.sqlx/query-36c30969226b767cf7d4eec854ccf1ff4c7fb7a158a0af6694c4cc3da0d66dd1.json b/litellm-rust/crates/db/.sqlx/query-36c30969226b767cf7d4eec854ccf1ff4c7fb7a158a0af6694c4cc3da0d66dd1.json new file mode 100644 index 00000000000..33901b021df --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-36c30969226b767cf7d4eec854ccf1ff4c7fb7a158a0af6694c4cc3da0d66dd1.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),\n 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM \"LiteLLM_OrganizationMembership\" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS \"data!: AuthData\" FROM \"LiteLLM_UserTable\" u WHERE u.sso_user_id = $1 ORDER BY u.user_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "36c30969226b767cf7d4eec854ccf1ff4c7fb7a158a0af6694c4cc3da0d66dd1" +} diff --git a/litellm-rust/crates/db/.sqlx/query-3b2157705b66ebcf6ee0c9805f15926a45b69c60677f527f227a2f8a58020b5b.json b/litellm-rust/crates/db/.sqlx/query-3b2157705b66ebcf6ee0c9805f15926a45b69c60677f527f227a2f8a58020b5b.json new file mode 100644 index 00000000000..c98477362f7 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-3b2157705b66ebcf6ee0c9805f15926a45b69c60677f527f227a2f8a58020b5b.json @@ -0,0 +1,25 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb),\n 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) AS \"data!: AuthData\" FROM \"LiteLLM_AgentIdentity\" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "3b2157705b66ebcf6ee0c9805f15926a45b69c60677f527f227a2f8a58020b5b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-498268b2d6edf244dfbe6d2c33327eae2c4d796d2891f6f3816289f8f71d7ea1.json b/litellm-rust/crates/db/.sqlx/query-498268b2d6edf244dfbe6d2c33327eae2c4d796d2891f6f3816289f8f71d7ea1.json new file mode 100644 index 00000000000..7bb00154919 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-498268b2d6edf244dfbe6d2c33327eae2c4d796d2891f6f3816289f8f71d7ea1.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_TagTable\" WHERE tag_name=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "498268b2d6edf244dfbe6d2c33327eae2c4d796d2891f6f3816289f8f71d7ea1" +} diff --git a/litellm-rust/crates/db/.sqlx/query-4ddd384812d4cc34d601873e73294c7e1bafada2968e53f4ecd16cd49d82cd5b.json b/litellm-rust/crates/db/.sqlx/query-4ddd384812d4cc34d601873e73294c7e1bafada2968e53f4ecd16cd49d82cd5b.json new file mode 100644 index 00000000000..a373a8246bf --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-4ddd384812d4cc34d601873e73294c7e1bafada2968e53f4ecd16cd49d82cd5b.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) AS \"data!: AuthData\"\nFROM \"LiteLLM_ObjectPermissionTable\" op\nWHERE op.object_permission_id = ANY($1::text[])\nORDER BY op.object_permission_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "4ddd384812d4cc34d601873e73294c7e1bafada2968e53f4ecd16cd49d82cd5b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-56a5d90aee114ff352a31fbec84ee9dd8dcb36772073095414df0d4493a3e363.json b/litellm-rust/crates/db/.sqlx/query-56a5d90aee114ff352a31fbec84ee9dd8dcb36772073095414df0d4493a3e363.json new file mode 100644 index 00000000000..26c9f609314 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-56a5d90aee114ff352a31fbec84ee9dd8dcb36772073095414df0d4493a3e363.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_ProjectTable\" WHERE project_id=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "56a5d90aee114ff352a31fbec84ee9dd8dcb36772073095414df0d4493a3e363" +} diff --git a/litellm-rust/crates/db/.sqlx/query-57d69d9d2ad5803d123717bd788223c5ce05000eaf8cb6888a66033a94fe7a09.json b/litellm-rust/crates/db/.sqlx/query-57d69d9d2ad5803d123717bd788223c5ce05000eaf8cb6888a66033a94fe7a09.json new file mode 100644 index 00000000000..d6587182619 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-57d69d9d2ad5803d123717bd788223c5ce05000eaf8cb6888a66033a94fe7a09.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_VerificationToken\" WHERE token=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "57d69d9d2ad5803d123717bd788223c5ce05000eaf8cb6888a66033a94fe7a09" +} diff --git a/litellm-rust/crates/db/.sqlx/query-65e9b4dc09101e63092902802978e012ff43436560342357afabbc0d5b961c44.json b/litellm-rust/crates/db/.sqlx/query-65e9b4dc09101e63092902802978e012ff43436560342357afabbc0d5b961c44.json new file mode 100644 index 00000000000..f18945c35ce --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-65e9b4dc09101e63092902802978e012ff43436560342357afabbc0d5b961c44.json @@ -0,0 +1,24 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = tm.budget_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_TeamMembership\" tm WHERE tm.user_id=$1 AND tm.team_id=$2\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "65e9b4dc09101e63092902802978e012ff43436560342357afabbc0d5b961c44" +} diff --git a/litellm-rust/crates/db/.sqlx/query-667556bca77708a5c027641bb3be14a6dcec24dbe6c5063ef9ef75cb196d886c.json b/litellm-rust/crates/db/.sqlx/query-667556bca77708a5c027641bb3be14a6dcec24dbe6c5063ef9ef75cb196d886c.json new file mode 100644 index 00000000000..46ac6b1664d --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-667556bca77708a5c027641bb3be14a6dcec24dbe6c5063ef9ef75cb196d886c.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(p) || jsonb_build_object('models', COALESCE(to_jsonb(p.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = p.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = p.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_ProjectTable\" p\nWHERE p.project_id = ANY($1::text[])\nORDER BY p.project_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "667556bca77708a5c027641bb3be14a6dcec24dbe6c5063ef9ef75cb196d886c" +} diff --git a/litellm-rust/crates/db/.sqlx/query-683dd0c1cdf39792c5b0c7ca482fee4c56e187d9f9e3d26d625fc3432e9d5652.json b/litellm-rust/crates/db/.sqlx/query-683dd0c1cdf39792c5b0c7ca482fee4c56e187d9f9e3d26d625fc3432e9d5652.json new file mode 100644 index 00000000000..affb7aae13b --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-683dd0c1cdf39792c5b0c7ca482fee4c56e187d9f9e3d26d625fc3432e9d5652.json @@ -0,0 +1,25 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('window_start', window_start, 'spend', spend) AS \"data!: AuthData\" FROM \"LiteLLM_BudgetWindowSpend\" WHERE entity_type=$1 AND entity_id=$2 AND window_duration=$3\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "683dd0c1cdf39792c5b0c7ca482fee4c56e187d9f9e3d26d625fc3432e9d5652" +} diff --git a/litellm-rust/crates/db/.sqlx/query-6af954d0490c3cea037cd7bd00993f983adaeaf5ab4166a90ab4a89ae39e2b0b.json b/litellm-rust/crates/db/.sqlx/query-6af954d0490c3cea037cd7bd00993f983adaeaf5ab4166a90ab4a89ae39e2b0b.json new file mode 100644 index 00000000000..5b386933a28 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-6af954d0490c3cea037cd7bd00993f983adaeaf5ab4166a90ab4a89ae39e2b0b.json @@ -0,0 +1,25 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(i) AS \"data!: AuthData\" FROM \"LiteLLM_RetiredAgentIdentity\" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "6af954d0490c3cea037cd7bd00993f983adaeaf5ab4166a90ab4a89ae39e2b0b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-786700d34873fa39e57ca9b0c1bec5df311bb3953f9756a5181821af9339cdba.json b/litellm-rust/crates/db/.sqlx/query-786700d34873fa39e57ca9b0c1bec5df311bb3953f9756a5181821af9339cdba.json new file mode 100644 index 00000000000..61e92bc7ae4 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-786700d34873fa39e57ca9b0c1bec5df311bb3953f9756a5181821af9339cdba.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),\n 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),\n 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),\n 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM \"LiteLLM_ModelTable\" m WHERE m.id=t.model_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = t.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_TeamTable\" t\nWHERE t.team_id = ANY($1::text[])\nORDER BY t.team_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "786700d34873fa39e57ca9b0c1bec5df311bb3953f9756a5181821af9339cdba" +} diff --git a/litellm-rust/crates/db/.sqlx/query-847628007ea76e788caf20d35219860f7a2da7d3d6d40c209322e0893b99ac69.json b/litellm-rust/crates/db/.sqlx/query-847628007ea76e788caf20d35219860f7a2da7d3d6d40c209322e0893b99ac69.json new file mode 100644 index 00000000000..c5612f85346 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-847628007ea76e788caf20d35219860f7a2da7d3d6d40c209322e0893b99ac69.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "WITH resolved AS (\n SELECT COALESCE(v.token, d.active_token_id) AS token\n FROM UNNEST($1::text[]) AS requested(token)\n LEFT JOIN \"LiteLLM_VerificationToken\" v ON v.token=requested.token\n LEFT JOIN \"LiteLLM_DeprecatedVerificationToken\" d ON d.token=requested.token AND d.revoke_at > (clock_timestamp() AT TIME ZONE 'UTC')\n)\nSELECT to_jsonb(view) AS \"data!: AuthData\" FROM (\nSELECT v.*,\nt.spend AS team_spend,\nt.max_budget AS team_max_budget,\nt.soft_budget AS team_soft_budget,\nt.model_max_budget AS team_model_max_budget,\nt.tpm_limit AS team_tpm_limit,\nt.rpm_limit AS team_rpm_limit,\nt.tpd_limit AS team_tpd_limit,\nt.models AS team_models,\nt.metadata AS team_metadata,\nt.blocked AS team_blocked,\nt.team_alias AS team_alias,\nt.members_with_roles AS team_members_with_roles,\nt.object_permission_id AS team_object_permission_id,\nt.organization_id AS org_id,\np.project_alias AS project_alias,\ntm.spend AS team_member_spend,\nb_tm.tpm_limit AS team_member_tpm_limit,\nb_tm.rpm_limit AS team_member_rpm_limit,\nm.aliases AS team_model_aliases,\nb.max_budget AS litellm_budget_table_max_budget,\nb.tpm_limit AS litellm_budget_table_tpm_limit,\nb.rpm_limit AS litellm_budget_table_rpm_limit,\nb.tpd_limit AS litellm_budget_table_tpd_limit,\nb.model_max_budget AS litellm_budget_table_model_max_budget,\nb.soft_budget AS litellm_budget_table_soft_budget,\no.metadata AS organization_metadata,\no.organization_alias AS organization_alias,\nb2.max_budget AS organization_max_budget,\nb2.tpm_limit AS organization_tpm_limit,\nb2.rpm_limit AS organization_rpm_limit\nFROM \"LiteLLM_VerificationToken\" v\nLEFT JOIN \"LiteLLM_TeamTable\" t ON v.team_id=t.team_id\nLEFT JOIN \"LiteLLM_TeamMembership\" tm ON v.team_id=tm.team_id AND tm.user_id=v.user_id\nLEFT JOIN \"LiteLLM_BudgetTable\" b_tm ON tm.budget_id=b_tm.budget_id\nLEFT JOIN \"LiteLLM_ModelTable\" m ON t.model_id=m.id\nLEFT JOIN \"LiteLLM_BudgetTable\" b ON v.budget_id=b.budget_id\nLEFT JOIN \"LiteLLM_ProjectTable\" p ON v.project_id=p.project_id\nLEFT JOIN \"LiteLLM_OrganizationTable\" o ON v.organization_id=o.organization_id\nLEFT JOIN \"LiteLLM_BudgetTable\" b2 ON o.budget_id=b2.budget_id\nWHERE v.token IN (SELECT token FROM resolved)\nORDER BY v.token\n) view\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "847628007ea76e788caf20d35219860f7a2da7d3d6d40c209322e0893b99ac69" +} diff --git a/litellm-rust/crates/db/.sqlx/query-8637bd849446309e631a4eaa47f7b51a9910fb27160356d9cc5dd03eb16985b2.json b/litellm-rust/crates/db/.sqlx/query-8637bd849446309e631a4eaa47f7b51a9910fb27160356d9cc5dd03eb16985b2.json new file mode 100644 index 00000000000..9479b08f1b8 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-8637bd849446309e631a4eaa47f7b51a9910fb27160356d9cc5dd03eb16985b2.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('tag_name', tag_name) AS \"data!: AuthData\" FROM \"LiteLLM_TagTable\" ORDER BY tag_name LIMIT $1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Int8" + ] + }, + "nullable": [ + null + ] + }, + "hash": "8637bd849446309e631a4eaa47f7b51a9910fb27160356d9cc5dd03eb16985b2" +} diff --git a/litellm-rust/crates/db/.sqlx/query-8d8ac844d36122a8e652c494519d5b7f5ea9ad7a158c5824ab8fc42e1e67dbdb.json b/litellm-rust/crates/db/.sqlx/query-8d8ac844d36122a8e652c494519d5b7f5ea9ad7a158c5824ab8fc42e1e67dbdb.json new file mode 100644 index 00000000000..1a8687649fc --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-8d8ac844d36122a8e652c494519d5b7f5ea9ad7a158c5824ab8fc42e1e67dbdb.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(a) || jsonb_build_object('access_model_names', COALESCE(to_jsonb(a.access_model_names), '[]'::jsonb),\n 'access_mcp_server_ids', COALESCE(to_jsonb(a.access_mcp_server_ids), '[]'::jsonb),\n 'access_agent_ids', COALESCE(to_jsonb(a.access_agent_ids), '[]'::jsonb),\n 'access_skill_names', COALESCE(to_jsonb(a.access_skill_names), '[]'::jsonb),\n 'assigned_team_ids', COALESCE(to_jsonb(a.assigned_team_ids), '[]'::jsonb),\n 'assigned_key_ids', COALESCE(to_jsonb(a.assigned_key_ids), '[]'::jsonb)) AS \"data!: AuthData\"\nFROM \"LiteLLM_AccessGroupTable\" a\nWHERE a.access_group_id = ANY($1::text[])\nORDER BY a.access_group_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "8d8ac844d36122a8e652c494519d5b7f5ea9ad7a158c5824ab8fc42e1e67dbdb" +} diff --git a/litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json b/litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json new file mode 100644 index 00000000000..a27f0fd3ea7 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(a) || jsonb_build_object('extra_headers', COALESCE(to_jsonb(a.extra_headers), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(a.agent_access_groups), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(a.access_group_ids), '[]'::jsonb), 'identity', (SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb),\n 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) FROM \"LiteLLM_AgentIdentity\" i WHERE i.agent_id=a.agent_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = a.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_AgentsTable\" a\nWHERE a.agent_name = ANY($1::text[])\nORDER BY a.agent_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db" +} diff --git a/litellm-rust/crates/db/.sqlx/query-ab4b895f3a1ee6ab9e86f13b4f02b5719d6fcbc961b5234cd5ad3100674b293b.json b/litellm-rust/crates/db/.sqlx/query-ab4b895f3a1ee6ab9e86f13b4f02b5719d6fcbc961b5234cd5ad3100674b293b.json new file mode 100644 index 00000000000..02e10380b08 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-ab4b895f3a1ee6ab9e86f13b4f02b5719d6fcbc961b5234cd5ad3100674b293b.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(t) || jsonb_build_object('models', COALESCE(to_jsonb(t.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = t.budget_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_TagTable\" t\nWHERE t.tag_name = ANY($1::text[])\nORDER BY t.tag_name\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "ab4b895f3a1ee6ab9e86f13b4f02b5719d6fcbc961b5234cd5ad3100674b293b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-b0d27eac417a4044836ce7ba4450b4bd94727d6f882992ff4b8419be945175ef.json b/litellm-rust/crates/db/.sqlx/query-b0d27eac417a4044836ce7ba4450b4bd94727d6f882992ff4b8419be945175ef.json new file mode 100644 index 00000000000..8abd18e6d11 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b0d27eac417a4044836ce7ba4450b4bd94727d6f882992ff4b8419be945175ef.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),\n 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM \"LiteLLM_OrganizationMembership\" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS \"data!: AuthData\" FROM \"LiteLLM_UserTable\" u WHERE LOWER(u.user_email) = LOWER($1) ORDER BY u.user_id LIMIT 1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "b0d27eac417a4044836ce7ba4450b4bd94727d6f882992ff4b8419be945175ef" +} diff --git a/litellm-rust/crates/db/.sqlx/query-b23d9fe187171500ebd8593571e9ba9059a84dbd20d9e148c72f23a62f52a981.json b/litellm-rust/crates/db/.sqlx/query-b23d9fe187171500ebd8593571e9ba9059a84dbd20d9e148c72f23a62f52a981.json new file mode 100644 index 00000000000..0bf8c2f0bf7 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b23d9fe187171500ebd8593571e9ba9059a84dbd20d9e148c72f23a62f52a981.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) AS \"data!: AuthData\"\nFROM \"LiteLLM_BudgetTable\" b\nWHERE b.budget_id = ANY($1::text[])\nORDER BY b.budget_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "b23d9fe187171500ebd8593571e9ba9059a84dbd20d9e148c72f23a62f52a981" +} diff --git a/litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json b/litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json new file mode 100644 index 00000000000..4e904c41350 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('agent_id', a.agent_id) AS \"data!: AuthData\"\nFROM \"LiteLLM_AgentsTable\" a\nWHERE a.agent_access_groups && $1::text[]\nORDER BY a.agent_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28" +} diff --git a/litellm-rust/crates/db/.sqlx/query-b840de2b7f4c4dfc7c2909062a4494d92ea3bf7d09bfeb6f6aa30a7bff7da7bf.json b/litellm-rust/crates/db/.sqlx/query-b840de2b7f4c4dfc7c2909062a4494d92ea3bf7d09bfeb6f6aa30a7bff7da7bf.json new file mode 100644 index 00000000000..f09242f49eb --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b840de2b7f4c4dfc7c2909062a4494d92ea3bf7d09bfeb6f6aa30a7bff7da7bf.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_UserTable\" WHERE user_id=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "b840de2b7f4c4dfc7c2909062a4494d92ea3bf7d09bfeb6f6aa30a7bff7da7bf" +} diff --git a/litellm-rust/crates/db/.sqlx/query-c265904d0f03ab912de3cfd5adae562f532616ee7d81288ffdcc9ba34d4762c0.json b/litellm-rust/crates/db/.sqlx/query-c265904d0f03ab912de3cfd5adae562f532616ee7d81288ffdcc9ba34d4762c0.json new file mode 100644 index 00000000000..eb35bccaf20 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-c265904d0f03ab912de3cfd5adae562f532616ee7d81288ffdcc9ba34d4762c0.json @@ -0,0 +1,24 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS \"data!: AuthData\" FROM \"LiteLLM_SpendLogs\" WHERE \"api_key\"=$1 AND \"startTime\" >= $2\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Timestamp" + ] + }, + "nullable": [ + null + ] + }, + "hash": "c265904d0f03ab912de3cfd5adae562f532616ee7d81288ffdcc9ba34d4762c0" +} diff --git a/litellm-rust/crates/db/.sqlx/query-c720cb3a555cd750cd929003591e8b674266007e4b11e388f722c9418f8ccea0.json b/litellm-rust/crates/db/.sqlx/query-c720cb3a555cd750cd929003591e8b674266007e4b11e388f722c9418f8ccea0.json new file mode 100644 index 00000000000..775ad5d1ce0 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-c720cb3a555cd750cd929003591e8b674266007e4b11e388f722c9418f8ccea0.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(a) || jsonb_build_object('extra_headers', COALESCE(to_jsonb(a.extra_headers), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(a.agent_access_groups), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(a.access_group_ids), '[]'::jsonb), 'identity', (SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb),\n 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) FROM \"LiteLLM_AgentIdentity\" i WHERE i.agent_id=a.agent_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = a.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_AgentsTable\" a\nWHERE a.agent_id = ANY($1::text[])\nORDER BY a.agent_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "c720cb3a555cd750cd929003591e8b674266007e4b11e388f722c9418f8ccea0" +} diff --git a/litellm-rust/crates/db/.sqlx/query-c879e82114e82685c3613fa4bc82f48ade878aae01105b6d7181f95cd5a16864.json b/litellm-rust/crates/db/.sqlx/query-c879e82114e82685c3613fa4bc82f48ade878aae01105b6d7181f95cd5a16864.json new file mode 100644 index 00000000000..81d3e59c212 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-c879e82114e82685c3613fa4bc82f48ade878aae01105b6d7181f95cd5a16864.json @@ -0,0 +1,25 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(i) AS \"data!: AuthData\" FROM \"LiteLLM_VerifiedSubject\" i WHERE i.issuer=$1 AND i.tenant_id=$2 AND i.oid=$3\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "c879e82114e82685c3613fa4bc82f48ade878aae01105b6d7181f95cd5a16864" +} diff --git a/litellm-rust/crates/db/.sqlx/query-cc5bac434a523f8deb0416b806412422927cedd9981b6b3617aa61b1086ab8f9.json b/litellm-rust/crates/db/.sqlx/query-cc5bac434a523f8deb0416b806412422927cedd9981b6b3617aa61b1086ab8f9.json new file mode 100644 index 00000000000..89ba50a6a96 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-cc5bac434a523f8deb0416b806412422927cedd9981b6b3617aa61b1086ab8f9.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('access_group_name', access_group_name) AS \"data!: AuthData\" FROM \"LiteLLM_ModelAccessGroupBudgetTable\" ORDER BY access_group_name LIMIT $1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Int8" + ] + }, + "nullable": [ + null + ] + }, + "hash": "cc5bac434a523f8deb0416b806412422927cedd9981b6b3617aa61b1086ab8f9" +} diff --git a/litellm-rust/crates/db/.sqlx/query-d804605eae54eaeaa01c7f511cb33c8443365b9262f2827e1c0b09f715b0227b.json b/litellm-rust/crates/db/.sqlx/query-d804605eae54eaeaa01c7f511cb33c8443365b9262f2827e1c0b09f715b0227b.json new file mode 100644 index 00000000000..cf7d30e8ddd --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-d804605eae54eaeaa01c7f511cb33c8443365b9262f2827e1c0b09f715b0227b.json @@ -0,0 +1,27 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object(\n'user_row', (SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),\n 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM \"LiteLLM_OrganizationMembership\" om WHERE om.user_id = u.user_id), '[]'::jsonb)) FROM \"LiteLLM_UserTable\" u WHERE u.user_id=$1),\n 'team_row', (SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),\n 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),\n 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),\n 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM \"LiteLLM_ModelTable\" m WHERE m.id=t.model_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = t.object_permission_id)) FROM \"LiteLLM_TeamTable\" t WHERE t.team_id=$2),\n'membership_row', (SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = tm.budget_id)) FROM \"LiteLLM_TeamMembership\" tm WHERE tm.user_id=$3 AND tm.team_id=$2),\n'organization_row', (SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = o.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = o.object_permission_id)) FROM \"LiteLLM_OrganizationTable\" o WHERE o.organization_id=$4),\n'project_row', (SELECT to_jsonb(p) || jsonb_build_object('models', COALESCE(to_jsonb(p.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = p.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = p.object_permission_id)) FROM \"LiteLLM_ProjectTable\" p WHERE p.project_id=$5)) AS \"data!: AuthData\"\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text", + "Text", + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "d804605eae54eaeaa01c7f511cb33c8443365b9262f2827e1c0b09f715b0227b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-dbe6bf2326dcb74823f6bfb90149e5d79069036df1c90481aa17b5ee7c6e4293.json b/litellm-rust/crates/db/.sqlx/query-dbe6bf2326dcb74823f6bfb90149e5d79069036df1c90481aa17b5ee7c6e4293.json new file mode 100644 index 00000000000..99c7b7da991 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-dbe6bf2326dcb74823f6bfb90149e5d79069036df1c90481aa17b5ee7c6e4293.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_EndUserTable\" WHERE user_id=$1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend!", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "dbe6bf2326dcb74823f6bfb90149e5d79069036df1c90481aa17b5ee7c6e4293" +} diff --git a/litellm-rust/crates/db/.sqlx/query-e6b5061b738fd353df83eaef9dc02f70cecce44f8e5d6673d2521a57de3513d8.json b/litellm-rust/crates/db/.sqlx/query-e6b5061b738fd353df83eaef9dc02f70cecce44f8e5d6673d2521a57de3513d8.json new file mode 100644 index 00000000000..cdeb9486129 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-e6b5061b738fd353df83eaef9dc02f70cecce44f8e5d6673d2521a57de3513d8.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),\n 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),\n 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),\n 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM \"LiteLLM_ModelTable\" m WHERE m.id=t.model_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = t.object_permission_id)) AS \"data!: AuthData\" FROM \"LiteLLM_TeamTable\" t WHERE t.team_alias = $1 ORDER BY t.team_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "e6b5061b738fd353df83eaef9dc02f70cecce44f8e5d6673d2521a57de3513d8" +} diff --git a/litellm-rust/crates/db/.sqlx/query-f6ed52b7b550c41a8d085f1a66bd1dcafc1f0cf69e33cb58ebe8642b75ee8d36.json b/litellm-rust/crates/db/.sqlx/query-f6ed52b7b550c41a8d085f1a66bd1dcafc1f0cf69e33cb58ebe8642b75ee8d36.json new file mode 100644 index 00000000000..4516ed84ca0 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-f6ed52b7b550c41a8d085f1a66bd1dcafc1f0cf69e33cb58ebe8642b75ee8d36.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('user_id', user_id) AS \"data!: AuthData\" FROM \"LiteLLM_EndUserTable\"\nWHERE blocked OR budget_id IS NOT NULL OR allowed_model_region IS NOT NULL OR default_model IS NOT NULL OR object_permission_id IS NOT NULL OR cardinality(models) > 0\nORDER BY user_id LIMIT $1\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Int8" + ] + }, + "nullable": [ + null + ] + }, + "hash": "f6ed52b7b550c41a8d085f1a66bd1dcafc1f0cf69e33cb58ebe8642b75ee8d36" +} diff --git a/litellm-rust/crates/db/Cargo.toml b/litellm-rust/crates/db/Cargo.toml index 4f8fe4e28e2..d377eca2086 100644 --- a/litellm-rust/crates/db/Cargo.toml +++ b/litellm-rust/crates/db/Cargo.toml @@ -6,7 +6,11 @@ license.workspace = true repository.workspace = true [dependencies] +veil.workspace = true +litellm-secret-value.workspace = true +tokio.workspace = true chrono.workspace = true +futures-util.workspace = true macro_rules_attribute.workspace = true serde.workspace = true serde_json.workspace = true @@ -16,4 +20,3 @@ thiserror.workspace = true [dev-dependencies] rstest.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/db/queries/auth/access_group.sql b/litellm-rust/crates/db/queries/auth/access_group.sql new file mode 100644 index 00000000000..32da5a7a38e --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/access_group.sql @@ -0,0 +1,9 @@ +SELECT to_jsonb(a) || jsonb_build_object('access_model_names', COALESCE(to_jsonb(a.access_model_names), '[]'::jsonb), + 'access_mcp_server_ids', COALESCE(to_jsonb(a.access_mcp_server_ids), '[]'::jsonb), + 'access_agent_ids', COALESCE(to_jsonb(a.access_agent_ids), '[]'::jsonb), + 'access_skill_names', COALESCE(to_jsonb(a.access_skill_names), '[]'::jsonb), + 'assigned_team_ids', COALESCE(to_jsonb(a.assigned_team_ids), '[]'::jsonb), + 'assigned_key_ids', COALESCE(to_jsonb(a.assigned_key_ids), '[]'::jsonb)) AS "data!: AuthData" +FROM "LiteLLM_AccessGroupTable" a +WHERE a.access_group_id = ANY($1::text[]) +ORDER BY a.access_group_id diff --git a/litellm-rust/crates/db/queries/auth/agent.sql b/litellm-rust/crates/db/queries/auth/agent.sql new file mode 100644 index 00000000000..67f07f85279 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/agent.sql @@ -0,0 +1,16 @@ +SELECT to_jsonb(a) || jsonb_build_object('extra_headers', COALESCE(to_jsonb(a.extra_headers), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(a.agent_access_groups), '[]'::jsonb), + 'access_group_ids', COALESCE(to_jsonb(a.access_group_ids), '[]'::jsonb), 'identity', (SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb), + 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) FROM "LiteLLM_AgentIdentity" i WHERE i.agent_id=a.agent_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = a.object_permission_id)) AS "data!: AuthData" +FROM "LiteLLM_AgentsTable" a +WHERE a.agent_id = ANY($1::text[]) +ORDER BY a.agent_id diff --git a/litellm-rust/crates/db/queries/auth/agent_access_group.sql b/litellm-rust/crates/db/queries/auth/agent_access_group.sql new file mode 100644 index 00000000000..750669b5b4c --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/agent_access_group.sql @@ -0,0 +1,4 @@ +SELECT jsonb_build_object('agent_id', a.agent_id) AS "data!: AuthData" +FROM "LiteLLM_AgentsTable" a +WHERE a.agent_access_groups && $1::text[] +ORDER BY a.agent_id diff --git a/litellm-rust/crates/db/queries/auth/agent_identity.sql b/litellm-rust/crates/db/queries/auth/agent_identity.sql new file mode 100644 index 00000000000..f148045f02b --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/agent_identity.sql @@ -0,0 +1,2 @@ +SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb), + 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) AS "data!: AuthData" FROM "LiteLLM_AgentIdentity" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3 diff --git a/litellm-rust/crates/db/queries/auth/agent_name.sql b/litellm-rust/crates/db/queries/auth/agent_name.sql new file mode 100644 index 00000000000..01776e8d95e --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/agent_name.sql @@ -0,0 +1,16 @@ +SELECT to_jsonb(a) || jsonb_build_object('extra_headers', COALESCE(to_jsonb(a.extra_headers), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(a.agent_access_groups), '[]'::jsonb), + 'access_group_ids', COALESCE(to_jsonb(a.access_group_ids), '[]'::jsonb), 'identity', (SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb), + 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) FROM "LiteLLM_AgentIdentity" i WHERE i.agent_id=a.agent_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = a.object_permission_id)) AS "data!: AuthData" +FROM "LiteLLM_AgentsTable" a +WHERE a.agent_name = ANY($1::text[]) +ORDER BY a.agent_id diff --git a/litellm-rust/crates/db/queries/auth/budget.sql b/litellm-rust/crates/db/queries/auth/budget.sql new file mode 100644 index 00000000000..47f7cf37189 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/budget.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) AS "data!: AuthData" +FROM "LiteLLM_BudgetTable" b +WHERE b.budget_id = ANY($1::text[]) +ORDER BY b.budget_id diff --git a/litellm-rust/crates/db/queries/auth/budget_window.sql b/litellm-rust/crates/db/queries/auth/budget_window.sql new file mode 100644 index 00000000000..5f327055c5a --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/budget_window.sql @@ -0,0 +1 @@ +SELECT jsonb_build_object('window_start', window_start, 'spend', spend) AS "data!: AuthData" FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type=$1 AND entity_id=$2 AND window_duration=$3 diff --git a/litellm-rust/crates/db/queries/auth/end_user.sql b/litellm-rust/crates/db/queries/auth/end_user.sql new file mode 100644 index 00000000000..578d1528f7c --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/end_user.sql @@ -0,0 +1,13 @@ +SELECT to_jsonb(e) || jsonb_build_object('models', COALESCE(to_jsonb(e.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = e.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = e.object_permission_id)) AS "data!: AuthData" +FROM "LiteLLM_EndUserTable" e +WHERE e.user_id = ANY($1::text[]) +ORDER BY e.user_id diff --git a/litellm-rust/crates/db/queries/auth/end_user_restricted_registry.sql b/litellm-rust/crates/db/queries/auth/end_user_restricted_registry.sql new file mode 100644 index 00000000000..ac3e5708b77 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/end_user_restricted_registry.sql @@ -0,0 +1,3 @@ +SELECT jsonb_build_object('user_id', user_id) AS "data!: AuthData" FROM "LiteLLM_EndUserTable" +WHERE blocked OR budget_id IS NOT NULL OR allowed_model_region IS NOT NULL OR default_model IS NOT NULL OR object_permission_id IS NOT NULL OR cardinality(models) > 0 +ORDER BY user_id LIMIT $1 diff --git a/litellm-rust/crates/db/queries/auth/fixture.sql b/litellm-rust/crates/db/queries/auth/fixture.sql new file mode 100644 index 00000000000..f50348acc1b --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/fixture.sql @@ -0,0 +1,58 @@ +INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, tpm_limit, rpm_limit, allowed_models, created_by, updated_by) +VALUES ('auth-b', 0, 22, 33, NULL, 'auth-test', 'auth-test'), + ('auth-b-key', 4.5, 44, 55, ARRAY['key-model'], 'auth-test', 'auth-test'), + ('auth-b-org', 13, 66, 77, ARRAY['org-model'], 'auth-test', 'auth-test'); +INSERT INTO "LiteLLM_ObjectPermissionTable" (object_permission_id, models, mcp_servers, agents) +VALUES ('auth-op', ARRAY['allowed-model'], NULL, ARRAY['auth-agent']); +INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, object_permission_id, models, spend, created_by, updated_by) +VALUES ('auth-org', 'Auth Org', 'auth-b-org', 'auth-op', NULL, 13.5, 'auth-test', 'auth-test'); +INSERT INTO "LiteLLM_ModelTable" (id, aliases, created_by, updated_by) +VALUES (-922, '{"alias":"actual"}', 'auth-test', 'auth-test'); +INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, organization_id, model_id, object_permission_id, models, admins, members, spend, max_budget, blocked, members_with_roles) +VALUES ('auth-team', 'Auth Team', 'auth-org', -922, 'auth-op', ARRAY['team-model'], NULL, NULL, 2.5, 25, true, '[{"user_id":"auth-u","role":"admin"}]'), + ('auth-team-empty', 'Auth Empty', NULL, NULL, NULL, NULL, NULL, NULL, 0, NULL, false, '[]'); +INSERT INTO "LiteLLM_UserTable" (user_id, user_email, sso_user_id, models, teams, spend) +VALUES ('auth-u', 'Test_%@Example.com', 'auth-sso', NULL, ARRAY['auth-team'], 3.5), + ('auth-other', 'test-xx@example.com', NULL, ARRAY['other-model'], NULL, 0); +INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role, budget_id) +VALUES ('auth-u', 'auth-org', 'org_admin', 'auth-b'); +INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, budget_id) +VALUES ('auth-u', 'auth-team', 6.5, 'auth-b'); +INSERT INTO "LiteLLM_ProjectTable" (project_id, project_alias, team_id, budget_id, object_permission_id, models, spend, created_by, updated_by) +VALUES ('auth-project', 'Auth Project', 'auth-team', 'auth-b', 'auth-op', NULL, 8.5, 'auth-test', 'auth-test'); +INSERT INTO "LiteLLM_VerificationToken" (token, user_id, team_id, organization_id, project_id, budget_id, spend, models) +VALUES (repeat('a', 64), 'auth-u', 'auth-team', 'auth-org', 'auth-project', 'auth-b-key', 4.5, ARRAY['key-model']), + (repeat('e', 64), NULL, NULL, NULL, NULL, NULL, 0, ARRAY[]::text[]); +INSERT INTO "LiteLLM_DeprecatedVerificationToken" (id, token, active_token_id, revoke_at) +VALUES ('auth-deprecated', repeat('b', 64), repeat('a', 64), now() + INTERVAL '1 hour'), + ('auth-expired', repeat('c', 64), repeat('a', 64), now() - INTERVAL '1 hour'), + ('auth-chained', repeat('d', 64), repeat('b', 64), now() + INTERVAL '1 hour'), + ('auth-shadowed', repeat('a', 64), repeat('e', 64), now() + INTERVAL '1 hour'); +INSERT INTO "LiteLLM_EndUserTable" (user_id, models, budget_id, object_permission_id, spend, blocked) +VALUES ('auth-end', NULL, 'auth-b', 'auth-op', 10.5, false), + ('auth-unrestricted', NULL, NULL, NULL, 0, false); +INSERT INTO "LiteLLM_TagTable" (tag_name, models, budget_id, spend) +VALUES ('auth-tag', NULL, 'auth-b', 11.5); +INSERT INTO "LiteLLM_ModelAccessGroupBudgetTable" (access_group_name, budget_id, spend) +VALUES ('auth-model-group', 'auth-b', 12.5); +INSERT INTO "LiteLLM_AccessGroupTable" (access_group_id, access_group_name, access_model_names, access_agent_ids, access_mcp_server_ids) +VALUES ('auth-access', 'Auth Access', ARRAY['group-model'], ARRAY['auth-agent'], NULL); +INSERT INTO "LiteLLM_AgentsTable" (agent_id, agent_name, agent_card_params, object_permission_id, identity_managed, enabled, extra_headers, agent_access_groups, created_by, updated_by) +VALUES ('auth-agent', 'Auth Agent', '{}', 'auth-op', true, false, NULL, ARRAY['auth-group'], 'auth-test', 'auth-test'); +INSERT INTO "LiteLLM_ManagedVectorStoresTable" (vector_store_id, custom_llm_provider, litellm_params, team_id, updated_at) +VALUES ('auth-vector', 'openai', '{"api_base":"https://example.com"}', 'auth-team', now()); +INSERT INTO "LiteLLM_AgentIdentity" (agent_id, provider, issuer, tenant_id, client_id, revision, active, required_roles, required_scopes) +VALUES ('auth-agent', 'microsoft_entra', 'auth-issuer', 'auth-tenant', 'auth-client', 'auth-revision', false, NULL, ARRAY['read']); +INSERT INTO "LiteLLM_RetiredAgentIdentity" (binding_id, provider, issuer, tenant_id, client_id) +VALUES ('auth-retired-binding', 'microsoft_entra', 'auth-issuer', 'auth-tenant', 'auth-retired-client'); +INSERT INTO "LiteLLM_RetiredAgent" (original_agent_id) VALUES ('auth-retired-agent'); +INSERT INTO "LiteLLM_VerifiedSubject" (subject_id, issuer, tenant_id, oid, user_id) +VALUES ('auth-subject', 'auth-issuer', 'auth-tenant', 'auth-oid', 'auth-u'); +INSERT INTO "LiteLLM_BudgetWindowSpend" (entity_type, entity_id, window_duration, window_start, spend, updated_at) +VALUES ('key', repeat('a',64), '1h', '2025-01-01 01:00:00', 2.5, now()), + ('team', 'auth-team', '1h', '2025-01-01 01:00:00', 2.5, now()); +INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, team_id, "startTime", "endTime", spend) +VALUES ('auth-log-before', 'completion', repeat('a',64), 'auth-team', '2025-01-01 00:59:59', '2025-01-01 01:00:01', 100), + ('auth-log-boundary', 'completion', repeat('a',64), 'auth-team', '2025-01-01 01:00:00', '2025-01-01 01:00:01', 2.5), + ('auth-log-after', 'completion', repeat('a',64), 'auth-team', '2025-01-01 01:00:01', '2025-01-01 01:00:02', 3.5), + ('auth-log-other', 'completion', repeat('e',64), 'auth-team-empty', '2025-01-01 01:00:00', '2025-01-01 01:00:01', 500); diff --git a/litellm-rust/crates/db/queries/auth/key.sql b/litellm-rust/crates/db/queries/auth/key.sql new file mode 100644 index 00000000000..c210fc6fcc4 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/key.sql @@ -0,0 +1,50 @@ +WITH resolved AS ( + SELECT COALESCE(v.token, d.active_token_id) AS token + FROM UNNEST($1::text[]) AS requested(token) + LEFT JOIN "LiteLLM_VerificationToken" v ON v.token=requested.token + LEFT JOIN "LiteLLM_DeprecatedVerificationToken" d ON d.token=requested.token AND d.revoke_at > (clock_timestamp() AT TIME ZONE 'UTC') +) +SELECT to_jsonb(view) AS "data!: AuthData" FROM ( +SELECT v.*, +t.spend AS team_spend, +t.max_budget AS team_max_budget, +t.soft_budget AS team_soft_budget, +t.model_max_budget AS team_model_max_budget, +t.tpm_limit AS team_tpm_limit, +t.rpm_limit AS team_rpm_limit, +t.tpd_limit AS team_tpd_limit, +t.models AS team_models, +t.metadata AS team_metadata, +t.blocked AS team_blocked, +t.team_alias AS team_alias, +t.members_with_roles AS team_members_with_roles, +t.object_permission_id AS team_object_permission_id, +t.organization_id AS org_id, +p.project_alias AS project_alias, +tm.spend AS team_member_spend, +b_tm.tpm_limit AS team_member_tpm_limit, +b_tm.rpm_limit AS team_member_rpm_limit, +m.aliases AS team_model_aliases, +b.max_budget AS litellm_budget_table_max_budget, +b.tpm_limit AS litellm_budget_table_tpm_limit, +b.rpm_limit AS litellm_budget_table_rpm_limit, +b.tpd_limit AS litellm_budget_table_tpd_limit, +b.model_max_budget AS litellm_budget_table_model_max_budget, +b.soft_budget AS litellm_budget_table_soft_budget, +o.metadata AS organization_metadata, +o.organization_alias AS organization_alias, +b2.max_budget AS organization_max_budget, +b2.tpm_limit AS organization_tpm_limit, +b2.rpm_limit AS organization_rpm_limit +FROM "LiteLLM_VerificationToken" v +LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id=t.team_id +LEFT JOIN "LiteLLM_TeamMembership" tm ON v.team_id=tm.team_id AND tm.user_id=v.user_id +LEFT JOIN "LiteLLM_BudgetTable" b_tm ON tm.budget_id=b_tm.budget_id +LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id=m.id +LEFT JOIN "LiteLLM_BudgetTable" b ON v.budget_id=b.budget_id +LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id=p.project_id +LEFT JOIN "LiteLLM_OrganizationTable" o ON v.organization_id=o.organization_id +LEFT JOIN "LiteLLM_BudgetTable" b2 ON o.budget_id=b2.budget_id +WHERE v.token IN (SELECT token FROM resolved) +ORDER BY v.token +) view diff --git a/litellm-rust/crates/db/queries/auth/managed_vector_store.sql b/litellm-rust/crates/db/queries/auth/managed_vector_store.sql new file mode 100644 index 00000000000..06765a45ac2 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/managed_vector_store.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(v) AS "data!: AuthData" +FROM "LiteLLM_ManagedVectorStoresTable" v +WHERE v.vector_store_id = ANY($1::text[]) +ORDER BY v.vector_store_id diff --git a/litellm-rust/crates/db/queries/auth/member.sql b/litellm-rust/crates/db/queries/auth/member.sql new file mode 100644 index 00000000000..d742f785c94 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/member.sql @@ -0,0 +1,2 @@ +SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = tm.budget_id)) AS "data!: AuthData" +FROM "LiteLLM_TeamMembership" tm WHERE tm.user_id=$1 AND tm.team_id=$2 diff --git a/litellm-rust/crates/db/queries/auth/model_access_group.sql b/litellm-rust/crates/db/queries/auth/model_access_group.sql new file mode 100644 index 00000000000..2d503c4d1ea --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/model_access_group.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(g) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = g.budget_id)) AS "data!: AuthData" +FROM "LiteLLM_ModelAccessGroupBudgetTable" g +WHERE g.access_group_name = ANY($1::text[]) +ORDER BY g.access_group_name diff --git a/litellm-rust/crates/db/queries/auth/model_access_group_registry.sql b/litellm-rust/crates/db/queries/auth/model_access_group_registry.sql new file mode 100644 index 00000000000..7bbc7c297b9 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/model_access_group_registry.sql @@ -0,0 +1 @@ +SELECT jsonb_build_object('access_group_name', access_group_name) AS "data!: AuthData" FROM "LiteLLM_ModelAccessGroupBudgetTable" ORDER BY access_group_name LIMIT $1 diff --git a/litellm-rust/crates/db/queries/auth/object_permission.sql b/litellm-rust/crates/db/queries/auth/object_permission.sql new file mode 100644 index 00000000000..e6617c76105 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/object_permission.sql @@ -0,0 +1,13 @@ +SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) AS "data!: AuthData" +FROM "LiteLLM_ObjectPermissionTable" op +WHERE op.object_permission_id = ANY($1::text[]) +ORDER BY op.object_permission_id diff --git a/litellm-rust/crates/db/queries/auth/org.sql b/litellm-rust/crates/db/queries/auth/org.sql new file mode 100644 index 00000000000..4ed266085d4 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/org.sql @@ -0,0 +1,13 @@ +SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = o.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = o.object_permission_id)) AS "data!: AuthData" +FROM "LiteLLM_OrganizationTable" o +WHERE o.organization_id = ANY($1::text[]) +ORDER BY o.organization_id diff --git a/litellm-rust/crates/db/queries/auth/org_alias.sql b/litellm-rust/crates/db/queries/auth/org_alias.sql new file mode 100644 index 00000000000..43bff1b3bd2 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/org_alias.sql @@ -0,0 +1,10 @@ +SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = o.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = o.object_permission_id)) AS "data!: AuthData" FROM "LiteLLM_OrganizationTable" o WHERE o.organization_alias = $1 ORDER BY o.organization_id diff --git a/litellm-rust/crates/db/queries/auth/prefetch.sql b/litellm-rust/crates/db/queries/auth/prefetch.sql new file mode 100644 index 00000000000..46ef3b36a11 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/prefetch.sql @@ -0,0 +1,42 @@ +SELECT jsonb_build_object( +'user_row', (SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb), + 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb), + 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb), + 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) FROM "LiteLLM_UserTable" u WHERE u.user_id=$1), + 'team_row', (SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb), + 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb), + 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb), + 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb), + 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb), + 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb), + 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM "LiteLLM_ModelTable" m WHERE m.id=t.model_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = t.object_permission_id)) FROM "LiteLLM_TeamTable" t WHERE t.team_id=$2), +'membership_row', (SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = tm.budget_id)) FROM "LiteLLM_TeamMembership" tm WHERE tm.user_id=$3 AND tm.team_id=$2), +'organization_row', (SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = o.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = o.object_permission_id)) FROM "LiteLLM_OrganizationTable" o WHERE o.organization_id=$4), +'project_row', (SELECT to_jsonb(p) || jsonb_build_object('models', COALESCE(to_jsonb(p.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = p.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = p.object_permission_id)) FROM "LiteLLM_ProjectTable" p WHERE p.project_id=$5)) AS "data!: AuthData" diff --git a/litellm-rust/crates/db/queries/auth/project.sql b/litellm-rust/crates/db/queries/auth/project.sql new file mode 100644 index 00000000000..563398011e7 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/project.sql @@ -0,0 +1,13 @@ +SELECT to_jsonb(p) || jsonb_build_object('models', COALESCE(to_jsonb(p.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = p.budget_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = p.object_permission_id)) AS "data!: AuthData" +FROM "LiteLLM_ProjectTable" p +WHERE p.project_id = ANY($1::text[]) +ORDER BY p.project_id diff --git a/litellm-rust/crates/db/queries/auth/retired_agent.sql b/litellm-rust/crates/db/queries/auth/retired_agent.sql new file mode 100644 index 00000000000..45ef5c36469 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/retired_agent.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(a) AS "data!: AuthData" +FROM "LiteLLM_RetiredAgent" a +WHERE a.original_agent_id = ANY($1::text[]) +ORDER BY a.original_agent_id diff --git a/litellm-rust/crates/db/queries/auth/retired_agent_identity.sql b/litellm-rust/crates/db/queries/auth/retired_agent_identity.sql new file mode 100644 index 00000000000..14847555364 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/retired_agent_identity.sql @@ -0,0 +1 @@ +SELECT to_jsonb(i) AS "data!: AuthData" FROM "LiteLLM_RetiredAgentIdentity" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3 diff --git a/litellm-rust/crates/db/queries/auth/spend_end_user.sql b/litellm-rust/crates/db/queries/auth/spend_end_user.sql new file mode 100644 index 00000000000..4321a4c4280 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_end_user.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_EndUserTable" WHERE user_id=$1 diff --git a/litellm-rust/crates/db/queries/auth/spend_key.sql b/litellm-rust/crates/db/queries/auth/spend_key.sql new file mode 100644 index 00000000000..42ff1986c5a --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_key.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_VerificationToken" WHERE token=$1 diff --git a/litellm-rust/crates/db/queries/auth/spend_member.sql b/litellm-rust/crates/db/queries/auth/spend_member.sql new file mode 100644 index 00000000000..885b0894d4c --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_member.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_TeamMembership" WHERE user_id=$1 AND team_id=$2 diff --git a/litellm-rust/crates/db/queries/auth/spend_model_access_group.sql b/litellm-rust/crates/db/queries/auth/spend_model_access_group.sql new file mode 100644 index 00000000000..cbc64093194 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_model_access_group.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_ModelAccessGroupBudgetTable" WHERE access_group_name=$1 diff --git a/litellm-rust/crates/db/queries/auth/spend_org.sql b/litellm-rust/crates/db/queries/auth/spend_org.sql new file mode 100644 index 00000000000..4a42cebf372 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_org.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_OrganizationTable" WHERE organization_id=$1 diff --git a/litellm-rust/crates/db/queries/auth/spend_project.sql b/litellm-rust/crates/db/queries/auth/spend_project.sql new file mode 100644 index 00000000000..ee1ad2e8d76 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_project.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_ProjectTable" WHERE project_id=$1 diff --git a/litellm-rust/crates/db/queries/auth/spend_tag.sql b/litellm-rust/crates/db/queries/auth/spend_tag.sql new file mode 100644 index 00000000000..a2b510baece --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_tag.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_TagTable" WHERE tag_name=$1 diff --git a/litellm-rust/crates/db/queries/auth/spend_team.sql b/litellm-rust/crates/db/queries/auth/spend_team.sql new file mode 100644 index 00000000000..c69d48520a4 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_team.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_TeamTable" WHERE team_id=$1 diff --git a/litellm-rust/crates/db/queries/auth/spend_user.sql b/litellm-rust/crates/db/queries/auth/spend_user.sql new file mode 100644 index 00000000000..1504e06c4a9 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/spend_user.sql @@ -0,0 +1 @@ +SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_UserTable" WHERE user_id=$1 diff --git a/litellm-rust/crates/db/queries/auth/tag.sql b/litellm-rust/crates/db/queries/auth/tag.sql new file mode 100644 index 00000000000..d056cf9d247 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/tag.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(t) || jsonb_build_object('models', COALESCE(to_jsonb(t.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = t.budget_id)) AS "data!: AuthData" +FROM "LiteLLM_TagTable" t +WHERE t.tag_name = ANY($1::text[]) +ORDER BY t.tag_name diff --git a/litellm-rust/crates/db/queries/auth/tag_registry.sql b/litellm-rust/crates/db/queries/auth/tag_registry.sql new file mode 100644 index 00000000000..430b4f3a273 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/tag_registry.sql @@ -0,0 +1 @@ +SELECT jsonb_build_object('tag_name', tag_name) AS "data!: AuthData" FROM "LiteLLM_TagTable" ORDER BY tag_name LIMIT $1 diff --git a/litellm-rust/crates/db/queries/auth/team.sql b/litellm-rust/crates/db/queries/auth/team.sql new file mode 100644 index 00000000000..9d0144bf5d3 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/team.sql @@ -0,0 +1,19 @@ +SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb), + 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb), + 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb), + 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb), + 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb), + 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb), + 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM "LiteLLM_ModelTable" m WHERE m.id=t.model_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = t.object_permission_id)) AS "data!: AuthData" +FROM "LiteLLM_TeamTable" t +WHERE t.team_id = ANY($1::text[]) +ORDER BY t.team_id diff --git a/litellm-rust/crates/db/queries/auth/team_alias.sql b/litellm-rust/crates/db/queries/auth/team_alias.sql new file mode 100644 index 00000000000..8ac1fa05bcb --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/team_alias.sql @@ -0,0 +1,16 @@ +SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb), + 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb), + 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb), + 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb), + 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb), + 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb), + 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM "LiteLLM_ModelTable" m WHERE m.id=t.model_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = t.object_permission_id)) AS "data!: AuthData" FROM "LiteLLM_TeamTable" t WHERE t.team_alias = $1 ORDER BY t.team_id diff --git a/litellm-rust/crates/db/queries/auth/user.sql b/litellm-rust/crates/db/queries/auth/user.sql new file mode 100644 index 00000000000..09132ecbd3d --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/user.sql @@ -0,0 +1,7 @@ +SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb), + 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb), + 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb), + 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS "data!: AuthData" +FROM "LiteLLM_UserTable" u +WHERE u.user_id = ANY($1::text[]) +ORDER BY u.user_id diff --git a/litellm-rust/crates/db/queries/auth/user_email.sql b/litellm-rust/crates/db/queries/auth/user_email.sql new file mode 100644 index 00000000000..cb5de48217b --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/user_email.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb), + 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb), + 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb), + 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS "data!: AuthData" FROM "LiteLLM_UserTable" u WHERE LOWER(u.user_email) = LOWER($1) ORDER BY u.user_id LIMIT 1 diff --git a/litellm-rust/crates/db/queries/auth/user_sso.sql b/litellm-rust/crates/db/queries/auth/user_sso.sql new file mode 100644 index 00000000000..28bf27d066a --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/user_sso.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb), + 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb), + 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb), + 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS "data!: AuthData" FROM "LiteLLM_UserTable" u WHERE u.sso_user_id = $1 ORDER BY u.user_id diff --git a/litellm-rust/crates/db/queries/auth/verified_subject.sql b/litellm-rust/crates/db/queries/auth/verified_subject.sql new file mode 100644 index 00000000000..ecf1440f820 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/verified_subject.sql @@ -0,0 +1 @@ +SELECT to_jsonb(i) AS "data!: AuthData" FROM "LiteLLM_VerifiedSubject" i WHERE i.issuer=$1 AND i.tenant_id=$2 AND i.oid=$3 diff --git a/litellm-rust/crates/db/queries/auth/window_key_spend.sql b/litellm-rust/crates/db/queries/auth/window_key_spend.sql new file mode 100644 index 00000000000..545abe8af8c --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/window_key_spend.sql @@ -0,0 +1 @@ +SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS "data!: AuthData" FROM "LiteLLM_SpendLogs" WHERE "api_key"=$1 AND "startTime" >= $2 diff --git a/litellm-rust/crates/db/queries/auth/window_team_spend.sql b/litellm-rust/crates/db/queries/auth/window_team_spend.sql new file mode 100644 index 00000000000..553873c857b --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/window_team_spend.sql @@ -0,0 +1 @@ +SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS "data!: AuthData" FROM "LiteLLM_SpendLogs" WHERE "team_id"=$1 AND "startTime" >= $2 diff --git a/litellm-rust/crates/db/src/auth.rs b/litellm-rust/crates/db/src/auth.rs new file mode 100644 index 00000000000..840b517120d --- /dev/null +++ b/litellm-rust/crates/db/src/auth.rs @@ -0,0 +1,814 @@ +use crate::{Error, executor::StatementPersistence}; +use sqlx::{Executor, Postgres}; + +#[derive(Clone, PartialEq, serde::Serialize, sqlx::Type, veil::Redact)] +#[serde(transparent)] +#[sqlx(transparent)] +pub struct AuthData(#[redact(with = "[REDACTED]")] pub serde_json::Value); + +#[macro_rules_attribute::apply(crate::db_row)] +pub struct AuthRow { + pub data: AuthData, +} + +#[derive(Clone, Copy, Debug, strum::EnumString, strum::VariantArray)] +pub enum AuthRead { + #[strum(serialize = "user")] + User, + #[strum(serialize = "team")] + Team, + #[strum(serialize = "org")] + Org, + #[strum(serialize = "project")] + Project, + #[strum(serialize = "end_user")] + EndUser, + #[strum(serialize = "budget")] + Budget, + #[strum(serialize = "object_permission")] + ObjectPermission, + #[strum(serialize = "access_group")] + AccessGroup, + #[strum(serialize = "managed_vector_store")] + ManagedVectorStore, + #[strum(serialize = "tag")] + Tag, + #[strum(serialize = "model_access_group")] + ModelAccessGroup, + #[strum(serialize = "agent")] + Agent, + #[strum(serialize = "agent_name")] + AgentName, + #[strum(serialize = "agent_access_group")] + AgentAccessGroup, + #[strum(serialize = "retired_agent")] + RetiredAgent, + #[strum(serialize = "key")] + Key, + #[strum(serialize = "member")] + Member, + #[strum(serialize = "agent_identity")] + AgentIdentity, + #[strum(serialize = "retired_agent_identity")] + RetiredAgentIdentity, + #[strum(serialize = "verified_subject")] + VerifiedSubject, + #[strum(serialize = "team_alias")] + TeamAlias, + #[strum(serialize = "org_alias")] + OrgAlias, + #[strum(serialize = "user_sso")] + UserSso, + #[strum(serialize = "user_email")] + UserEmail, + #[strum(serialize = "tag_registry")] + TagRegistry, + #[strum(serialize = "model_access_group_registry")] + ModelAccessGroupRegistry, + #[strum(serialize = "end_user_restricted_registry")] + EndUserRestrictedRegistry, + #[strum(serialize = "prefetch")] + Prefetch, + #[strum(serialize = "budget_window")] + BudgetWindow, + #[strum(serialize = "window_spend")] + WindowSpend, +} + +#[derive(Clone, Copy, Debug, strum::EnumString, strum::VariantArray)] +pub enum AuthSpend { + #[strum(serialize = "key")] + Key, + #[strum(serialize = "user")] + User, + #[strum(serialize = "team")] + Team, + #[strum(serialize = "org")] + Org, + #[strum(serialize = "project")] + Project, + #[strum(serialize = "end_user")] + EndUser, + #[strum(serialize = "tag")] + Tag, + #[strum(serialize = "model_access_group")] + ModelAccessGroup, + #[strum(serialize = "member")] + Member, +} + +fn exact(ids: &[String], count: usize) -> Result<(), Error> { + if ids.len() == count { + Ok(()) + } else { + Err(Error::InvalidAuthArguments) + } +} + +fn hashed_key(id: &str) -> Result<(), Error> { + if id.len() == 64 && id.bytes().all(|byte| byte.is_ascii_hexdigit()) { + Ok(()) + } else { + Err(Error::InvalidAuthArguments) + } +} + +fn entity(value: &str) -> Result<&'static str, Error> { + match value { + "Key" => Ok("key"), + "Team" => Ok("team"), + _ => Err(Error::InvalidAuthArguments), + } +} + +pub async fn auth_read<'e>( + executor: impl Executor<'e, Database = Postgres>, + kind: AuthRead, + ids: &[String], + persistent: bool, +) -> Result, Error> { + let executor = StatementPersistence { + executor, + persistent, + }; + if ids.len() > 1000 || ids.iter().any(|id| id.len() > 65536) { + return Err(Error::InvalidAuthArguments); + } + let rows = match kind { + AuthRead::User => { + sqlx::query_file_as!(AuthRow, "queries/auth/user.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Team => { + sqlx::query_file_as!(AuthRow, "queries/auth/team.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Org => { + sqlx::query_file_as!(AuthRow, "queries/auth/org.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Project => { + sqlx::query_file_as!(AuthRow, "queries/auth/project.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::EndUser => { + sqlx::query_file_as!(AuthRow, "queries/auth/end_user.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Budget => { + sqlx::query_file_as!(AuthRow, "queries/auth/budget.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::ObjectPermission => { + sqlx::query_file_as!(AuthRow, "queries/auth/object_permission.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::AccessGroup => { + sqlx::query_file_as!(AuthRow, "queries/auth/access_group.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::ManagedVectorStore => { + sqlx::query_file_as!(AuthRow, "queries/auth/managed_vector_store.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Tag => { + sqlx::query_file_as!(AuthRow, "queries/auth/tag.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::ModelAccessGroup => { + sqlx::query_file_as!(AuthRow, "queries/auth/model_access_group.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Agent => { + sqlx::query_file_as!(AuthRow, "queries/auth/agent.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::AgentName => { + sqlx::query_file_as!(AuthRow, "queries/auth/agent_name.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::AgentAccessGroup => { + sqlx::query_file_as!(AuthRow, "queries/auth/agent_access_group.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::RetiredAgent => { + sqlx::query_file_as!(AuthRow, "queries/auth/retired_agent.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Key => { + ids.iter().try_for_each(|id| hashed_key(id))?; + sqlx::query_file_as!(AuthRow, "queries/auth/key.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::Member => { + exact(ids, 2)?; + sqlx::query_file_as!(AuthRow, "queries/auth/member.sql", &ids[0], &ids[1]) + .fetch_all(executor) + .await? + } + AuthRead::AgentIdentity => { + exact(ids, 3)?; + sqlx::query_file_as!( + AuthRow, + "queries/auth/agent_identity.sql", + &ids[0], + &ids[1], + &ids[2] + ) + .fetch_all(executor) + .await? + } + AuthRead::RetiredAgentIdentity => { + exact(ids, 3)?; + sqlx::query_file_as!( + AuthRow, + "queries/auth/retired_agent_identity.sql", + &ids[0], + &ids[1], + &ids[2] + ) + .fetch_all(executor) + .await? + } + AuthRead::VerifiedSubject => { + exact(ids, 3)?; + sqlx::query_file_as!( + AuthRow, + "queries/auth/verified_subject.sql", + &ids[0], + &ids[1], + &ids[2] + ) + .fetch_all(executor) + .await? + } + AuthRead::TeamAlias => { + exact(ids, 1)?; + sqlx::query_file_as!(AuthRow, "queries/auth/team_alias.sql", &ids[0]) + .fetch_all(executor) + .await? + } + AuthRead::OrgAlias => { + exact(ids, 1)?; + sqlx::query_file_as!(AuthRow, "queries/auth/org_alias.sql", &ids[0]) + .fetch_all(executor) + .await? + } + AuthRead::UserSso => { + exact(ids, 1)?; + sqlx::query_file_as!(AuthRow, "queries/auth/user_sso.sql", &ids[0]) + .fetch_all(executor) + .await? + } + AuthRead::UserEmail => { + exact(ids, 1)?; + sqlx::query_file_as!(AuthRow, "queries/auth/user_email.sql", &ids[0]) + .fetch_all(executor) + .await? + } + AuthRead::TagRegistry => { + exact(ids, 1)?; + let limit: i64 = ids[0] + .parse() + .ok() + .filter(|value| (1..=100_001).contains(value)) + .ok_or(Error::InvalidAuthArguments)?; + sqlx::query_file_as!(AuthRow, "queries/auth/tag_registry.sql", limit) + .fetch_all(executor) + .await? + } + AuthRead::ModelAccessGroupRegistry => { + exact(ids, 1)?; + let limit: i64 = ids[0] + .parse() + .ok() + .filter(|value| (1..=100_001).contains(value)) + .ok_or(Error::InvalidAuthArguments)?; + sqlx::query_file_as!( + AuthRow, + "queries/auth/model_access_group_registry.sql", + limit + ) + .fetch_all(executor) + .await? + } + AuthRead::EndUserRestrictedRegistry => { + exact(ids, 1)?; + let limit: i64 = ids[0] + .parse() + .ok() + .filter(|value| (1..=100_001).contains(value)) + .ok_or(Error::InvalidAuthArguments)?; + sqlx::query_file_as!( + AuthRow, + "queries/auth/end_user_restricted_registry.sql", + limit + ) + .fetch_all(executor) + .await? + } + AuthRead::Prefetch => { + exact(ids, 1)?; + let refs: [Option; 5] = + serde_json::from_str(&ids[0]).map_err(|_| Error::InvalidAuthArguments)?; + sqlx::query_file_as!( + AuthRow, + "queries/auth/prefetch.sql", + refs[0].as_deref(), + refs[1].as_deref(), + refs[2].as_deref(), + refs[3].as_deref(), + refs[4].as_deref() + ) + .fetch_all(executor) + .await? + } + AuthRead::BudgetWindow => { + exact(ids, 3)?; + let entity = entity(&ids[0])?; + if entity == "key" { + hashed_key(&ids[1])?; + } + sqlx::query_file_as!( + AuthRow, + "queries/auth/budget_window.sql", + entity, + &ids[1], + &ids[2] + ) + .fetch_all(executor) + .await? + } + AuthRead::WindowSpend => { + exact(ids, 3)?; + let start = chrono::DateTime::parse_from_rfc3339(&ids[2]) + .map(|value| value.naive_utc()) + .or_else(|_| chrono::NaiveDateTime::parse_from_str(&ids[2], "%Y-%m-%dT%H:%M:%S%.f")) + .map_err(|_| Error::InvalidAuthArguments)?; + match entity(&ids[0])? { + "key" => { + hashed_key(&ids[1])?; + sqlx::query_file_as!( + AuthRow, + "queries/auth/window_key_spend.sql", + &ids[1], + start + ) + .fetch_all(executor) + .await? + } + _ => { + sqlx::query_file_as!( + AuthRow, + "queries/auth/window_team_spend.sql", + &ids[1], + start + ) + .fetch_all(executor) + .await? + } + } + } + }; + Ok(rows) +} + +pub async fn auth_spend<'e>( + executor: impl Executor<'e, Database = Postgres>, + kind: AuthSpend, + id: &str, + persistent: bool, +) -> Result, Error> { + let executor = StatementPersistence { + executor, + persistent, + }; + if id.len() > 65536 { + return Err(Error::InvalidAuthArguments); + } + let row = match kind { + AuthSpend::Key => { + hashed_key(id)?; + sqlx::query_file!("queries/auth/spend_key.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend) + } + AuthSpend::User => sqlx::query_file!("queries/auth/spend_user.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend), + AuthSpend::Team => sqlx::query_file!("queries/auth/spend_team.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend), + AuthSpend::Org => sqlx::query_file!("queries/auth/spend_org.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend), + AuthSpend::Project => sqlx::query_file!("queries/auth/spend_project.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend), + AuthSpend::EndUser => sqlx::query_file!("queries/auth/spend_end_user.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend), + AuthSpend::Tag => sqlx::query_file!("queries/auth/spend_tag.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend), + AuthSpend::ModelAccessGroup => { + sqlx::query_file!("queries/auth/spend_model_access_group.sql", id) + .fetch_optional(executor) + .await? + .map(|row| row.spend) + } + AuthSpend::Member => { + let (user_id, team_id) = id.rsplit_once(':').ok_or(Error::InvalidAuthArguments)?; + sqlx::query_file!("queries/auth/spend_member.sql", user_id, team_id) + .fetch_optional(executor) + .await? + .map(|row| row.spend) + } + }; + Ok(row) +} + +#[cfg(test)] +mod tests { + use super::{AuthRead, AuthSpend, auth_read, auth_spend}; + use crate::test_support::database_pool; + use rstest::rstest; + use serde_json::{Value, json}; + use sqlx::{PgPool, Postgres, Transaction}; + use std::error::Error as StdError; + + async fn seeded(pool: &PgPool) -> Result, sqlx::Error> { + let mut tx = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "Transaction-local test fixture outside the production query set" + )] + sqlx::raw_sql(include_str!("../queries/auth/fixture.sql")) + .execute(&mut *tx) + .await?; + Ok(tx) + } + + fn strings(values: &[&str]) -> Vec { + values.iter().map(|value| (*value).to_owned()).collect() + } + + #[rstest] + #[case::vector_store(AuthRead::ManagedVectorStore, &["auth-vector"], "/litellm_params", json!({"api_base":"https://example.com"}))] + #[case::vector_store_nullable(AuthRead::ManagedVectorStore, &["auth-vector"], "/vector_store_metadata", json!(null))] + #[case::user(AuthRead::User, &["auth-u"], "/organization_memberships/0/user_role", json!("org_admin"))] + #[case::team(AuthRead::Team, &["auth-team"], "/litellm_model_table/model_aliases/alias", json!("actual"))] + #[case::team_null_lists(AuthRead::Team, &["auth-team-empty"], "/models", json!([]))] + #[case::team_missing_relation(AuthRead::Team, &["auth-team-empty"], "/object_permission", json!(null))] + #[case::org(AuthRead::Org, &["auth-org"], "/litellm_budget_table/max_budget", json!(13.0))] + #[case::project(AuthRead::Project, &["auth-project"], "/litellm_budget_table/max_budget", json!(0.0))] + #[case::end_user(AuthRead::EndUser, &["auth-end"], "/models", json!([]))] + #[case::budget(AuthRead::Budget, &["auth-b"], "/allowed_models", json!([]))] + #[case::object_permission(AuthRead::ObjectPermission, &["auth-op"], "/mcp_servers", json!([]))] + #[case::access_group(AuthRead::AccessGroup, &["auth-access"], "/access_model_names", json!(["group-model"]))] + #[case::tag(AuthRead::Tag, &["auth-tag"], "/models", json!([]))] + #[case::model_access_group(AuthRead::ModelAccessGroup, &["auth-model-group"], "/litellm_budget_table/max_budget", json!(0.0))] + #[case::member(AuthRead::Member, &["auth-u","auth-team"], "/litellm_budget_table/rpm_limit", json!(33))] + #[case::agent(AuthRead::Agent, &["auth-agent"], "/identity/active", json!(false))] + #[case::agent_name(AuthRead::AgentName, &["Auth Agent"], "/identity/active", json!(false))] + #[case::agent_group(AuthRead::AgentAccessGroup, &["other-group", "auth-group"], "", json!({"agent_id":"auth-agent"}))] + #[case::retired_agent(AuthRead::RetiredAgent, &["auth-retired-agent"], "/original_agent_id", json!("auth-retired-agent"))] + #[case::agent_identity(AuthRead::AgentIdentity, &["microsoft_entra","auth-tenant","auth-client"], "/required_roles", json!([]))] + #[case::retired_agent_identity(AuthRead::RetiredAgentIdentity, &["microsoft_entra","auth-tenant","auth-retired-client"], "/binding_id", json!("auth-retired-binding"))] + #[case::verified_subject(AuthRead::VerifiedSubject, &["auth-issuer","auth-tenant","auth-oid"], "/user_id", json!("auth-u"))] + #[case::team_alias(AuthRead::TeamAlias, &["Auth Team"], "/team_id", json!("auth-team"))] + #[case::org_alias(AuthRead::OrgAlias, &["Auth Org"], "/organization_id", json!("auth-org"))] + #[case::sso(AuthRead::UserSso, &["auth-sso"], "/user_id", json!("auth-u"))] + #[case::email_literal_insensitive(AuthRead::UserEmail, &["test_%@example.COM"], "/user_id", json!("auth-u"))] + #[case::tags_registry(AuthRead::TagRegistry, &["5001"], "/tag_name", json!("auth-tag"))] + #[case::model_registry(AuthRead::ModelAccessGroupRegistry, &["5001"], "/access_group_name", json!("auth-model-group"))] + #[case::restricted_registry(AuthRead::EndUserRestrictedRegistry, &["5001"], "/user_id", json!("auth-end"))] + #[case::budget_window(AuthRead::BudgetWindow, &["Team","auth-team","1h"], "/spend", json!(2.5))] + #[case::window_spend_offset(AuthRead::WindowSpend, &["Team","auth-team","2025-01-01T02:00:00+01:00"], "/spend", json!(6.0))] + #[case::window_spend_naive(AuthRead::WindowSpend, &["Team","auth-team","2025-01-01T01:00:00"], "/spend", json!(6.0))] + #[case::window_spend_no_rows(AuthRead::WindowSpend, &["Team","absent","2025-01-01T01:00:00Z"], "/spend", json!(0.0))] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn projects_auth_rows( + #[case] kind: AuthRead, + #[case] ids: &[&str], + #[case] pointer: &str, + #[case] expected: Value, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = seeded(&pool).await?; + let rows = auth_read(&mut *tx, kind, &strings(ids), false).await?; + assert_eq!(rows.len(), 1); + let actual = rows[0].data.0.pointer(pointer).expect("projected field"); + if expected.is_number() { + assert_eq!(actual.as_f64(), expected.as_f64()); + } else { + assert_eq!(actual, &expected); + } + Ok(()) + } + + #[rstest] + #[case::active('a', true)] + #[case::deprecated('b', true)] + #[case::expired('c', false)] + #[case::no_chaining('d', false)] + #[case::missing('f', false)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn combined_key_join_and_deprecation( + #[case] token: char, + #[case] found: bool, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = seeded(&pool).await?; + let rows = auth_read( + &mut *tx, + AuthRead::Key, + &[token.to_string().repeat(64)], + false, + ) + .await?; + assert_eq!(rows.len(), usize::from(found)); + if !found { + return Ok(()); + } + let row = &rows[0].data.0; + assert_eq!(row["token"], "a".repeat(64)); + assert_eq!(row["team_spend"], 2.5); + assert_eq!(row["team_blocked"], true); + assert_eq!(row["org_id"], "auth-org"); + assert_eq!(row["team_models"], json!(["team-model"])); + assert_eq!(row["team_model_aliases"], json!({"alias":"actual"})); + assert_eq!(row["team_member_spend"], 6.5); + assert_eq!(row["team_member_tpm_limit"], 22); + assert_eq!(row["litellm_budget_table_max_budget"], 4.5); + assert_eq!(row["organization_max_budget"], 13.0); + assert_eq!(row["project_alias"], "Auth Project"); + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn combined_key_keeps_null_join_values( + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = seeded(&pool).await?; + let rows = auth_read(&mut *tx, AuthRead::Key, &["e".repeat(64)], false).await?; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].data.0["team_models"], Value::Null); + assert_eq!(rows[0].data.0["team_blocked"], Value::Null); + assert_eq!(rows[0].data.0["organization_max_budget"], Value::Null); + assert_eq!(rows[0].data.0["spend"], 0.0); + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn prefetch_has_same_nested_projections( + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = seeded(&pool).await?; + let refs = json!(["auth-u", "auth-team", "auth-u", "auth-org", "auth-project"]).to_string(); + let rows = auth_read(&mut *tx, AuthRead::Prefetch, &[refs], false).await?; + let user = auth_read(&mut *tx, AuthRead::User, &strings(&["auth-u"]), false).await?; + let team = auth_read(&mut *tx, AuthRead::Team, &strings(&["auth-team"]), false).await?; + let member = auth_read( + &mut *tx, + AuthRead::Member, + &strings(&["auth-u", "auth-team"]), + false, + ) + .await?; + let org = auth_read(&mut *tx, AuthRead::Org, &strings(&["auth-org"]), false).await?; + let project = auth_read( + &mut *tx, + AuthRead::Project, + &strings(&["auth-project"]), + false, + ) + .await?; + assert_eq!(rows.len(), 1); + assert_eq!( + rows[0].data.0, + json!({ + "user_row": user[0].data.0, "team_row": team[0].data.0, + "membership_row": member[0].data.0, "organization_row": org[0].data.0, + "project_row": project[0].data.0 + }) + ); + let absent = auth_read( + &mut *tx, + AuthRead::Prefetch, + &[json!([null, null, null, null, null]).to_string()], + false, + ) + .await?; + assert_eq!( + absent[0].data.0, + json!({"user_row":null,"team_row":null,"membership_row":null,"organization_row":null,"project_row":null}) + ); + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn batch_ids_preserve_tenant_boundaries_and_missing_rows( + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = seeded(&pool).await?; + let rows = auth_read( + &mut *tx, + AuthRead::User, + &strings(&["auth-u", "missing", "auth-u"]), + false, + ) + .await?; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].data.0["user_id"], "auth-u"); + assert!( + auth_read(&mut *tx, AuthRead::User, &[], false) + .await? + .is_empty() + ); + assert!( + auth_read( + &mut *tx, + AuthRead::Member, + &strings(&["auth-other", "auth-team"]), + false + ) + .await? + .is_empty() + ); + assert!( + auth_read( + &mut *tx, + AuthRead::AgentIdentity, + &strings(&["microsoft_entra", "other-tenant", "auth-client"]), + false + ) + .await? + .is_empty() + ); + assert!( + auth_read( + &mut *tx, + AuthRead::UserEmail, + &strings(&["test%@example.com"]), + false + ) + .await? + .is_empty() + ); + assert!( + auth_read(&mut *tx, AuthRead::User, &strings(&["' OR 1=1 --"]), false) + .await? + .is_empty() + ); + assert!( + auth_read( + &mut *tx, + AuthRead::AgentAccessGroup, + &strings(&["other-group"]), + false + ) + .await? + .is_empty() + ); + Ok(()) + } + + #[rstest] + #[case::key(AuthSpend::Key, "a".repeat(64), 4.5)] + #[case::user(AuthSpend::User, "auth-u".into(), 3.5)] + #[case::zero(AuthSpend::User, "auth-other".into(), 0.0)] + #[case::team(AuthSpend::Team, "auth-team".into(), 2.5)] + #[case::org(AuthSpend::Org, "auth-org".into(), 13.5)] + #[case::project(AuthSpend::Project, "auth-project".into(), 8.5)] + #[case::end_user(AuthSpend::EndUser, "auth-end".into(), 10.5)] + #[case::tag(AuthSpend::Tag, "auth-tag".into(), 11.5)] + #[case::model_group(AuthSpend::ModelAccessGroup, "auth-model-group".into(), 12.5)] + #[case::member(AuthSpend::Member, "auth-u:auth-team".into(), 6.5)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn spend_reads( + #[case] kind: AuthSpend, + #[case] id: String, + #[case] expected: f64, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = seeded(&pool).await?; + assert_eq!( + auth_spend(&mut *tx, kind, &id, false).await?, + Some(expected) + ); + let absent = match kind { + AuthSpend::Key => "f".repeat(64), + AuthSpend::Member => "missing:missing".into(), + _ => "missing".into(), + }; + assert_eq!(auth_spend(&mut *tx, kind, &absent, false).await?, None); + Ok(()) + } + + #[rstest] + #[case::key_window(AuthRead::BudgetWindow, vec!["Key".into(), "a".repeat(64), "1h".into()], 2.5)] + #[case::key_logs(AuthRead::WindowSpend, vec!["Key".into(), "a".repeat(64), "2025-01-01T01:00:00Z".into()], 6.0)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn key_window_spend( + #[case] kind: AuthRead, + #[case] ids: Vec, + #[case] expected: f64, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = seeded(&pool).await?; + let rows = auth_read(&mut *tx, kind, &ids, false).await?; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].data.0["spend"], expected); + Ok(()) + } + + #[rstest] + #[case::member(AuthRead::Member, vec!["one".into()])] + #[case::registry(AuthRead::TagRegistry, vec!["0".into()])] + #[case::huge_registry(AuthRead::TagRegistry, vec!["100002".into()])] + #[case::raw_token(AuthRead::Key, vec!["sk-raw-secret".into()])] + #[case::prefetch_shape(AuthRead::Prefetch, vec!["[]".into()])] + #[case::unknown_window(AuthRead::BudgetWindow, strings(&["Unknown","id","1h"]))] + #[tokio::test] + async fn invalid_arguments_fail_before_queries( + #[case] kind: AuthRead, + #[case] ids: Vec, + ) { + let pool = sqlx::postgres::PgPoolOptions::new() + .connect_lazy("postgresql://localhost/unused") + .unwrap(); + pool.close().await; + assert!(matches!( + auth_read(&pool, kind, &ids, false).await, + Err(crate::Error::InvalidAuthArguments) + )); + } + + #[rstest] + #[tokio::test] + async fn database_errors_are_not_misses() { + let pool = sqlx::postgres::PgPoolOptions::new() + .connect_lazy("postgresql://localhost/unused") + .unwrap(); + pool.close().await; + assert!(matches!( + auth_read(&pool, AuthRead::User, &strings(&["id"]), false).await, + Err(crate::Error::Query(sqlx::Error::PoolClosed)) + )); + assert!(matches!( + auth_spend(&pool, AuthSpend::User, "id", false).await, + Err(crate::Error::Query(sqlx::Error::PoolClosed)) + )); + } + + #[rstest] + fn projected_credentials_are_redacted() { + let row = super::AuthRow { + data: super::AuthData(json!({"password":"credential-value"})), + }; + let debug = format!("{row:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); + } +} diff --git a/litellm-rust/crates/db/src/error.rs b/litellm-rust/crates/db/src/error.rs index af0ee030896..5c5bdee7382 100644 --- a/litellm-rust/crates/db/src/error.rs +++ b/litellm-rust/crates/db/src/error.rs @@ -1,5 +1,13 @@ #[derive(Debug, thiserror::Error)] pub enum Error { + #[error("invalid database setting {0}")] + InvalidSetting(&'static str), + #[error("invalid auth read arguments")] + InvalidAuthArguments, + #[error("native database is not configured")] + NotConfigured, + #[error("database connection timed out")] + ConnectTimeout(#[from] tokio::time::error::Elapsed), #[error("database query failed")] Query(#[from] sqlx::Error), #[error("spend batch has {ids} ids and {costs} costs")] @@ -9,6 +17,25 @@ pub enum Error { } impl Error { + pub fn is_database_failure(&self) -> bool { + match self { + Self::ConnectTimeout(_) => true, + Self::Query(error) => matches!( + error, + sqlx::Error::Database(_) + | sqlx::Error::Io(_) + | sqlx::Error::Tls(_) + | sqlx::Error::PoolTimedOut + | sqlx::Error::PoolClosed + ), + Self::InvalidSetting(_) + | Self::InvalidAuthArguments + | Self::NotConfigured + | Self::InvalidSpendBatch { .. } + | Self::InvalidTimestamp(_) => false, + } + } + pub fn is_retryable(&self) -> bool { match self { Self::Query(sqlx::Error::Database(error)) => { @@ -42,6 +69,24 @@ mod tests { assert_eq!(Error::from(error).is_retryable(), retryable); } + #[rstest] + #[case::closed_pool(sqlx::Error::PoolClosed, true)] + #[case::pool_timeout(sqlx::Error::PoolTimedOut, true)] + #[case::io(sqlx::Error::Io(std::io::ErrorKind::ConnectionReset.into()), true)] + #[case::tls( + sqlx::Error::Tls(Box::new(std::io::Error::from(std::io::ErrorKind::ConnectionReset))), + true + )] + #[case::decode( + sqlx::Error::Decode(Box::new(std::io::Error::from(std::io::ErrorKind::InvalidData))), + false + )] + #[case::missing_column(sqlx::Error::ColumnNotFound("spend".into()), false)] + #[case::missing_row(sqlx::Error::RowNotFound, false)] + fn distinguishes_io_from_native_invariants(#[case] error: sqlx::Error, #[case] io: bool) { + assert_eq!(Error::from(error).is_database_failure(), io); + } + #[rstest] #[case::connection_capacity("53300", true)] #[case::deadlock("40P01", true)] @@ -72,7 +117,9 @@ mod tests { error.as_database_error().and_then(|error| error.code()), Some(code.into()) ); - assert_eq!(Error::from(error).is_retryable(), retryable); + let error = Error::from(error); + assert!(error.is_database_failure()); + assert_eq!(error.is_retryable(), retryable); tx.rollback().await?; Ok(()) } diff --git a/litellm-rust/crates/db/src/executor.rs b/litellm-rust/crates/db/src/executor.rs new file mode 100644 index 00000000000..a656c53e408 --- /dev/null +++ b/litellm-rust/crates/db/src/executor.rs @@ -0,0 +1,84 @@ +use futures_util::{future::BoxFuture, stream::BoxStream}; +use sqlx::{ + Describe, Either, Execute, Executor, Postgres, SqlStr, + postgres::{PgArguments, PgQueryResult, PgRow, PgStatement, PgTypeInfo}, +}; + +#[derive(Debug)] +pub(crate) struct StatementPersistence { + pub executor: E, + pub persistent: bool, +} + +struct QueryPersistence { + query: Q, + persistent: bool, +} + +impl<'q, Q: Execute<'q, Postgres>> Execute<'q, Postgres> for QueryPersistence { + fn sql(self) -> SqlStr { + self.query.sql() + } + + fn statement(&self) -> Option<&PgStatement> { + self.query.statement() + } + + fn take_arguments(&mut self) -> Result, sqlx::error::BoxDynError> { + self.query.take_arguments() + } + + fn persistent(&self) -> bool { + self.persistent + } +} + +impl<'c, E: Executor<'c, Database = Postgres>> Executor<'c> for StatementPersistence { + type Database = Postgres; + + fn fetch_many<'e, 'q: 'e, Q>( + self, + query: Q, + ) -> BoxStream<'e, Result, sqlx::Error>> + where + 'c: 'e, + Q: 'q + Execute<'q, Postgres>, + { + self.executor.fetch_many(QueryPersistence { + query, + persistent: self.persistent, + }) + } + + fn fetch_optional<'e, 'q: 'e, Q>( + self, + query: Q, + ) -> BoxFuture<'e, Result, sqlx::Error>> + where + 'c: 'e, + Q: 'q + Execute<'q, Postgres>, + { + self.executor.fetch_optional(QueryPersistence { + query, + persistent: self.persistent, + }) + } + + fn prepare_with<'e>( + self, + sql: SqlStr, + parameters: &'e [PgTypeInfo], + ) -> BoxFuture<'e, Result> + where + 'c: 'e, + { + self.executor.prepare_with(sql, parameters) + } + + fn describe<'e>(self, sql: SqlStr) -> BoxFuture<'e, Result, sqlx::Error>> + where + 'c: 'e, + { + self.executor.describe(sql) + } +} diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 541f97c7db9..3264e3a8592 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -6,18 +6,23 @@ macro_rules_attribute::attribute_alias! { #[derive(Clone, Debug, PartialEq, sqlx::FromRow)]; } +mod auth; mod daily_spend; mod end_user; mod error; +mod executor; mod failure; +mod pool; mod spend; #[cfg(test)] mod test_support; +pub use auth::{AuthRead, AuthRow, AuthSpend, auth_read, auth_spend}; pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow}; pub use end_user::{EndUser, RateLimits, find_end_user}; pub use error::Error; pub use failure::DbFailure; +pub use pool::{DatabasePools, DatabaseSettings}; pub use spend::{ DailySpend, DailyTable, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend, add_end_user_spend, add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, diff --git a/litellm-rust/crates/db/src/pool.rs b/litellm-rust/crates/db/src/pool.rs new file mode 100644 index 00000000000..74a48e2cd5d --- /dev/null +++ b/litellm-rust/crates/db/src/pool.rs @@ -0,0 +1,495 @@ +use std::time::Duration; + +use litellm_secret_value::SecretValue; +use sqlx::{ + ConnectOptions, PgPool, + postgres::{PgConnectOptions, PgPoolOptions, PgSslMode}, +}; + +use crate::Error; + +#[derive(Clone, Debug)] +pub struct DatabaseSettings { + pub host: String, + pub port: u16, + pub username: String, + pub password: Option, + pub database: String, + pub schema: String, + pub ssl_mode: PgSslMode, + pub ssl_root_cert: Option, + pub max_connections: u32, + pub acquire_timeout: Duration, + pub connect_timeout: Duration, + pub idle_timeout: Option, + pub statement_cache_capacity: usize, +} + +impl DatabaseSettings { + fn pool(&self) -> Result { + if self.max_connections == 0 { + return Err(Error::InvalidSetting("max_connections")); + } + if self.host.is_empty() + || self.username.is_empty() + || self.database.is_empty() + || self.schema.is_empty() + { + return Err(Error::InvalidSetting("connection identity")); + } + let options = PgConnectOptions::new_without_pgpass() + .host(&self.host) + .port(self.port) + .username(&self.username) + .password(self.password.as_ref().map_or("", SecretValue::expose)) + .database(&self.database) + .application_name("native-auth") + .ssl_mode(self.ssl_mode) + .statement_cache_capacity(self.statement_cache_capacity) + .options([( + "search_path", + format!("\"{}\"", self.schema.replace('"', "\"\"")), + )]) + .disable_statement_logging(); + let options = match &self.ssl_root_cert { + Some(cert) => options.ssl_root_cert(cert), + None => options, + }; + Ok(PgPoolOptions::new() + .max_connections(self.max_connections) + .acquire_timeout(if self.acquire_timeout.is_zero() { + Duration::from_secs(30) + } else { + self.acquire_timeout + }) + .idle_timeout(self.idle_timeout) + .connect_lazy_with(options)) + } +} + +#[derive(Clone, Debug)] +pub struct DatabasePools { + writer: PgPool, + reader: Option, + reader_available: bool, + writer_unbounded: bool, + reader_unbounded: bool, + writer_persistent: bool, + reader_persistent: bool, +} + +async fn acquire( + pool: &PgPool, + unbounded: bool, +) -> Result, sqlx::Error> { + loop { + match pool.acquire().await { + Err(sqlx::Error::PoolTimedOut) if unbounded => continue, + result => return result, + } + } +} + +async fn connect_pool(pool: &PgPool, settings: &DatabaseSettings) -> Result<(), Error> { + let connection = acquire(pool, settings.acquire_timeout.is_zero()); + if settings.connect_timeout.is_zero() { + connection.await?; + } else { + tokio::time::timeout(settings.connect_timeout, connection).await??; + } + Ok(()) +} + +impl DatabasePools { + pub async fn connect( + writer: DatabaseSettings, + reader: Option, + ) -> Result { + let writer_pool = writer.pool()?; + let reader_pool = reader.as_ref().map(DatabaseSettings::pool).transpose()?; + let writer_result = connect_pool(&writer_pool, &writer).await; + let reader_result = match (&reader_pool, &reader) { + (Some(pool), Some(settings)) => Some(connect_pool(pool, settings).await), + _ => None, + }; + let reader_available = reader_result.as_ref().is_some_and(Result::is_ok); + drop(reader_result); + if let Err(error) = writer_result + && !reader_available + { + writer_pool.close().await; + if let Some(pool) = reader_pool { + pool.close().await; + } + return Err(error); + } + Ok(Self { + writer: writer_pool, + reader: reader_pool, + reader_available, + writer_unbounded: writer.acquire_timeout.is_zero(), + reader_unbounded: reader + .as_ref() + .is_some_and(|settings| settings.acquire_timeout.is_zero()), + writer_persistent: writer.statement_cache_capacity > 0, + reader_persistent: reader + .as_ref() + .is_some_and(|settings| settings.statement_cache_capacity > 0), + }) + } + + pub fn reader_available(&self) -> bool { + self.reader_available + } + + pub fn persistent(&self, writer: bool) -> bool { + if !writer && self.reader_available { + self.reader_persistent + } else { + self.writer_persistent + } + } + + pub fn pool(&self, writer: bool) -> &PgPool { + match (&self.reader, writer, self.reader_available) { + (Some(pool), false, true) => pool, + _ => &self.writer, + } + } + + pub async fn acquire( + &self, + writer: bool, + ) -> Result, Error> { + let unbounded = if !writer && self.reader_available { + self.reader_unbounded + } else { + self.writer_unbounded + }; + Ok(acquire(self.pool(writer), unbounded).await?) + } + + pub fn set_passwords( + &self, + writer: Option, + reader: Option, + ) -> Result<(), Error> { + if self.writer.is_closed() { + return Err(sqlx::Error::PoolClosed.into()); + } + if reader.is_some() && self.reader.is_none() { + return Err(Error::InvalidSetting("reader_password")); + } + if let Some(password) = writer { + self.writer.set_connect_options( + (*self.writer.connect_options()) + .clone() + .password(password.expose()), + ); + } + if let (Some(pool), Some(password)) = (&self.reader, reader) { + pool.set_connect_options( + (*pool.connect_options()) + .clone() + .password(password.expose()), + ); + } + Ok(()) + } + + pub async fn close(&self) { + if let Some(reader) = &self.reader { + tokio::join!(self.writer.close(), reader.close()); + } else { + self.writer.close().await; + } + } +} + +#[cfg(test)] +mod tests { + use std::{error::Error as StdError, time::Duration}; + + use litellm_secret_value::SecretValue; + use rstest::{fixture, rstest}; + use sqlx::{ConnectOptions, PgPool, postgres::PgPoolOptions}; + + use super::{DatabasePools, DatabaseSettings, acquire}; + use crate::test_support::database_pool; + + #[fixture] + async fn settings() -> Result> { + let pool = database_pool().await?; + let options = pool.connect_options(); + let settings = DatabaseSettings { + host: options.get_socket().map_or_else( + || options.get_host().to_owned(), + |socket| socket.to_string_lossy().into_owned(), + ), + port: options.get_port(), + username: options.get_username().into(), + password: options.to_url_lossy().password().map(SecretValue::from), + database: options + .get_database() + .ok_or("database must be explicit")? + .into(), + schema: "public".into(), + ssl_mode: options.get_ssl_mode(), + ssl_root_cert: None, + max_connections: 1, + acquire_timeout: Duration::from_secs(2), + connect_timeout: Duration::from_secs(2), + idle_timeout: None, + statement_cache_capacity: 0, + }; + pool.close().await; + Ok(settings) + } + + #[rstest] + #[case::writer_only(false, false)] + #[case::reader_healthy(true, false)] + #[case::reader_unavailable(true, true)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn reader_selection_and_shutdown( + #[case] configured: bool, + #[case] unavailable: bool, + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = DatabaseSettings { + statement_cache_capacity: 1, + ..settings? + }; + let reader = configured.then(|| DatabaseSettings { + database: if unavailable { + "native_auth_missing_database".into() + } else { + settings.database.clone() + }, + connect_timeout: Duration::from_millis(100), + statement_cache_capacity: 0, + ..settings.clone() + }); + let pools = DatabasePools::connect(settings, reader).await?; + assert_eq!(pools.reader_available(), configured && !unavailable); + assert!(pools.persistent(true)); + assert_eq!(pools.persistent(false), !pools.reader_available()); + let mut writer = pools.acquire(true).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test verifies connection provenance rather than a production query" + )] + let (writer_pid, app, schema): (i32, String, String) = sqlx::query_as( + "SELECT pg_backend_pid(), current_setting('application_name'), current_schema()", + ) + .fetch_one(&mut *writer) + .await?; + assert_eq!(app, "native-auth"); + assert_eq!(schema, "public"); + drop(writer); + let mut reader = pools.acquire(false).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test verifies pool selection by backend identity" + )] + let reader_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *reader) + .await?; + assert_eq!(writer_pid == reader_pid, !configured || unavailable); + drop(reader); + let clone = pools.clone(); + pools.close().await; + clone.close().await; + assert!(matches!( + clone.acquire(false).await, + Err(crate::Error::Query(sqlx::Error::PoolClosed)) + )); + assert!( + clone + .set_passwords(Some(SecretValue::from("next")), None) + .is_err() + ); + Ok(()) + } + + #[rstest] + #[case::uncached(0, 0)] + #[case::cached(4, 2)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn respects_statement_persistence( + #[case] capacity: usize, + #[case] expected: i64, + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = DatabaseSettings { + statement_cache_capacity: capacity, + ..settings? + }; + let pools = DatabasePools::connect(settings, None).await?; + let mut connection = pools.acquire(false).await?; + crate::auth_read( + &mut *connection, + crate::AuthRead::User, + &[], + pools.persistent(false), + ) + .await?; + crate::auth_spend( + &mut *connection, + crate::AuthSpend::User, + "missing", + pools.persistent(false), + ) + .await?; + #[expect( + clippy::disallowed_methods, + reason = "Test observes server-side prepared statements on the selected connection" + )] + let prepared: i64 = sqlx::query_scalar("SELECT count(*) FROM pg_prepared_statements") + .persistent(false) + .fetch_one(&mut *connection) + .await?; + assert_eq!(prepared, expected); + drop(connection); + pools.close().await; + Ok(()) + } + + #[rstest] + #[case::reader_survives(true)] + #[case::both_unavailable(false)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn unavailable_writer( + #[case] healthy_reader: bool, + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = settings?; + let writer = DatabaseSettings { + database: "native_auth_missing_database".into(), + connect_timeout: Duration::from_millis(100), + acquire_timeout: Duration::from_millis(100), + ..settings.clone() + }; + let result = DatabasePools::connect(writer, healthy_reader.then_some(settings)).await; + if !healthy_reader { + assert!(result.is_err()); + return Ok(()); + } + let pools = result?; + assert!(pools.reader_available()); + assert!(pools.acquire(false).await.is_ok()); + assert!(pools.acquire(true).await.is_err()); + pools.close().await; + Ok(()) + } + + #[rstest] + #[case::bounded(false)] + #[case::unbounded(true)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn acquisition_deadline_and_cancellation( + #[case] unbounded: bool, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let source = database_pool?; + let pool = PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(Duration::from_millis(20)) + .connect_with((*source.connect_options()).clone()) + .await?; + let held = pool.acquire().await?; + let result = + tokio::time::timeout(Duration::from_millis(150), acquire(&pool, unbounded)).await; + if unbounded { + assert!(result.is_err()); + } else { + assert!(matches!(result, Ok(Err(sqlx::Error::PoolTimedOut)))); + } + drop(held); + let recovered = + tokio::time::timeout(Duration::from_secs(2), acquire(&pool, unbounded)).await??; + drop(recovered); + pool.close().await; + assert!(matches!( + acquire(&pool, unbounded).await, + Err(sqlx::Error::PoolClosed) + )); + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn query_error_does_not_replay_on_writer( + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = settings?; + let reader = DatabaseSettings { + schema: "native_auth_missing_schema".into(), + ..settings.clone() + }; + let pools = DatabasePools::connect(settings, Some(reader)).await?; + let mut read = pools.acquire(false).await?; + assert!( + crate::auth_read( + &mut *read, + crate::AuthRead::User, + &[], + pools.persistent(false) + ) + .await + .is_err() + ); + drop(read); + let mut write = pools.acquire(true).await?; + assert!( + crate::auth_read( + &mut *write, + crate::AuthRead::User, + &[], + pools.persistent(true) + ) + .await? + .is_empty() + ); + drop(write); + pools.close().await; + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn zero_timeouts_and_secret_redaction( + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = DatabaseSettings { + acquire_timeout: Duration::ZERO, + connect_timeout: Duration::ZERO, + ..settings? + }; + let debug = format!( + "{:?}", + DatabaseSettings { + password: Some(SecretValue::from("sensitive-test-password")), + ..settings.clone() + } + ); + assert!(!debug.contains("sensitive-test-password")); + assert!(debug.contains("[REDACTED]")); + let pools = DatabasePools::connect(settings, None).await?; + assert!(!pools.reader_available()); + assert!(pools.acquire(false).await.is_ok()); + assert!( + pools + .set_passwords(None, Some(SecretValue::from("reader"))) + .is_err() + ); + pools.close().await; + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6ecfebf0096..ffd2579999b 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,8 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +litellm-db.workspace = true +litellm-request-auth-core.workspace = true litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/db.rs b/litellm-rust/crates/python-bridge/src/db.rs index 03fb30b986f..4f9ed6eec16 100644 --- a/litellm-rust/crates/python-bridge/src/db.rs +++ b/litellm-rust/crates/python-bridge/src/db.rs @@ -1,19 +1,174 @@ -use litellm_host_python::ready_future; -use pyo3::{prelude::*, types::PyNone}; +use std::time::Duration; + +use litellm_db::{AuthRead, AuthSpend, DatabasePools, DatabaseSettings}; +use litellm_host_python::{Pythonized, ready_future, run_async_value}; +use litellm_secret_value::SecretValue; +use pyo3::{ + exceptions::{PyOSError, PyRuntimeError, PyValueError}, + prelude::*, + types::PyNone, +}; #[pyclass(frozen)] -pub(crate) struct NativeDatabase {} +pub(crate) struct NativeDatabase { + pools: Option, +} + +fn db_error(error: litellm_db::Error) -> PyErr { + if error.is_database_failure() { + return PyOSError::new_err(litellm_error::Report::new(&error).to_string()); + } + match error { + litellm_db::Error::InvalidSetting(_) | litellm_db::Error::InvalidAuthArguments => { + PyValueError::new_err(litellm_error::Report::new(&error).to_string()) + } + _ => PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string()), + } +} + +fn duration(seconds: f64, name: &'static str) -> PyResult { + Duration::try_from_secs_f64(seconds) + .map_err(|_| db_error(litellm_db::Error::InvalidSetting(name))) +} + +fn settings(value: &Bound<'_, PyAny>) -> PyResult { + let ssl_mode: String = value.getattr("ssl_mode")?.extract()?; + let idle_timeout: Option = value.getattr("idle_timeout")?.extract()?; + Ok(DatabaseSettings { + host: value.getattr("host")?.extract()?, + port: value.getattr("port")?.extract()?, + username: value.getattr("username")?.extract()?, + password: value + .getattr("password")? + .extract::>()? + .map(SecretValue::from), + database: value.getattr("database")?.extract()?, + schema: value.getattr("schema")?.extract()?, + ssl_mode: ssl_mode + .parse() + .map_err(|_| db_error(litellm_db::Error::InvalidSetting("ssl_mode")))?, + ssl_root_cert: value.getattr("ssl_root_cert")?.extract()?, + max_connections: value.getattr("max_connections")?.extract()?, + acquire_timeout: duration( + value.getattr("acquire_timeout")?.extract()?, + "acquire_timeout", + )?, + connect_timeout: duration( + value.getattr("connect_timeout")?.extract()?, + "connect_timeout", + )?, + idle_timeout: idle_timeout + .map(|seconds| duration(seconds, "idle_timeout")) + .transpose()?, + statement_cache_capacity: value.getattr("statement_cache_capacity")?.extract()?, + }) +} -// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first -// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens #[pymethods] impl NativeDatabase { #[staticmethod] - fn connect(py: Python<'_>) -> PyResult> { - ready_future(py, Bound::new(py, NativeDatabase {})?.as_any()) + #[pyo3(signature = (writer_settings=None, reader_settings=None))] + fn connect<'py>( + py: Python<'py>, + writer_settings: Option<&Bound<'py, PyAny>>, + reader_settings: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let writer = writer_settings.map(settings).transpose()?; + let reader = reader_settings.map(settings).transpose()?; + let Some(writer) = writer else { + if reader.is_some() { + return Err(db_error(litellm_db::Error::InvalidSetting( + "writer_settings", + ))); + } + return ready_future(py, Bound::new(py, Self { pools: None })?.as_any()); + }; + run_async_value(py, async move { + let pools = DatabasePools::connect(writer, reader) + .await + .map_err(db_error)?; + Ok(Self { pools: Some(pools) }) + }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { - ready_future(py, PyNone::get(py).as_any()) + let Some(pools) = self.pools.clone() else { + return ready_future(py, PyNone::get(py).as_any()); + }; + run_async_value(py, async move { + pools.close().await; + Ok(()) + }) + } + + #[getter] + fn reader_available(&self) -> bool { + self.pools + .as_ref() + .is_some_and(DatabasePools::reader_available) + } + + #[pyo3(signature = (writer_password=None, reader_password=None))] + fn set_passwords( + &self, + writer_password: Option, + reader_password: Option, + ) -> PyResult<()> { + self.pools + .as_ref() + .ok_or_else(|| db_error(litellm_db::Error::NotConfigured))? + .set_passwords( + writer_password.map(SecretValue::from), + reader_password.map(SecretValue::from), + ) + .map_err(db_error) + } + + #[pyo3(signature = (kind, ids, writer=false))] + fn auth_read<'py>( + &self, + py: Python<'py>, + kind: &str, + ids: Vec, + writer: bool, + ) -> PyResult> { + let kind: AuthRead = kind + .parse() + .map_err(|_| db_error(litellm_db::Error::InvalidAuthArguments))?; + let pools = self + .pools + .clone() + .ok_or_else(|| db_error(litellm_db::Error::NotConfigured))?; + run_async_value(py, async move { + let mut connection = pools.acquire(writer).await.map_err(db_error)?; + let rows = + litellm_db::auth_read(&mut *connection, kind, &ids, pools.persistent(writer)) + .await + .map_err(db_error)?; + Ok(Pythonized( + rows.into_iter().map(|row| row.data).collect::>(), + )) + }) + } + + fn auth_spend<'py>( + &self, + py: Python<'py>, + kind: &str, + id: String, + ) -> PyResult> { + let kind: AuthSpend = kind + .parse() + .map_err(|_| db_error(litellm_db::Error::InvalidAuthArguments))?; + let pools = self + .pools + .clone() + .ok_or_else(|| db_error(litellm_db::Error::NotConfigured))?; + run_async_value(py, async move { + let mut connection = pools.acquire(false).await.map_err(db_error)?; + litellm_db::auth_spend(&mut *connection, kind, &id, pools.persistent(false)) + .await + .map_err(db_error) + }) } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 09b787b2c94..953140011d9 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,6 +13,7 @@ mod model_registry; mod preflight; mod python_settings; mod redis; +mod request_auth; mod router; mod routes; mod secrets; @@ -42,7 +43,7 @@ mod _native { #[pymodule_export] use crate::redis::client::NativeRedisClient; #[pymodule_export] - use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes}; + use crate::redis::store::auth_cache::{NativeAuthCache, NativeAuthSubscription, PyLocalSizes}; #[pymodule_export] use crate::redis::store::rate_limits::{ NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired, @@ -61,6 +62,48 @@ mod _native { NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns, }; #[pymodule_export] + use crate::request_auth::{ + NativeAuthContinuation, request_auth_access, request_auth_admission, + request_auth_agent_access, request_auth_agent_ceiling, request_auth_agent_ceiling_source, + request_auth_agent_grants, request_auth_agent_group_load, request_auth_agent_invocation, + request_auth_budget_alert, request_auth_budget_alert_plan, request_auth_budget_exemption, + request_auth_budget_price, request_auth_budget_route, request_auth_budget_window_start, + request_auth_budgets, request_auth_claude_membership, request_auth_claude_normalization, + request_auth_codex_model, request_auth_codex_normalization, request_auth_codex_turn, + request_auth_common_budget, request_auth_common_budget_scopes, request_auth_common_flow, + request_auth_common_read, request_auth_context, request_auth_context_end_user, + request_auth_context_member, request_auth_context_models, request_auth_context_percentage, + request_auth_context_temporary, request_auth_context_throttle, request_auth_context_token, + request_auth_custom_flow, request_auth_end_user_apply, request_auth_end_user_assignments, + request_auth_end_user_default_id, request_auth_end_user_lookup, + request_auth_end_user_normalize_id, request_auth_end_user_registry, + request_auth_end_user_validation, request_auth_expired, request_auth_fallback_access, + request_auth_fallback_candidate, request_auth_key_model_access, + request_auth_litellm_header_token, request_auth_llm_route, request_auth_managed, + request_auth_managed_model_grants, request_auth_managed_models, request_auth_member_limit, + request_auth_member_models, request_auth_model_access, request_auth_model_attribution, + request_auth_model_budget_aliases, request_auth_model_budget_check, + request_auth_model_budget_query, request_auth_model_budget_resolve, + request_auth_model_budget_scope, request_auth_model_budget_team_exclusions, + request_auth_model_groups, request_auth_model_helpers, request_auth_model_projection, + request_auth_org_admin, request_auth_org_admin_raw, request_auth_org_route, + request_auth_owner_teams, request_auth_personal_budget, + request_auth_personal_model_listing, request_auth_personal_model_request, + request_auth_personal_models, request_auth_personal_owner, request_auth_reservation, + request_auth_resolved_models, request_auth_route_access, request_auth_route_enable, + request_auth_route_role, request_auth_route_role_api, request_auth_route_role_expand, + request_auth_route_role_inside, request_auth_route_role_jwt, + request_auth_route_role_listed, request_auth_route_scope, request_auth_scheme_credential, + request_auth_select_api_key, request_auth_select_websocket_api_key, + request_auth_spend_counter_policy, request_auth_start, request_auth_team_aliases, + request_auth_team_override, request_auth_unpriced_model, + }; + #[pymodule_export] + use crate::request_auth::{ + request_auth_budget_lookup_emails, request_auth_budget_lookup_member, + request_auth_budget_lookup_organization, request_auth_budget_lookup_soft, + }; + #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; @@ -115,6 +158,95 @@ mod tests { Python::initialize(); Python::attach(|py| { let mut expected = vec![ + "NativeAuthContinuation", + "request_auth_access", + "request_auth_admission", + "request_auth_agent_access", + "request_auth_agent_ceiling", + "request_auth_agent_ceiling_source", + "request_auth_agent_grants", + "request_auth_agent_group_load", + "request_auth_agent_invocation", + "request_auth_budget_alert", + "request_auth_budget_alert_plan", + "request_auth_budget_exemption", + "request_auth_budget_lookup_emails", + "request_auth_budget_lookup_member", + "request_auth_budget_lookup_organization", + "request_auth_budget_lookup_soft", + "request_auth_budget_price", + "request_auth_budget_route", + "request_auth_budget_window_start", + "request_auth_budgets", + "request_auth_common_budget", + "request_auth_common_budget_scopes", + "request_auth_common_flow", + "request_auth_common_read", + "request_auth_context", + "request_auth_context_end_user", + "request_auth_context_member", + "request_auth_context_models", + "request_auth_context_percentage", + "request_auth_context_temporary", + "request_auth_context_throttle", + "request_auth_context_token", + "request_auth_custom_flow", + "request_auth_end_user_apply", + "request_auth_end_user_assignments", + "request_auth_end_user_default_id", + "request_auth_end_user_lookup", + "request_auth_end_user_normalize_id", + "request_auth_end_user_registry", + "request_auth_end_user_validation", + "request_auth_expired", + "request_auth_fallback_access", + "request_auth_fallback_candidate", + "request_auth_key_model_access", + "request_auth_litellm_header_token", + "request_auth_llm_route", + "request_auth_managed", + "request_auth_managed_model_grants", + "request_auth_managed_models", + "request_auth_member_limit", + "request_auth_member_models", + "request_auth_model_access", + "request_auth_model_attribution", + "request_auth_model_budget_check", + "request_auth_model_budget_aliases", + "request_auth_model_budget_scope", + "request_auth_model_budget_team_exclusions", + "request_auth_model_budget_query", + "request_auth_model_budget_resolve", + "request_auth_model_groups", + "request_auth_model_helpers", + "request_auth_model_projection", + "request_auth_org_admin", + "request_auth_org_admin_raw", + "request_auth_org_route", + "request_auth_owner_teams", + "request_auth_personal_budget", + "request_auth_personal_model_listing", + "request_auth_personal_model_request", + "request_auth_personal_models", + "request_auth_personal_owner", + "request_auth_reservation", + "request_auth_resolved_models", + "request_auth_route_access", + "request_auth_route_enable", + "request_auth_route_role", + "request_auth_route_role_api", + "request_auth_route_role_expand", + "request_auth_route_role_inside", + "request_auth_route_role_jwt", + "request_auth_route_role_listed", + "request_auth_scheme_credential", + "request_auth_select_api_key", + "request_auth_select_websocket_api_key", + "request_auth_spend_counter_policy", + "request_auth_start", + "request_auth_team_aliases", + "request_auth_team_override", + "request_auth_unpriced_model", "RustBridgeDeclined", "RustUpstreamError", "ForkedAfterNativeRuntimeStarted", @@ -141,6 +273,7 @@ mod tests { "NativeRequestRedis", "NativeSend", "NativeAuthCache", + "NativeAuthSubscription", "NativeRateLimits", "WindowCount", "AtomicApplied", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs index 8d7788b136e..0d13dda3d3b 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs @@ -1,5 +1,9 @@ use std::sync::Arc; +use futures_util::{StreamExt, stream::BoxStream}; +use litellm_redis_client::SubscriptionEvent; +use tokio::sync::Mutex; + use litellm_python_compat::Value; use litellm_redis_store::auth_cache::{AuthCache, Settings}; use pyo3::prelude::*; @@ -44,6 +48,35 @@ impl NativeAuthCache { self.0.has_redis() } + fn attach(&self, client: &NativeRedisClient) -> bool { + self.0.attach(remote(client)) + } + + fn current_generation(&self) -> u64 { + self.0.current_generation() + } + + #[pyo3(signature = (key, value, generation, ttl_seconds = None))] + fn set_if_generation( + &self, + key: &str, + value: PyValue, + generation: u64, + ttl_seconds: Option, + ) -> PyResult { + Ok(self + .0 + .set_if_generation(key, value.0, generation, ttl(ttl_seconds)?)) + } + + fn subscribe(&self, client: &NativeRedisClient, channel: &str) -> NativeAuthSubscription { + self.0.subscriber_connected(false); + NativeAuthSubscription { + cache: self.0.clone(), + stream: Arc::new(Mutex::new(Some(client.client.subscribe_events(channel)))), + } + } + /// `request` answers a key its pre-call reads already found absent without another `GET`. #[pyo3(signature = (key, request = None))] fn get<'py>( @@ -245,6 +278,65 @@ impl NativeAuthCache { } } +#[pyclass(frozen)] +pub(crate) struct NativeAuthSubscription { + cache: Arc, + stream: Arc>>>, +} + +#[pymethods] +impl NativeAuthSubscription { + fn is_connected(&self) -> bool { + self.cache.local_ready() + } + + fn next<'py>(&self, py: Python<'py>) -> PyResult> { + let cache = self.cache.clone(); + let stream = self.stream.clone(); + litellm_host_python::run_async_value(py, async move { + let mut guard = stream.lock().await; + let Some(stream) = guard.as_mut() else { + return Ok(None::>); + }; + match stream.next().await { + Some(SubscriptionEvent::Connected) => cache.subscriber_connected(true), + Some(SubscriptionEvent::Message(payload)) => { + if let Ok(message) = serde_json::from_slice::(&payload) + && let Some(key) = + message.get("cache_key").and_then(serde_json::Value::as_str) + && message + .get("new_value") + .and_then(serde_json::Value::as_f64) + .is_none() + { + cache.delete_local(key).map_err(kind_error)?; + } + return Ok(Some(Python::attach(|py| { + pyo3::types::PyBytes::new(py, &payload).unbind() + }))); + } + _ => cache.subscriber_connected(false), + } + Ok(None) + }) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + self.cache.subscriber_connected(false); + let stream = self.stream.clone(); + litellm_host_python::run_async_value(py, async move { + stream.lock().await.take(); + Ok(()) + }) + } +} + +impl Drop for NativeAuthSubscription { + fn drop(&mut self) { + self.cache.subscriber_connected(false); + } +} + /// Entries the management-object and the user-key partitions hold in memory. #[pyclass(frozen, eq, get_all, name = "LocalSizes")] #[derive(PartialEq)] diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs new file mode 100644 index 00000000000..2ee093d8611 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -0,0 +1,1086 @@ +use std::sync::Mutex; + +use litellm_host_python::{from_py, to_py}; +use litellm_request_auth_core::{self as core, credential, protocol::*}; +use pyo3::{ + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, +}; + +type ProjectedStep = (Py, Option); + +#[pyfunction] +pub(crate) fn request_auth_spend_counter_policy( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::spend_counter_policy::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::spend_counter_policy::decide(facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_common_read( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::ReadFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::read(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_reservation( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::reservation::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::reservation::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_common_flow( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::plan(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_common_budget_scopes( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::BudgetFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::budgets(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_custom_flow( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::CustomFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::custom(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_expired( + py: Python<'_>, + expires: &Bound<'_, PyAny>, + now: &Bound<'_, PyAny>, +) -> PyResult { + let expires: core::protocol::DateTime = from_py(expires)?; + let now: core::protocol::DateTime = from_py(now)?; + py.detach(move || expires.before(&now)).map_err(map_error) +} + +#[pyfunction] +pub(crate) fn request_auth_context( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::auth_context::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::auth_context::plan(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_grants( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_grants::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::agent_grants::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_managed( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed::ManagedFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::managed::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_access( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_access::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::agent_access::decide(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_invocation( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_access::Invocation = from_py(facts)?; + to_py( + py, + &py.detach(move || core::agent_access::invocation(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_admission( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::admission::AdmissionFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::admission::plan(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_access(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::access::AccessFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::access::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_soft( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult { + let facts: core::budget_lookup::SoftFacts = from_py(facts)?; + Ok(py.detach(move || core::budget_lookup::read_soft_emails(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_organization( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::OrganizationFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_lookup::organization(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_member( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::MemberFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_lookup::member(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_emails( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::EmailFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_lookup::emails(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_window_start( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_budgets::WindowFacts = from_py(facts)?; + Ok(py.detach(move || core::common_budgets::window_start(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_resolve( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_data::ResolveFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::resolve(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_aliases( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_aliases::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_aliases::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_scope( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_aliases::ScopeFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_aliases::scope(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_team_exclusions( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: Vec = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::team_exclusions(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_query( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_data::QueryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::query(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_check( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_data::Snapshot = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::check(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_price( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_pricing::ModelPriceFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_pricing::price(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_route( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_pricing::RouteBudgetFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_pricing::route(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_common_budget( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_budgets::BudgetFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_budgets::check(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_alert( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult { + let facts: core::common_budgets::AlertFacts = from_py(facts)?; + Ok(py.detach(move || core::common_budgets::alert(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_member_limit( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_budgets::MemberLimitFacts = from_py(facts)?; + py.detach(move || core::common_budgets::member_limit(&facts)) + .map_err(map_error) +} + +fn map_error(error: core::Error) -> PyErr { + let message = litellm_error::Report::new(&error).to_string(); + match error { + core::Error::IncompatibleDateTimes => PyTypeError::new_err(message), + core::Error::InvalidAnswer => PyRuntimeError::new_err(message), + core::Error::Pattern(_) => PyRuntimeError::new_err(message), + } +} + +#[pyfunction] +pub(crate) fn request_auth_codex_normalization( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::codex_normalization::Facts = from_py(facts)?; + let decision = py + .detach(move || core::codex_normalization::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_claude_membership( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::claude_membership::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::claude_membership::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_claude_normalization( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::claude_normalization::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::claude_normalization::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_codex_turn( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::codex_turn::Request = from_py(facts)?; + let decision = py + .detach(move || core::codex_turn::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_scope( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_scope::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::route_scope::decide(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_codex_model( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::codex_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::codex_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +fn project(py: Python<'_>, step: Result) -> PyResult { + let step = step.map_err(map_error)?; + Ok(( + to_py(py, &step.node)?, + step.resume.map(|resume| NativeAuthContinuation { + resume: Mutex::new(Some(resume)), + }), + )) +} + +#[pyclass(frozen, module = "litellm.rust_bridge._native")] +pub(crate) struct NativeAuthContinuation { + resume: Mutex>, +} + +#[pymethods] +impl NativeAuthContinuation { + fn resume(&self, py: Python<'_>, answer: &Bound<'_, PyAny>) -> PyResult { + let resume = self + .resume + .lock() + .map_err(|_| PyRuntimeError::new_err("auth continuation lock poisoned"))? + .take() + .ok_or_else(|| PyRuntimeError::new_err("auth continuation already consumed"))?; + let answer: Answer = from_py(answer)?; + project(py, py.detach(move || resume(answer))) + } +} + +#[pyfunction] +pub(crate) fn request_auth_start( + py: Python<'_>, + facts: &Bound<'_, PyAny>, + settings: &Bound<'_, PyAny>, + now: &Bound<'_, PyAny>, + resolved: bool, +) -> PyResult { + let facts: RequestFacts = from_py(facts)?; + let settings: Settings = from_py(settings)?; + let now: DateTime = from_py(now)?; + project( + py, + py.detach(move || { + if resolved { + core::start_resolved_key(facts, settings, now) + } else { + core::start(facts, settings, now) + } + }), + ) +} + +#[derive(serde::Deserialize, strum::VariantArray)] +enum BudgetEntry { + Key, + User, + TeamMember, +} + +#[pyfunction] +pub(crate) fn request_auth_budgets( + py: Python<'_>, + entry: &Bound<'_, PyAny>, +) -> PyResult { + let entry: BudgetEntry = from_py(entry)?; + project( + py, + py.detach(move || match entry { + BudgetEntry::Key => core::check_model_budgets(), + BudgetEntry::User => core::check_user_model_budgets(), + BudgetEntry::TeamMember => core::check_team_member_model_budgets(), + }), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_select_api_key( + py: Python<'_>, + keys: &Bound<'_, PyAny>, +) -> PyResult> { + let keys: credential::PresentedKeys = from_py(keys)?; + to_py(py, &py.detach(move || credential::select_api_key(keys))) +} + +#[pyfunction] +pub(crate) fn request_auth_select_websocket_api_key( + py: Python<'_>, + keys: &Bound<'_, PyAny>, +) -> PyResult> { + let keys: credential::WebSocketKeys = from_py(keys)?; + to_py( + py, + &py.detach(move || credential::select_websocket_api_key(keys)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_scheme_credential(value: &str) -> Option { + credential::scheme_credential(value) +} + +#[pyfunction] +pub(crate) fn request_auth_litellm_header_token(value: &str) -> String { + credential::litellm_header_token(value) +} + +#[pyfunction] +pub(crate) fn request_auth_route_access( + py: Python<'_>, + request: &Bound<'_, PyAny>, +) -> PyResult> { + let request: core::route_access::Request = from_py(request)?; + let decision = py + .detach(move || core::route_access::decide(request)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_org_route( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::org_routes::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::org_routes::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_org_admin( + py: Python<'_>, + memberships: &Bound<'_, PyAny>, + organizations: Vec, +) -> PyResult { + let memberships: Vec = from_py(memberships)?; + Ok(py.detach(move || core::org_routes::user_is_admin(&memberships, &organizations))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_access( + py: Python<'_>, + request: &Bound<'_, PyAny>, +) -> PyResult> { + let request: core::model_access::Request = from_py(request)?; + let decision = py + .detach(move || core::model_access::decide(request)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_model_groups( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_groups::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::model_groups::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_fallback_access( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::fallback_access::Facts = from_py(facts)?; + let decision = py + .detach(move || core::fallback_access::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_exemption( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_pricing::ExemptionFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_pricing::exemption(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_key_model_access( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::key_model_access::Facts = from_py(facts)?; + let decision = py + .detach(move || core::key_model_access::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_unpriced_model( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::unpriced_models::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::unpriced_models::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_alert_plan( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_alerts::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_alerts::plan(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_attribution( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_attribution::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_attribution::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_fallback_candidate( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::fallback_candidate::Facts = from_py(facts)?; + let decision = py + .detach(move || core::fallback_candidate::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_budget( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_budget::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_budget::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_models::Request = from_py(facts)?; + let decision = py + .detach(move || core::personal_models::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_owner_teams( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::owner_teams::Facts = from_py(facts)?; + let decision = py + .detach(move || core::owner_teams::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_context_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::ModelCheckFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::model_checks(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_end_user( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::EndUserFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::end_user_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_token( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::TokenDefaultsFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::token_assignments(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_member( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::MemberFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::member_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_temporary( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::TemporaryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::temporary(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_throttle( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::ThrottleFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::throttle(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_percentage( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::NumericFacts = from_py(facts)?; + Ok(py.detach(move || core::budget_context::percentage(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_team_override( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::team_override::Facts = from_py(facts)?; + let decision = py + .detach(move || core::team_override::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_model_helpers( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_helpers::Facts = from_py(facts)?; + let decision = py + .detach(move || core::model_helpers::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_owner( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_owner::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_owner::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_model_request( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_model_request::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_model_request::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_member_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::member_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::member_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_resolved_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::resolved_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::resolved_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_normalize_id( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::IdFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::normalize_id(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_default_id( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::DefaultFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::default_id(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_apply( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::ApplyFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::apply_default(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_registry( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::RegistryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::registry(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_assignments( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::AssignmentFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::assignments(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_lookup( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::LookupFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::lookup(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_validation( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::ValidationFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::validation(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_model_listing( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_model_listing::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_model_listing::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_ceiling( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_ceilings::Facts = from_py(facts)?; + let decision = py + .detach(move || core::agent_ceilings::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_group_load( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_ceilings::LoadFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::agent_ceilings::load(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_team_aliases( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::team_aliases::Facts = from_py(facts)?; + let decision = py + .detach(move || core::team_aliases::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_managed_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::managed_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_ceiling_source( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed_models::CeilingSource = from_py(facts)?; + to_py( + py, + &py.detach(move || core::managed_models::ceiling_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_managed_model_grants( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed_models::GrantFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::managed_models::grants(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::RoleFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::role(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_api( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::ApiFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::api(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_jwt( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::JwtFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::jwt(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_inside( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::InsideFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::inside(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_expand( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::Entry = from_py(facts)?; + let decision = py.detach(move || core::route_roles::expand(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_listed( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::ListedFacts = from_py(facts)?; + let decision = py + .detach(move || core::route_roles::listed(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_org_admin_raw( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::org_routes::AdminFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::org_routes::admin(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_projection( + py: Python<'_>, + facts: &Bound<'_, PyAny>, + reads: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_access::Facts = from_py(facts)?; + let reads: Vec = from_py(reads)?; + let decision = py + .detach(move || core::model_projection::decide(&facts, &reads)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_enable( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_enable::Facts = from_py(facts)?; + let decision = py + .detach(move || core::route_enable::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +#[pyo3(signature = (context, configured=None))] +pub(crate) fn request_auth_llm_route( + py: Python<'_>, + context: &Bound<'_, PyAny>, + configured: Option, +) -> PyResult> { + let context: core::route_access::RouteContext = from_py(context)?; + let decision = py + .detach(move || core::route_enable::llm_route(&context, configured)) + .map_err(map_error)?; + to_py(py, &decision) +} diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index d9a8b3af7e8..4f477c3ee3c 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -155,6 +155,10 @@ impl RedisClient { self.0.subscriber.clone().messages(channel.to_owned()) } + pub fn subscribe_events(&self, channel: &str) -> BoxStream<'static, crate::SubscriptionEvent> { + self.0.subscriber.clone().events(channel.to_owned()) + } + pub fn pipeline(&self) -> Pipeline { Pipeline::new(self.0.namespace.clone()) } diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 2754666c329..1a995334ced 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -23,6 +23,7 @@ pub use litellm_redis_types::{ }; pub use namespace::Namespace; pub use pipeline::Pipeline; +pub use pubsub::SubscriptionEvent; pub use redis::{ BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value, cluster_routing::Slot, diff --git a/litellm-rust/crates/redis-client/src/pubsub.rs b/litellm-rust/crates/redis-client/src/pubsub.rs index b6b18d3e481..a0004a239b6 100644 --- a/litellm-rust/crates/redis-client/src/pubsub.rs +++ b/litellm-rust/crates/redis-client/src/pubsub.rs @@ -1,7 +1,10 @@ use std::{sync::Arc, time::Duration}; use futures_util::{StreamExt, stream::BoxStream}; -use redis::{ConnectionInfo, RedisResult, aio::PubSubStream}; +use redis::{ + ConnectionInfo, RedisResult, + aio::{PubSubSink, PubSubStream}, +}; use crate::{ ConnectError, Credentials, Standalone, Target, Timeouts, config, @@ -11,6 +14,18 @@ use crate::{ // litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS const BACKOFF_INITIAL: Duration = Duration::from_secs(5); const BACKOFF_MAX: Duration = Duration::from_secs(60); +const IDLE_PROBE: Duration = Duration::from_secs(15); + +pub enum SubscriptionEvent { + Connected, + Disconnected, + Message(Vec), +} + +struct Live { + sink: PubSubSink, + stream: PubSubStream, +} /// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends /// `SUBSCRIBE`, authenticated once when it connects. @@ -66,43 +81,76 @@ impl Subscriber { } pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec> { + self.events(channel) + .filter_map(|event| async move { + match event { + SubscriptionEvent::Message(payload) => Some(payload), + _ => None, + } + }) + .boxed() + } + + pub(crate) fn events(self, channel: String) -> BoxStream<'static, SubscriptionEvent> { futures_util::stream::unfold( - (self, channel, None::, BACKOFF_INITIAL), - |(subscriber, channel, live, backoff)| async move { + (self, channel, None::, BACKOFF_INITIAL, false), + |(subscriber, channel, live, backoff, retry)| async move { let mut live = live; let mut backoff = backoff; + if retry { + tokio::time::sleep(backoff).await; + } loop { let Some(current) = live.as_mut() else { - match subscriber.subscribe(&channel).await { + let event = match subscriber.subscribe(&channel).await { Ok(subscribed) => { live = Some(subscribed); backoff = BACKOFF_INITIAL; + SubscriptionEvent::Connected } - Err(error) => { - tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying"); - tokio::time::sleep(backoff).await; - backoff = (backoff * 2).min(BACKOFF_MAX); + Err(_) => { + if retry { + backoff = (backoff * 2).min(BACKOFF_MAX); + } + SubscriptionEvent::Disconnected } - } - continue; + }; + let retry = live.is_none(); + return Some((event, (subscriber, channel, live, backoff, retry))); }; - match current.next().await { - Some(message) => { + match tokio::time::timeout(IDLE_PROBE, current.stream.next()).await { + Ok(Some(message)) => { let payload = message.get_payload_bytes().to_vec(); - return Some((payload, (subscriber, channel, live, backoff))); + return Some(( + SubscriptionEvent::Message(payload), + (subscriber, channel, live, backoff, false), + )); } - None => live = None, + Err(_) => { + if matches!( + tokio::time::timeout( + subscriber.connect_timeout, + current.sink.ping::() + ) + .await, + Ok(Ok(_)) + ) { + continue; + } + } + Ok(None) => {} } - tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing"); - tokio::time::sleep(backoff).await; - backoff = (backoff * 2).min(BACKOFF_MAX); + return Some(( + SubscriptionEvent::Disconnected, + (subscriber, channel, None, backoff, true), + )); } }, ) .boxed() } - async fn subscribe(&self, channel: &str) -> RedisResult { + async fn subscribe(&self, channel: &str) -> RedisResult { match &self.source { Source::Sentinel(sentinel) => { let master = sentinel.master().await?; @@ -122,7 +170,7 @@ impl Subscriber { } } - async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { + async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { let settings = config::authenticated(info.redis_settings().clone(), &self.credentials).await?; let client = redis::Client::open(info.set_redis_settings(settings))?; @@ -135,7 +183,8 @@ impl Subscriber { )) })??; subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?; - Ok(pubsub.into_on_message()) + let (sink, stream) = pubsub.split(); + Ok(Live { sink, stream }) } } diff --git a/litellm-rust/crates/redis-client/tests/pubsub.rs b/litellm-rust/crates/redis-client/tests/pubsub.rs index 2ad7cb38a5a..e3fda48d765 100644 --- a/litellm-rust/crates/redis-client/tests/pubsub.rs +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -7,7 +7,7 @@ use litellm_circuit_breaker::{Breaker, Config}; use litellm_redis_client::commands; use litellm_redis_client::{ Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, - Target, Timeouts, + SubscriptionEvent, Target, Timeouts, }; use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually}; use rstest::rstest; @@ -16,6 +16,85 @@ use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation"; +#[rstest] +#[tokio::test] +async fn a_silent_dead_subscription_is_detected_and_reconnected( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let proxy = support::proxied(&server).await; + let client = RedisClient::new(RedisConfig { + target: proxy.target.clone(), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let mut events = client.subscribe_events(CHANNEL); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); + proxy.blackhole_open_connections(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(20), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Disconnected) + )); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(10), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); +} + +#[rstest] +#[tokio::test] +async fn lifecycle_reports_disconnect_before_reconnect( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let mut events = client.subscribe_events(CHANNEL); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); + server.restart().await.unwrap(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Disconnected) + )); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(15), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); + eventually(async || { + client + .send(commands::publish(CHANNEL, b"after".as_slice())) + .await + .ok() + }) + .await; + assert!( + matches!(tokio::time::timeout(Duration::from_secs(5), events.next()).await.unwrap(), Some(SubscriptionEvent::Message(bytes)) if bytes == b"after") + ); +} + fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver> { let (sender, receiver) = unbounded_channel(); let mut messages = client.subscribe(channel); diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml index 1dd61d76d3c..1cd752abbef 100644 --- a/litellm-rust/crates/redis-store/Cargo.toml +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -11,6 +11,7 @@ litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-python-compat.workspace = true litellm-redis-client.workspace = true +moka = { workspace = true, features = ["sync"] } strum.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } diff --git a/litellm-rust/crates/redis-store/src/auth_cache/local.rs b/litellm-rust/crates/redis-store/src/auth_cache/local.rs new file mode 100644 index 00000000000..ed533fd2580 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/local.rs @@ -0,0 +1,108 @@ +use std::time::Duration; + +use litellm_python_compat::Value; +use moka::sync::Cache; + +use super::{Settings, is_key_object}; + +const MAX_AUTH_TTL: Duration = Duration::from_secs(60); + +#[derive(Clone)] +struct Entry { + value: Value, + expires: Duration, +} + +pub(super) struct Local { + objects: Cache, + keys: Cache, + pub generation: u64, + pub healthy: bool, +} + +impl Local { + pub fn new(settings: Settings) -> Self { + let partition = || { + Cache::builder() + .max_capacity(settings.memory_entries as u64) + .build() + }; + Self { + objects: partition(), + keys: partition(), + generation: 0, + healthy: true, + } + } + + fn partition(&self, key: &str) -> &Cache { + if is_key_object(key) { + &self.keys + } else { + &self.objects + } + } + + pub fn get(&self, key: &str, now: Duration) -> Option { + if !self.healthy { + return None; + } + let entry = self.partition(key).get(key)?; + if now >= entry.expires { + self.partition(key).invalidate(key); + return None; + } + Some(entry.value) + } + + pub fn set(&self, key: &str, value: Value, ttl: Option, now: Duration) { + if !self.healthy || key.starts_with("sk-") { + return; + } + let ttl = ttl.unwrap_or(Duration::from_secs(600)); + let ttl = if is_key_object(key) { + ttl.min(MAX_AUTH_TTL) + } else { + ttl + }; + self.partition(key).insert( + key.to_owned(), + Entry { + value, + expires: now.saturating_add(ttl), + }, + ); + } + + pub fn delete(&mut self, key: &str) { + self.generation = self.generation.wrapping_add(1); + self.partition(key).invalidate(key); + } + + pub fn flush(&mut self) { + self.generation = self.generation.wrapping_add(1); + self.keys.invalidate_all(); + self.objects.invalidate_all(); + } + + pub fn reconfigure(&mut self, settings: Settings) { + let replacement = Self::new(settings); + for (key, value) in &self.keys { + replacement.keys.insert((*key).clone(), value); + } + for (key, value) in &self.objects { + replacement.objects.insert((*key).clone(), value); + } + self.keys = replacement.keys; + self.objects = replacement.objects; + } + + pub fn len(&self) -> (usize, usize) { + self.objects.run_pending_tasks(); + self.keys.run_pending_tasks(); + ( + self.objects.entry_count() as usize, + self.keys.entry_count() as usize, + ) + } +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/mod.rs b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs index b5a83b5aec6..4ddb68df710 100644 --- a/litellm-rust/crates/redis-store/src/auth_cache/mod.rs +++ b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs @@ -1,4 +1,5 @@ mod keys; +mod local; mod types; use std::{ @@ -12,23 +13,19 @@ use litellm_redis_client::{Declared, RedisBatch, commands}; pub use self::keys::is_key_object; pub use self::types::Settings; use crate::{ - Error, Number, Pending, RequestRedis, Send, - dual_cache::DualCache, + Error, Number, Pending, RequestRedis, Send, codec, + failure_log::redis_failure, memory::{Clock, lock, system_clock}, remote::{Remote, RemoteSlot}, - request::ready, + request::{pending, ready}, }; /// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it. const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"; -/// The proxy's auth object cache with `enable_redis_auth_cache` on: `UserApiKeyCache`'s two -/// `DualCache` partitions, so churn in management objects cannot evict user-key objects, in front of -/// one Redis shared by every pod. Keys are the logical cache keys `UserApiKeyCache` callers already -/// use; values are what `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`. pub struct AuthCache { - key_objects: DualCache, - objects: DualCache, + local: Mutex, + now: Clock, redis: Arc, settings: Mutex, } @@ -40,10 +37,9 @@ impl AuthCache { pub fn with_clock(remote: Option, settings: Settings, now: Clock) -> Self { let redis = Arc::new(RemoteSlot::new(remote)); - let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone()); Self { - key_objects: partition(), - objects: partition(), + local: Mutex::new(local::Local::new(settings)), + now, redis, settings: Mutex::new(settings), } @@ -53,6 +49,40 @@ impl AuthCache { self.redis.get().is_some() } + pub fn attach(&self, remote: Remote) -> bool { + self.redis.attach(remote) + } + + pub fn current_generation(&self) -> u64 { + lock(&self.local).generation + } + + pub fn local_ready(&self) -> bool { + lock(&self.local).healthy + } + + pub fn subscriber_connected(&self, connected: bool) { + let mut local = lock(&self.local); + local.flush(); + local.healthy = connected; + } + + pub fn set_if_generation( + &self, + key: &str, + value: Value, + generation: u64, + ttl: Option, + ) -> bool { + let ttl = ttl.or(self.settings().memory_ttl); + let local = lock(&self.local); + if local.generation != generation || !local.healthy || key.starts_with("sk-") { + return false; + } + local.set(key, value, ttl, (self.now)()); + true + } + pub fn settings(&self) -> Settings { *lock(&self.settings) } @@ -61,37 +91,33 @@ impl AuthCache { /// keep what they hold. pub fn reconfigure(&self, settings: Settings) { *lock(&self.settings) = settings; - self.key_objects.reconfigure(settings.memory()); - self.objects.reconfigure(settings.memory()); + lock(&self.local).reconfigure(settings); } pub fn get_local(&self, key: &str) -> Result, Error> { - Ok(self.partition(key).memory().get(key)?) + Ok(lock(&self.local).get(key, (self.now)())) } /// A memory write with `ttl`, or `default_in_memory_ttl` when it names none. pub fn set_local(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { - let partition = self.partition(key); - Ok(partition - .memory() - .set(key, value, ttl.or(partition.default_in_memory_ttl()))?) + let ttl = ttl.or(self.settings().memory_ttl); + lock(&self.local).set(key, value, ttl, (self.now)()); + Ok(()) } pub fn delete_local(&self, key: &str) -> Result<(), Error> { - Ok(self.partition(key).memory().delete(key)?) + lock(&self.local).delete(key); + Ok(()) } pub fn flush_local(&self) -> Result<(), Error> { - self.key_objects.memory().flush()?; - Ok(self.objects.memory().flush()?) + lock(&self.local).flush(); + Ok(()) } /// Entries held by the management-object and the user-key partitions. pub fn len_local(&self) -> Result<(usize, usize), Error> { - Ok(( - self.objects.memory().len()?, - self.key_objects.memory().len()?, - )) + Ok(lock(&self.local).len()) } /// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read @@ -101,12 +127,42 @@ impl AuthCache { request: Option<&RequestRedis>, key: &str, ) -> Result, Error> { - self.partition(key).get(request, key, None).await + let generation = self.current_generation(); + if let Some(value) = self.get_local(key)? { + return Ok(Some(value)); + } + if !lock(&self.local).healthy || key.starts_with("sk-") { + return Ok(None); + } + let Some(client) = self.redis.client() else { + return Ok(None); + }; + if request.is_some_and(|request| request.misses(client).contains(key)) { + return Ok(None); + } + let value = client + .send(commands::get::>(key)) + .await + .map_err(Error::from) + .and_then(|raw| Ok(types::decode(raw)?)) + .unwrap_or_else(|error| { + redis_failure!(&error, "auth cache read failed"); + None + }); + if let Some(value) = &value + && !self.set_if_generation(key, value.clone(), generation, None) + { + return Ok(None); + } + Ok(value) } /// `DualCache.async_set_cache`, writing `json.dumps(value)`. pub async fn set(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { - self.partition(key).set(key, value, ttl).await + if let Err(error) = self.declare_set(Send::Now, key, value, ttl)?.await { + redis_failure!(&error, "auth cache write failed"); + } + Ok(()) } /// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`. @@ -116,7 +172,14 @@ impl AuthCache { value: Value, ttl: Option, ) -> Result<(), Error> { - self.partition(key).set_str(key, value, ttl).await + let encoded = codec::str(value.clone()); + if let Err(error) = self + .declare_encoded(Send::Now, key, value, ttl, Ok(encoded))? + .await + { + redis_failure!(&error, "auth cache write failed"); + } + Ok(()) } /// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks. @@ -127,12 +190,56 @@ impl AuthCache { value: Value, ttl: Option, ) -> Result, Error> { - self.partition(key).declare_set(send, key, value, ttl) + let encoded = codec::dumps(value.clone()); + self.declare_encoded(send, key, value, ttl, encoded) + } + + fn declare_encoded( + &self, + send: Send<'_>, + key: &str, + value: Value, + ttl: Option, + encoded: Result, + ) -> Result, Error> { + let ttl = ttl.or(self.settings().memory_ttl); + self.set_local(key, value, ttl)?; + if key.starts_with("sk-") { + return Ok(ready(())); + } + let Some(remote) = self.redis.get() else { + return Ok(ready(())); + }; + let encoded = match encoded { + Ok(encoded) => encoded, + Err(error) => { + redis_failure!(&error, "auth cache value did not encode"); + return Ok(ready(())); + } + }; + if let Some(misses) = send.misses(&remote.client) { + misses.forget(key); + } + Ok(pending(send.batch(&remote.client).add(commands::set( + key, + encoded, + ttl.unwrap_or(remote.default_ttl), + )))) } /// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks. pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { - self.partition(key).declare_delete(send, key) + self.delete_local(key)?; + let Some(client) = self.redis.client() else { + return Ok(ready(())); + }; + if let Some(misses) = send.misses(client) { + misses.note([key.to_owned()]); + } + let deleted = send.batch(client).add(commands::del(&[key])); + Ok(Box::pin( + async move { Ok(deleted.resolve().await.map(drop)?) }, + )) } /// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then @@ -142,10 +249,30 @@ impl AuthCache { entries: Vec<(String, Value)>, ttl: Option, ) -> Result<(), Error> { - let (key_objects, others): (Vec<_>, Vec<_>) = + let (keys, objects): (Vec<_>, Vec<_>) = entries.into_iter().partition(|(key, _)| is_key_object(key)); - self.key_objects.set_pipeline(key_objects, ttl).await?; - self.objects.set_pipeline(others, ttl).await + self.set_partition(keys, ttl).await?; + self.set_partition(objects, ttl).await + } + + async fn set_partition( + &self, + entries: Vec<(String, Value)>, + ttl: Option, + ) -> Result<(), Error> { + let ttl = ttl.or(self.settings().memory_ttl); + let entries = entries + .into_iter() + .map(|(key, value)| (key, value, ttl)) + .collect(); + let result = match self.write_back(Send::Now, entries) { + Ok(written) => written.await, + Err(error) => Err(error), + }; + if let Err(error) = result { + redis_failure!(&error, "auth cache pipeline write failed"); + } + Ok(()) } /// The auth prefetch write-back: memory first, an entry without a TTL taking @@ -158,6 +285,10 @@ impl AuthCache { send: Send<'_>, entries: Vec<(String, Value, Option)>, ) -> Result, Error> { + let entries: Vec<_> = entries + .into_iter() + .filter(|(key, _, _)| !key.starts_with("sk-")) + .collect(); let encoded = self .redis .get() @@ -199,13 +330,30 @@ impl AuthCache { /// `UserApiKeyCache.async_delete_cache_keys`: the key-object partition and the rest as two /// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails. pub async fn delete(&self, keys: &[String]) -> Result<(), Error> { - let (key_objects, others): (Vec, Vec) = + for key in keys { + self.delete_local(key)?; + } + let (keys, objects): (Vec<_>, Vec<_>) = keys.iter().cloned().partition(|key| is_key_object(key)); - let (key_objects, others) = tokio::join!( - self.key_objects.delete_keys(&key_objects), - self.objects.delete_keys(&others) + let (keys, objects) = tokio::join!( + self.delete_partition(&keys), + self.delete_partition(&objects) ); - key_objects.and(others) + keys.and(objects) + } + + async fn delete_partition(&self, keys: &[String]) -> Result<(), Error> { + let Some(client) = self.redis.client() else { + return Ok(()); + }; + for chunk in keys.chunks(1000) { + client + .send(commands::del( + &chunk.iter().map(String::as_str).collect::>(), + )) + .await?; + } + Ok(()) } /// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on @@ -247,12 +395,67 @@ impl AuthCache { .send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount)) .await?) } +} - fn partition(&self, key: &str) -> &DualCache { - if is_key_object(key) { - &self.key_objects - } else { - &self.objects - } +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use std::sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }; + + const KEY: &str = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"; + + #[rstest] + fn invalidation_fences_an_in_flight_fill_and_allows_a_new_one() { + let cache = AuthCache::new(None, Settings::default()); + let before = cache.current_generation(); + cache.delete_local(KEY).unwrap(); + assert!(!cache.set_if_generation(KEY, Value::from("revoked"), before, None)); + assert_eq!(cache.get_local(KEY).unwrap(), None); + let after = cache.current_generation(); + assert!(cache.set_if_generation(KEY, Value::from("fresh"), after, None)); + assert_eq!(cache.get_local(KEY).unwrap(), Some(Value::from("fresh"))); + } + + #[rstest] + fn disconnect_bypasses_writes_and_reconnect_fences_old_fills() { + let cache = AuthCache::new(None, Settings::default()); + cache.set_local(KEY, Value::Bool(true), None).unwrap(); + cache.subscriber_connected(false); + assert_eq!(cache.get_local(KEY).unwrap(), None); + cache.set_local(KEY, Value::Bool(true), None).unwrap(); + let disconnected = cache.current_generation(); + assert!(!cache.set_if_generation(KEY, Value::Bool(true), disconnected, None)); + cache.subscriber_connected(true); + assert_eq!(cache.get_local(KEY).unwrap(), None); + assert!(!cache.set_if_generation(KEY, Value::Bool(true), disconnected, None)); + } + + #[rstest] + fn auth_entries_expire_even_without_a_configured_ttl() { + let seconds = Arc::new(AtomicU64::new(100)); + let clock = seconds.clone(); + let cache = AuthCache::with_clock( + None, + Settings::default(), + Arc::new(move || Duration::from_secs(clock.load(Ordering::SeqCst))), + ); + cache.set_local(KEY, Value::Bool(true), None).unwrap(); + assert_eq!(cache.get_local(KEY).unwrap(), Some(Value::Bool(true))); + seconds.store(160, Ordering::SeqCst); + assert_eq!(cache.get_local(KEY).unwrap(), None); + } + + #[rstest] + fn raw_credentials_are_not_cached() { + let cache = AuthCache::new(None, Settings::default()); + cache + .set_local("sk-do-not-cache", Value::Bool(true), None) + .unwrap(); + assert_eq!(cache.get_local("sk-do-not-cache").unwrap(), None); + assert_eq!(cache.len_local().unwrap(), (0, 0)); } } diff --git a/litellm-rust/crates/redis-store/src/auth_cache/types.rs b/litellm-rust/crates/redis-store/src/auth_cache/types.rs index e933ac19af9..d66ae858e44 100644 --- a/litellm-rust/crates/redis-store/src/auth_cache/types.rs +++ b/litellm-rust/crates/redis-store/src/auth_cache/types.rs @@ -2,7 +2,7 @@ use std::time::Duration; use litellm_python_compat::Value; -use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings}; +use crate::{DecodeError, codec::python_value}; /// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -21,15 +21,6 @@ impl Default for Settings { } } -impl Settings { - pub(super) fn memory(self) -> MemorySettings { - MemorySettings { - entries: Some(self.memory_entries), - ttl: self.memory_ttl, - } - } -} - pub(super) fn decode(raw: Option>) -> Result, DecodeError> { raw.as_deref().map(python_value).transpose() } diff --git a/litellm-rust/crates/redis-store/src/dual_cache/mod.rs b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs index b2c277a118e..4158653be4b 100644 --- a/litellm-rust/crates/redis-store/src/dual_cache/mod.rs +++ b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs @@ -88,6 +88,7 @@ impl DualCache { /// `DualCache.update_cache_ttl` and `update_in_memory_max_size` together: both apply in place, /// and the entries memory already holds stay. + #[cfg(test)] pub(crate) fn reconfigure(&self, memory: MemorySettings) { *lock(&self.default_in_memory_ttl) = memory.ttl; self.memory.update_max_size(memory.entries); diff --git a/litellm-rust/crates/redis-store/src/dual_cache/write.rs b/litellm-rust/crates/redis-store/src/dual_cache/write.rs index 6fcc6f32741..79d843783df 100644 --- a/litellm-rust/crates/redis-store/src/dual_cache/write.rs +++ b/litellm-rust/crates/redis-store/src/dual_cache/write.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_redis_client::{Declared, RedisBatch, commands}; +use litellm_redis_client::commands; use super::DualCache; use crate::{ @@ -35,69 +35,6 @@ impl DualCache { .await } - /// `DualCache._set_on_batch`: memory now, the `SET` on the batch `send` picks. Memory and Redis - /// both take `ttl`, else `default_in_memory_ttl`, and Redis falls back to its own default. - pub(crate) fn declare_set( - &self, - send: Send<'_>, - key: &str, - value: V, - ttl: Option, - ) -> Result, Error> { - self.declare_encoded(send, key, value, ttl, codec::dumps) - } - - /// `DualCache.async_set_cache_pipeline`: memory, then one pipeline of `SET`s. A value - /// `json.dumps` refuses aborts the whole pipeline, and any failure only logs, as `RedisCache` - /// does. - pub(crate) async fn set_pipeline( - &self, - entries: Vec<(String, V)>, - ttl: Option, - ) -> Result<(), Error> { - let ttl = ttl.or(self.default_in_memory_ttl()); - let encoded = self.redis.get().map(|_| { - entries - .iter() - .map(|(_, value)| codec::dumps(value.clone())) - .collect::, Error>>() - }); - let keys = entries - .into_iter() - .map(|(key, value)| { - self.memory.set(&key, value, ttl)?; - Ok(key) - }) - .collect::, Error>>()?; - let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { - return Ok(()); - }; - let failure = match encoded { - Err(error) => Some(error), - Ok(values) => { - let pipeline = RedisBatch::pipelined(remote.client.clone()); - let ttl = ttl.unwrap_or(remote.default_ttl); - let written = keys - .iter() - .zip(values) - .map(|(key, value)| pipeline.add(commands::set(key, value, ttl))) - .collect::>(); - Declared::all(written) - .await - .into_iter() - .find_map(Result::err) - .map(Error::from) - } - }; - if let Some(error) = failure { - redis_failure!( - &error, - "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", - ); - } - Ok(()) - } - async fn set_now( &self, key: &str, @@ -147,22 +84,6 @@ impl DualCache { } impl DualCache { - /// `DualCache.async_delete_cache_pre_call`: memory now, the `DEL` on the batch `send` picks, - /// after which the request reads the key as missing. - pub(crate) fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { - self.memory.delete(key)?; - let Some(client) = self.redis.client() else { - return Ok(ready(())); - }; - if let Some(misses) = send.misses(client) { - misses.note([key.to_owned()]); - } - let deleted = send.batch(client).add(commands::del(&[key])); - Ok(Box::pin( - async move { Ok(deleted.resolve().await.map(drop)?) }, - )) - } - /// `DualCache.async_delete_cache_keys`: memory, then one multi-key `DEL` per chunk, in order. pub(crate) async fn delete_keys(&self, keys: &[String]) -> Result<(), Error> { keys.iter().try_for_each(|key| self.memory.delete(key))?; diff --git a/litellm-rust/crates/redis-store/src/memory.rs b/litellm-rust/crates/redis-store/src/memory.rs index f89d5926760..af6ac9bb0e1 100644 --- a/litellm-rust/crates/redis-store/src/memory.rs +++ b/litellm-rust/crates/redis-store/src/memory.rs @@ -50,10 +50,6 @@ impl Memory { self.cache.flush_cache() } - pub fn len(&self) -> Result { - self.cache.len() - } - pub fn expires_at(&self, key: &str) -> Result, Error> { self.cache.expires_at(key) } @@ -63,6 +59,7 @@ impl Memory { } /// `DualCache.update_in_memory_max_size`: `None` restores the 200-entry default. + #[cfg(test)] pub fn update_max_size(&self, entries: Option) { self.cache.update_max_size_in_memory(entries); } diff --git a/litellm-rust/crates/request-auth-core/Cargo.toml b/litellm-rust/crates/request-auth-core/Cargo.toml new file mode 100644 index 00000000000..0b51939bb9c --- /dev/null +++ b/litellm-rust/crates/request-auth-core/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-request-auth-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true + +[dependencies] +fancy-regex.workspace = true +litellm-secret-value.workspace = true +serde.workspace = true +sha2.workspace = true +strum.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/request-auth-core/src/access.rs b/litellm-rust/crates/request-auth-core/src/access.rs new file mode 100644 index 00000000000..cbe8c19815b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/access.rs @@ -0,0 +1,714 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum AccessFacts { + ConcurrentFailures { + failures: Vec, + terminal: Vec, + }, + UserParamRequest { + enabled: bool, + }, + ClientTagRequest { + enabled: bool, + llm_route: Option, + }, + GuardrailRequest { + requested_keys: Vec, + }, + TraceRequest { + required: bool, + }, + Search { + name: String, + allowed: Vec, + strict: bool, + }, + TraceGate { + required: bool, + method: String, + llm_route: bool, + passthrough: bool, + telemetry: bool, + }, + TraceStep { + source: TraceSource, + value: Option, + present: bool, + headers_only: bool, + }, + Blocked { + blocked: bool, + }, + UserParam { + enabled: bool, + method: Option, + llm_route: bool, + mcp_route: bool, + has_user: bool, + }, + ClientTags { + enabled: bool, + llm_route: bool, + has_tags: bool, + }, + Guardrails { + requested_keys: Vec, + modify_guardrails: Option, + }, + AgentTrace { + required: bool, + header: Option, + }, + TeamTrace { + required: bool, + method: String, + llm_route: bool, + passthrough: bool, + telemetry: bool, + headers_only: bool, + header: Option, + body_trace: Option, + body_trace_present: bool, + traceparent_id: Option, + selected_trace: Option, + selected_trace_present: bool, + other_trace: Option, + }, + ToolRequest { + token_present: bool, + call_types: Option>, + capable_types: Vec, + }, + ToolAllowlist { + key_nonempty_list: bool, + team_nonempty_list: Option, + requested_count: Option, + }, + Tools { + requested: Vec, + key_allowed: Option>, + team_allowed: Option>, + }, + GrantLayers { + deny_by_default: bool, + identity_present: bool, + master_key: bool, + dashboard: bool, + virtual_key: bool, + session_token: bool, + team_present: bool, + }, + VectorRequest { + client_present: bool, + strict: bool, + ids_read: bool, + unreadable: bool, + grants: Option>>>, + requested_count: Option, + }, + VectorStores { + requested: Vec, + allowed: Option>, + strict: bool, + }, + TeamLookup { + dashboard: bool, + missing: bool, + recorded_models: Vec, + allow_unavailable: Option, + }, +} + +#[derive(Clone, Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum Denial { + Search, + Blocked, + UserParam, + ClientTags, + Guardrails, + AgentTrace, + TeamTrace, + Tools, + VectorStores, + MissingVectorGrant, + MalformedVectorIds, + TeamLookup, +} + +#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, strum::VariantArray)] +pub enum TraceSource { + Gate, + Header, + Body, + Traceparent, + Selected, + Other, +} + +#[derive(Debug, Default, PartialEq, Serialize)] +pub struct AccessDecision { + pub read: Option, + pub denial: Option, + pub rejected: Vec, + pub layers: Vec, + pub rejected_index: Option, + pub next_trace: Option, + pub failure_index: Option, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum AccessRead { + UserParam, + ClientTagRoute, + ClientTags, + GuardrailPermission, + UnavailablePolicy, + ToolRoute, + ToolAllowlist, + TeamToolList, + ToolNames, + KeyToolGrant, + TeamToolGrant, + StrictVectorIds, + LegacyVectorIds, + LegacyVectorGrants, + RagVectorIds, + VectorGrants, +} + +fn modifies_guardrails(keys: &[String]) -> bool { + keys.iter().any(|key| { + matches!( + key.as_str(), + "guardrails" + | "disable_global_guardrails" + | "disable_global_guardrail" + | "opted_out_global_guardrails" + ) + }) +} + +fn deny_if(deny: bool, denial: Denial) -> AccessDecision { + AccessDecision { + denial: deny.then_some(denial), + ..AccessDecision::default() + } +} + +fn nonempty(value: &Option) -> bool { + value.as_ref().is_some_and(|value| !value.is_empty()) +} + +pub fn decide(facts: AccessFacts) -> AccessDecision { + match facts { + AccessFacts::VectorRequest { + client_present, + strict, + ids_read, + unreadable, + grants, + requested_count, + } => { + if !client_present { + return AccessDecision::default(); + } + let read = if !ids_read { + Some(if strict { + AccessRead::StrictVectorIds + } else { + AccessRead::LegacyVectorIds + }) + } else if unreadable { + let Some(grants) = grants else { + return AccessDecision { + read: Some(AccessRead::LegacyVectorGrants), + ..AccessDecision::default() + }; + }; + return AccessDecision { + denial: grants + .iter() + .flatten() + .any(|grant| !grant.is_empty()) + .then_some(Denial::MalformedVectorIds), + ..AccessDecision::default() + }; + } else { + match requested_count { + None => Some(AccessRead::RagVectorIds), + Some(0) => None, + Some(_) => Some(AccessRead::VectorGrants), + } + }; + AccessDecision { + read, + ..AccessDecision::default() + } + } + AccessFacts::ToolRequest { + token_present, + call_types, + capable_types, + } => AccessDecision { + read: if !token_present { + None + } else { + match call_types { + None => Some(AccessRead::ToolRoute), + Some(types) => types + .iter() + .any(|kind| capable_types.contains(kind)) + .then_some(AccessRead::ToolAllowlist), + } + }, + ..AccessDecision::default() + }, + AccessFacts::ToolAllowlist { + key_nonempty_list, + team_nonempty_list, + requested_count, + } => AccessDecision { + read: if !key_nonempty_list && team_nonempty_list.is_none() { + Some(AccessRead::TeamToolList) + } else if !key_nonempty_list && team_nonempty_list != Some(true) { + None + } else { + match requested_count { + None => Some(AccessRead::ToolNames), + Some(0) => None, + Some(_) if key_nonempty_list => Some(AccessRead::KeyToolGrant), + Some(_) => Some(AccessRead::TeamToolGrant), + } + }, + ..AccessDecision::default() + }, + AccessFacts::ConcurrentFailures { failures, terminal } => AccessDecision { + failure_index: terminal + .iter() + .position(|failed| *failed) + .or_else(|| failures.iter().position(|failed| *failed)), + ..AccessDecision::default() + }, + AccessFacts::UserParamRequest { enabled } => AccessDecision { + read: enabled.then_some(AccessRead::UserParam), + ..AccessDecision::default() + }, + AccessFacts::ClientTagRequest { enabled, llm_route } => AccessDecision { + read: match (enabled, llm_route) { + (false, _) | (_, Some(false)) => None, + (_, None) => Some(AccessRead::ClientTagRoute), + (_, Some(true)) => Some(AccessRead::ClientTags), + }, + ..AccessDecision::default() + }, + AccessFacts::GuardrailRequest { requested_keys } => AccessDecision { + read: modifies_guardrails(&requested_keys).then_some(AccessRead::GuardrailPermission), + ..AccessDecision::default() + }, + AccessFacts::TraceRequest { required } => AccessDecision { + next_trace: required.then_some(TraceSource::Gate), + ..AccessDecision::default() + }, + AccessFacts::Search { + name, + allowed, + strict, + } => deny_if( + (strict || !allowed.is_empty()) && !allowed.contains(&name), + Denial::Search, + ), + AccessFacts::TraceGate { + required, + method, + llm_route, + passthrough, + telemetry, + } => AccessDecision { + next_trace: (required + && !matches!(method.as_str(), "GET" | "HEAD" | "OPTIONS") + && (llm_route || passthrough) + && !telemetry) + .then_some(TraceSource::Header), + ..AccessDecision::default() + }, + AccessFacts::TraceStep { + source, + value, + present, + headers_only, + } => { + if nonempty(&value) { + return AccessDecision::default(); + } + let next_trace = match source { + TraceSource::Gate => Some(TraceSource::Header), + TraceSource::Header if headers_only => None, + TraceSource::Header => Some(TraceSource::Body), + TraceSource::Body if !present => Some(TraceSource::Traceparent), + TraceSource::Body | TraceSource::Traceparent => Some(TraceSource::Selected), + TraceSource::Selected if !present => Some(TraceSource::Other), + TraceSource::Selected | TraceSource::Other => None, + }; + AccessDecision { + denial: next_trace.is_none().then_some(Denial::TeamTrace), + next_trace, + ..AccessDecision::default() + } + } + AccessFacts::Blocked { blocked } => deny_if(blocked, Denial::Blocked), + AccessFacts::UserParam { + enabled, + method, + llm_route, + mcp_route, + has_user, + } => deny_if( + enabled + && method.is_some_and(|value| value.to_uppercase() == "POST") + && llm_route + && !mcp_route + && !has_user, + Denial::UserParam, + ), + AccessFacts::ClientTags { + enabled, + llm_route, + has_tags, + } => deny_if(enabled && llm_route && has_tags, Denial::ClientTags), + AccessFacts::Guardrails { + requested_keys, + modify_guardrails, + } => deny_if( + modify_guardrails == Some(false) + && requested_keys.iter().any(|key| { + matches!( + key.as_str(), + "guardrails" + | "disable_global_guardrails" + | "disable_global_guardrail" + | "opted_out_global_guardrails" + ) + }), + Denial::Guardrails, + ), + AccessFacts::AgentTrace { required, header } => { + deny_if(required && !nonempty(&header), Denial::AgentTrace) + } + AccessFacts::TeamTrace { + required, + method, + llm_route, + passthrough, + telemetry, + headers_only, + header, + body_trace, + body_trace_present, + traceparent_id, + selected_trace, + selected_trace_present, + other_trace, + } => { + let gate = required + && !matches!(method.as_str(), "GET" | "HEAD" | "OPTIONS") + && (llm_route || passthrough) + && !telemetry; + let body = nonempty(&body_trace) + || (!body_trace_present && nonempty(&traceparent_id)) + || if selected_trace_present { + nonempty(&selected_trace) + } else { + nonempty(&other_trace) + }; + deny_if( + gate && !nonempty(&header) && (headers_only || !body), + Denial::TeamTrace, + ) + } + AccessFacts::Tools { + requested, + key_allowed, + team_allowed, + } => { + let allowed = key_allowed + .filter(|allowed| !allowed.is_empty()) + .or(team_allowed); + let rejected: Vec<_> = match allowed { + Some(allowed) if !allowed.is_empty() => requested + .into_iter() + .filter(|name| !allowed.contains(name)) + .collect(), + _ => Vec::new(), + }; + AccessDecision { + denial: (!rejected.is_empty()).then_some(Denial::Tools), + rejected, + ..AccessDecision::default() + } + } + AccessFacts::GrantLayers { + deny_by_default, + identity_present, + master_key, + dashboard, + virtual_key, + session_token, + team_present, + } => { + let layers = if !deny_by_default || !identity_present || master_key || dashboard { + Vec::new() + } else if !(virtual_key && !session_token) && !team_present { + vec!["user".into()] + } else { + [ + ("key", virtual_key && !session_token), + ("team", team_present), + ] + .into_iter() + .filter_map(|(name, required)| required.then_some(name.into())) + .collect() + }; + AccessDecision { + layers, + ..AccessDecision::default() + } + } + AccessFacts::VectorStores { + requested, + allowed, + strict, + } => match allowed { + None => deny_if(strict, Denial::MissingVectorGrant), + Some(allowed) if allowed.is_empty() => deny_if(strict, Denial::MissingVectorGrant), + Some(allowed) => { + let rejected_index = requested.iter().position(|name| !allowed.contains(name)); + AccessDecision { + denial: rejected_index.map(|_| Denial::VectorStores), + rejected_index, + ..AccessDecision::default() + } + } + }, + AccessFacts::TeamLookup { + dashboard, + missing, + recorded_models, + allow_unavailable, + } => { + if !dashboard && !missing && recorded_models.is_empty() && allow_unavailable.is_none() { + return AccessDecision { + read: Some(AccessRead::UnavailablePolicy), + ..AccessDecision::default() + }; + } + deny_if( + !dashboard + && (missing || (recorded_models.is_empty() && allow_unavailable != Some(true))), + Denial::TeamLookup, + ) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::no_database(false, true, false, false, None, None, (None, None))] + #[case::strict_ids( + true, + true, + false, + false, + None, + None, + (Some(AccessRead::StrictVectorIds), None) + )] + #[case::legacy_ids( + true, + false, + false, + false, + None, + None, + (Some(AccessRead::LegacyVectorIds), None) + )] + #[case::unreadable_grants( + true, + false, + true, + true, + None, + None, + (Some(AccessRead::LegacyVectorGrants), None) + )] + #[case::unrestricted_unreadable(true, false, true, true, Some(vec![None, Some(vec![])]), None, (None, None))] + #[case::restricted_unreadable(true, false, true, true, Some(vec![Some(vec!["a".into()])]), None, (None, Some(Denial::MalformedVectorIds)))] + #[case::rag_ids( + true, + true, + true, + false, + None, + None, + (Some(AccessRead::RagVectorIds), None) + )] + #[case::no_ids(true, true, true, false, None, Some(0), (None, None))] + #[case::identity_grants( + true, + true, + true, + false, + None, + Some(2), + (Some(AccessRead::VectorGrants), None) + )] + fn vector_read_order( + #[case] client_present: bool, + #[case] strict: bool, + #[case] ids_read: bool, + #[case] unreadable: bool, + #[case] grants: Option>>>, + #[case] requested_count: Option, + #[case] expected: (Option, Option), + ) { + assert_eq!( + decide(AccessFacts::VectorRequest { + client_present, + strict, + ids_read, + unreadable, + grants, + requested_count, + }), + AccessDecision { + read: expected.0, + denial: expected.1, + ..AccessDecision::default() + } + ); + } + + #[rstest] + #[case::no_token(false, None, None)] + #[case::route(true, None, Some(AccessRead::ToolRoute))] + #[case::no_call_type(true, Some(vec![]), None)] + #[case::non_tool_route(true, Some(vec!["embedding".into()]), None)] + #[case::tool_route(true, Some(vec!["embedding".into(), "completion".into()]), Some(AccessRead::ToolAllowlist))] + fn tool_route_reads( + #[case] token_present: bool, + #[case] call_types: Option>, + #[case] read: Option, + ) { + assert_eq!( + decide(AccessFacts::ToolRequest { + token_present, + call_types, + capable_types: vec!["completion".into()], + }) + .read, + read + ); + } + + #[rstest] + #[case::key_override(true, None, None, Some(AccessRead::ToolNames))] + #[case::team_fallback(false, None, None, Some(AccessRead::TeamToolList))] + #[case::no_allowlist(false, Some(false), None, None)] + #[case::team_enabled(false, Some(true), None, Some(AccessRead::ToolNames))] + #[case::no_tools(true, None, Some(0), None)] + #[case::key_grant(true, None, Some(1), Some(AccessRead::KeyToolGrant))] + #[case::team_grant(false, Some(true), Some(1), Some(AccessRead::TeamToolGrant))] + fn tool_grant_reads( + #[case] key_nonempty_list: bool, + #[case] team_nonempty_list: Option, + #[case] requested_count: Option, + #[case] read: Option, + ) { + assert_eq!( + decide(AccessFacts::ToolAllowlist { + key_nonempty_list, + team_nonempty_list, + requested_count, + }) + .read, + read + ); + } + + #[rstest] + #[case::disabled(false, None, None)] + #[case::read_route(true, None, Some(AccessRead::ClientTagRoute))] + #[case::non_llm(true, Some(false), None)] + #[case::read_tags(true, Some(true), Some(AccessRead::ClientTags))] + fn client_tag_reads( + #[case] enabled: bool, + #[case] llm_route: Option, + #[case] read: Option, + ) { + assert_eq!( + decide(AccessFacts::ClientTagRequest { enabled, llm_route }).read, + read + ); + } + + #[rstest] + #[case::allow(vec![false, false], vec![false, false], None)] + #[case::scope_priority(vec![true, true], vec![false, false], Some(0))] + #[case::second_scope(vec![false, true], vec![false, false], Some(1))] + #[case::terminal_fault(vec![true, true], vec![false, true], Some(1))] + fn gathered_failures_preserve_scope_order_unless_native_failed( + #[case] failures: Vec, + #[case] terminal: Vec, + #[case] expected: Option, + ) { + assert_eq!( + decide(AccessFacts::ConcurrentFailures { failures, terminal }).failure_index, + expected + ); + } + + #[rstest] + #[case::key_overrides_team(Some(vec!["key".into()]), Some(vec!["team".into()]), vec!["team".into()])] + #[case::empty_key_inherits(Some(vec![]), Some(vec!["team".into()]), vec!["key".into()])] + #[case::unrestricted(None, Some(vec![]), vec![])] + fn tools( + #[case] key: Option>, + #[case] team: Option>, + #[case] rejected: Vec, + ) { + let result = decide(AccessFacts::Tools { + requested: vec!["key".into(), "team".into()], + key_allowed: key, + team_allowed: team, + }); + assert_eq!( + result.denial, + (!rejected.is_empty()).then_some(Denial::Tools) + ); + assert_eq!(result.rejected, rejected); + } + + #[rstest] + #[case::missing_strict(None, true, Some(Denial::MissingVectorGrant))] + #[case::empty_legacy(Some(vec![]), false, None)] + #[case::empty_strict(Some(vec![]), true, Some(Denial::MissingVectorGrant))] + #[case::denied(Some(vec!["other".into()]), false, Some(Denial::VectorStores))] + #[case::granted(Some(vec!["store".into()]), true, None)] + fn vectors( + #[case] allowed: Option>, + #[case] strict: bool, + #[case] denial: Option, + ) { + let result = decide(AccessFacts::VectorStores { + requested: vec!["store".into()], + allowed, + strict, + }); + assert_eq!(result.denial, denial); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/admission.rs b/litellm-rust/crates/request-auth-core/src/admission.rs new file mode 100644 index 00000000000..8c23d47ef8e --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/admission.rs @@ -0,0 +1,183 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum AdmissionFacts { + CommonGate { + force_virtual_key_checks: bool, + public_route: bool, + additional_public_route: bool, + unauthenticated_passthrough: bool, + master_key_present: bool, + jwt_enabled: bool, + oauth_enabled: bool, + oauth_proxy_enabled: bool, + custom_auth: bool, + custom_auth_common_checks: bool, + }, + Reservation { + skip_budget_checks: bool, + disabled: bool, + apply_user_budget_to_team_keys: bool, + fail_closed: bool, + billing_agent: bool, + }, + Ownership { + virtual_key: bool, + custom_auth: bool, + jwt: bool, + oauth: bool, + }, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum CommonSkip { + PublicRoute, + UnauthenticatedPassthrough, + NoAuthDevelopment, + CustomAuth, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum Owner { + Python, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum Candidate { + VirtualKey, + CustomAuth, + Validator, + Other, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum AdmissionPlan { + Common { + skip: Option, + }, + Reserve { + reserve: bool, + apply_user_budget_to_team_keys: bool, + fail_closed: bool, + }, + Observe { + owner: Owner, + candidate: Candidate, + }, +} + +pub fn plan(facts: AdmissionFacts) -> AdmissionPlan { + match facts { + AdmissionFacts::CommonGate { + force_virtual_key_checks, + public_route, + additional_public_route, + unauthenticated_passthrough, + master_key_present, + jwt_enabled, + oauth_enabled, + oauth_proxy_enabled, + custom_auth, + custom_auth_common_checks, + } => { + let skip = if force_virtual_key_checks { + None + } else if public_route || additional_public_route { + Some(CommonSkip::PublicRoute) + } else if unauthenticated_passthrough { + Some(CommonSkip::UnauthenticatedPassthrough) + } else if !master_key_present && !jwt_enabled && !oauth_enabled && !oauth_proxy_enabled + { + Some(CommonSkip::NoAuthDevelopment) + } else if custom_auth && !custom_auth_common_checks { + Some(CommonSkip::CustomAuth) + } else { + None + }; + AdmissionPlan::Common { skip } + } + AdmissionFacts::Reservation { + skip_budget_checks, + disabled, + apply_user_budget_to_team_keys, + fail_closed, + billing_agent, + } => AdmissionPlan::Reserve { + reserve: !skip_budget_checks && !disabled, + apply_user_budget_to_team_keys, + fail_closed: fail_closed || billing_agent, + }, + AdmissionFacts::Ownership { + virtual_key, + custom_auth, + jwt, + oauth, + } => { + let candidate = if custom_auth { + Candidate::CustomAuth + } else if jwt || oauth { + Candidate::Validator + } else if virtual_key { + Candidate::VirtualKey + } else { + Candidate::Other + }; + AdmissionPlan::Observe { + owner: Owner::Python, + candidate, + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::normal(false, false, false, true, false)] + #[case::free(true, false, false, false, false)] + #[case::disabled(false, true, false, false, false)] + #[case::managed(false, false, true, true, true)] + fn reservation( + #[case] skip: bool, + #[case] disabled: bool, + #[case] billing: bool, + #[case] reserve: bool, + #[case] closed: bool, + ) { + assert_eq!( + plan(AdmissionFacts::Reservation { + skip_budget_checks: skip, + disabled, + billing_agent: billing, + fail_closed: false, + apply_user_budget_to_team_keys: true + }), + AdmissionPlan::Reserve { + reserve, + apply_user_budget_to_team_keys: true, + fail_closed: closed + } + ); + } + + #[rstest] + fn ownership_never_transfers_request() { + assert_eq!( + plan(AdmissionFacts::Ownership { + virtual_key: true, + custom_auth: false, + jwt: false, + oauth: false + }), + AdmissionPlan::Observe { + owner: Owner::Python, + candidate: Candidate::VirtualKey + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/agent_access.rs b/litellm-rust/crates/request-auth-core/src/agent_access.rs new file mode 100644 index 00000000000..555014741fc --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/agent_access.rs @@ -0,0 +1,269 @@ +use serde::{Deserialize, Serialize}; + +use crate::managed::Agent; + +#[derive(Clone, Debug, Deserialize)] +pub struct Facts { + pub registry_managed: bool, + pub database_present: bool, + pub target_loaded: bool, + pub target: Option, + pub subject_present: bool, + pub key_hash_present: bool, + pub subject_managed: bool, + pub session: bool, + pub custom: bool, + pub authority_loaded: bool, + pub explicit_grants: Option>, + pub legacy_loaded: bool, + pub legacy_grants: Option>, + pub target_ids: Option>, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Decision { + ReadTarget, + ReadAuthority, + ReadExplicitGrants, + ReadLegacyGrants, + ReadTargetIds, + Allow, + Deny, +} + +pub fn decide(f: &Facts) -> Decision { + if f.registry_managed || f.database_present { + if !f.target_loaded { + return Decision::ReadTarget; + } + if f.target.is_none() && f.registry_managed { + return Decision::Deny; + } + if let Some(target) = f.target.as_ref().filter(|target| target.identity_managed) { + if !target.enabled + || !target + .identity + .as_ref() + .is_some_and(|identity| identity.active) + || !f.subject_present + { + return Decision::Deny; + } + if f.key_hash_present + && !f.subject_managed + && !f.session + && !f.custom + && !f.authority_loaded + { + return Decision::ReadAuthority; + } + return match f.explicit_grants.as_ref() { + None => Decision::ReadExplicitGrants, + Some(grants) if grants.contains(&target.agent_id) => Decision::Allow, + Some(_) => Decision::Deny, + }; + } + } + if !f.legacy_loaded { + return Decision::ReadLegacyGrants; + } + let Some(grants) = f.legacy_grants.as_ref() else { + return Decision::Allow; + }; + match f.target_ids.as_ref() { + None => Decision::ReadTargetIds, + Some(ids) if ids.iter().any(|id| grants.contains(id)) => Decision::Allow, + Some(_) => Decision::Deny, + } +} + +#[derive(Deserialize)] +pub struct Invocation { + pub registered: Option, + pub has_store: bool, + pub target_loaded: bool, + pub target: Option, + pub actor_managed: bool, + pub actor_id_present: bool, + pub permissions: Option, + pub assigned: bool, + pub fee_loaded: bool, + pub fee: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum InvocationAction { + Complete, + ReadTarget, + ReadPermissions, + Assign, + ReadFee, + Deny, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum InvocationDenial { + DatabaseRequired, + Removed, + CallerNotPermitted, + InvalidFee, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct InvocationPlan { + pub action: InvocationAction, + pub denial: Option, + pub use_target: bool, + pub assign_billing: bool, +} + +pub fn invocation(f: &Invocation) -> InvocationPlan { + let action = |action, denial| InvocationPlan { + action, + denial, + use_target: f.target.is_some(), + assign_billing: !f.actor_id_present + && f.target + .as_ref() + .or(f.registered.as_ref()) + .is_some_and(|agent| agent.identity_managed), + }; + let Some(registered) = f.registered.as_ref() else { + return action(InvocationAction::Complete, None); + }; + let registered_managed = registered.identity_managed || registered.identity.is_some(); + if !f.has_store && registered_managed { + return action( + InvocationAction::Deny, + Some(InvocationDenial::DatabaseRequired), + ); + } + if !f.target_loaded { + return action(InvocationAction::ReadTarget, None); + } + if f.target.is_none() && registered_managed { + return action(InvocationAction::Deny, Some(InvocationDenial::Removed)); + } + let effective = f.target.as_ref().unwrap_or(registered); + if !effective.identity_managed && !f.actor_managed { + return action(InvocationAction::Complete, None); + } + match f.permissions.as_ref().map(decide) { + None => return action(InvocationAction::ReadPermissions, None), + Some(Decision::Allow) => (), + Some(_) => { + return action( + InvocationAction::Deny, + Some(InvocationDenial::CallerNotPermitted), + ); + } + } + if !f.assigned { + return action(InvocationAction::Assign, None); + } + if !f.fee_loaded { + return action(InvocationAction::ReadFee, None); + } + if !f.fee.is_some_and(|fee| fee.is_finite() && fee >= 0.0) { + return action(InvocationAction::Deny, Some(InvocationDenial::InvalidFee)); + } + action(InvocationAction::Complete, None) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::free(Some(0.0), None)] + #[case::priced(Some(1.0), None)] + #[case::missing(None, Some(InvocationDenial::InvalidFee))] + #[case::negative(Some(-1.0), Some(InvocationDenial::InvalidFee))] + #[case::nan(Some(f64::NAN), Some(InvocationDenial::InvalidFee))] + #[case::infinity(Some(f64::INFINITY), Some(InvocationDenial::InvalidFee))] + fn invocation_fee_is_checked_after_assignment( + #[case] fee: Option, + #[case] denial: Option, + ) { + let facts = Invocation { + registered: Some(Agent { + agent_id: "agent".into(), + enabled: true, + identity_managed: true, + identity: None, + execution_mode: "both".into(), + }), + has_store: true, + target_loaded: true, + target: Some(Agent { + agent_id: "agent".into(), + enabled: true, + identity_managed: true, + identity: None, + execution_mode: "both".into(), + }), + actor_managed: false, + actor_id_present: false, + permissions: Some(Facts { + registry_managed: false, + database_present: false, + target_loaded: false, + target: None, + subject_present: true, + key_hash_present: false, + subject_managed: false, + session: false, + custom: false, + authority_loaded: false, + explicit_grants: None, + legacy_loaded: true, + legacy_grants: None, + target_ids: None, + }), + assigned: true, + fee_loaded: true, + fee, + }; + assert_eq!(invocation(&facts).denial, denial); + assert!(invocation(&facts).assign_billing); + assert_eq!( + invocation(&Invocation { + assigned: false, + ..facts + }) + .action, + InvocationAction::Assign + ); + } + + #[rstest] + #[case::no_database(false, false, false, Decision::ReadLegacyGrants)] + #[case::database(false, true, false, Decision::ReadTarget)] + #[case::managed_missing(true, true, true, Decision::Deny)] + #[case::ordinary_missing(false, true, true, Decision::ReadLegacyGrants)] + fn target_read_precedes_grants( + #[case] registry_managed: bool, + #[case] database_present: bool, + #[case] target_loaded: bool, + #[case] expected: Decision, + ) { + let facts = Facts { + registry_managed, + database_present, + target_loaded, + target: None, + subject_present: false, + key_hash_present: false, + subject_managed: false, + session: false, + custom: false, + authority_loaded: false, + explicit_grants: None, + legacy_loaded: false, + legacy_grants: None, + target_ids: None, + }; + assert_eq!(decide(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs b/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs new file mode 100644 index 00000000000..e270e291078 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs @@ -0,0 +1,151 @@ +use crate::Error; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Status { + Found, + Missing, + Unavailable, +} + +#[derive(Clone, Deserialize)] +pub struct Group { + pub status: Status, + pub models: Vec, + pub mcp_server_ids: Vec, + pub agent_ids: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub ids: Option>, + pub groups: Option>, + pub check_db_only: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadIds, + ReadGroups { + authoritative: bool, + }, + Complete { + present: bool, + ids: Vec, + models: BTreeSet, + mcp_server_ids: BTreeSet, + agent_ids: BTreeSet, + indices: Vec, + }, +} + +pub fn decide(facts: &Facts) -> Result { + let Some(ids) = &facts.ids else { + return Ok(Decision::ReadIds); + }; + let empty = |present| Decision::Complete { + present, + ids: vec![], + models: BTreeSet::new(), + mcp_server_ids: BTreeSet::new(), + agent_ids: BTreeSet::new(), + indices: vec![], + }; + if ids.is_empty() { + return Ok(empty(false)); + } + let Some(groups) = &facts.groups else { + return Ok(Decision::ReadGroups { + authoritative: facts.check_db_only, + }); + }; + if groups.len() != ids.len() { + return Err(Error::InvalidAnswer); + } + let indices: Vec<_> = groups + .iter() + .enumerate() + .filter_map(|(index, group)| (group.status == Status::Found).then_some(index)) + .collect(); + Ok(Decision::Complete { + present: true, + ids: ids + .iter() + .zip(groups) + .filter(|(_, group)| group.status != Status::Missing) + .map(|(id, _)| id.clone()) + .collect(), + models: indices + .iter() + .flat_map(|index| groups[*index].models.iter().cloned()) + .collect(), + mcp_server_ids: indices + .iter() + .flat_map(|index| groups[*index].mcp_server_ids.iter().cloned()) + .collect(), + agent_ids: indices + .iter() + .flat_map(|index| groups[*index].agent_ids.iter().cloned()) + .collect(), + indices, + }) +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum ReadStatus { + Success, + NotFound, + HttpError, +} + +#[derive(Clone, Deserialize)] +pub struct LoadFacts { + pub database_present: bool, + pub check_db_only: bool, + pub status: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum LoadDecision { + Read, + Return, + DatabaseUnavailable, + HttpUnavailable, + Raise, +} + +pub fn load(facts: &LoadFacts) -> LoadDecision { + if !facts.database_present { + return LoadDecision::DatabaseUnavailable; + } + match facts.status { + None => LoadDecision::Read, + Some(ReadStatus::Success) => LoadDecision::Return, + Some(ReadStatus::NotFound) => LoadDecision::Raise, + Some(ReadStatus::HttpError) if facts.check_db_only => LoadDecision::Raise, + Some(ReadStatus::HttpError) => LoadDecision::HttpUnavailable, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::soft(false, LoadDecision::HttpUnavailable)] + #[case::authoritative(true, LoadDecision::Raise)] + fn http_lookup_policy(#[case] check_db_only: bool, #[case] expected: LoadDecision) { + assert_eq!( + load(&LoadFacts { + database_present: true, + check_db_only, + status: Some(ReadStatus::HttpError) + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/agent_grants.rs b/litellm-rust/crates/request-auth-core/src/agent_grants.rs new file mode 100644 index 00000000000..5956c7c8390 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/agent_grants.rs @@ -0,0 +1,408 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Source { + pub team_id: Option, + pub grants: Option>, + pub unavailable: bool, +} + +#[derive(Deserialize)] +pub struct ConfiguredAgent { + pub agent_id: String, + pub groups: Vec, +} + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Resolve { + managed: bool, + strict: bool, + role: Option, + key_team_loaded: bool, + key_team: Option>, + default_deny: Option, + caller_loaded: bool, + own: Option>, + ceiling_loaded: bool, + ceiling: Option>, + }, + Failure { + strict: bool, + closed: bool, + }, + OwnerTeam { + strict: bool, + permission_id_present: bool, + permission_present: bool, + }, + OwnerUnion { + teams: Vec, + withheld: Vec, + unreadable: bool, + }, + BillingTeam { + sources: Vec, + agent_id: String, + }, + Human { + sources: Vec, + team_id: Option, + }, + Delegated { + sources: Vec, + agent_id: String, + team_id: Option, + explicit: bool, + allowed: Option>, + }, + TeamLookup { + subject: bool, + team_id: Option, + personal: Option, + database: bool, + loaded: bool, + row_present: bool, + strict: bool, + }, + GroupMatches { + agents: Vec, + groups: Vec, + }, +} + +#[derive(Default, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Action { + #[default] + Complete, + ReadManaged, + ReadKeyTeam, + ReadDefault, + ReadCaller, + ReadCeiling, + ReadPersonalEligibility, + ReadOwnerTeams, + ReadTeam, + ReadDeclared, + Unavailable, + SelectTeam, +} + +#[derive(Default, Debug, PartialEq, Eq, Serialize)] +pub struct Decision { + pub action: Action, + pub grants: Option>, + pub team_id: Option, +} + +fn action(action: Action) -> Decision { + Decision { + action, + ..Decision::default() + } +} +fn grants(grants: Option>) -> Decision { + Decision { + grants, + ..Decision::default() + } +} +fn intersection(left: Option>, right: Option>) -> Option> { + match (left, right) { + (None, right) => right, + (left, None) => left, + (Some(left), Some(right)) => { + Some(left.into_iter().filter(|id| right.contains(id)).collect()) + } + } +} + +pub fn decide(facts: Facts) -> Decision { + match facts { + Facts::Resolve { + managed, + strict, + role, + key_team_loaded, + key_team, + default_deny, + caller_loaded, + own, + ceiling_loaded, + ceiling, + } => { + if managed { + return action(Action::ReadManaged); + } + if !key_team_loaded { + return action(Action::ReadKeyTeam); + } + let admin = role.as_deref() == Some("proxy_admin"); + if !strict && !admin && default_deny.is_none() { + return action(Action::ReadDefault); + } + if (strict || (!admin && default_deny == Some(true))) && key_team.is_none() { + return grants(Some(vec![])); + } + if !caller_loaded { + return action(Action::ReadCaller); + } + if !ceiling_loaded { + return action(Action::ReadCeiling); + } + grants(intersection(own, ceiling)) + } + Facts::Failure { strict, closed } => { + if strict { + action(Action::Unavailable) + } else { + grants(closed.then(Vec::new)) + } + } + Facts::OwnerTeam { + strict, + permission_id_present, + permission_present, + } => { + if permission_id_present && !permission_present { + if strict { + action(Action::Unavailable) + } else { + grants(Some(vec![])) + } + } else { + action(Action::ReadDeclared) + } + } + Facts::OwnerUnion { + teams, + withheld, + unreadable, + } => { + if teams.iter().any(|source| source.unavailable) { + return action(Action::Unavailable); + } + let restricted = teams.is_empty() + || unreadable + || teams.iter().any(|source| source.grants.is_some()) + || withheld + .iter() + .any(|source| source.unavailable || source.grants.is_some()); + grants(restricted.then(|| { + teams + .into_iter() + .flat_map(|source| source.grants.unwrap_or_default()) + .collect::>() + .into_iter() + .collect() + })) + } + Facts::BillingTeam { sources, agent_id } => { + let selected = sources + .iter() + .filter(|source| { + source + .grants + .as_ref() + .is_some_and(|ids| ids.contains(&agent_id)) + }) + .filter_map(|source| source.team_id.clone()) + .min(); + let team_id = selected.or_else(|| { + sources + .iter() + .all(|source| source.grants.is_none() && !source.unavailable) + .then(|| { + sources + .iter() + .filter_map(|source| source.team_id.clone()) + .min() + }) + .flatten() + }); + Decision { + team_id, + ..Decision::default() + } + } + Facts::Human { sources, team_id } => grants(Some( + sources + .into_iter() + .filter(|source| source.team_id.is_none() || source.team_id == team_id) + .flat_map(|source| source.grants.unwrap_or_default()) + .collect::>() + .into_iter() + .collect(), + )), + Facts::Delegated { + sources, + agent_id, + team_id, + explicit, + allowed, + } => { + if sources.iter().any(|source| { + source.team_id.is_none() + && source + .grants + .as_ref() + .is_some_and(|ids| ids.contains(&agent_id)) + }) { + return Decision { + team_id, + ..Decision::default() + }; + } + let teams: BTreeSet<_> = sources + .iter() + .filter(|source| { + source + .grants + .as_ref() + .is_some_and(|ids| ids.contains(&agent_id)) + }) + .filter_map(|source| source.team_id.as_ref()) + .filter(|team| { + allowed + .as_ref() + .is_none_or(|allowed| allowed.contains(team)) + }) + .cloned() + .collect(); + if team_id.as_ref().is_some_and(|team| teams.contains(team)) { + return Decision { + team_id, + ..Decision::default() + }; + } + if !explicit && let Some(first) = teams.first() { + return Decision { + team_id: Some(first.clone()), + ..Decision::default() + }; + } + action(Action::SelectTeam) + } + Facts::TeamLookup { + subject, + team_id, + personal, + database, + loaded, + row_present, + strict, + } => { + if !subject { + return grants(None); + } + if team_id.is_none() { + return match personal { + None => action(Action::ReadPersonalEligibility), + Some(true) => action(Action::ReadOwnerTeams), + Some(false) => grants(None), + }; + } + if !database || (loaded && !row_present) { + return grants(strict.then(Vec::new)); + } + if !loaded { + action(Action::ReadTeam) + } else { + action(Action::ReadDeclared) + } + } + Facts::GroupMatches { agents, groups } => grants(Some( + agents + .into_iter() + .filter(|agent| agent.groups.iter().any(|id| groups.contains(id))) + .map(|agent| agent.agent_id) + .collect(), + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::legacy(false, false, Action::Complete, None)] + #[case::closed(false, true, Action::Complete, Some(vec![]))] + #[case::strict(true, false, Action::Unavailable, None)] + #[case::strict_closed(true, true, Action::Unavailable, None)] + fn lookup_failure( + #[case] strict: bool, + #[case] closed: bool, + #[case] expected: Action, + #[case] expected_grants: Option>, + ) { + assert_eq!( + decide(Facts::Failure { strict, closed }), + Decision { + action: expected, + grants: expected_grants, + team_id: None + } + ); + } + + #[rstest] + #[case::unrestricted(None, false, None)] + #[case::withheld(Some(vec!["a".into()]), false, Some(vec![]))] + #[case::unreadable(None, true, Some(vec![]))] + fn withheld_never_widens( + #[case] withheld: Option>, + #[case] unreadable: bool, + #[case] expected: Option>, + ) { + assert_eq!( + decide(Facts::OwnerUnion { + teams: vec![Source { + team_id: Some("team".into()), + grants: None, + unavailable: false + }], + withheld: vec![Source { + team_id: Some("held".into()), + grants: withheld, + unavailable: false + }], + unreadable + }) + .grants, + expected + ); + } + + #[rstest] + #[case::implicit(false, Some("a".into()), Action::Complete)] + #[case::explicit(true, None, Action::SelectTeam)] + fn delegated_requires_explicit_grant( + #[case] explicit: bool, + #[case] team_id: Option, + #[case] expected: Action, + ) { + assert_eq!( + decide(Facts::Delegated { + sources: vec![Source { + team_id: Some("a".into()), + grants: Some(vec!["agent".into()]), + unavailable: false + }], + agent_id: "agent".into(), + team_id: Some("other".into()), + explicit, + allowed: None + }), + Decision { + action: expected, + team_id, + grants: None + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/auth_context.rs b/litellm-rust/crates/request-auth-core/src/auth_context.rs new file mode 100644 index 00000000000..005a1ec57c1 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/auth_context.rs @@ -0,0 +1,169 @@ +use serde::{Deserialize, Serialize}; + +use crate::access::{self, AccessFacts, AccessRead}; + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Fetch { + team_id: Option, + ui_team_id: String, + user_id: Option, + project_id: Option, + end_user_id: Option, + }, + Team { + team_id: Option, + ui_team_id: String, + failed: bool, + missing: bool, + recorded_models: Vec, + allow_unavailable: Option, + }, + Admin { + role: Option, + user_id: Option, + admin_id: String, + user_row: bool, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum TeamSource { + None, + Token, + Loaded, + ReadUnavailablePolicy, + Failure, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Plan { + Fetch { + team: bool, + user: bool, + project: bool, + end_user: bool, + }, + Team { + source: TeamSource, + }, + Admin { + force_admin: bool, + preserve_budget: bool, + }, +} + +pub fn plan(facts: Facts) -> Plan { + match facts { + Facts::Fetch { + team_id, + ui_team_id, + user_id, + project_id, + end_user_id, + } => Plan::Fetch { + team: team_id.as_ref().is_some_and(|id| id != &ui_team_id), + user: user_id.is_some(), + project: project_id.is_some(), + end_user: end_user_id.is_some_and(|id| !id.is_empty()), + }, + Facts::Team { + team_id, + ui_team_id, + failed, + missing, + recorded_models, + allow_unavailable, + } => { + let dashboard = team_id.as_ref() == Some(&ui_team_id); + let source = if !failed { + if dashboard { + TeamSource::Token + } else { + TeamSource::Loaded + } + } else if team_id.is_none() { + TeamSource::None + } else { + let decision = access::decide(AccessFacts::TeamLookup { + dashboard, + missing, + recorded_models, + allow_unavailable, + }); + if decision.read == Some(AccessRead::UnavailablePolicy) { + TeamSource::ReadUnavailablePolicy + } else if decision.denial.is_some() { + TeamSource::Failure + } else { + TeamSource::Token + } + }; + Plan::Team { source } + } + Facts::Admin { + role, + user_id, + admin_id, + user_row, + } => Plan::Admin { + force_admin: role.as_deref() == Some("proxy_admin"), + preserve_budget: user_row && user_id.as_ref() != Some(&admin_id), + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(None, false)] + #[case::dashboard(Some("ui"), false)] + #[case::empty(Some(""), true)] + #[case::team(Some("team"), true)] + fn fetch_team(#[case] team: Option<&str>, #[case] expected: bool) { + assert_eq!( + plan(Facts::Fetch { + team_id: team.map(str::to_owned), + ui_team_id: "ui".into(), + user_id: Some("".into()), + project_id: Some("".into()), + end_user_id: Some("".into()) + }), + Plan::Fetch { + team: expected, + user: true, + project: true, + end_user: false + } + ); + } + + #[rstest] + #[case::master(Some("admin"), true, false)] + #[case::personal(Some("user"), true, true)] + #[case::no_row(Some("user"), false, false)] + #[case::missing_id(None, true, true)] + fn admin_budget_projection( + #[case] user: Option<&str>, + #[case] user_row: bool, + #[case] preserve_budget: bool, + ) { + assert_eq!( + plan(Facts::Admin { + role: Some("proxy_admin".into()), + user_id: user.map(str::to_owned), + admin_id: "admin".into(), + user_row + }), + Plan::Admin { + force_admin: true, + preserve_budget + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_alerts.rs b/litellm-rust/crates/request-auth-core/src/budget_alerts.rs new file mode 100644 index 00000000000..e6b6ad2a274 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_alerts.rs @@ -0,0 +1,233 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] +pub struct Failure { + pub kind: String, + pub message: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ThresholdFacts { + pub key: String, + pub digit: bool, + pub digits: Option, + pub ratio: Option, + pub integer_failure: Option, + pub ratio_failure: Option, + pub text: Option, + pub items: Option>>, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ConfigFacts { + pub global_entries: Vec, + pub local_entries: Vec, + pub global_failure: Option, + pub local_failure: Option, + pub order: Option>, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Facts { + pub member: bool, + pub spend: Option, + pub limit: Option, + pub default_fraction: f64, + pub config: Option, +} + +#[derive(Clone, Debug, Serialize, PartialEq)] +pub struct Threshold { + pub key: String, + pub recipients: Vec, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct Decision { + pub read_config: bool, + pub read_order: bool, + pub global_keys: Vec, + pub local_keys: Vec, + pub emit: bool, + pub thresholds: Option>, + pub failure: Option, +} + +fn member_entry(entry: &ThresholdFacts) -> Result { + if !entry.digit || entry.key.chars().count() > 3 { + return Ok(false); + } + if let Some(error) = &entry.integer_failure { + return Err(error.clone()); + } + Ok(entry + .digits + .as_deref() + .and_then(|digits| digits.parse::().ok()) + .is_some_and(|n| (1..=100).contains(&n))) +} + +fn eligible(entries: &[ThresholdFacts], member: bool) -> Result, Failure> { + entries.iter().try_fold(Vec::new(), |mut selected, entry| { + if !member || member_entry(entry)? { + selected.push(entry); + } + Ok(selected) + }) +} + +fn recipients(entry: &ThresholdFacts) -> Vec { + let strings: Vec<&str> = if let Some(text) = &entry.text { + text.split(',').collect() + } else { + entry + .items + .as_deref() + .unwrap_or_default() + .iter() + .filter_map(Option::as_deref) + .collect() + }; + strings + .into_iter() + .map(|s| s.trim_matches(|c: char| c.is_whitespace() || matches!(c, '\u{1c}'..='\u{1f}'))) + .filter(|s| !s.is_empty()) + .map(str::to_owned) + .collect() +} + +fn merge(key: &str, global: &[&ThresholdFacts], local: &[&ThresholdFacts]) -> Threshold { + let recipients = global + .iter() + .chain(local) + .filter(|entry| entry.key == key) + .flat_map(|entry| recipients(entry)) + .fold(Vec::new(), |mut merged, email| { + if !merged.contains(&email) { + merged.push(email); + } + merged + }); + Threshold { + key: key.to_owned(), + recipients, + } +} + +fn evaluate(facts: &Facts) -> Result { + if !facts.member && !(facts.limit.is_some() && facts.spend.is_some_and(|spend| spend > 0.0)) { + return Ok(Decision::default()); + } + let Some(config) = &facts.config else { + return Ok(Decision { + read_config: true, + ..Decision::default() + }); + }; + if let Some(failure) = &config.global_failure { + return Err(failure.clone()); + } + if !facts.member + && let Some(failure) = &config.local_failure + { + return Err(failure.clone()); + } + let global = eligible(&config.global_entries, false)?; + let local = if facts.member && config.local_failure.is_some() { + vec![] + } else { + eligible(&config.local_entries, facts.member)? + }; + if global.is_empty() && local.is_empty() { + let emit = !facts.member + && matches!((facts.spend, facts.limit), (Some(spend), Some(limit)) if spend >= limit * facts.default_fraction && spend < limit); + return Ok(Decision { + emit, + ..Decision::default() + }); + } + let Some(order) = &config.order else { + return Ok(Decision { + read_order: true, + global_keys: global.iter().map(|entry| entry.key.clone()).collect(), + local_keys: local.iter().map(|entry| entry.key.clone()).collect(), + ..Decision::default() + }); + }; + let thresholds: Vec = order + .iter() + .map(|key| merge(key, &global, &local)) + .collect(); + if facts.member && facts.spend.is_some_and(|spend| spend <= 0.0) { + return Ok(Decision::default()); + } + let mut minimum: Option<&ThresholdFacts> = None; + for key in order { + let Some(entry) = global.iter().chain(&local).find(|entry| entry.key == *key) else { + continue; + }; + if !entry.digit { + continue; + } + if let Some(failure) = &entry.integer_failure { + return Err(failure.clone()); + } + let Some(digits) = entry.digits.as_ref() else { + continue; + }; + if minimum + .and_then(|entry| entry.digits.as_ref()) + .is_none_or(|current| (digits.len(), digits) < (current.len(), current)) + { + minimum = Some(entry); + } + } + let Some(minimum) = minimum else { + return Ok(Decision::default()); + }; + if let Some(failure) = &minimum.ratio_failure { + return Err(failure.clone()); + } + let emit = matches!((facts.spend, facts.limit, minimum.ratio), (Some(spend), Some(limit), Some(ratio)) if spend.partial_cmp(&(limit * ratio)) != Some(std::cmp::Ordering::Less)); + Ok(Decision { + emit, + thresholds: Some(thresholds), + ..Decision::default() + }) +} + +pub fn plan(facts: &Facts) -> Decision { + evaluate(facts).unwrap_or_else(|failure| Decision { + failure: Some(failure), + ..Decision::default() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::before(7.9, false)] + #[case::at(8.0, true)] + #[case::at_limit(10.0, false)] + #[case::nan(f64::NAN, false)] + #[case::zero(0.0, false)] + fn default_key_threshold(#[case] spend: f64, #[case] emit: bool) { + let facts = Facts { + member: false, + spend: Some(spend), + limit: Some(10.0), + default_fraction: 0.8, + config: Some(ConfigFacts { + global_entries: vec![], + local_entries: vec![], + global_failure: None, + local_failure: None, + order: None, + }), + }; + assert_eq!(plan(&facts).emit, emit); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_context.rs b/litellm-rust/crates/request-auth-core/src/budget_context.rs new file mode 100644 index 00000000000..5162edd0836 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_context.rs @@ -0,0 +1,484 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum ModelScope { + Team, + EndUser, + CustomEndUser, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelCheckFacts { + pub scope: ModelScope, + pub entity_id: Option, + pub budget_count: usize, + pub model_count: usize, +} + +pub fn model_checks(facts: &ModelCheckFacts) -> Vec { + if facts.budget_count == 0 + || facts.entity_id.is_none() + || (facts.scope == ModelScope::EndUser && facts.entity_id.as_deref() == Some("")) + { + return vec![]; + } + (0..facts.model_count).collect() +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct BudgetFields { + pub max_budget: Option, + pub tpm_limit: Option, + pub rpm_limit: Option, + pub tpd_limit: Option, + pub model_budget_count: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum EndUserMode { + Models, + Token, + Default, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct EndUserFacts { + pub mode: EndUserMode, + pub object_present: bool, + pub end_user_id: Option, + pub prisma_present: bool, + pub key_budget_id: Option, + pub key_read: bool, + pub key_budget: Option, + pub global_config_read: bool, + pub global_budget_id: Option, + pub global_read: bool, + pub global_budget: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum EndUserSource { + Own, + ReadKey, + Key, + ReadGlobalConfig, + ReadGlobal, + Global, + None, +} + +pub fn end_user_source(facts: &EndUserFacts) -> EndUserSource { + use EndUserSource::*; + if facts.mode != EndUserMode::Default && facts.object_present { + return Own; + } + if facts.mode == EndUserMode::Models + && (!facts.prisma_present || facts.end_user_id.as_deref().is_none_or(str::is_empty)) + { + return None; + } + if facts.key_budget_id.is_some() { + if !facts.key_read { + return ReadKey; + } + if facts.key_budget.is_some() { + return Key; + } + } + if !facts.global_config_read { + return ReadGlobalConfig; + } + if facts.global_budget_id.is_none() { + return None; + } + if !facts.global_read { + return ReadGlobal; + } + if facts.global_budget.is_some() { + Global + } else { + None + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TokenDefaultsFacts { + pub current: BudgetFields, + pub keep_token_limits: bool, + pub default_present: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum TokenField { + MaxBudget, + Tpm, + Rpm, + Tpd, + Models, +} + +pub fn token_assignments(facts: &TokenDefaultsFacts) -> Vec { + if !facts.default_present { + return vec![]; + } + [ + (TokenField::MaxBudget, facts.current.max_budget.is_none()), + (TokenField::Tpm, facts.current.tpm_limit.is_none()), + (TokenField::Rpm, facts.current.rpm_limit.is_none()), + (TokenField::Tpd, facts.current.tpd_limit.is_none()), + ( + TokenField::Models, + facts.current.model_budget_count.is_none(), + ), + ] + .into_iter() + .filter_map(|(field, missing)| (!facts.keep_token_limits || missing).then_some(field)) + .collect() +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MembershipSnapshot { + pub budget_id: Option, + pub model_budget_count: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, strum::VariantArray)] +pub enum LookupFailure { + Membership, + Default, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MemberFacts { + pub user_id: Option, + pub team_id: Option, + pub prisma_present: bool, + pub default_id: Option, + pub skip_budget_checks: bool, + pub membership_read: bool, + pub membership: Option, + pub default_read: bool, + pub default_budget: Option, + pub lookup_failure: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum MemberSource { + Done, + ReadMembership, + ReadDefault, + Member, + Default, + Clear, + RaiseLookup, + SkipLookup, +} + +pub fn member_source(facts: &MemberFacts) -> MemberSource { + use MemberSource::*; + if facts.user_id.is_none() || facts.team_id.is_none() || !facts.prisma_present { + return Done; + } + if facts.lookup_failure.is_some() { + return if facts.skip_budget_checks { + SkipLookup + } else { + RaiseLookup + }; + } + if !facts.membership_read { + return ReadMembership; + } + if facts.membership.as_ref().is_some_and(|member| { + member.budget_id.is_some() + && member.budget_id != facts.default_id + && member.model_budget_count.is_some_and(|count| count > 0) + }) { + return Member; + } + if facts.default_id.is_none() { + return Done; + } + if !facts.default_read { + return ReadDefault; + } + if facts + .default_budget + .as_ref() + .is_some_and(|budget| budget.model_budget_count.is_some()) + { + Default + } else { + Clear + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct IncreaseFacts { + pub number: Option, + pub truthy: bool, + pub type_name: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TemporaryFacts { + pub max_budget: Option, + pub metadata_read: bool, + pub has_increase: bool, + pub has_expiry: bool, + pub expires_at: Option, + pub now: Option, + pub increase: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum TemporaryRead { + Metadata, + Expiry, + Now, + Increase, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct TemporaryDecision { + pub read: Option, + pub copy: bool, + pub max_budget: Option, + pub invalid_increase_type: Option, +} + +fn temp_read(read: TemporaryRead) -> TemporaryDecision { + TemporaryDecision { + read: Some(read), + ..TemporaryDecision::default() + } +} + +pub fn temporary(facts: &TemporaryFacts) -> TemporaryDecision { + let Some(limit) = facts.max_budget else { + return TemporaryDecision::default(); + }; + if !facts.metadata_read { + return temp_read(TemporaryRead::Metadata); + } + if !facts.has_increase || !facts.has_expiry { + return TemporaryDecision::default(); + } + let Some(expiry) = facts.expires_at else { + return temp_read(TemporaryRead::Expiry); + }; + let Some(now) = facts.now else { + return temp_read(TemporaryRead::Now); + }; + if expiry <= now { + return TemporaryDecision::default(); + } + let Some(increase) = &facts.increase else { + return temp_read(TemporaryRead::Increase); + }; + if !increase.truthy { + return TemporaryDecision::default(); + } + match increase.number { + Some(amount) => TemporaryDecision { + copy: true, + max_budget: Some(limit + amount), + ..TemporaryDecision::default() + }, + None => TemporaryDecision { + invalid_increase_type: Some(increase.type_name.clone()), + ..TemporaryDecision::default() + }, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct NumericFacts { + pub number: Option, + pub boolean: bool, +} + +pub fn percentage(facts: &NumericFacts) -> Option { + if facts.boolean { + return None; + } + facts.number.filter(|pct| *pct > 0.0 && *pct <= 1.0) +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ThrottleFacts { + pub opt_in: bool, + pub tpm_limit: Option, + pub rpm_limit: Option, + pub percentage: Option, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct ThrottleDecision { + pub read_percentage: bool, + pub percentage: Option, +} + +pub fn throttle(facts: &ThrottleFacts) -> ThrottleDecision { + if !facts.opt_in || (facts.tpm_limit.is_none() && facts.rpm_limit.is_none()) { + return ThrottleDecision::default(); + } + match &facts.percentage { + None => ThrottleDecision { + read_percentage: true, + ..ThrottleDecision::default() + }, + Some(raw) => ThrottleDecision { + percentage: percentage(raw), + ..ThrottleDecision::default() + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::team_empty(ModelScope::Team, Some(""), vec![0, 1])] + #[case::team_none(ModelScope::Team, None, vec![])] + #[case::enduser_empty(ModelScope::EndUser, Some(""), vec![])] + #[case::enduser(ModelScope::EndUser, Some("user"), vec![0, 1])] + #[case::custom_empty(ModelScope::CustomEndUser, Some(""), vec![0, 1])] + #[case::custom_absent(ModelScope::CustomEndUser, None, vec![])] + fn model_identity_gates( + #[case] scope: ModelScope, + #[case] identity: Option<&str>, + #[case] expected: Vec, + ) { + assert_eq!( + model_checks(&ModelCheckFacts { + scope, + entity_id: identity.map(str::to_owned), + budget_count: 1, + model_count: 2 + }), + expected + ); + } + + #[rstest] + #[case::zero(0.0, false, None)] + #[case::negative(-0.5, false, None)] + #[case::nan(f64::NAN, false, None)] + #[case::infinite(f64::INFINITY, false, None)] + #[case::one(1.0, false, Some(1.0))] + #[case::boolean(1.0, true, None)] + fn numeric_percentages( + #[case] number: f64, + #[case] boolean: bool, + #[case] expected: Option, + ) { + assert_eq!( + percentage(&NumericFacts { + number: Some(number), + boolean + }), + expected + ); + } + + #[rstest] + #[case::expired(10, 10, 3.0, false, None)] + #[case::before(9, 10, 3.0, false, None)] + #[case::active(11, 10, 3.0, true, Some(13.0))] + #[case::negative(11, 10, -3.0, true, Some(7.0))] + #[case::zero(11, 10, 0.0, false, None)] + fn temporary_copy_boundary( + #[case] expiry: i64, + #[case] now: i64, + #[case] amount: f64, + #[case] copy: bool, + #[case] expected: Option, + ) { + let facts = TemporaryFacts { + max_budget: Some(10.0), + metadata_read: true, + has_increase: true, + has_expiry: true, + expires_at: Some(expiry), + now: Some(now), + increase: Some(IncreaseFacts { + number: Some(amount), + truthy: amount != 0.0, + type_name: "float".to_owned(), + }), + }; + let result = temporary(&facts); + assert_eq!( + (result.copy, result.max_budget, result.read), + (copy, expected, None) + ); + } + + #[rstest] + #[case::member_override(Some("member"), Some(1), Some("default"), false, MemberSource::Member)] + #[case::no_override(None, Some(1), None, false, MemberSource::Done)] + #[case::empty_override( + Some("member"), + Some(0), + Some("default"), + false, + MemberSource::ReadDefault + )] + #[case::current_default( + Some("default"), + Some(1), + Some("default"), + false, + MemberSource::ReadDefault + )] + #[case::missing_default(None, None, Some("default"), true, MemberSource::Clear)] + fn member_default_staging( + #[case] member_id: Option<&str>, + #[case] count: Option, + #[case] default_id: Option<&str>, + #[case] default_read: bool, + #[case] expected: MemberSource, + ) { + let facts = MemberFacts { + user_id: Some(String::new()), + team_id: Some(String::new()), + prisma_present: true, + default_id: default_id.map(str::to_owned), + skip_budget_checks: false, + membership_read: true, + membership: Some(MembershipSnapshot { + budget_id: member_id.map(str::to_owned), + model_budget_count: count, + }), + default_read, + default_budget: None, + lookup_failure: None, + }; + assert_eq!(member_source(&facts), expected); + } + + #[rstest] + #[case::replace(false, true, vec![TokenField::MaxBudget, TokenField::Tpm, TokenField::Rpm, TokenField::Tpd, TokenField::Models])] + #[case::keep_empty(true, true, vec![TokenField::Rpm, TokenField::Tpd])] + #[case::no_default(false, false, vec![])] + fn token_assignment_order( + #[case] keep: bool, + #[case] present: bool, + #[case] expected: Vec, + ) { + let facts = TokenDefaultsFacts { + current: BudgetFields { + max_budget: Some(0.0), + tpm_limit: Some(0), + model_budget_count: Some(0), + ..BudgetFields::default() + }, + keep_token_limits: keep, + default_present: present, + }; + assert_eq!(token_assignments(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_lookup.rs b/litellm-rust/crates/request-auth-core/src/budget_lookup.rs new file mode 100644 index 00000000000..c5c150d2bdb --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_lookup.rs @@ -0,0 +1,284 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct OrganizationFacts { + pub token_present: bool, + pub client_present: bool, + pub team_present: bool, + pub token_id: Option, + pub team_id: Option, + pub row_present: Option, + pub lookup_failed: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum OrganizationAction { + Skip, + ReadTokenId, + ReadTeamId, + ReadOrganization, + UseOrganization, +} + +pub fn organization(facts: &OrganizationFacts) -> OrganizationAction { + use OrganizationAction::*; + if !facts.token_present || !facts.client_present { + return Skip; + } + match facts.token_id { + None => return ReadTokenId, + Some(false) if !facts.team_present => return Skip, + Some(false) => match facts.team_id { + None => return ReadTeamId, + Some(false) => return Skip, + Some(true) => {} + }, + Some(true) => {} + } + if facts.lookup_failed { + return Skip; + } + match facts.row_present { + None => ReadOrganization, + Some(false) => Skip, + Some(true) => UseOrganization, + } +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct MemberFacts { + pub team_present: bool, + pub token_present: bool, + pub team_id: Option, + pub user_id: Option, + pub membership_loaded: bool, + pub limits_read: bool, + pub maximum_present: bool, + pub default_id_string: bool, + pub default_read: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum MemberAction { + Skip, + ReadTeamId, + ReadUserId, + ReadMembership, + ReadLimits, + ReadDefault, + UseLimits, +} + +pub fn member(facts: &MemberFacts) -> MemberAction { + use MemberAction::*; + if !facts.team_present { + return Skip; + } + match facts.team_id { + None => return ReadTeamId, + Some(false) => return Skip, + Some(true) => {} + } + if !facts.token_present { + return Skip; + } + match facts.user_id { + None => return ReadUserId, + Some(false) => return Skip, + Some(true) => {} + } + if !facts.membership_loaded { + return ReadMembership; + } + if !facts.limits_read { + return ReadLimits; + } + if !facts.maximum_present && facts.default_id_string && !facts.default_read { + return ReadDefault; + } + UseLimits +} + +#[derive(Clone, Debug, Deserialize)] +pub enum EmailFacts { + Text(String), + List(Vec>), + Other, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct SoftFacts { + pub limit: Option, + pub spend: Option, + pub token_present: bool, +} + +pub fn read_soft_emails(facts: &SoftFacts) -> bool { + facts.token_present + && facts + .limit + .zip(facts.spend) + .is_some_and(|(limit, spend)| spend >= limit) +} + +fn stripped(value: &str) -> &str { + value.trim_matches(|character: char| { + character.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&character) + }) +} + +pub fn emails(facts: &EmailFacts) -> Vec { + match facts { + EmailFacts::Text(value) => value + .split(',') + .map(stripped) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + .collect(), + EmailFacts::List(values) => values + .iter() + .flatten() + .filter(|value| !stripped(value).is_empty()) + .cloned() + .collect(), + EmailFacts::Other => Vec::new(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::below(Some(1.0), Some(0.0), true, false)] + #[case::at_limit(Some(1.0), Some(1.0), true, true)] + #[case::no_token(Some(1.0), Some(2.0), false, false)] + #[case::no_limit(None, Some(2.0), true, false)] + #[case::nan(Some(f64::NAN), Some(2.0), true, false)] + fn lazy_soft_metadata( + #[case] limit: Option, + #[case] spend: Option, + #[case] token_present: bool, + #[case] expected: bool, + ) { + assert_eq!( + read_soft_emails(&SoftFacts { + limit, + spend, + token_present + }), + expected + ); + } + + #[rstest] + fn failed_organization_lookup_is_skipped() { + assert_eq!( + organization(&OrganizationFacts { + token_present: true, + client_present: true, + token_id: Some(true), + lookup_failed: true, + ..Default::default() + }), + OrganizationAction::Skip + ); + } + + #[rstest] + #[case::missing_token(false, true, None, None, OrganizationAction::Skip)] + #[case::missing_client(true, false, None, None, OrganizationAction::Skip)] + #[case::token_first(true, true, None, None, OrganizationAction::ReadTokenId)] + #[case::token_selected(true, true, Some(true), None, OrganizationAction::ReadOrganization)] + #[case::missing_row(true, true, Some(true), Some(false), OrganizationAction::Skip)] + #[case::loaded_row( + true, + true, + Some(true), + Some(true), + OrganizationAction::UseOrganization + )] + fn organization_stages( + #[case] token_present: bool, + #[case] client_present: bool, + #[case] token_id: Option, + #[case] row_present: Option, + #[case] expected: OrganizationAction, + ) { + assert_eq!( + organization(&OrganizationFacts { + token_present, + client_present, + token_id, + row_present, + ..Default::default() + }), + expected + ); + } + + #[rstest] + #[case::no_team(false, None, OrganizationAction::Skip)] + #[case::read_team(true, None, OrganizationAction::ReadTeamId)] + #[case::no_team_id(true, Some(false), OrganizationAction::Skip)] + #[case::team_selected(true, Some(true), OrganizationAction::ReadOrganization)] + fn organization_fallback( + #[case] team_present: bool, + #[case] team_id: Option, + #[case] expected: OrganizationAction, + ) { + assert_eq!( + organization(&OrganizationFacts { + token_present: true, + client_present: true, + token_id: Some(false), + team_present, + team_id, + row_present: None, + lookup_failed: false + }), + expected + ); + } + + #[rstest] + #[case::load_membership(false, false, false, false, false, MemberAction::ReadMembership)] + #[case::read_limits(true, false, false, false, false, MemberAction::ReadLimits)] + #[case::own_maximum(true, true, true, true, false, MemberAction::UseLimits)] + #[case::default(true, true, false, true, false, MemberAction::ReadDefault)] + #[case::invalid_default(true, true, false, false, false, MemberAction::UseLimits)] + #[case::loaded_default(true, true, false, true, true, MemberAction::UseLimits)] + fn member_stages( + #[case] membership_loaded: bool, + #[case] limits_read: bool, + #[case] maximum_present: bool, + #[case] default_id_string: bool, + #[case] default_read: bool, + #[case] expected: MemberAction, + ) { + assert_eq!( + member(&MemberFacts { + team_present: true, + token_present: true, + team_id: Some(true), + user_id: Some(true), + membership_loaded, + limits_read, + maximum_present, + default_id_string, + default_read + }), + expected + ); + } + + #[rstest] + #[case::comma_list(EmailFacts::Text(" a, ,b,a ".into()), vec!["a", "b", "a"])] + #[case::list_keeps_space(EmailFacts::List(vec![Some(" a ".into()), None, Some("\u{1c}".into())]), vec![" a "])] + #[case::python_whitespace(EmailFacts::Text("\u{1c}a\u{1f},\u{85}b\u{a0}".into()), vec!["a", "b"])] + #[case::other(EmailFacts::Other, vec![])] + fn email_normalization(#[case] facts: EmailFacts, #[case] expected: Vec<&str>) { + assert_eq!(emails(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_pricing.rs b/litellm-rust/crates/request-auth-core/src/budget_pricing.rs new file mode 100644 index 00000000000..29329d5c0ea --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_pricing.rs @@ -0,0 +1,261 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize)] +pub struct DeploymentPrice { + pub has_input_price: bool, + pub has_output_price: bool, + pub input_price: Option, + pub output_price: Option, + pub query_price: Option, + pub ptu_count: Option, + pub cost_per_ptu_per_hour: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelPriceFacts { + pub cached: Option, + pub lookup_failed: bool, + pub group_found: bool, + pub input_price: Option, + pub output_price: Option, + pub deployments: Option>, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct PriceDecision { + pub zero_cost: bool, + pub cache_value: Option, + pub read_deployments: bool, +} + +pub fn price(facts: &ModelPriceFacts) -> PriceDecision { + if let Some(cached) = facts.cached { + return PriceDecision { + zero_cost: cached, + cache_value: None, + read_deployments: false, + }; + } + if facts.lookup_failed { + return PriceDecision { + zero_cost: false, + cache_value: None, + read_deployments: false, + }; + } + let group_zero = facts.group_found + && matches!((facts.input_price, facts.output_price), (Some(input), Some(output)) if !(input > 0.0 || output > 0.0)); + if group_zero && facts.deployments.is_none() { + return PriceDecision { + zero_cost: false, + cache_value: None, + read_deployments: true, + }; + } + let deployments = facts.deployments.as_deref().unwrap_or_default(); + let zero_cost = group_zero + && deployments + .iter() + .any(|d| d.has_input_price || d.has_output_price) + && !deployments.iter().any(|d| { + d.input_price.is_some_and(|p| p > 0.0) + || d.output_price.is_some_and(|p| p > 0.0) + || d.query_price.is_some_and(|p| p > 0.0) + || d.ptu_count.is_some() && d.cost_per_ptu_per_hour.is_some() + }); + PriceDecision { + zero_cost, + cache_value: Some(zero_cost), + read_deployments: false, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct RouteBudgetFacts { + pub side_effect_route: bool, + pub discovery_route: bool, + pub llm_route: bool, + pub registered_pass_through: Option, + pub route: String, + pub method: Option, + pub body_present: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum RouteBudgetDecision { + ReadPassThrough, + ReadPricing, + Skip, +} + +pub fn route(facts: &RouteBudgetFacts) -> RouteBudgetDecision { + use RouteBudgetDecision::*; + if !facts.side_effect_route { + if facts.discovery_route { + return Skip; + } + if !facts.llm_route { + match facts.registered_pass_through { + None => return ReadPassThrough, + Some(false) => return Skip, + Some(true) => (), + } + } + } + let mcp_discovery = facts.body_present + && (matches!( + facts.route.as_str(), + "/mcp-rest/tools/list" | "/v1/mcp/tools" + ) || !matches!( + facts.route.as_str(), + "/mcp/tools/call" | "/mcp-rest/tools/call" + ) && (facts.route == "/mcp" || facts.route.starts_with("/mcp/")) + && matches!( + facts.method.as_deref(), + Some("initialize" | "notifications/initialized" | "ping" | "tools/list") + )); + if mcp_discovery { Skip } else { ReadPricing } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ExemptionFacts { + pub route: RouteBudgetFacts, + pub model_input_present: bool, + pub router_present: bool, + pub model_count: usize, + #[serde(default)] + pub cache_read: bool, + pub prices: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct ExemptionDecision { + pub read_pass_through: bool, + pub read_price_cache: bool, + pub read_model: Option, + pub exempt: bool, +} + +pub fn exemption(facts: &ExemptionFacts) -> ExemptionDecision { + let done = |exempt| ExemptionDecision { + read_pass_through: false, + read_price_cache: false, + read_model: None, + exempt, + }; + match route(&facts.route) { + RouteBudgetDecision::ReadPassThrough => { + return ExemptionDecision { + read_pass_through: true, + ..done(false) + }; + } + RouteBudgetDecision::Skip => return done(true), + RouteBudgetDecision::ReadPricing => (), + } + if !facts.model_input_present || !facts.router_present { + return done(false); + } + if !facts.cache_read { + return ExemptionDecision { + read_price_cache: true, + ..done(false) + }; + } + if facts.prices.iter().any(|facts| !price(facts).zero_cost) { + return done(false); + } + if facts.prices.len() < facts.model_count { + return ExemptionDecision { + read_model: Some(facts.prices.len()), + ..done(false) + }; + } + done(true) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::zero(Some(0.0), true)] + #[case::negative(Some(-1.0), true)] + #[case::nan(Some(f64::NAN), true)] + #[case::missing(None, false)] + #[case::positive(Some(1.0), false)] + fn preserves_group_comparison(#[case] input: Option, #[case] zero_cost: bool) { + let facts = ModelPriceFacts { + cached: None, + lookup_failed: false, + group_found: true, + input_price: input, + output_price: Some(0.0), + deployments: Some(vec![DeploymentPrice { + has_input_price: true, + has_output_price: false, + input_price: None, + output_price: None, + query_price: None, + ptu_count: None, + cost_per_ptu_per_hour: None, + }]), + }; + assert_eq!( + price(&facts), + PriceDecision { + zero_cost, + cache_value: Some(zero_cost), + read_deployments: false, + } + ); + } + #[rstest] + #[case::read_first(vec![], Some(0), false)] + #[case::read_second(vec![true], Some(1), false)] + #[case::priced_first(vec![false], None, false)] + #[case::all_free(vec![true, true], None, true)] + #[case::priced_second(vec![true, false], None, false)] + fn exemption_reads_models_in_order( + #[case] cached: Vec, + #[case] read: Option, + #[case] exempt: bool, + ) { + let facts = ExemptionFacts { + route: RouteBudgetFacts { + side_effect_route: false, + discovery_route: false, + llm_route: true, + route: "/chat/completions".to_owned(), + method: None, + body_present: false, + registered_pass_through: None, + }, + model_input_present: true, + router_present: true, + model_count: 2, + cache_read: true, + prices: cached + .into_iter() + .map(|value| ModelPriceFacts { + cached: Some(value), + lookup_failed: false, + group_found: false, + input_price: None, + output_price: None, + deployments: None, + }) + .collect(), + }; + assert_eq!( + exemption(&facts), + ExemptionDecision { + read_pass_through: false, + read_price_cache: false, + read_model: read, + exempt + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/claude_membership.rs b/litellm-rust/crates/request-auth-core/src/claude_membership.rs new file mode 100644 index 00000000000..2840e4aaa9b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/claude_membership.rs @@ -0,0 +1,100 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] +pub struct Facts { + pub name: String, + pub router_present: bool, + pub alias_count: usize, + pub global_alias_exists: Option, + pub alias_entries: Vec>, + pub registry_loaded: bool, + pub registry_alias_exists: Option, + pub deployment: Option, + pub names: Option>, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { present: bool }, + ReadGlobal, + ReadAlias { index: usize }, + ReadRegistry, + ReadAliases, + ReadDeployment, + ReadNames, +} + +pub fn decide(facts: &Facts) -> Decision { + let Some(global) = &facts.global_alias_exists else { + return Decision::ReadGlobal; + }; + if *global { + return Decision::Complete { present: true }; + } + for index in 0..facts.alias_count { + let Some(aliases) = facts.alias_entries.get(index) else { + return Decision::ReadAlias { index }; + }; + if *aliases == Some(true) { + return Decision::Complete { present: true }; + } + } + if !facts.router_present { + return Decision::Complete { present: false }; + } + if !facts.registry_loaded { + return Decision::ReadRegistry; + } + let Some(aliases) = &facts.registry_alias_exists else { + return Decision::ReadAliases; + }; + if *aliases { + return Decision::Complete { present: true }; + } + match facts.deployment { + None => return Decision::ReadDeployment, + Some(true) => return Decision::Complete { present: true }, + Some(false) => (), + } + let Some(names) = &facts.names else { + return Decision::ReadNames; + }; + Decision::Complete { + present: !names.is_empty(), + } +} + +#[cfg(test)] +mod tests { + use super::{Decision, Facts, decide}; + use rstest::rstest; + + #[rstest] + #[case::global(Some(true), Decision::Complete { present: true })] + #[case::missing(None, Decision::ReadGlobal)] + #[case::next_alias(Some(false), Decision::ReadAlias { index: 0 })] + fn alias_short_circuit(#[case] global_alias_exists: Option, #[case] expected: Decision) { + let facts = Facts { + name: "model".into(), + router_present: true, + alias_count: 3, + global_alias_exists, + ..Facts::default() + }; + assert_eq!(decide(&facts), expected); + } + + #[rstest] + #[case::no_router(false, Decision::Complete { present: false })] + #[case::registry(true, Decision::ReadRegistry)] + fn registry_after_all_aliases(#[case] router_present: bool, #[case] expected: Decision) { + let facts = Facts { + name: "model".into(), + router_present, + global_alias_exists: Some(false), + ..Facts::default() + }; + assert_eq!(decide(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/claude_normalization.rs b/litellm-rust/crates/request-auth-core/src/claude_normalization.rs new file mode 100644 index 00000000000..7a853e5ae2d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/claude_normalization.rs @@ -0,0 +1,176 @@ +use serde::{Deserialize, Serialize}; + +use crate::claude_membership; + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct Facts { + pub route: String, + pub request_present: bool, + pub router_present: Option, + pub normalized: Option, + pub model_loaded: bool, + pub requested: Option, + pub body_loaded: bool, + pub body_model: Option, + pub settings_loaded: bool, + pub membership: Vec, + pub marked: bool, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete, + ReadRouter, + ReadNormalized, + ReadModel, + ReadBodyModel, + ReadSettings, + ReadMembership { name: String, index: usize }, + MarkNormalized, + Rewrite { target: String }, +} + +fn unmarked(name: &str) -> &str { + if name.to_ascii_lowercase().ends_with("[1m]") { + &name[..name.len() - 4] + } else { + name + } +} + +fn decoded(name: &str) -> Option { + let encoded = unmarked(name).strip_prefix("claude-router-")?; + if !encoded.len().is_multiple_of(2) + || !encoded + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + { + return None; + } + let bytes: Option> = encoded + .as_bytes() + .as_chunks::<2>() + .0 + .iter() + .map(|pair| u8::from_str_radix(std::str::from_utf8(pair).ok()?, 16).ok()) + .collect(); + String::from_utf8(bytes?) + .ok() + .filter(|value| !value.is_empty()) +} + +fn membership(facts: &Facts, index: usize, name: &str) -> Result { + if let Some(read) = facts.membership.get(index).filter(|read| read.name == name) + && let claude_membership::Decision::Complete { present } = claude_membership::decide(read) + { + return Ok(present); + } + Err(Decision::ReadMembership { + name: name.into(), + index, + }) +} + +fn target(facts: &Facts, name: &str) -> Result, Decision> { + if membership(facts, 0, name)? { + return Ok(None); + } + let plain = unmarked(name); + let marked = plain != name; + if marked && membership(facts, 1, plain)? { + return Ok(Some(plain.into())); + } + let Some(model) = decoded(name) else { + return Ok(None); + }; + Ok(membership(facts, if marked { 2 } else { 1 }, &model)?.then_some(model)) +} + +pub fn decide(facts: &Facts) -> Decision { + if !matches!( + facts.route.as_str(), + "/messages" + | "/v1/messages" + | "/chat/completions" + | "/v1/chat/completions" + | "/responses" + | "/v1/responses" + ) { + return Decision::Complete; + } + match facts.router_present { + None => return Decision::ReadRouter, + Some(false) => return Decision::Complete, + Some(true) => (), + } + if facts.request_present { + match facts.normalized { + None => return Decision::ReadNormalized, + Some(true) => return Decision::Complete, + Some(false) => (), + } + } + if !facts.model_loaded { + return Decision::ReadModel; + } + let Some(requested) = facts.requested.as_deref() else { + return Decision::Complete; + }; + if !facts.body_loaded { + return Decision::ReadBodyModel; + } + if Some(requested) != facts.body_model.as_deref() { + return Decision::Complete; + } + if !requested.starts_with("claude-router-") && !requested.to_ascii_lowercase().ends_with("[1m]") + { + return Decision::Complete; + } + if !facts.settings_loaded { + return Decision::ReadSettings; + } + let result = match target(facts, requested) { + Err(read) => return read, + Ok(value) => value, + }; + if facts.request_present && !facts.marked { + return Decision::MarkNormalized; + } + match result { + Some(target) => Decision::Rewrite { target }, + None => Decision::Complete, + } +} + +#[cfg(test)] +mod tests { + use super::{Decision, Facts, decide, decoded}; + use rstest::rstest; + + #[rstest] + #[case::canonical("claude-router-6d6f64656c", Some("model"))] + #[case::marked("claude-router-6d6f64656c[1M]", Some("model"))] + #[case::uppercase("claude-router-6D6F64656C", None)] + #[case::whitespace("claude-router-6d 6f64656c", None)] + #[case::invalid_utf8("claude-router-ff", None)] + #[case::empty("claude-router-", None)] + #[case::odd("claude-router-1", None)] + #[case::unicode("claude-router-e6a8a1e59e8b", Some("模型"))] + fn canonical_picker_decoding(#[case] name: &str, #[case] expected: Option<&str>) { + assert_eq!(decoded(name).as_deref(), expected); + } + + #[rstest] + #[case::irrelevant("/models", Decision::Complete)] + #[case::eligible("/v1/messages", Decision::ReadRouter)] + fn route_before_reads(#[case] route: &str, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + route: route.into(), + ..Facts::default() + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/codex_models.rs b/litellm-rust/crates/request-auth-core/src/codex_models.rs new file mode 100644 index 00000000000..1dbcbbc902c --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/codex_models.rs @@ -0,0 +1,427 @@ +use crate::{Error, model_helpers}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Routes, + NativeRoutes, + Chatgpt, + Grants, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Candidate { + pub provider_loaded: bool, + pub provider: Option, + pub model: Option, + pub name: Option, + pub public_loaded: bool, + pub public_name: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Grants { + pub key: Vec, + pub team: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub kind: Kind, + pub model: Option, + pub router_present: bool, + pub team_id: Option, + pub bundled: Option>, + pub priced: Option, + pub candidates: Option>, + pub alias_loaded: bool, + pub alias: Option, + pub grants: Option, + pub prefixed_grants: Option, + pub bare: Vec, + pub prefixed: Vec, + pub route: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum Field { + Provider, + Model, + Name, + PublicName, +} + +#[derive(Debug, Serialize, PartialEq, Eq)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { + allowed: bool, + target: Option, + }, + ReadBundled, + ReadCost { + model: String, + }, + ReadCandidates { + model: String, + }, + ReadAlias { + model: String, + }, + ReadCandidate { + index: usize, + field: Field, + }, + ReadGrants { + prefixed: bool, + }, + ReadHelper { + index: usize, + model: String, + prefixed: bool, + }, +} + +fn complete(allowed: bool) -> Decision { + Decision::Complete { + allowed, + target: None, + } +} + +fn grants(facts: &Facts, model: &str, prefixed: bool) -> Result, Error> { + let selected = if prefixed { + &facts.prefixed_grants + } else { + &facts.grants + }; + let Some(grants) = selected else { + return Ok(Err(Decision::ReadGrants { prefixed })); + }; + let scopes = if prefixed { + &facts.prefixed + } else { + &facts.bare + }; + for (index, expected) in [&grants.key, &grants.team].into_iter().enumerate() { + let Some(scope) = scopes.get(index) else { + return Ok(Err(Decision::ReadHelper { + index, + model: model.to_owned(), + prefixed, + })); + }; + if scope.model != model + || scope.grants.as_ref() != Some(expected) + || scope.kind != model_helpers::Kind::Helper + { + return Err(Error::InvalidAnswer); + } + match model_helpers::decide(scope)? { + model_helpers::Decision::Complete { allowed: false } => return Ok(Ok(false)), + model_helpers::Decision::Complete { allowed: true } => (), + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Ok(true)) +} + +fn candidate( + candidate: &Candidate, + index: usize, + facts: &Facts, + model: &str, +) -> Result { + let read = |field| Decision::ReadCandidate { index, field }; + if !candidate.provider_loaded { + return Err(read(Field::Provider)); + } + let chatgpt = if candidate.provider.as_deref() == Some("chatgpt") { + true + } else { + let backend = candidate + .model + .as_deref() + .ok_or_else(|| read(Field::Model))?; + backend.split('/').next() == Some("chatgpt") + }; + if facts.kind == Kind::Chatgpt { + return Ok(chatgpt); + } + if chatgpt { + return Ok(false); + } + let alias = facts + .alias + .as_deref() + .filter(|alias| !alias.is_empty()) + .unwrap_or(model); + let name = candidate.name.as_deref().ok_or_else(|| read(Field::Name))?; + if name == model || name == alias { + return Ok(true); + } + if !candidate.public_loaded { + return Err(read(Field::PublicName)); + } + Ok(candidate + .public_name + .as_deref() + .is_some_and(|name| name == model || name == alias)) +} + +pub fn decide(facts: &Facts) -> Result { + if facts.kind == Kind::NativeRoutes + && facts.route.as_deref() != Some("/backend-api/codex/responses") + { + return Ok(complete(false)); + } + let Some(model) = facts.model.as_deref().filter(|_| facts.router_present) else { + return Ok(complete(false)); + }; + if facts.kind == Kind::Grants { + return Ok(match grants(facts, model, false)? { + Ok(allowed) => complete(allowed), + Err(read) => read, + }); + } + if facts.kind == Kind::Chatgpt { + if model.starts_with("chatgpt/") { + return Ok(complete(false)); + } + match grants(facts, model, false)? { + Ok(true) => return Ok(complete(false)), + Err(read) => return Ok(read), + Ok(false) => (), + } + } else { + let Some(bundled) = &facts.bundled else { + return Ok(Decision::ReadBundled); + }; + if bundled.iter().any(|name| name == model) { + return Ok(complete(false)); + } + match facts.priced { + None => { + return Ok(Decision::ReadCost { + model: format!("chatgpt/{model}"), + }); + } + Some(true) => return Ok(complete(false)), + Some(false) => (), + } + } + let target = if facts.kind == Kind::Chatgpt { + format!("chatgpt/{model}") + } else { + model.to_owned() + }; + let Some(candidates) = &facts.candidates else { + return Ok(Decision::ReadCandidates { model: target }); + }; + if matches!(facts.kind, Kind::Routes | Kind::NativeRoutes) && !facts.alias_loaded { + return Ok(Decision::ReadAlias { + model: model.to_owned(), + }); + } + if candidates.is_empty() { + return Ok(complete(false)); + } + for (index, value) in candidates.iter().enumerate() { + match candidate(value, index, facts, model) { + Ok(false) => return Ok(complete(false)), + Err(read) => return Ok(read), + Ok(true) => (), + } + } + if matches!(facts.kind, Kind::Routes | Kind::NativeRoutes) { + return Ok(complete(true)); + } + Ok(match grants(facts, &target, true)? { + Ok(true) => Decision::Complete { + allowed: true, + target: Some(target), + }, + Ok(false) => complete(false), + Err(read) => read, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + fn helper(model: &str, grants: &[String]) -> model_helpers::Facts { + model_helpers::Facts { + kind: model_helpers::Kind::Helper, + model: model.into(), + grants: Some(grants.to_owned()), + team_id: None, + team_aliases: vec![], + router_present: false, + groups: None, + names: None, + provider_loaded: false, + provider: None, + provider_reads: None, + } + } + + #[rstest] + #[case::all_candidates_both_grants(true, true, true)] + #[case::mixed_candidates(false, true, false)] + #[case::team_denies(true, false, false)] + #[case::mixed_and_team_denies(false, false, false)] + fn prefix_requires_all_chatgpt_deployments_and_both_grants( + base: Facts, + #[case] all_chatgpt: bool, + #[case] team_allows: bool, + #[case] allowed: bool, + ) { + let grants = Grants { + key: vec!["chatgpt/model".into()], + team: vec![ + if team_allows { + "chatgpt/model" + } else { + "other" + } + .into(), + ], + }; + let first = Candidate { + provider_loaded: true, + provider: Some("chatgpt".into()), + model: Some("chatgpt/model".into()), + name: Some("chatgpt/model".into()), + public_loaded: false, + public_name: None, + }; + let second = Candidate { + provider: Some(if all_chatgpt { "chatgpt" } else { "openai" }.into()), + model: Some( + if all_chatgpt { + "chatgpt/model" + } else { + "openai/model" + } + .into(), + ), + ..first.clone() + }; + let facts = Facts { + kind: Kind::Chatgpt, + candidates: Some(vec![first, second]), + bare: vec![helper("model", &grants.key)], + prefixed: vec![ + helper("chatgpt/model", &grants.key), + helper("chatgpt/model", &grants.team), + ], + grants: Some(grants.clone()), + prefixed_grants: Some(grants), + ..base + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::Complete { + allowed, + target: allowed.then(|| "chatgpt/model".into()) + } + ); + } + + #[fixture] + fn base() -> Facts { + Facts { + kind: Kind::Routes, + model: Some("model".into()), + router_present: true, + team_id: None, + bundled: None, + priced: None, + candidates: None, + alias_loaded: false, + alias: None, + grants: None, + prefixed_grants: None, + bare: vec![], + prefixed: vec![], + route: None, + } + } + + #[rstest] + #[case::no_model(None, true, false)] + #[case::no_router(Some("model"), false, false)] + #[case::needs_bundled_catalog(Some("model"), true, true)] + fn routing_starts_lazily( + base: Facts, + #[case] model: Option<&str>, + #[case] router: bool, + #[case] read: bool, + ) { + let facts = Facts { + model: model.map(str::to_owned), + router_present: router, + ..base + }; + let expected = if read { + Decision::ReadBundled + } else { + Decision::Complete { + allowed: false, + target: None, + } + }; + assert_eq!(decide(&facts).unwrap(), expected); + } + + #[rstest] + #[case::bundled("model", false)] + #[case::other("other", true)] + fn bundled_route_stops_before_cost(base: Facts, #[case] bundled: &str, #[case] read: bool) { + let facts = Facts { + bundled: Some(vec![bundled.into()]), + ..base + }; + let expected = if read { + Decision::ReadCost { + model: "chatgpt/model".into(), + } + } else { + Decision::Complete { + allowed: false, + target: None, + } + }; + assert_eq!(decide(&facts).unwrap(), expected); + } + + #[rstest] + #[case::alias_pending(false, true)] + #[case::empty_candidates(true, false)] + fn empty_candidates_still_read_alias( + base: Facts, + #[case] alias_loaded: bool, + #[case] read: bool, + ) { + let facts = Facts { + bundled: Some(vec![]), + priced: Some(false), + candidates: Some(vec![]), + alias_loaded, + ..base + }; + let expected = if read { + Decision::ReadAlias { + model: "model".into(), + } + } else { + Decision::Complete { + allowed: false, + target: None, + } + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/codex_normalization.rs b/litellm-rust/crates/request-auth-core/src/codex_normalization.rs new file mode 100644 index 00000000000..041c0cca72d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/codex_normalization.rs @@ -0,0 +1,201 @@ +use crate::{Error, codex_models, codex_turn, route_access::RouteContext, route_enable}; +use serde::{Deserialize, Serialize}; + +const NATIVE_ROUTE: &str = "/backend-api/codex/responses"; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Alias, + Prefix, +} + +#[derive(Deserialize)] +pub struct Facts { + pub kind: Kind, + pub route: String, + pub request_present: bool, + pub router_present: Option, + pub context: Option, + pub eu_jev: Option, + pub resolved: Option, + pub marked: bool, + pub model_loaded: bool, + pub model: Option, + pub explicit: Option, + pub provider: Option, + pub routing: Option, + pub body_present: Option, + pub settings_loaded: bool, + pub target: Option, + pub provider_after: Option, + pub turn: Option, + pub prefix: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete, + ReadRouter, + ReadContext, + ReadEuJev, + ReadResolved, + MarkResolved, + ReadModel, + ReadExplicit, + ReadProvider, + ReadRouting, + ReadBody, + ReadSettings, + ReadProviderAfter, + ReadTurn, + ReadPrefix, + Rewrite { target: String }, +} + +fn routed(facts: &codex_models::Facts) -> Result { + match codex_models::decide(facts)? { + codex_models::Decision::Complete { allowed, .. } => Ok(allowed), + _ => Err(Error::InvalidAnswer), + } +} + +pub fn decide(facts: &Facts) -> Result { + if !facts.request_present { + return Ok(Decision::Complete); + } + if facts.kind == Kind::Prefix && (facts.route != NATIVE_ROUTE || facts.model.is_none()) { + return Ok(Decision::Complete); + } + match facts.router_present { + None => return Ok(Decision::ReadRouter), + Some(false) => return Ok(Decision::Complete), + Some(true) => (), + } + if facts.kind == Kind::Prefix { + let Some(prefix) = &facts.prefix else { + return Ok(Decision::ReadPrefix); + }; + return match codex_models::decide(prefix)? { + codex_models::Decision::Complete { + target: Some(target), + .. + } => Ok(Decision::Rewrite { target }), + codex_models::Decision::Complete { .. } => Ok(Decision::Complete), + _ => Err(Error::InvalidAnswer), + }; + } + let Some(context) = &facts.context else { + return Ok(Decision::ReadContext); + }; + match route_enable::llm_route(context, facts.eu_jev)? { + route_enable::LlmDecision::Deny => return Ok(Decision::Complete), + route_enable::LlmDecision::ReadEuJev => return Ok(Decision::ReadEuJev), + route_enable::LlmDecision::Allow => (), + } + match facts.resolved { + None => return Ok(Decision::ReadResolved), + Some(true) => return Ok(Decision::Complete), + Some(false) => (), + } + if !facts.marked { + return Ok(Decision::MarkResolved); + } + if !facts.model_loaded { + return Ok(Decision::ReadModel); + } + match facts.explicit { + None => return Ok(Decision::ReadExplicit), + Some(true) => return Ok(Decision::Complete), + Some(false) => (), + } + let Some(provider) = facts.provider else { + return Ok(Decision::ReadProvider); + }; + if provider { + if facts.route != NATIVE_ROUTE { + return Ok(Decision::Complete); + } + let Some(routing) = &facts.routing else { + return Ok(Decision::ReadRouting); + }; + if !routed(routing)? { + return Ok(Decision::Complete); + } + } + let Some(model) = &facts.model else { + return Ok(Decision::Complete); + }; + match facts.body_present { + None => return Ok(Decision::ReadBody), + Some(false) => return Ok(Decision::Complete), + Some(true) => (), + } + if !facts.settings_loaded { + return Ok(Decision::ReadSettings); + } + let Some(target) = facts + .target + .as_ref() + .filter(|target| !target.is_empty() && *target != model) + else { + return Ok(Decision::Complete); + }; + let Some(provider_after) = facts.provider_after else { + return Ok(Decision::ReadProviderAfter); + }; + if provider_after { + let Some(turn) = &facts.turn else { + return Ok(Decision::ReadTurn); + }; + match codex_turn::turn(turn)? { + codex_turn::Decision::Complete { allowed: true } => (), + codex_turn::Decision::Complete { allowed: false } => return Ok(Decision::Complete), + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Rewrite { + target: target.clone(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing_request(false, NATIVE_ROUTE, Some("model"), Decision::Complete)] + #[case::wrong_route(true, "/v1/responses", Some("model"), Decision::Complete)] + #[case::non_string(true, NATIVE_ROUTE, None, Decision::Complete)] + #[case::native(true, NATIVE_ROUTE, Some("model"), Decision::ReadRouter)] + fn prefix_does_not_read_router_on_early_exit( + #[case] request: bool, + #[case] route: &str, + #[case] model: Option<&str>, + #[case] expected: Decision, + ) { + let facts = Facts { + kind: Kind::Prefix, + route: route.into(), + request_present: request, + router_present: None, + context: None, + eu_jev: None, + resolved: None, + marked: false, + model_loaded: true, + model: model.map(str::to_owned), + explicit: None, + provider: None, + routing: None, + body_present: None, + settings_loaded: false, + target: None, + provider_after: None, + turn: None, + prefix: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/codex_turn.rs b/litellm-rust/crates/request-auth-core/src/codex_turn.rs new file mode 100644 index 00000000000..4215b762193 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/codex_turn.rs @@ -0,0 +1,230 @@ +use crate::{Error, codex_models, route_access, route_enable, route_scope}; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +pub struct KeyFacts { + pub proxy: Option, + pub enabled: Option, + pub virtual_route: Option, + pub role_loaded: bool, + pub role: Option, + pub admin: Option, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum JwtRead { + Returned, + Exception, +} + +#[derive(Deserialize)] +pub struct AccessFacts { + pub key: KeyFacts, + pub jwt_token_present: bool, + pub jwt: Option, +} + +#[derive(Deserialize)] +pub struct Credentials { + pub bearer_present: bool, + pub gateway_present: bool, + pub equal: bool, +} + +#[derive(Deserialize)] +pub struct Facts { + pub model: Option, + pub native: bool, + pub credentials: Option, + pub access: Option, + pub routing: Option, + pub catalog_loaded: bool, + pub slugs: Option>, +} + +#[derive(Deserialize)] +pub enum Request { + Key(Box), + Access(Box), + Turn(Box), + Jwt(JwtRead), + JwtToken { + claims_present: bool, + token_present: bool, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool }, + Failure { message: String }, + ReadProxy, + ReadEnabled, + ReadVirtual, + ReadRole, + ReadAdmin, + ReadCredentials, + ReadAccess, + ReadRouting, + ReadJwt, + ReadCatalog, +} + +fn complete(allowed: bool) -> Decision { + Decision::Complete { allowed } +} + +fn scope(facts: &route_scope::Facts) -> Result, Error> { + match route_scope::decide(facts) { + route_scope::Decision::Allow => Ok(Ok(true)), + route_scope::Decision::Deny => Ok(Ok(false)), + route_scope::Decision::Failure { message } => Ok(Err(message)), + _ => Err(Error::InvalidAnswer), + } +} + +pub fn key(facts: &KeyFacts) -> Result { + let Some(proxy) = &facts.proxy else { + return Ok(Decision::ReadProxy); + }; + match scope(proxy)? { + Ok(true) => (), + Ok(false) => return Ok(complete(false)), + Err(message) => return Ok(Decision::Failure { message }), + } + let Some(enabled) = &facts.enabled else { + return Ok(Decision::ReadEnabled); + }; + match route_enable::decide(enabled)? { + route_enable::Decision::Allow => (), + route_enable::Decision::Deny { .. } | route_enable::Decision::PropagateHttp => { + return Ok(complete(false)); + } + _ => return Err(Error::InvalidAnswer), + } + let Some(virtual_route) = &facts.virtual_route else { + return Ok(Decision::ReadVirtual); + }; + if !route_access::virtual_route(virtual_route)?.allowed { + return Ok(complete(false)); + } + if !facts.role_loaded { + return Ok(Decision::ReadRole); + } + if facts.role.as_deref() == Some("proxy_admin") { + return Ok(complete(true)); + } + let Some(admin) = &facts.admin else { + return Ok(Decision::ReadAdmin); + }; + Ok(match scope(admin)? { + Ok(allowed) => complete(allowed), + Err(message) => Decision::Failure { message }, + }) +} + +pub fn access(facts: &AccessFacts) -> Result { + match key(&facts.key)? { + Decision::Complete { allowed: true } => (), + Decision::Complete { allowed: false } => return Ok(complete(false)), + _ => return Err(Error::InvalidAnswer), + } + if !facts.jwt_token_present { + return Ok(complete(true)); + } + Ok(match facts.jwt { + None => Decision::ReadJwt, + Some(read) => jwt(read), + }) +} + +fn jwt(read: JwtRead) -> Decision { + complete(matches!(read, JwtRead::Returned)) +} + +pub fn turn(facts: &Facts) -> Result { + if facts.native { + let Some(credentials) = &facts.credentials else { + return Ok(Decision::ReadCredentials); + }; + if !credentials.bearer_present || !credentials.gateway_present || credentials.equal { + return Ok(complete(false)); + } + } + let Some(access_facts) = &facts.access else { + return Ok(Decision::ReadAccess); + }; + let Some(model) = &facts.model else { + return Ok(complete(false)); + }; + let Some(routing) = &facts.routing else { + return Ok(Decision::ReadRouting); + }; + match codex_models::decide(routing)? { + codex_models::Decision::Complete { allowed: true, .. } => (), + codex_models::Decision::Complete { allowed: false, .. } => return Ok(complete(false)), + _ => return Err(Error::InvalidAnswer), + } + match access(access_facts)? { + Decision::Complete { allowed: true } => (), + other => return Ok(other), + } + if !facts.catalog_loaded { + return Ok(Decision::ReadCatalog); + } + Ok(complete( + facts + .slugs + .as_ref() + .is_some_and(|slugs| !slugs.contains(model)), + )) +} + +pub fn decide(request: Request) -> Result { + match request { + Request::Key(facts) => key(&facts), + Request::Access(facts) => access(&facts), + Request::Turn(facts) => turn(&facts), + Request::Jwt(read) => Ok(jwt(read)), + Request::JwtToken { + claims_present, + token_present, + } => Ok(complete(claims_present && !token_present)), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(None, Decision::ReadCredentials)] + #[case::no_bearer(Some(Credentials { bearer_present: false, gateway_present: true, equal: false }), complete(false))] + #[case::no_gateway(Some(Credentials { bearer_present: true, gateway_present: false, equal: false }), complete(false))] + #[case::equal(Some(Credentials { bearer_present: true, gateway_present: true, equal: true }), complete(false))] + #[case::separate(Some(Credentials { bearer_present: true, gateway_present: true, equal: false }), Decision::ReadAccess)] + fn credentials_precede_route_reads( + #[case] credentials: Option, + #[case] expected: Decision, + ) { + let facts = Facts { + model: Some("model".into()), + native: true, + credentials, + access: None, + routing: None, + catalog_loaded: false, + slugs: None, + }; + assert_eq!(turn(&facts).unwrap(), expected); + } + + #[rstest] + #[case::returned(JwtRead::Returned, true)] + #[case::exception(JwtRead::Exception, false)] + fn jwt_wrapper_uses_only_validator_outcome(#[case] read: JwtRead, #[case] allowed: bool) { + assert_eq!(jwt(read), complete(allowed)); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/common_budgets.rs b/litellm-rust/crates/request-auth-core/src/common_budgets.rs new file mode 100644 index 00000000000..9f3463fe5c1 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/common_budgets.rs @@ -0,0 +1,230 @@ +use std::cmp::Ordering::{Equal, Greater, Less}; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Scope { + Team, + Key, + TeamWindow, + KeyWindow, + Project, + Organization, + Tag, + AccessGroup, + User, + TeamMember, + EndUser, + Proxy, + KeySoft, + TeamSoft, + ProjectSoft, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct BudgetFacts { + pub scope: Scope, + pub limit: Option, + pub spend: Option, + pub token_present: bool, + pub team_key: bool, + pub apply_user_budget_to_team_keys: bool, + pub info_route: bool, + pub llm_route: bool, + pub model_list_route: bool, + pub skip: bool, + pub alert_recipients: Vec, + pub throttle_percentage: Option, + pub throttle_opt_in: bool, + pub tpm_limit: Option, + pub rpm_limit: Option, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize)] +pub struct BudgetDecision { + pub read_spend: bool, + pub denied: bool, + pub alert: Option<&'static str>, + pub throttle_percentage: Option, +} + +pub fn check(facts: &BudgetFacts) -> BudgetDecision { + use Scope::*; + let Some(limit) = facts.limit else { + return BudgetDecision::default(); + }; + if matches!(facts.scope, User) && facts.team_key && !facts.apply_user_budget_to_team_keys + || matches!(facts.scope, EndUser) && facts.info_route + || matches!(facts.scope, Project) && !limit.is_finite() + || matches!(facts.scope, AccessGroup) && limit <= 0.0 + || matches!(facts.scope, Proxy) + && (limit.partial_cmp(&0.0) != Some(Greater) + || facts.skip + || !facts.llm_route + || facts.model_list_route) + { + return BudgetDecision::default(); + } + let soft = matches!(facts.scope, KeySoft | TeamSoft | ProjectSoft); + let Some(spend) = facts.spend else { + return BudgetDecision { + read_spend: !soft && facts.scope != Proxy, + ..BudgetDecision::default() + }; + }; + if soft { + let alert = spend >= limit + && match facts.scope { + KeySoft => limit != 0.0, + TeamSoft => facts.token_present && !facts.alert_recipients.is_empty(), + ProjectSoft => facts.token_present, + _ => false, + }; + return BudgetDecision { + alert: alert.then_some("soft_budget"), + ..BudgetDecision::default() + }; + } + let exceeded = match facts.scope { + Team | Proxy => limit.is_finite() && spend > limit, + EndUser => spend > limit, + Tag => !matches!(spend.partial_cmp(&limit), Some(Less | Equal)), + AccessGroup => spend.partial_cmp(&limit) != Some(Less), + _ => limit.is_finite() && spend >= limit, + }; + let alert = match facts.scope { + Key => Some("token_budget"), + User => Some("user_budget"), + Team if exceeded && facts.token_present => Some("team_budget"), + Project if exceeded && facts.token_present => Some("project_budget"), + Organization if exceeded => Some("organization_budget"), + TeamMember if limit.is_finite() => Some("team_member_threshold"), + _ => None, + }; + let throttle_percentage = if exceeded + && facts.scope == Key + && facts.throttle_opt_in + && (facts.tpm_limit.is_some() || facts.rpm_limit.is_some()) + { + facts + .throttle_percentage + .filter(|pct| *pct > 0.0 && *pct <= 1.0) + } else { + None + }; + BudgetDecision { + read_spend: false, + denied: exceeded && throttle_percentage.is_none(), + alert, + throttle_percentage, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct AlertFacts { + pub spend: Option, + pub max_budget: Option, + pub configured: bool, + pub minimum_percent: Option, + pub default_ratio: f64, + pub team_member: bool, +} + +pub fn alert(facts: &AlertFacts) -> bool { + let (Some(spend), Some(limit)) = (facts.spend, facts.max_budget) else { + return false; + }; + if facts.team_member { + if !facts.configured || spend <= 0.0 { + return false; + } + } else if spend.partial_cmp(&0.0) != Some(Greater) { + return false; + } + if facts.configured { + return facts + .minimum_percent + .is_some_and(|pct| spend.partial_cmp(&(limit * (pct / 100.0))) != Some(Less)); + } + spend >= limit * facts.default_ratio && spend < limit +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MemberLimitFacts { + pub member_limit: Option, + pub default_limit: Option, + pub temporary_increase: Option, + pub expires_at: Option, + pub now: f64, + #[serde(default)] + pub now_naive: bool, +} + +pub fn member_limit(facts: &MemberLimitFacts) -> Result, crate::Error> { + let base = facts.member_limit.or_else(|| { + facts + .default_limit + .filter(|limit| !matches!(limit.partial_cmp(&0.0), Some(Less | Equal))) + }); + let Some(base) = base else { + return Ok(None); + }; + let increase = match (facts.temporary_increase, facts.expires_at) { + (Some(_), Some(_)) if facts.now_naive => return Err(crate::Error::IncompatibleDateTimes), + (Some(increase), Some(expiry)) if expiry > facts.now => increase, + _ => 0.0, + }; + Ok(Some(base + increase)) +} + +#[derive(Clone, Debug, Deserialize)] +pub struct WindowFacts { + pub duration_seconds: Option, + pub reset_at: Option, + pub now: f64, +} + +pub fn window_start(facts: &WindowFacts) -> Option { + facts + .duration_seconds + .map(|seconds| facts.reset_at.unwrap_or(facts.now) - seconds) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::team_equality(Scope::Team, 1.0, 1.0, false)] + #[case::key_equality(Scope::Key, 1.0, 1.0, true)] + #[case::tag_nan(Scope::Tag, f64::NAN, 1.0, true)] + #[case::group_nan_spend(Scope::AccessGroup, 1.0, f64::NAN, true)] + #[case::key_nan(Scope::Key, f64::NAN, 1.0, false)] + #[case::end_user_infinity(Scope::EndUser, 1.0, f64::INFINITY, true)] + fn boundary( + #[case] scope: Scope, + #[case] limit: f64, + #[case] spend: f64, + #[case] denied: bool, + ) { + let facts = BudgetFacts { + scope, + limit: Some(limit), + spend: Some(spend), + token_present: false, + team_key: false, + apply_user_budget_to_team_keys: false, + info_route: false, + llm_route: true, + model_list_route: false, + skip: false, + alert_recipients: vec![], + throttle_percentage: None, + throttle_opt_in: false, + tpm_limit: None, + rpm_limit: None, + }; + assert_eq!(check(&facts).denied, denied); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/common_flow.rs b/litellm-rust/crates/request-auth-core/src/common_flow.rs new file mode 100644 index 00000000000..2f6a4e9873d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/common_flow.rs @@ -0,0 +1,358 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, strum::VariantArray)] +pub enum ModelBudgetScope { + Key, + User, + Member, +} + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum ReadFacts { + ModelBudgets { + scope: ModelBudgetScope, + entity_present: Option, + budget_present: Option, + mapping: Option, + team_present: Option, + }, + FallbackAllowed { + router: bool, + disabled: Option, + }, + BudgetSkip { + skip: bool, + exempt: Option, + }, + Pricing { + enabled: bool, + llm_route: Option, + }, + Unpriced { + count: usize, + }, + MemberFallback { + router: bool, + multipart: Option, + resource_bound: Option, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum ReadPlan { + CheckModelBudgets, + ReadModelBudgetPresence, + ReadModelBudgetShape, + ReadModelBudgetToken, + ReadModelBudgetUser, + ReadModelBudgetTeam, + AllowFallback, + ReadFallbackDisabled, + ReadExemption, + SkipBudgets, + CheckBudgets, + ReadRoute, + ReadUnpriced, + RejectUnpriced, + ReadContentType, + ReadResourceBinding, + ReadFallbackSettings, + Continue, +} + +pub fn read(facts: ReadFacts) -> ReadPlan { + match facts { + ReadFacts::ModelBudgets { + scope, + entity_present, + budget_present, + mapping, + team_present, + } => model_budgets(scope, entity_present, budget_present, mapping, team_present), + ReadFacts::FallbackAllowed { + router: true, + disabled: Some(false), + } => ReadPlan::AllowFallback, + ReadFacts::FallbackAllowed { + router: true, + disabled: None, + } => ReadPlan::ReadFallbackDisabled, + ReadFacts::FallbackAllowed { .. } => ReadPlan::Continue, + ReadFacts::BudgetSkip { skip: true, .. } + | ReadFacts::BudgetSkip { + exempt: Some(true), .. + } => ReadPlan::SkipBudgets, + ReadFacts::BudgetSkip { + exempt: Some(false), + .. + } => ReadPlan::CheckBudgets, + ReadFacts::BudgetSkip { exempt: None, .. } => ReadPlan::ReadExemption, + ReadFacts::Pricing { enabled: false, .. } + | ReadFacts::Pricing { + llm_route: Some(false), + .. + } => ReadPlan::Continue, + ReadFacts::Pricing { + llm_route: None, .. + } => ReadPlan::ReadRoute, + ReadFacts::Pricing { + llm_route: Some(true), + .. + } => ReadPlan::ReadUnpriced, + ReadFacts::Unpriced { count: 0 } => ReadPlan::Continue, + ReadFacts::Unpriced { .. } => ReadPlan::RejectUnpriced, + ReadFacts::MemberFallback { router: false, .. } + | ReadFacts::MemberFallback { + multipart: Some(true), + .. + } + | ReadFacts::MemberFallback { + resource_bound: Some(true), + .. + } => ReadPlan::Continue, + ReadFacts::MemberFallback { + multipart: None, .. + } => ReadPlan::ReadContentType, + ReadFacts::MemberFallback { + resource_bound: None, + .. + } => ReadPlan::ReadResourceBinding, + ReadFacts::MemberFallback { .. } => ReadPlan::ReadFallbackSettings, + } +} + +fn model_budgets( + scope: ModelBudgetScope, + entity: Option, + budget: Option, + mapping: Option, + team: Option, +) -> ReadPlan { + use ReadPlan as P; + let stages = match scope { + ModelBudgetScope::Key => vec![ + (budget, P::ReadModelBudgetPresence), + (entity, P::ReadModelBudgetToken), + ], + ModelBudgetScope::User => vec![ + (entity, P::ReadModelBudgetUser), + (mapping, P::ReadModelBudgetShape), + (budget, P::ReadModelBudgetPresence), + ], + ModelBudgetScope::Member => vec![ + (entity, P::ReadModelBudgetUser), + (team, P::ReadModelBudgetTeam), + (mapping, P::ReadModelBudgetShape), + (budget, P::ReadModelBudgetPresence), + ], + }; + for (present, missing) in stages { + match present { + None => return missing, + Some(false) => return P::Continue, + Some(true) => {} + } + } + P::CheckModelBudgets +} + +#[derive(Deserialize)] +pub struct Facts { + pub model_present: bool, + pub skip_budgets: bool, + pub team_present: bool, + pub user_present: bool, + pub token_present: bool, + pub user_id: Option, + pub agent_id: Option, + pub customer_present: bool, + pub customer_models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct Plan { + pub membership: bool, + pub team_models: bool, + pub member_models: bool, + pub user_models: bool, + pub customer_models: bool, + pub agent_trace: bool, +} + +pub fn plan(facts: Facts) -> Plan { + Plan { + membership: facts.token_present + && (facts.model_present || !facts.skip_budgets) + && facts.team_present + && facts.user_id.is_some(), + team_models: facts.model_present && facts.team_present, + member_models: facts.model_present + && facts.team_present + && facts.token_present + && facts.user_id.is_some_and(|id| !id.is_empty()), + user_models: facts.model_present && !facts.team_present && facts.user_present, + customer_models: facts.customer_present && !facts.customer_models.is_empty(), + agent_trace: facts.token_present && facts.agent_id.is_some_and(|id| !id.is_empty()), + } +} + +#[derive(Deserialize)] +pub struct BudgetFacts { + pub token_present: bool, + pub group_count: usize, + pub customer_present: bool, + pub customer_budget_present: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum BudgetScope { + Team, + TeamWindows, + KeyWindows, + TeamSoft, + Organization, + Tags, + Groups, + User, + Member, + Customer, +} + +pub fn budgets(facts: BudgetFacts) -> Vec { + use BudgetScope as S; + [ + Some(S::Team), + Some(S::TeamWindows), + facts.token_present.then_some(S::KeyWindows), + Some(S::TeamSoft), + Some(S::Organization), + Some(S::Tags), + (facts.group_count > 0).then_some(S::Groups), + Some(S::User), + Some(S::Member), + (facts.customer_present && facts.customer_budget_present).then_some(S::Customer), + ] + .into_iter() + .flatten() + .collect() +} + +#[derive(Deserialize)] +pub struct CustomFacts { + pub end_user_id: Option, + pub end_user_present: bool, + pub common_checks: bool, + pub model_present: bool, + pub router_present: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct CustomPlan { + pub read_end_user: bool, + pub check_end_user: bool, + pub check_model_access: bool, + pub read_pricing: bool, +} + +pub fn custom(facts: CustomFacts) -> CustomPlan { + CustomPlan { + read_end_user: facts.end_user_id.is_some(), + check_end_user: facts.end_user_present && !facts.common_checks, + check_model_access: facts.common_checks, + read_pricing: facts.model_present && facts.router_present, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::available(true, true, true, true, ReadPlan::CheckModelBudgets)] + #[case::missing_identity(false, true, true, true, ReadPlan::Continue)] + #[case::empty_budget(true, false, true, true, ReadPlan::Continue)] + #[case::malformed_member_budget(true, true, false, true, ReadPlan::Continue)] + #[case::no_team(true, true, true, false, ReadPlan::Continue)] + fn model_budget_availability( + #[case] entity_present: bool, + #[case] budget_present: bool, + #[case] mapping: bool, + #[case] team_present: bool, + #[case] expected: ReadPlan, + ) { + assert_eq!( + read(ReadFacts::ModelBudgets { + scope: ModelBudgetScope::Member, + entity_present: Some(entity_present), + budget_present: Some(budget_present), + mapping: Some(mapping), + team_present: Some(team_present) + }), + expected + ); + } + + #[rstest] + #[case::explicit_skip(ReadFacts::BudgetSkip {skip:true, exempt:None}, ReadPlan::SkipBudgets)] + #[case::exempt(ReadFacts::BudgetSkip {skip:false, exempt:Some(true)}, ReadPlan::SkipBudgets)] + #[case::check(ReadFacts::BudgetSkip {skip:false, exempt:Some(false)}, ReadPlan::CheckBudgets)] + #[case::read_exemption(ReadFacts::BudgetSkip {skip:false, exempt:None}, ReadPlan::ReadExemption)] + #[case::pricing_disabled(ReadFacts::Pricing {enabled:false, llm_route:None}, ReadPlan::Continue)] + #[case::pricing_route(ReadFacts::Pricing {enabled:true, llm_route:None}, ReadPlan::ReadRoute)] + #[case::pricing_skip(ReadFacts::Pricing {enabled:true, llm_route:Some(false)}, ReadPlan::Continue)] + #[case::pricing_check(ReadFacts::Pricing {enabled:true, llm_route:Some(true)}, ReadPlan::ReadUnpriced)] + #[case::priced(ReadFacts::Unpriced {count:0}, ReadPlan::Continue)] + #[case::unpriced(ReadFacts::Unpriced {count:1}, ReadPlan::RejectUnpriced)] + #[case::no_router(ReadFacts::MemberFallback {router:false, multipart:None, resource_bound:None}, ReadPlan::Continue)] + #[case::content_type(ReadFacts::MemberFallback {router:true, multipart:None, resource_bound:None}, ReadPlan::ReadContentType)] + #[case::multipart(ReadFacts::MemberFallback {router:true, multipart:Some(true), resource_bound:None}, ReadPlan::Continue)] + #[case::resource_read(ReadFacts::MemberFallback {router:true, multipart:Some(false), resource_bound:None}, ReadPlan::ReadResourceBinding)] + #[case::bound_resource(ReadFacts::MemberFallback {router:true, multipart:Some(false), resource_bound:Some(true)}, ReadPlan::Continue)] + #[case::settings(ReadFacts::MemberFallback {router:true, multipart:Some(false), resource_bound:Some(false)}, ReadPlan::ReadFallbackSettings)] + fn selected_reads(#[case] facts: ReadFacts, #[case] expected: ReadPlan) { + assert_eq!(read(facts), expected); + } + + #[rstest] + #[case::absent(None, false, false)] + #[case::empty(Some(""), true, false)] + #[case::user(Some("user"), true, true)] + fn member_identity(#[case] user: Option<&str>, #[case] membership: bool, #[case] models: bool) { + let actual = plan(Facts { + model_present: true, + skip_budgets: false, + team_present: true, + user_present: false, + token_present: true, + user_id: user.map(str::to_owned), + agent_id: None, + customer_present: false, + customer_models: vec![], + }); + assert_eq!(actual.membership, membership); + assert_eq!(actual.member_models, models); + } + + #[rstest] + #[case::empty(false, 0, false, false, vec![BudgetScope::Team, BudgetScope::TeamWindows, BudgetScope::TeamSoft, BudgetScope::Organization, BudgetScope::Tags, BudgetScope::User, BudgetScope::Member])] + #[case::all(true, 1, true, true, vec![BudgetScope::Team, BudgetScope::TeamWindows, BudgetScope::KeyWindows, BudgetScope::TeamSoft, BudgetScope::Organization, BudgetScope::Tags, BudgetScope::Groups, BudgetScope::User, BudgetScope::Member, BudgetScope::Customer])] + fn budget_order( + #[case] token_present: bool, + #[case] group_count: usize, + #[case] customer_present: bool, + #[case] customer_budget_present: bool, + #[case] expected: Vec, + ) { + assert_eq!( + budgets(BudgetFacts { + token_present, + group_count, + customer_present, + customer_budget_present + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/credential.rs b/litellm-rust/crates/request-auth-core/src/credential.rs new file mode 100644 index 00000000000..a246b36e9db --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/credential.rs @@ -0,0 +1,225 @@ +use litellm_secret_value::SecretValue; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Default)] +pub struct PresentedKeys { + pub custom_header: Option, + pub litellm_header: Option, + pub authorization: Option, + pub api_key: Option, + pub x_api_key: Option, + pub google: Option, + pub apim: Option, + pub query: Option, + pub pass_through: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum KeySource { + #[serde(rename = "custom_header")] + CustomHeader, + #[serde(rename = "x-litellm-api-key")] + LitellmHeader, + #[serde(rename = "authorization")] + Authorization, + #[serde(rename = "api-key")] + ApiKey, + #[serde(rename = "x-api-key")] + XApiKey, + #[serde(rename = "x-goog-api-key")] + Google, + #[serde(rename = "ocp-apim-subscription-key")] + Apim, + #[serde(rename = "key")] + Query, + #[serde(rename = "pass_through_endpoint")] + PassThrough, +} + +#[derive(veil::Redact, Serialize)] +pub struct SelectedKey { + pub source: KeySource, + #[redact(with = "[REDACTED]")] + pub presented: String, + #[redact(with = "[REDACTED]")] + pub token: String, +} + +#[derive(Debug, Deserialize, Default)] +pub struct WebSocketKeys { + pub authorization: Option, + pub api_key: Option, + pub subprotocols: Option, +} + +#[derive(Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum WebSocketDenial { + #[serde(rename = "no_api_key")] + NoApiKey, + #[serde(rename = "invalid_authorization")] + InvalidAuthorization, +} + +#[derive(veil::Redact, Serialize)] +#[serde(tag = "kind", content = "value")] +pub enum WebSocketResult { + Token(#[redact(with = "[REDACTED]")] String), + Denial(WebSocketDenial), +} + +fn whitespace(ch: char) -> bool { + ch.is_whitespace() || matches!(ch, '\u{1c}'..='\u{1f}') +} + +fn credential_part(value: &str) -> Option<&str> { + let end = value + .find(|ch| ch == '/' || ch == ',' || whitespace(ch)) + .unwrap_or(value.len()); + (end != 0).then_some(&value[..end]) +} + +pub fn scheme_credential(value: &str) -> Option { + if value.starts_with("AWS4-HMAC-SHA256") { + let bearer = value + .match_indices("Credential=Bearer") + .find_map(|(index, prefix)| { + let tail = &value[index + prefix.len()..]; + if !tail.starts_with(whitespace) { + return None; + } + credential_part(tail.trim_start_matches(whitespace)) + }); + return bearer + .or_else(|| { + value + .match_indices("Credential=") + .find_map(|(index, prefix)| credential_part(&value[index + prefix.len()..])) + }) + .map(str::to_owned); + } + let (scheme, token) = value.split_once(' ')?; + matches!(scheme.to_lowercase().as_str(), "bearer" | "basic").then(|| token.to_owned()) +} + +pub fn litellm_header_token(value: &str) -> String { + scheme_credential(value).unwrap_or_else(|| value.to_owned()) +} + +pub fn select_api_key(keys: PresentedKeys) -> Option { + if let Some(value) = keys.custom_header { + return Some(SelectedKey { + source: KeySource::CustomHeader, + presented: value.expose().to_owned(), + token: scheme_credential(value.expose()).unwrap_or_default(), + }); + } + if let Some(value) = keys.litellm_header { + return Some(SelectedKey { + source: KeySource::LitellmHeader, + presented: value.expose().to_owned(), + token: litellm_header_token(value.expose()), + }); + } + if let Some(value) = keys + .authorization + .filter(|value| !value.expose().is_empty()) + { + return Some(SelectedKey { + source: KeySource::Authorization, + presented: value.expose().to_owned(), + token: scheme_credential(value.expose()).unwrap_or_default(), + }); + } + [ + (KeySource::ApiKey, keys.api_key), + (KeySource::XApiKey, keys.x_api_key), + (KeySource::Google, keys.google), + (KeySource::Apim, keys.apim), + ( + KeySource::Query, + keys.query.filter(|value| !value.expose().is_empty()), + ), + (KeySource::PassThrough, keys.pass_through), + ] + .into_iter() + .find_map(|(source, value)| { + value.map(|value| SelectedKey { + source, + presented: value.expose().to_owned(), + token: value.expose().to_owned(), + }) + }) +} + +pub fn select_websocket_api_key(keys: WebSocketKeys) -> WebSocketResult { + if let Some(value) = keys + .authorization + .filter(|value| !value.expose().is_empty()) + { + let Some((scheme, token)) = value.expose().split_once(' ') else { + return WebSocketResult::Denial(WebSocketDenial::InvalidAuthorization); + }; + if scheme.to_lowercase() != "bearer" { + return WebSocketResult::Denial(WebSocketDenial::InvalidAuthorization); + } + return WebSocketResult::Token(token.trim_matches(whitespace).to_owned()); + } + if let Some(value) = keys.api_key.filter(|value| !value.expose().is_empty()) { + return WebSocketResult::Token(value.expose().to_owned()); + } + let token = keys.subprotocols.as_ref().and_then(|value| { + value.expose().split(',').find_map(|protocol| { + protocol + .trim_matches(whitespace) + .strip_prefix("openai-insecure-api-key.") + }) + }); + match token { + Some(token) if !token.is_empty() => WebSocketResult::Token(token.to_owned()), + _ => WebSocketResult::Denial(WebSocketDenial::NoApiKey), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::bearer("Bearer token", Some(" token"))] + #[case::empty("Bearer ", Some(""))] + #[case::basic("bAsIc token", Some("token"))] + #[case::unsupported("Token token", None)] + #[case::aws("AWS4-HMAC-SHA256 Credential=Bearer\u{1c}token/date", Some("token"))] + #[case::aws_precedence( + "AWS4-HMAC-SHA256 Credential=first/date, Credential=Bearer second/date", + Some("second") + )] + fn schemes(#[case] value: &str, #[case] expected: Option<&str>) { + assert_eq!(scheme_credential(value).as_deref(), expected); + } + + #[rstest] + fn empty_custom_header_still_wins() { + let selected = select_api_key(PresentedKeys { + custom_header: Some("".into()), + api_key: Some("other".into()), + ..Default::default() + }) + .unwrap(); + assert_eq!(selected.source, KeySource::CustomHeader); + assert_eq!(selected.token, ""); + } + + #[rstest] + fn secrets_are_redacted() { + let keys = PresentedKeys { + api_key: Some("credential-value".into()), + ..Default::default() + }; + let keys_debug = format!("{keys:?}"); + let debug = format!("{keys_debug} {:?}", select_api_key(keys).unwrap()); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("REDACTED")); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/end_user_context.rs b/litellm-rust/crates/request-auth-core/src/end_user_context.rs new file mode 100644 index 00000000000..ef7a5084bc2 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/end_user_context.rs @@ -0,0 +1,428 @@ +use serde::{Deserialize, Serialize}; +use strum::VariantArray; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct IdFacts { + pub value: Option, +} + +pub fn normalize_id(facts: &IdFacts) -> Option { + facts.value.clone().filter(|value| !value.is_empty()) +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DefaultFacts { + pub prisma_present: bool, + pub explicit_id: Option, + pub global_read: bool, + pub global_id: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DefaultDecision { + pub read_global: bool, + pub budget_id: Option, +} + +pub fn default_id(facts: &DefaultFacts) -> DefaultDecision { + if facts.explicit_id.is_none() && !facts.global_read { + return DefaultDecision { + read_global: true, + budget_id: None, + }; + } + if !facts.prisma_present { + return DefaultDecision { + read_global: false, + budget_id: None, + }; + } + DefaultDecision { + read_global: false, + budget_id: facts + .explicit_id + .clone() + .or_else(|| facts.global_id.clone()), + } +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ApplyFacts { + pub own_budget_id: Option, + pub own_budget_present: bool, + pub key_id: Option, + pub global_read: bool, + pub global_id: Option, + pub default_read: bool, + pub default_present: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ApplyDecision { + Keep, + ReadGlobal, + ReadDefault, + Copy, +} + +pub fn apply_default(facts: &ApplyFacts) -> ApplyDecision { + if facts.own_budget_id.is_some() && facts.own_budget_present { + return ApplyDecision::Keep; + } + if facts.key_id.is_none() { + if !facts.global_read { + return ApplyDecision::ReadGlobal; + } + if facts.global_id.is_none() { + return ApplyDecision::Keep; + } + } + if !facts.default_read { + return ApplyDecision::ReadDefault; + } + if facts.default_present { + ApplyDecision::Copy + } else { + ApplyDecision::Keep + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct RegistryFacts { + pub end_user_id: String, + pub global_id: Option, + pub key_id: Option, + pub validate_id: bool, + pub token_max_present: bool, + pub registry_read: bool, + pub registry: Option>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum RegistryDecision { + Read, + Restricted, + Unrestricted, +} + +pub fn registry(facts: &RegistryFacts) -> RegistryDecision { + if facts.global_id.is_some() + || facts.key_id.is_some() + || facts.validate_id + || facts.token_max_present + { + return RegistryDecision::Restricted; + } + if !facts.registry_read { + return RegistryDecision::Read; + } + if facts + .registry + .as_ref() + .is_some_and(|ids| !ids.contains(&facts.end_user_id)) + { + RegistryDecision::Unrestricted + } else { + RegistryDecision::Restricted + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum Field { + EndUserId, + Tpm, + Rpm, + Tpd, + MaxBudget, + Models, + Region, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct Fields { + pub tpm: bool, + pub rpm: bool, + pub tpd: bool, + pub max_budget: bool, + pub models: bool, + pub region: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum AssignmentMode { + Params, + Token, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct AssignmentFacts { + pub mode: AssignmentMode, + pub fields: Fields, +} + +pub fn assignments(facts: &AssignmentFacts) -> Vec { + let fields = &facts.fields; + match facts.mode { + AssignmentMode::Params => [ + (Field::Tpm, fields.tpm), + (Field::Rpm, fields.rpm), + (Field::Tpd, fields.tpd), + (Field::MaxBudget, fields.max_budget), + (Field::Models, fields.models), + ] + .into_iter() + .filter_map(|(field, present)| present.then_some(field)) + .collect(), + AssignmentMode::Token => [ + (Field::EndUserId, true), + (Field::Tpm, fields.tpm), + (Field::Rpm, fields.rpm), + (Field::Tpd, fields.tpd), + (Field::Region, fields.region), + (Field::Models, fields.models), + ] + .into_iter() + .filter_map(|(field, present)| present.then_some(field)) + .collect(), + } +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct LookupFacts { + pub own_budget_present: bool, + pub key_id: Option, + pub end_user_read: bool, + pub end_user_present: bool, + pub global_read: bool, + pub global_id: Option, + pub default_read: bool, + pub default_present: bool, + pub token_max_present: bool, + pub applied: bool, + pub read_failed: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum LookupAction { + ReadEndUser, + ReadGlobal, + ReadDefault, + ApplyOwn, + ApplyDefault, + Done, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct LookupDecision { + pub action: LookupAction, + pub set_max_budget: bool, + pub project_budget: bool, +} + +pub fn lookup(facts: &LookupFacts) -> LookupDecision { + let action = if facts.read_failed || facts.applied { + LookupAction::Done + } else if !facts.end_user_read { + LookupAction::ReadEndUser + } else if facts.end_user_present { + LookupAction::ApplyOwn + } else if facts.key_id.is_none() && !facts.global_read { + LookupAction::ReadGlobal + } else if facts.key_id.is_none() && facts.global_id.is_none() { + LookupAction::Done + } else if !facts.default_read { + LookupAction::ReadDefault + } else if facts.default_present { + LookupAction::ApplyDefault + } else { + LookupAction::Done + }; + LookupDecision { + project_budget: action == LookupAction::ApplyDefault + || (action == LookupAction::ApplyOwn && facts.own_budget_present), + set_max_budget: facts.applied + && !facts.read_failed + && !facts.end_user_present + && facts.default_present + && !facts.token_max_present, + action, + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ValidationMode { + #[default] + Resolve, + Exists, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ValidationFacts { + pub mode: ValidationMode, + pub end_user_id: Option, + pub enabled: bool, + pub prisma_present: bool, + pub global_id: Option, + pub key_id: Option, + pub cache_read: bool, + pub cached: Option, + pub end_user_read: bool, + pub end_user_present: bool, + pub user_read: bool, + pub user_present: bool, + pub cache_written: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ValidationRead { + Cache, + EndUser, + User, + WriteCache, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ValidationDecision { + pub read: Option, + pub end_user_id: Option, + pub valid: bool, + pub cache_value: Option, + pub cache_ttl: u64, + pub user_email: Option, +} + +fn validation_result(facts: &ValidationFacts, valid: bool) -> ValidationDecision { + let retain = valid + || facts.key_id.is_some() + || facts.global_id.as_ref().is_some_and(|id| !id.is_empty()); + ValidationDecision { + end_user_id: facts.end_user_id.clone().filter(|_| retain), + valid, + ..ValidationDecision::default() + } +} + +pub fn validation(facts: &ValidationFacts) -> ValidationDecision { + let Some(id) = &facts.end_user_id else { + return ValidationDecision::default(); + }; + if facts.mode == ValidationMode::Resolve { + if !facts.enabled || !facts.prisma_present { + return ValidationDecision { + end_user_id: Some(id.clone()), + ..ValidationDecision::default() + }; + } + if !facts.cache_read { + return ValidationDecision { + read: Some(ValidationRead::Cache), + ..ValidationDecision::default() + }; + } + match facts.cached.as_deref() { + Some("valid") => return validation_result(facts, true), + Some("invalid") => return validation_result(facts, false), + _ => {} + } + } + if !facts.end_user_read { + return ValidationDecision { + read: Some(ValidationRead::EndUser), + ..ValidationDecision::default() + }; + } + if !facts.end_user_present && !facts.user_read { + return ValidationDecision { + read: Some(ValidationRead::User), + user_email: id.contains('@').then(|| id.clone()), + ..ValidationDecision::default() + }; + } + let valid = facts.end_user_present || facts.user_present; + if facts.mode == ValidationMode::Resolve && !facts.cache_written { + return ValidationDecision { + read: Some(ValidationRead::WriteCache), + valid, + cache_value: Some(if valid { "valid" } else { "invalid" }.to_owned()), + cache_ttl: if valid { 300 } else { 60 }, + ..ValidationDecision::default() + }; + } + validation_result(facts, valid) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(None, None)] + #[case::empty(Some(""), None)] + #[case::space(Some(" "), Some(" "))] + #[case::id(Some("budget"), Some("budget"))] + fn normalizes_metadata_id(#[case] value: Option<&str>, #[case] expected: Option<&str>) { + assert_eq!( + normalize_id(&IdFacts { + value: value.map(str::to_owned) + }) + .as_deref(), + expected + ); + } + + #[rstest] + #[case::explicit(Some("key"), Some("global"), Some("key"))] + #[case::empty_explicit(Some(""), Some("global"), Some(""))] + #[case::global(None, Some("global"), Some("global"))] + #[case::absent(None, None, None)] + fn selects_default( + #[case] explicit: Option<&str>, + #[case] global: Option<&str>, + #[case] expected: Option<&str>, + ) { + let facts = DefaultFacts { + prisma_present: true, + explicit_id: explicit.map(str::to_owned), + global_read: true, + global_id: global.map(str::to_owned), + }; + assert_eq!(default_id(&facts).budget_id.as_deref(), expected); + } + + #[rstest] + #[case::none(None, None, None)] + #[case::empty_global(Some(""), None, None)] + #[case::global(Some("budget"), None, Some("missing"))] + #[case::key(None, Some(""), Some("missing"))] + fn retains_unknown_ids( + #[case] global: Option<&str>, + #[case] key: Option<&str>, + #[case] expected: Option<&str>, + ) { + let facts = ValidationFacts { + end_user_id: Some("missing".to_owned()), + enabled: true, + prisma_present: true, + global_id: global.map(str::to_owned), + key_id: key.map(str::to_owned), + cache_read: true, + cached: Some("invalid".to_owned()), + ..ValidationFacts::default() + }; + assert_eq!(validation(&facts).end_user_id.as_deref(), expected); + } + + #[rstest] + #[case::missing(AssignmentMode::Params, vec![])] + #[case::token_identity(AssignmentMode::Token, vec![Field::EndUserId])] + fn absent_fields(#[case] mode: AssignmentMode, #[case] expected: Vec) { + assert_eq!( + assignments(&AssignmentFacts { + mode, + fields: Fields::default() + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/error.rs b/litellm-rust/crates/request-auth-core/src/error.rs new file mode 100644 index 00000000000..0f75826b4f7 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/error.rs @@ -0,0 +1,9 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid auth pattern")] + Pattern(#[source] fancy_regex::Error), + #[error("invalid auth effect answer")] + InvalidAnswer, + #[error("can't compare offset-naive and offset-aware datetimes")] + IncompatibleDateTimes, +} diff --git a/litellm-rust/crates/request-auth-core/src/fallback_access.rs b/litellm-rust/crates/request-auth-core/src/fallback_access.rs new file mode 100644 index 00000000000..6c080e12ab2 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/fallback_access.rs @@ -0,0 +1,145 @@ +use serde::{Deserialize, Serialize}; + +use crate::{ + Error, + model_access::{self, Decision as ModelDecision, Facts as ModelFacts, Request, Response}, +}; + +#[derive(Clone, Deserialize)] +pub struct GroupedScope { + pub facts: ModelFacts, + pub group_ids: Vec, + pub group_models: Option>, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub refuse_agent_keys: bool, + pub agent_id: Option, + pub key: Option, + pub team: Option, + pub embedded_team: Option, + pub assigned_models: Option>, + pub assigned_scope: Option, + pub member: Option, + pub project: Option, + pub team_present: bool, + pub embedded_team_grants: Vec, + pub member_grants: Vec, + pub project_grants: Vec, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { denial: Option }, + ReadKey, + ReadTeam, + ReadEmbeddedTeam, + ReadMember, + ReadProject, + ReadKeyGroups, + ReadTeamGroups, + ReadAssignedGroups, + ReadAssignedScope, +} + +pub fn grouped(scope: GroupedScope, merge: bool) -> Result { + let Response::Decision(decision) = model_access::decide(Request::Grouped { + facts: scope.facts, + group_ids: scope.group_ids, + group_models: scope.group_models, + merge, + })? + else { + return Err(Error::InvalidAnswer); + }; + Ok(decision) +} + +pub fn decide(facts: Facts) -> Result { + if facts.refuse_agent_keys && facts.agent_id.is_some() { + return Ok(Decision::Deny { denial: None }); + } + let Some(key) = facts.key else { + return Ok(Decision::ReadKey); + }; + let key = grouped(key, false)?; + if matches!( + key, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadKeyGroups); + } + if !key.allowed() { + return Ok(Decision::Deny { denial: Some(key) }); + } + if facts.team_present && facts.team.is_none() { + return Ok(Decision::ReadTeam); + } + if let Some(team) = facts.team { + let decision = grouped(team, true)?; + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadTeamGroups); + } + if !decision.allowed() { + if !matches!(decision, ModelDecision::Deny { .. }) { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + let Some(grants) = facts.assigned_models else { + return Ok(Decision::ReadAssignedGroups); + }; + if grants.is_empty() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + let Some(assigned) = facts.assigned_scope else { + return Ok(Decision::ReadAssignedScope); + }; + if !model_access::check(&assigned)?.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + } else if facts.embedded_team.is_none() && !facts.embedded_team_grants.is_empty() { + return Ok(Decision::ReadEmbeddedTeam); + } else if let Some(team) = facts.embedded_team { + let decision = model_access::check(&team)?; + if !decision.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + if facts.member.is_none() && !facts.member_grants.is_empty() { + return Ok(Decision::ReadMember); + } + if let Some(scope) = facts.member { + let decision = model_access::check(&scope)?; + if !decision.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + if facts.project.is_none() && !facts.project_grants.is_empty() { + return Ok(Decision::ReadProject); + } + if let Some(scope) = facts.project { + let decision = model_access::check(&scope)?; + if !decision.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + Ok(Decision::Allow) +} diff --git a/litellm-rust/crates/request-auth-core/src/fallback_candidate.rs b/litellm-rust/crates/request-auth-core/src/fallback_candidate.rs new file mode 100644 index 00000000000..214c9bd5ded --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/fallback_candidate.rs @@ -0,0 +1,322 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; +use strum::VariantArray; + +use crate::{ + Error, budget_pricing, common_budgets, fallback_access, model_access, model_attribution, + model_budget_data, route_access, unpriced_models, +}; + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum Entity { + Key, + User, + EndUser, + Team, + TeamMember, +} + +#[derive(Clone, Copy, Default, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Membership { + #[default] + Unread, + Loaded, + Failed, +} + +#[derive(Clone, Deserialize)] +pub struct BudgetScope { + pub entity: Entity, + pub identity_present: bool, + pub configured: bool, +} + +#[derive(Clone, Deserialize)] +pub struct BudgetRead { + pub model_index: usize, + pub entity: Entity, + pub snapshot: model_budget_data::Snapshot, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub router_present: bool, + pub user_present: bool, + pub team_present: bool, + pub membership: Membership, + pub block_unpriced: bool, + pub route: Option, + pub unpriced: Option, + pub customer_present: bool, + pub customer_grants: Option>, + pub customer: Option, + pub price: Option, + pub budget_models: Option>, + pub budget_scopes: Option>, + pub budgets: Vec, + pub authorization: Option, + pub attribution_reset: bool, + pub attribution: Option, + pub attribution_stamped: bool, + pub group_budgets_read: bool, + pub group_budgets: Vec, + pub io_failed: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { usable: bool, log_failure: bool }, + ReadMembership, + ReadRoute, + ReadUnpriced, + ReadCustomerGrants, + ReadCustomer, + ReadPrice, + ReadBudgetModels, + ReadBudgetScopes, + ReadBudget { model_index: usize, entity: Entity }, + ReadAuthorization, + ResetAttribution, + ReadAttribution, + StampAttribution { groups: Vec }, + ReadGroupBudgets { groups: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let complete = |usable, log_failure| { + Ok(Decision::Complete { + usable, + log_failure, + }) + }; + if !facts.router_present + || !facts.user_present + || !facts.team_present + || facts.membership == Membership::Failed + { + return complete(false, false); + } + if facts.membership == Membership::Unread { + return Ok(Decision::ReadMembership); + } + if facts.io_failed { + return complete(false, false); + } + if facts.block_unpriced { + let Some(route) = &facts.route else { + return Ok(Decision::ReadRoute); + }; + if route_access::decide(route_access::Request::Classify { + context: route.clone(), + kind: route_access::Classification::Llm, + })? + .allowed + { + let Some(unpriced) = &facts.unpriced else { + return Ok(Decision::ReadUnpriced); + }; + match unpriced_models::decide(unpriced) { + unpriced_models::Decision::Complete { unpriced: true } => { + return complete(false, false); + } + unpriced_models::Decision::Complete { unpriced: false } => (), + _ => return Err(Error::InvalidAnswer), + } + } + } + if facts.customer_present { + let Some(grants) = &facts.customer_grants else { + return Ok(Decision::ReadCustomerGrants); + }; + if !grants.is_empty() { + let Some(customer) = &facts.customer else { + return Ok(Decision::ReadCustomer); + }; + match model_access::decide(model_access::Request::Customer(customer.clone()))? { + model_access::Response::Decision(decision) if decision.allowed() => (), + model_access::Response::Decision(_) => return complete(false, true), + _ => return Err(Error::InvalidAnswer), + } + } + } + let Some(price) = &facts.price else { + return Ok(Decision::ReadPrice); + }; + let priced = budget_pricing::price(price); + if priced.read_deployments { + return Err(Error::InvalidAnswer); + } + if !priced.zero_cost { + let Some(models) = &facts.budget_models else { + return Ok(Decision::ReadBudgetModels); + }; + let Some(scopes) = &facts.budget_scopes else { + return Ok(Decision::ReadBudgetScopes); + }; + let mut seen = BTreeSet::new(); + for (model_index, model) in models.iter().enumerate() { + if !seen.insert(model) { + continue; + } + for entity in Entity::VARIANTS { + if !scopes.iter().any(|scope| { + scope.entity == *entity && scope.identity_present && scope.configured + }) { + continue; + } + let Some(reading) = facts + .budgets + .iter() + .find(|read| read.model_index == model_index && read.entity == *entity) + else { + return Ok(Decision::ReadBudget { + model_index, + entity: *entity, + }); + }; + let budget = model_budget_data::check(&reading.snapshot); + if budget.read.is_some() { + return Err(Error::InvalidAnswer); + } + if budget.denied { + return complete(false, false); + } + } + } + } + let Some(authorization) = &facts.authorization else { + return Ok(Decision::ReadAuthorization); + }; + match fallback_access::decide(authorization.clone())? { + fallback_access::Decision::Allow => (), + fallback_access::Decision::Deny { .. } => return complete(false, false), + _ => return Err(Error::InvalidAnswer), + } + if !facts.attribution_reset { + return Ok(Decision::ResetAttribution); + } + let Some(attribution) = &facts.attribution else { + return Ok(Decision::ReadAttribution); + }; + let model_attribution::Decision::Complete { groups, .. } = + model_attribution::decide(attribution) + else { + return Err(Error::InvalidAnswer); + }; + if !facts.attribution_stamped { + return Ok(Decision::StampAttribution { groups }); + } + if !groups.is_empty() { + if !facts.group_budgets_read { + return Ok(Decision::ReadGroupBudgets { groups }); + } + for facts in &facts.group_budgets { + let budget = common_budgets::check(facts); + if budget.read_spend { + return Err(Error::InvalidAnswer); + } + if budget.denied { + return complete(false, true); + } + } + } + complete(true, false) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + #[fixture] + fn facts() -> Facts { + Facts { + router_present: true, + user_present: true, + team_present: true, + ..Facts::default() + } + } + + #[rstest] + #[case::router(false, true, true, Decision::Complete { usable: false, log_failure: false })] + #[case::user(true, false, true, Decision::Complete { usable: false, log_failure: false })] + #[case::team(true, true, false, Decision::Complete { usable: false, log_failure: false })] + #[case::ready(true, true, true, Decision::ReadMembership)] + fn prerequisites( + facts: Facts, + #[case] router: bool, + #[case] user: bool, + #[case] team: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(&Facts { + router_present: router, + user_present: user, + team_present: team, + ..facts + }) + .unwrap(), + expected + ); + } + + #[rstest] + #[case::key(0, Entity::Key)] + #[case::user(1, Entity::User)] + #[case::enduser(2, Entity::EndUser)] + #[case::team(3, Entity::Team)] + #[case::member(4, Entity::TeamMember)] + fn budgets(facts: Facts, #[case] count: usize, #[case] expected: Entity) { + let facts = Facts { + membership: Membership::Loaded, + price: Some(budget_pricing::ModelPriceFacts { + cached: Some(false), + lookup_failed: false, + group_found: false, + input_price: None, + output_price: None, + deployments: None, + }), + budget_models: Some(vec!["model".into()]), + budget_scopes: Some( + Entity::VARIANTS + .iter() + .map(|entity| BudgetScope { + entity: *entity, + identity_present: true, + configured: true, + }) + .collect(), + ), + budgets: Entity::VARIANTS + .iter() + .take(count) + .map(|entity| BudgetRead { + model_index: 0, + entity: *entity, + snapshot: model_budget_data::Snapshot { + max_budget: None, + local_spend: None, + legacy_spend: None, + shared_spend: None, + shared_started_at: None, + shared_window_seconds: None, + now: 0.0, + }, + }) + .collect(), + ..facts + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::ReadBudget { + model_index: 0, + entity: expected + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/key.rs b/litellm-rust/crates/request-auth-core/src/key.rs new file mode 100644 index 00000000000..00eb6a82a97 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/key.rs @@ -0,0 +1,442 @@ +use litellm_secret_value::SecretValue; +use sha2::{Digest, Sha256}; + +use crate::{Error, model_budgets, protocol::*}; + +struct KeyCheck { + settings: Settings, + now: DateTime, + is_llm_api_route: bool, + key_label: SecretValue, +} + +pub fn start(facts: RequestFacts, settings: Settings, now: DateTime) -> Result { + let Some(master_key) = &settings.master_key else { + return Step::allow(Grant::NoMasterKey { + token: facts.token.map(|token| token.expose().to_owned()), + }); + }; + let Some(token) = facts.token.clone() else { + return Step::deny(Denial::MissingKey); + }; + if token.expose().is_empty() { + return Step::deny(Denial::MalformedKey); + } + if facts.route == "/user/auth" { + return if settings.allow_user_auth { + Step::allow(Grant::UserAuthRoute) + } else { + Step::deny(Denial::UserAuthDisabled) + }; + } + if &token == master_key { + return Step::effect( + Effect::ResolveEndUser { + prefetch_key: false, + }, + |(): ()| Step::allow(Grant::MasterKey), + ); + } + let hashed_token = format!("{:x}", Sha256::digest(token.expose().as_bytes())); + Step::effect( + Effect::ResolveEndUser { prefetch_key: true }, + move |(): ()| { + Step::effect( + Effect::ResolvePresentedKey { + hashed_token: hashed_token.clone(), + }, + move |presented: Option| { + after_presented_key(presented, token, hashed_token, facts, settings, now) + }, + ) + }, + ) +} + +pub fn start_resolved_key( + facts: RequestFacts, + settings: Settings, + now: DateTime, +) -> Result { + check_key(KeyCheck { + settings, + now, + is_llm_api_route: facts.is_llm_api_route, + key_label: facts.token.unwrap_or_else(|| "".into()), + }) +} + +fn after_presented_key( + presented: Option, + token: SecretValue, + hashed_token: String, + facts: RequestFacts, + settings: Settings, + now: DateTime, +) -> Result { + if let Some(PresentedKey::AdminKey { expires }) = &presented { + if let Some(expires) = expires + && expires.before(&now)? + { + let expires = expires.clone(); + return Step::effect(Effect::EvictPresentedKey { hashed_token }, move |(): ()| { + Step::deny(Denial::ExpiredKey { + expires, + now, + key_label: token.expose().to_owned(), + }) + }); + } + return Step::allow(Grant::CachedAdminKey); + } + if settings.master_key_only_routes.contains(&facts.route) { + return Step::deny(Denial::MasterKeyOnlyRoute); + } + if !settings.has_db { + return Step::deny(Denial::NoDatabase); + } + if presented.is_some() { + let check = KeyCheck { + settings, + now, + is_llm_api_route: facts.is_llm_api_route, + key_label: token, + }; + return Step::effect( + Effect::PrepareKey { + loaded_from_db: false, + }, + move |(): ()| check_key(check), + ); + } + if !token.expose().starts_with("sk-") { + return Step::deny(Denial::NotAVirtualKey); + } + let check = KeyCheck { + settings, + now, + is_llm_api_route: facts.is_llm_api_route, + key_label: hashed_token.clone().into(), + }; + Step::effect(Effect::LoadKey { hashed_token }, move |(): ()| { + Step::effect( + Effect::PrepareKey { + loaded_from_db: true, + }, + move |(): ()| check_key(check), + ) + }) +} + +fn check_key(check: KeyCheck) -> Result { + Step::effect(Effect::InspectKey, move |key: KeyFacts| { + if key.blocked { + return Step::deny(Denial::KeyBlocked); + } + Step::effect( + Effect::EnforceModelAccess, + move |facts: crate::key_model_access::Facts| match crate::key_model_access::decide( + facts, + )? { + crate::key_model_access::Decision::Allow => { + Step::effect(Effect::PrefetchAuthObjects, move |(): ()| { + check_owner(check, key) + }) + } + crate::key_model_access::Decision::Deny { denial } => { + Step::deny(Denial::ModelAccessDenied { decision: denial }) + } + _ => Err(Error::InvalidAnswer), + }, + ) + }) +} + +fn check_owner(check: KeyCheck, key: KeyFacts) -> Result { + let Some(user_id) = key.user_id.clone() else { + return check_budget_exemption(check, key); + }; + Step::effect( + Effect::LoadKeyOwner { + user_id: user_id.clone(), + }, + move |owner: Option| { + if owner.is_some_and(|owner| owner.deactivated) { + return Step::deny(Denial::OwnerDeactivated { user_id }); + } + check_budget_exemption(check, key) + }, + ) +} + +fn check_budget_exemption(check: KeyCheck, key: KeyFacts) -> Result { + Step::effect( + Effect::CheckBudgetExemption, + move |facts: crate::budget_pricing::ExemptionFacts| { + let decision = crate::budget_pricing::exemption(&facts); + if decision.read_pass_through + || decision.read_price_cache + || decision.read_model.is_some() + { + return Err(Error::InvalidAnswer); + } + let exempt = decision.exempt; + if exempt || key.team_member_spend.is_none() || !check.settings.has_db { + return check_expiry(check, key, exempt); + } + let (Some(user_id), Some(team_id)) = (key.user_id.clone(), key.team_id.clone()) else { + return check_expiry(check, key, exempt); + }; + Step::effect( + Effect::ReadTeamMemberBudget { + user_id: user_id.clone(), + team_id: team_id.clone(), + }, + move |budget: Option| { + let Some(budget) = budget else { + return check_expiry(check, key, false); + }; + if budget.spend < budget.max_budget { + return check_expiry(check, key, false); + } + Step::effect( + Effect::AlertTeamMemberBudget { + user_id: user_id.clone(), + team_id: team_id.clone(), + budget: budget.clone(), + }, + move |(): ()| { + Step::deny(Denial::TeamMemberBudgetExceeded { + user_id, + team_id, + spend: budget.spend, + max_budget: budget.max_budget, + }) + }, + ) + }, + ) + }, + ) +} + +fn check_expiry(check: KeyCheck, key: KeyFacts, exempt: bool) -> Result { + if let Some(expires) = &key.expires + && expires.before(&check.now)? + { + return Step::deny(Denial::ExpiredKey { + expires: expires.clone(), + now: check.now, + key_label: check.key_label.expose().to_owned(), + }); + } + if exempt { + return load_context(check, key, None); + } + Step::effect(Effect::SendKeyBudgetAlerts, move |(): ()| { + check_key_budget(check, key) + }) +} + +fn check_key_budget(check: KeyCheck, key: KeyFacts) -> Result { + let Some(max_budget) = key.max_budget.filter(|_| check.is_llm_api_route) else { + return check_soft_budget(check, key, None); + }; + Step::effect(Effect::ReadKeySpend { max_budget }, move |spend: f64| { + Step::effect(Effect::AlertKeyBudget { spend }, move |(): ()| { + if !max_budget.is_finite() || spend < max_budget { + return check_soft_budget(check, key, None); + } + if let Some(throttle) = key.budget_throttle_pct { + return check_soft_budget(check, key, Some(throttle)); + } + Step::deny(Denial::KeyBudgetExceeded { spend, max_budget }) + }) + }) +} + +fn check_soft_budget(check: KeyCheck, key: KeyFacts, throttle: Option) -> Result { + let alert = key + .soft_budget + .is_some_and(|budget| budget != 0.0 && key.spend >= budget); + let then = + move || model_budgets::key_budgets(Box::new(move || load_context(check, key, throttle))); + if alert { + return Step::effect(Effect::AlertKeySoftBudget, move |(): ()| then()); + } + then() +} + +fn load_context(check: KeyCheck, key: KeyFacts, throttle: Option) -> Result { + Step::effect(Effect::LoadKeyContext, move |(): ()| { + if !(check.settings.proxy_max_budget > 0.0 && check.settings.has_db) { + return grant(key, throttle); + } + Step::effect( + Effect::ReadProxySpend, + move |spend: Option| match spend { + None => grant(key, throttle), + Some(spend) => Step::effect(Effect::AlertProxyBudget { spend }, move |(): ()| { + grant(key, throttle) + }), + }, + ) + }) +} + +fn grant(key: KeyFacts, throttle: Option) -> Result { + if !key.has_token { + return Step::deny(Denial::KeyWithoutToken); + } + Step::allow(Grant::VirtualKey { + budget_throttle_pct: throttle, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + #[fixture] + fn check() -> KeyCheck { + KeyCheck { + settings: Settings { + master_key: Some("sk-master".into()), + allow_user_auth: false, + has_db: true, + master_key_only_routes: ["/admin".into()].into(), + proxy_max_budget: 0.0, + }, + now: DateTime { + wall_microseconds: 100, + offset_microseconds: Some(0), + timezone_id: 1, + reference: 0, + }, + is_llm_api_route: true, + key_label: "credential-value".into(), + } + } + + #[fixture] + fn key() -> KeyFacts { + KeyFacts { + has_token: true, + blocked: false, + expires: None, + spend: 0.0, + soft_budget: None, + max_budget: None, + team_member_spend: None, + user_id: None, + team_id: None, + budget_throttle_pct: None, + } + } + + #[rstest] + fn blocked_key_stops_before_model_access(check: KeyCheck, key: KeyFacts) { + let step = check_key(check).unwrap(); + assert!(matches!(step.node, Node::Effect(Effect::InspectKey))); + let result = (step.resume.unwrap())(Answer::Key(KeyFacts { + blocked: true, + ..key + })) + .unwrap(); + assert!(matches!( + result.node, + Node::Outcome(Outcome::Deny { + denial: Denial::KeyBlocked + }) + )); + assert!(result.resume.is_none()); + } + + #[rstest] + #[case::at_limit(1.0, 1.0, true)] + #[case::below_limit(0.9, 1.0, false)] + #[case::nan_limit(1.0, f64::NAN, false)] + #[case::infinite_limit(1.0, f64::INFINITY, false)] + #[case::nan_spend(f64::NAN, 1.0, true)] + fn budget_alert_precedes_the_decision( + check: KeyCheck, + key: KeyFacts, + #[case] spend: f64, + #[case] max_budget: f64, + #[case] denied: bool, + ) { + let read = check_key_budget( + check, + KeyFacts { + max_budget: Some(max_budget), + ..key + }, + ) + .unwrap(); + assert!(matches!( + read.node, + Node::Effect(Effect::ReadKeySpend { .. }) + )); + let alert = (read.resume.unwrap())(Answer::Spend(spend)).unwrap(); + assert!(matches!( + alert.node, + Node::Effect(Effect::AlertKeyBudget { .. }) + )); + let result = (alert.resume.unwrap())(Answer::Unit).unwrap(); + if denied { + assert!(matches!( + result.node, + Node::Outcome(Outcome::Deny { + denial: Denial::KeyBudgetExceeded { .. } + }) + )); + } else { + assert!(matches!( + result.node, + Node::Effect(Effect::ReadModelBudgetScope) + )); + } + } + + #[rstest] + fn cached_admin_skips_route_and_database_checks(check: KeyCheck) { + let settings = Settings { + has_db: false, + ..check.settings + }; + let step = start( + RequestFacts { + route: "/admin".into(), + token: Some("other".into()), + is_llm_api_route: false, + }, + settings, + check.now, + ) + .unwrap(); + let lookup = (step.resume.unwrap())(Answer::Unit).unwrap(); + let result = (lookup.resume.unwrap())(Answer::Presented(Some(PresentedKey::AdminKey { + expires: None, + }))) + .unwrap(); + assert!(matches!( + result.node, + Node::Outcome(Outcome::Allow { + grant: Grant::CachedAdminKey + }) + )); + } + + #[rstest] + fn public_expiry_payload_does_not_leak_into_debug(check: KeyCheck) { + let denial = Denial::ExpiredKey { + expires: check.now.clone(), + now: check.now, + key_label: check.key_label.expose().to_owned(), + }; + let debug = format!("{denial:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("REDACTED")); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/key_model_access.rs b/litellm-rust/crates/request-auth-core/src/key_model_access.rs new file mode 100644 index 00000000000..12fffc55528 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/key_model_access.rs @@ -0,0 +1,88 @@ +use serde::{Deserialize, Serialize}; + +use crate::{ + Error, + fallback_access::{GroupedScope, grouped}, + model_access::Decision as ModelDecision, +}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub config_nonempty: bool, + pub key_models: Vec, + pub primary_loaded: bool, + pub primary: Option, + pub fallback_names: Option>, + pub fallback_scopes: Vec, + pub routing: Vec, +} + +#[derive(Clone, Deserialize)] +pub enum Routing { + Routed, + Failed, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { denial: ModelDecision }, + ReadPrimary, + ReadPrimaryGroups, + ReadFallbackNames, + ReadFallbackScope { index: usize }, + ReadFallbackGroups { index: usize }, + RouteFallback { index: usize }, + RoutingFailure { index: usize }, +} + +pub fn decide(facts: Facts) -> Result { + if facts.config_nonempty + || facts + .key_models + .iter() + .any(|model| model == "all-team-models") + { + return Ok(Decision::Allow); + } + if !facts.primary_loaded { + return Ok(Decision::ReadPrimary); + } + if let Some(primary) = facts.primary { + let decision = grouped(primary, false)?; + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadPrimaryGroups); + } + if !decision.allowed() { + return Ok(Decision::Deny { denial: decision }); + } + } + let Some(names) = facts.fallback_names else { + return Ok(Decision::ReadFallbackNames); + }; + for index in 0..names.len() { + let Some(scope) = facts.fallback_scopes.get(index) else { + return Ok(Decision::ReadFallbackScope { index }); + }; + let decision = grouped(scope.clone(), false)?; + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadFallbackGroups { index }); + } + if !decision.allowed() { + return Ok(Decision::Deny { denial: decision }); + } + match facts.routing.get(index) { + None => return Ok(Decision::RouteFallback { index }), + Some(Routing::Failed) => return Ok(Decision::RoutingFailure { index }), + Some(Routing::Routed) => {} + } + } + Ok(Decision::Allow) +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs new file mode 100644 index 00000000000..2e7ae845911 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -0,0 +1,62 @@ +#![forbid(unsafe_code)] + +pub mod access; +pub mod admission; +pub mod agent_access; +pub mod agent_ceilings; +pub mod agent_grants; +pub mod auth_context; +pub mod budget_alerts; +pub mod budget_context; +pub mod budget_lookup; +pub mod budget_pricing; +pub mod claude_membership; +pub mod claude_normalization; +pub mod codex_models; +pub mod codex_normalization; +pub mod codex_turn; +pub mod common_budgets; +pub mod common_flow; +pub mod credential; +pub mod end_user_context; +mod error; +pub mod fallback_access; +pub mod fallback_candidate; +mod key; +pub mod key_model_access; +pub mod managed; +pub mod managed_models; +pub mod member_models; +pub mod model_access; +pub mod model_attribution; +pub mod model_budget_aliases; +pub mod model_budget_data; +mod model_budgets; +pub mod model_groups; +pub mod model_helpers; +pub mod model_projection; +pub mod org_routes; +pub mod owner_teams; +pub mod patterns; +pub mod personal_budget; +pub mod personal_model_listing; +pub mod personal_model_request; +pub mod personal_models; +pub mod personal_owner; +pub mod protocol; +pub mod reservation; +pub mod resolved_models; +pub mod route_access; +pub mod route_enable; +pub mod route_roles; +pub mod route_scope; +pub mod spend_counter_policy; +pub mod team_aliases; +pub mod team_override; +pub mod unpriced_models; + +pub use error::Error; +pub use key::{start, start_resolved_key}; +pub use model_budgets::{ + check_model_budgets, check_team_member_model_budgets, check_user_model_budgets, +}; diff --git a/litellm-rust/crates/request-auth-core/src/managed.rs b/litellm-rust/crates/request-auth-core/src/managed.rs new file mode 100644 index 00000000000..deb7354f601 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/managed.rs @@ -0,0 +1,348 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize)] +pub struct BoundIdentity { + pub active: bool, + pub revision: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Agent { + pub agent_id: String, + pub enabled: bool, + pub identity_managed: bool, + pub identity: Option, + pub execution_mode: String, +} + +#[derive(Debug, Deserialize)] +pub struct Context { + pub agent_id: String, + pub binding_revision: Option, + pub mode: String, + pub user_id: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum ManagedFacts { + Actor { + agent: Agent, + context: Option, + }, + ActorLookup { + agent_id: Option, + has_database: bool, + registry_loaded: bool, + registry_managed: bool, + agent_loaded: bool, + agent: Option, + retired: Option, + context: Option, + jwt_claims: bool, + delegation_verified: bool, + human_grants: Option>, + policy_assigned: bool, + }, + UnionGrants { + grants: Vec>>, + }, + IntersectGrants { + left: Option>, + right: Option>, + }, + DeclaredGrants { + direct: Vec, + declared_groups: Vec, + unified_groups: Vec, + group_agents: Option>, + unified_agents: Option>, + }, + ManagedAgentGrants { + permission_present: bool, + own: Option>, + ceilings: Option>>, + caller_loaded: bool, + caller: Option>, + mode: Option, + user_present: bool, + human: Option>, + }, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum ManagedDenial { + Disabled, + MissingContext, + ChangedIdentity, + ExecutionMode, + MissingHuman, + DatabaseRequired, + Removed, + VerifiedIdentityRequired, + HumanNotPermitted, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum ManagedAction { + Allow, + Deny, + ReadRegistry, + ReadAgent, + ReadRetired, + AssignPolicy, + ReadHumanGrants, + Grants, + ReadDeclaredGroups, + ReadUnifiedGroups, + ReadOwn, + ReadCeilings, + ReadCaller, +} + +#[derive(Debug, PartialEq, Serialize)] +pub struct ManagedDecision { + pub action: ManagedAction, + pub denial: Option, + pub grants: Option>, +} + +fn action(action: ManagedAction) -> ManagedDecision { + ManagedDecision { + action, + denial: None, + grants: None, + } +} + +fn denied(denial: ManagedDenial) -> ManagedDecision { + ManagedDecision { + action: ManagedAction::Deny, + denial: Some(denial), + grants: None, + } +} + +fn actor(agent: &Agent, context: Option<&Context>) -> ManagedDecision { + if !agent.enabled + || !agent + .identity + .as_ref() + .is_some_and(|identity| identity.active) + { + return denied(ManagedDenial::Disabled); + } + let Some(context) = context else { + return denied(ManagedDenial::MissingContext); + }; + if context.agent_id != agent.agent_id + || agent + .identity + .as_ref() + .is_some_and(|identity| context.binding_revision.as_ref() != Some(&identity.revision)) + { + return denied(ManagedDenial::ChangedIdentity); + } + if agent.execution_mode != context.mode && agent.execution_mode != "both" { + return denied(ManagedDenial::ExecutionMode); + } + if context.mode == "delegated" + && !context + .user_id + .as_ref() + .is_some_and(|user| !user.is_empty()) + { + return denied(ManagedDenial::MissingHuman); + } + action(ManagedAction::Allow) +} + +pub fn decide(facts: ManagedFacts) -> ManagedDecision { + match facts { + ManagedFacts::Actor { agent, context } => actor(&agent, context.as_ref()), + ManagedFacts::ActorLookup { + agent_id, + has_database, + registry_loaded, + registry_managed, + agent_loaded, + agent, + retired, + context, + jwt_claims, + delegation_verified, + human_grants, + policy_assigned, + } => { + if agent_id.is_none() { + return action(ManagedAction::Allow); + } + if !has_database { + if !registry_loaded { + return action(ManagedAction::ReadRegistry); + } + return if context.is_some() || registry_managed { + denied(ManagedDenial::DatabaseRequired) + } else { + action(ManagedAction::Allow) + }; + } + if !agent_loaded { + return action(ManagedAction::ReadAgent); + } + let Some(agent) = agent else { + return match retired { + None => action(ManagedAction::ReadRetired), + Some(retired) if retired || context.is_some() => denied(ManagedDenial::Removed), + Some(_) => action(ManagedAction::Allow), + }; + }; + if !agent.identity_managed { + return action(ManagedAction::Allow); + } + if jwt_claims && context.is_none() { + return denied(ManagedDenial::VerifiedIdentityRequired); + } + let checked = actor(&agent, context.as_ref()); + if checked.denial.is_some() { + return checked; + } + if !policy_assigned { + return action(ManagedAction::AssignPolicy); + } + if context + .as_ref() + .is_some_and(|context| context.mode == "delegated") + && !delegation_verified + { + return match human_grants { + None => action(ManagedAction::ReadHumanGrants), + Some(grants) if !grants.contains(&agent.agent_id) => { + denied(ManagedDenial::HumanNotPermitted) + } + Some(_) => action(ManagedAction::Allow), + }; + } + action(ManagedAction::Allow) + } + ManagedFacts::UnionGrants { grants } => { + let grants = if !grants.is_empty() && grants.iter().all(Option::is_none) { + None + } else { + Some(grants.into_iter().flatten().flatten().collect()) + }; + ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants, + } + } + ManagedFacts::IntersectGrants { left, right } => { + let grants = match (left, right) { + (None, right) => right, + (left, None) => left, + (Some(left), Some(right)) => { + Some(left.into_iter().filter(|id| right.contains(id)).collect()) + } + }; + ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants, + } + } + ManagedFacts::DeclaredGrants { + direct, + declared_groups, + unified_groups, + group_agents, + unified_agents, + } => { + if !declared_groups.is_empty() && group_agents.is_none() { + return action(ManagedAction::ReadDeclaredGroups); + } + if !unified_groups.is_empty() && unified_agents.is_none() { + return action(ManagedAction::ReadUnifiedGroups); + } + let grants = + if direct.is_empty() && declared_groups.is_empty() && unified_groups.is_empty() { + None + } else { + Some( + direct + .into_iter() + .chain(group_agents.into_iter().flatten()) + .chain(unified_agents.into_iter().flatten()) + .collect(), + ) + }; + ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants, + } + } + ManagedFacts::ManagedAgentGrants { + permission_present, + own, + ceilings, + caller_loaded, + caller, + mode, + user_present, + human, + } => { + let grants = |ids| ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants: Some(ids), + }; + if !permission_present { + return grants(vec![]); + } + let Some(own) = own else { + return action(ManagedAction::ReadOwn); + }; + let Some(ceilings) = ceilings else { + return action(ManagedAction::ReadCeilings); + }; + if !caller_loaded { + return action(ManagedAction::ReadCaller); + } + let capped: Vec<_> = own + .into_iter() + .filter(|id| { + ceilings.iter().all(|ceiling| ceiling.contains(id)) + && caller.as_ref().is_none_or(|ids| ids.contains(id)) + }) + .collect(); + if mode.is_none() || mode.as_deref() == Some("autonomous") { + return grants(capped); + } + if !user_present { + return grants(vec![]); + } + let Some(human) = human else { + return action(ManagedAction::ReadHumanGrants); + }; + grants(capped.into_iter().filter(|id| human.contains(id)).collect()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::empty(vec![], Some(vec![]))] + #[case::unrestricted(vec![None], None)] + #[case::mixed(vec![None, Some(vec!["a".into()])], Some(vec!["a".into()]))] + fn union(#[case] grants: Vec>>, #[case] expected: Option>) { + assert_eq!( + decide(ManagedFacts::UnionGrants { grants }).grants, + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/managed_models.rs b/litellm-rust/crates/request-auth-core/src/managed_models.rs new file mode 100644 index 00000000000..f618b38b371 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/managed_models.rs @@ -0,0 +1,253 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + member_models, + model_access::{self, Decision as ModelDecision}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Default, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Purpose { + #[default] + All, + Ceilings, + Caller, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub purpose: Purpose, + pub model_present: bool, + pub token_present: bool, + pub managed_loaded: bool, + pub managed_present: bool, + pub managed_models_valid: bool, + pub managed_models_count: usize, + pub managed: Option, + pub ceilings: Option>>, + pub aliases_loaded: bool, + pub scopes: Vec, + pub caller_loaded: bool, + pub caller_present: bool, + pub team_loaded: bool, + pub team_present: bool, + pub team: Option, + pub member: Option, + pub user_loaded: bool, + pub user_present: bool, + pub user_grants: Vec, + pub user: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: ModelDecision, + member_models: Option>, + }, + ReadManaged, + ReadManagedScope, + ReadCeilings, + ReadAliases, + ReadCeilingScope { + index: usize, + }, + ReadCaller, + ReadTeam, + ReadTeamScope, + ReadTeamGroups, + ReadMember, + ReadUser, + ReadUserScope, +} + +pub fn decide(facts: &Facts) -> Result { + let deny = |denial| { + Ok(Decision::Deny { + denial, + member_models: None, + }) + }; + if facts.purpose == Purpose::All && !facts.managed_loaded { + return Ok(Decision::ReadManaged); + } + if !facts.model_present || !facts.token_present { + return Ok(Decision::Allow); + } + if facts.purpose == Purpose::All && facts.managed_present { + if !facts.managed_models_valid || facts.managed_models_count == 0 { + return deny(ModelDecision::NoAgentModels); + } + let Some(scope) = &facts.managed else { + return Ok(Decision::ReadManagedScope); + }; + let result = model_access::check(scope)?; + if !result.allowed() { + return deny(result); + } + } + if facts.purpose != Purpose::Caller { + let Some(ceilings) = &facts.ceilings else { + return Ok(Decision::ReadCeilings); + }; + if !facts.aliases_loaded { + return Ok(Decision::ReadAliases); + } + for (index, models) in ceilings.iter().enumerate() { + if models.is_empty() { + return deny(ModelDecision::EmptyAgentCeiling); + } + let Some(scope) = facts.scopes.get(index) else { + return Ok(Decision::ReadCeilingScope { index }); + }; + let result = model_access::check(scope)?; + if !result.allowed() { + return deny(result); + } + } + if facts.purpose == Purpose::Ceilings { + return Ok(Decision::Allow); + } + } + if !facts.caller_loaded { + return Ok(Decision::ReadCaller); + } + if !facts.caller_present { + return Ok(Decision::Allow); + } + if !facts.team_loaded { + return Ok(Decision::ReadTeam); + } + if facts.team_present { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeamScope); + }; + let result = fallback_access::grouped(team.clone(), true)?; + if matches!( + result, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadTeamGroups); + } + if !result.allowed() { + return deny(result); + } + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember); + }; + return match member_models::decide(member)? { + member_models::Decision::Allow => Ok(Decision::Allow), + member_models::Decision::Deny { + denial, + member_models, + } => Ok(Decision::Deny { + denial, + member_models, + }), + _ => Err(Error::InvalidAnswer), + }; + } + if !facts.user_loaded { + return Ok(Decision::ReadUser); + } + let precheck = model_access::decide(model_access::Request::UserPrecheck { + present: facts.user_present, + grants: facts.user_grants.clone(), + })?; + let model_access::Response::Decision(precheck) = precheck else { + return Err(Error::InvalidAnswer); + }; + if !matches!(precheck, ModelDecision::ReadScope) { + return if precheck.allowed() { + Ok(Decision::Allow) + } else { + deny(precheck) + }; + } + let Some(user) = &facts.user else { + return Ok(Decision::ReadUserScope); + }; + let result = model_access::check(user)?; + if result.allowed() { + Ok(Decision::Allow) + } else { + deny(result) + } +} + +#[derive(Clone, Default, Deserialize)] +pub struct CeilingSource { + pub agent_present: bool, + pub managed_loaded: bool, + pub managed_present: bool, + pub loaded: bool, + pub count: usize, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum SourceDecision { + ReadManaged, + ReadManagedCeilings, + ReadUnmanagedCeiling, + Complete { indices: Vec }, +} + +pub fn ceiling_source(facts: &CeilingSource) -> SourceDecision { + if !facts.agent_present { + return SourceDecision::Complete { indices: vec![] }; + } + if !facts.managed_loaded { + return SourceDecision::ReadManaged; + } + if facts.loaded { + return SourceDecision::Complete { + indices: (0..facts.count).collect(), + }; + } + if facts.managed_present { + SourceDecision::ReadManagedCeilings + } else { + SourceDecision::ReadUnmanagedCeiling + } +} + +#[derive(Clone, Deserialize)] +pub struct GrantFacts { + pub models: Option>>, + pub ceilings: Vec>, +} + +pub fn grants(facts: &GrantFacts) -> std::collections::BTreeSet { + facts + .models + .iter() + .flatten() + .flatten() + .chain(facts.ceilings.iter().flatten()) + .cloned() + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::main(Purpose::All, Decision::ReadManaged)] + #[case::groups(Purpose::Ceilings, Decision::Allow)] + #[case::caller(Purpose::Caller, Decision::Allow)] + fn empty_model_read_order(#[case] purpose: Purpose, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + purpose, + ..Facts::default() + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/member_models.rs b/litellm-rust/crates/request-auth-core/src/member_models.rs new file mode 100644 index 00000000000..a4cacc8959b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/member_models.rs @@ -0,0 +1,77 @@ +use crate::{Error, model_access}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub team_present: bool, + pub membership_loaded: bool, + pub grants: Option>, + pub scope: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: model_access::Decision, + member_models: Option>, + }, + ReadMembership, + ReadGrants, + ReadScope, +} + +pub fn decide(facts: &Facts) -> Result { + if !facts.user_present || !facts.team_present { + return Ok(Decision::Allow); + } + if !facts.membership_loaded { + return Ok(Decision::ReadMembership); + } + let Some(grants) = &facts.grants else { + return Ok(Decision::ReadGrants); + }; + if grants.is_empty() { + return Ok(Decision::Allow); + } + let Some(scope) = &facts.scope else { + return Ok(Decision::ReadScope); + }; + let denial = model_access::check(scope)?; + if denial.allowed() { + return Ok(Decision::Allow); + } + let member_models = + matches!(denial, model_access::Decision::Deny { .. }).then(|| grants.clone()); + Ok(Decision::Deny { + denial, + member_models, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(false, false, Decision::Allow)] + #[case::lookup(true, false, Decision::ReadMembership)] + #[case::preloaded(true, true, Decision::ReadGrants)] + fn membership_read( + #[case] user_present: bool, + #[case] membership_loaded: bool, + #[case] expected: Decision, + ) { + let facts = Facts { + user_present, + team_present: true, + membership_loaded, + grants: None, + scope: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_access.rs b/litellm-rust/crates/request-auth-core/src/model_access.rs new file mode 100644 index 00000000000..8c9e7cb1b87 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_access.rs @@ -0,0 +1,726 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use serde::{Deserialize, Serialize}; + +use crate::{Error, patterns}; + +#[derive(Clone, Default, Deserialize)] +pub struct Provider { + pub stripped_model: String, + pub provider: String, +} + +#[derive(Clone, Default, Deserialize)] +pub struct ModelRecord { + pub access_groups: Vec, + pub provider: Option, + pub deployment_team: Option, + pub deployment_public_name: Option, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Registry { + pub present: bool, + pub names: Vec, + pub aliases: BTreeMap, + pub records: BTreeMap, + pub compaction_parents: BTreeMap, + pub helper_present: Option, + pub evaluations: Option>, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Aliases { + pub team: BTreeMap, + pub team_order: Option>, + pub key: BTreeMap, + #[serde(rename = "global_")] + pub global: BTreeMap, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub grants: Vec, + pub team_id: Option, + pub aliases: Aliases, + pub registry: Registry, + pub object_type: ObjectType, + pub fallback_depth: usize, + pub max_depth: usize, + pub is_list: bool, +} + +#[derive(Clone, Copy, Deserialize, Serialize, Debug, PartialEq, Eq, strum::VariantArray)] +#[serde(rename_all = "snake_case")] +pub enum ObjectType { + User, + Customer, + Team, + Key, + Org, + Project, + Agent, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + model: String, + grants: Vec, + object_type: ObjectType, + }, + Depth { + model: Option, + }, + NoDefaultModels, + NoAgentModels, + EmptyAgentCeiling, + PersonalTeam { + model: String, + }, + ReadGroups, + ReadScope, + ReadGroupScope { + grants: Vec, + }, +} + +impl Decision { + pub fn allowed(&self) -> bool { + matches!(self, Self::Allow) + } +} + +fn unique(values: impl Iterator) -> Vec { + let mut seen = BTreeSet::new(); + values.filter(|value| seen.insert(value.clone())).collect() +} + +pub fn resolve_key_grants( + grants: &[String], + team_id: Option<&str>, + team_grants: &[String], +) -> Vec { + if grants.iter().any(|grant| grant == "all-team-models") { + return if team_id.is_none() { + Vec::new() + } else { + team_grants.to_vec() + }; + } + grants.to_vec() +} + +fn expanded_grants(facts: &Facts) -> Vec { + if !facts.grants.iter().any(|grant| grant == "all-team-models") + || facts.team_id.is_none() + || !facts.registry.present + { + return facts.grants.clone(); + } + let others = facts + .grants + .iter() + .filter(|grant| grant.as_str() != "all-team-models") + .cloned(); + if facts.registry.names.is_empty() { + let values: Vec<_> = others.collect(); + return if values.is_empty() { + facts.grants.clone() + } else { + values + }; + } + unique(others.chain(facts.registry.names.iter().cloned())) +} + +fn helper(model: &str, facts: &Facts, aliases: bool) -> Result { + let grants = expanded_grants(facts); + let record = facts.registry.records.get(model); + let groups = record.map_or(&[][..], |record| record.access_groups.as_slice()); + if grants.iter().any(|grant| groups.contains(grant)) { + return Ok(true); + } + let filtered: Vec<_> = grants + .iter() + .filter(|grant| !groups.contains(grant)) + .collect(); + if aliases && facts.aliases.team.contains_key(model) { + return Ok(true); + } + for grant in &filtered { + if patterns::model_pattern(model, grant)? { + return Ok(true); + } + } + if let Some(provider) = record.and_then(|record| record.provider.as_ref()) + && (provider.stripped_model != model || !model.contains('/')) + { + for grant in &filtered { + if patterns::model_pattern( + &format!("{}/{}", provider.provider, provider.stripped_model), + grant, + )? { + return Ok(true); + } + } + } + Ok((filtered.is_empty() && grants.is_empty()) + || filtered.iter().any(|grant| { + matches!(grant.as_str(), "*" | "all-proxy-models") || grant.as_str() == model + })) +} + +pub fn candidate_names( + model: &str, + aliases: &Aliases, + registry: &Registry, + team_id: Option<&str>, +) -> (Vec, Vec) { + let after_team = aliases.team.get(model).map_or(model, String::as_str); + let after_key = aliases + .key + .get(after_team) + .map_or(after_team, String::as_str); + let after_global = aliases + .global + .get(after_key) + .map_or(after_key, String::as_str); + let dispatched = aliases + .key + .get(after_global) + .map_or(after_global, String::as_str); + if after_key != after_team || dispatched != after_global { + return (vec![dispatched.to_owned()], vec![]); + } + let candidates = std::iter::once(model.to_owned()) + .chain(registry.compaction_parents.get(model).cloned()) + .chain( + aliases + .global + .get(model) + .or_else(|| registry.aliases.get(model)) + .filter(|target| !target.is_empty()) + .cloned(), + ) + .collect(); + let Some(team) = team_id else { + return (candidates, vec![]); + }; + if !registry.present || !model.starts_with(&format!("model_name_{team}_")) { + return (candidates, vec![]); + } + let Some(record) = registry + .records + .get(model) + .filter(|record| record.deployment_team.as_deref() == Some(team)) + else { + return (candidates, vec![]); + }; + let public = record + .deployment_public_name + .clone() + .filter(|name| !name.is_empty()); + let ordered: Vec<&String> = aliases.team_order.as_ref().map_or_else( + || aliases.team.keys().collect(), + |names| names.iter().collect(), + ); + let legacy = ordered + .into_iter() + .filter(|alias| { + aliases + .team + .get(*alias) + .is_some_and(|target| target == model) + }) + .cloned(); + (candidates, public.into_iter().chain(legacy).collect()) +} + +pub fn live_team_target(model: &str, facts: &Facts) -> String { + let Some(target) = facts.aliases.team.get(model) else { + return model.to_owned(); + }; + let team = facts.team_id.as_deref().unwrap_or("None"); + if facts.registry.present + && target.starts_with(&format!("model_name_{team}_")) + && !facts.registry.names.contains(target) + { + return model.to_owned(); + } + target.clone() +} + +pub fn check(facts: &Facts) -> Result { + check_at(facts, 0) +} + +fn check_at(facts: &Facts, skip: usize) -> Result { + if let Some(evaluations) = &facts.registry.evaluations { + let Some(evaluation) = evaluations + .iter() + .skip(skip) + .find(|evaluation| evaluation.grants == facts.grants) + else { + return Ok(Decision::ReadGroupScope { + grants: facts.grants.clone(), + }); + }; + let scoped = Facts { + aliases: evaluation + .aliases + .clone() + .unwrap_or_else(|| facts.aliases.clone()), + ..facts.clone() + }; + return match crate::model_projection::decide(&scoped, &evaluation.reads)? { + crate::model_projection::Decision::Complete { decision } => Ok(decision), + _ => Err(Error::InvalidAnswer), + }; + } + if facts.fallback_depth >= facts.max_depth { + return Ok(Decision::Depth { model: None }); + } + for model in &facts.requested { + if facts.is_list && facts.fallback_depth + 1 >= facts.max_depth { + return Ok(Decision::Depth { + model: Some(model.clone()), + }); + } + let (candidates, owning_names) = candidate_names( + model, + &facts.aliases, + &facts.registry, + facts.team_id.as_deref(), + ); + let mut allowed = false; + for candidate in candidates { + if helper(&candidate, facts, true)? { + allowed = true; + break; + } + } + if !allowed { + for candidate in owning_names { + if helper(&candidate, facts, false)? { + allowed = true; + break; + } + } + } + if !allowed { + return Ok(Decision::Deny { + model: model.clone(), + grants: facts.grants.clone(), + object_type: facts.object_type, + }); + } + } + Ok(Decision::Allow) +} + +#[derive(Deserialize)] +pub enum Request { + Object(Facts), + Grouped { + facts: Facts, + group_ids: Vec, + group_models: Option>, + merge: bool, + }, + Helper { + model: String, + facts: Facts, + }, + User { + facts: Facts, + present: bool, + }, + UserPrecheck { + present: bool, + grants: Vec, + }, + Customer(Facts), + Agent { + facts: Facts, + managed_grants: Option>, + ceilings: Vec>, + }, + Personal { + requested: Vec, + teams: Vec, + }, + KeyGrants { + grants: Vec, + team_id: Option, + team_grants: Vec, + }, + Candidates(Facts), + PersonalApplies { + enabled: bool, + token_present: bool, + team_id: Option, + user_id: Option, + master: bool, + role: Option, + api_key_present: bool, + inheritance: bool, + }, + AssignedGroups { + groups: Vec, + team_id: Option, + key_id: Option, + }, +} + +#[derive(Deserialize)] +pub struct AssignedGroup { + pub assigned_team_ids: Vec, + pub assigned_key_ids: Vec, + pub models: Vec, +} + +#[derive(Deserialize)] +pub struct PersonalTeam { + pub blocked: bool, + pub available: bool, + pub team: Facts, + pub member: Option, +} + +#[derive(Serialize)] +#[serde(untagged)] +pub enum Response { + Decision(Decision), + Models(Vec), + Predicate(bool), +} + +pub fn decide(request: Request) -> Result { + let decision = match request { + Request::Object(facts) => check(&facts)?, + Request::Grouped { + facts, + group_ids, + group_models, + merge, + } => { + let primary = check(&facts)?; + if primary.allowed() + || !matches!(primary, Decision::Deny { .. }) + || group_ids.is_empty() + { + primary + } else if let Some(models) = group_models { + if models.is_empty() { + primary + } else { + check_at( + &Facts { + grants: if merge { + unique(facts.grants.iter().cloned().chain(models)) + } else { + models + }, + ..facts + }, + 1, + )? + } + } else { + Decision::ReadGroups + } + } + Request::Helper { model, facts } => { + return Ok(Response::Predicate(helper(&model, &facts, true)?)); + } + Request::User { facts, present } => { + if !present { + Decision::Allow + } else if facts + .grants + .iter() + .any(|grant| grant == "no-default-models") + { + Decision::NoDefaultModels + } else { + check(&facts)? + } + } + Request::UserPrecheck { present, grants } => { + if !present { + Decision::Allow + } else if grants.iter().any(|grant| grant == "no-default-models") { + Decision::NoDefaultModels + } else { + Decision::ReadScope + } + } + Request::Customer(facts) => { + if facts.registry.evaluations.is_some() { + return Ok(Response::Decision(check(&facts)?)); + } + for name in &facts.requested { + let target = live_team_target(name, &facts); + if target != *name && facts.grants.contains(name) { + continue; + } + let customer = Facts { + requested: vec![target], + is_list: false, + team_id: None, + aliases: Aliases { + key: facts.aliases.key.clone(), + global: facts.aliases.global.clone(), + ..Default::default() + }, + ..facts.clone() + }; + let result = check(&customer)?; + if !result.allowed() { + return Ok(Response::Decision(result)); + } + } + Decision::Allow + } + Request::Agent { + facts, + managed_grants, + ceilings, + } => { + if facts.requested.is_empty() + || (facts.requested.len() == 1 && !facts.is_list && facts.requested[0].is_empty()) + { + return Ok(Response::Decision(Decision::Allow)); + } + let dispatched = facts + .requested + .iter() + .map(|name| live_team_target(name, &facts)) + .collect(); + let scope = Facts { + requested: dispatched, + aliases: Aliases { + key: facts.aliases.key.clone(), + global: facts.aliases.global.clone(), + ..Default::default() + }, + object_type: ObjectType::Agent, + ..facts + }; + if let Some(grants) = managed_grants { + if grants.is_empty() { + return Ok(Response::Decision(Decision::NoAgentModels)); + } + let result = check(&Facts { + grants, + ..scope.clone() + })?; + if !result.allowed() { + return Ok(Response::Decision(result)); + } + } + for grants in ceilings { + if grants.is_empty() { + return Ok(Response::Decision(Decision::EmptyAgentCeiling)); + } + let result = check(&Facts { + grants, + ..scope.clone() + })?; + if !result.allowed() { + return Ok(Response::Decision(result)); + } + } + Decision::Allow + } + Request::Personal { requested, teams } => { + for name in requested { + let mut allowed = false; + for team in &teams { + if team.blocked || !team.available { + continue; + } + if !check(&Facts { + requested: vec![name.clone()], + is_list: false, + ..team.team.clone() + })? + .allowed() + { + continue; + } + if let Some(member) = &team.member + && !check(&Facts { + requested: vec![name.clone()], + is_list: false, + ..member.clone() + })? + .allowed() + { + continue; + } + allowed = true; + break; + } + if !allowed { + return Ok(Response::Decision(Decision::PersonalTeam { model: name })); + } + } + Decision::Allow + } + Request::KeyGrants { + grants, + team_id, + team_grants, + } => { + return Ok(Response::Models(resolve_key_grants( + &grants, + team_id.as_deref(), + &team_grants, + ))); + } + Request::Candidates(facts) => { + return Ok(Response::Models(unique(facts.requested.iter().flat_map( + |model| { + let (candidates, owned) = candidate_names( + model, + &facts.aliases, + &facts.registry, + facts.team_id.as_deref(), + ); + candidates.into_iter().chain(owned) + }, + )))); + } + Request::PersonalApplies { + enabled, + token_present, + team_id, + user_id, + master, + role, + api_key_present, + inheritance, + } => { + return Ok(Response::Predicate( + enabled + && token_present + && team_id.is_none() + && user_id.is_some() + && !master + && role.as_deref() != Some("proxy_admin") + && (!inheritance || api_key_present), + )); + } + Request::AssignedGroups { + groups, + team_id, + key_id, + } => { + return Ok(Response::Models(unique( + groups + .into_iter() + .filter(|group| { + team_id.as_ref().is_some_and(|id| { + !id.is_empty() && group.assigned_team_ids.contains(id) + }) || key_id + .as_ref() + .is_some_and(|id| !id.is_empty() && group.assigned_key_ids.contains(id)) + }) + .flat_map(|group| group.models), + ))); + } + }; + Ok(Response::Decision(decision)) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(None, false)] + #[case::foreign(Some("other"), false)] + #[case::owned(Some("team"), true)] + fn ownership_aliases_require_matching_deployment( + #[case] owner: Option<&str>, + #[case] owned: bool, + ) { + let model = "model_name_team_private"; + let aliases = Aliases { + team: BTreeMap::from([ + ("z-last".into(), model.into()), + ("a-first".into(), model.into()), + ]), + team_order: Some(vec!["z-last".into(), "a-first".into()]), + ..Default::default() + }; + let records = owner + .map(|team| { + BTreeMap::from([( + model.into(), + ModelRecord { + deployment_team: Some(team.into()), + deployment_public_name: Some("public".into()), + ..Default::default() + }, + )]) + }) + .unwrap_or_default(); + let registry = Registry { + present: true, + records, + ..Default::default() + }; + let (_, names) = candidate_names(model, &aliases, ®istry, Some("team")); + let expected: Vec = if owned { + vec!["public".into(), "z-last".into(), "a-first".into()] + } else { + vec![] + }; + assert_eq!(names, expected); + } + + #[rstest] + #[case::absent(false, vec!["no-default-models"], Decision::Allow)] + #[case::restricted(true, vec!["no-default-models"], Decision::NoDefaultModels)] + #[case::empty(true, vec![], Decision::ReadScope)] + #[case::grants(true, vec!["model"], Decision::ReadScope)] + fn user_precheck(#[case] present: bool, #[case] grants: Vec<&str>, #[case] expected: Decision) { + let response = decide(Request::UserPrecheck { + present, + grants: grants.into_iter().map(str::to_owned).collect(), + }) + .unwrap(); + let Response::Decision(decision) = response else { + panic!("expected decision") + }; + assert_eq!(decision, expected); + } + + #[rstest] + #[case::empty(vec![], true)] + #[case::wildcard(vec!["*"], true)] + #[case::proxy(vec!["all-proxy-models"], true)] + #[case::other(vec!["other"], false)] + #[case::exact(vec!["model"], true)] + fn model_grants(#[case] grants: Vec<&str>, #[case] expected: bool) { + let facts = Facts { + requested: vec!["model".into()], + grants: grants.into_iter().map(str::to_owned).collect(), + team_id: None, + aliases: Aliases::default(), + registry: Registry::default(), + object_type: ObjectType::Key, + fallback_depth: 0, + max_depth: 10, + is_list: false, + }; + assert_eq!(check(&facts).unwrap().allowed(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_attribution.rs b/litellm-rust/crates/request-auth-core/src/model_attribution.rs new file mode 100644 index 00000000000..7cea0a344b0 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_attribution.rs @@ -0,0 +1,147 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub model_present: bool, + pub token_present: bool, + pub router_present: bool, + pub database_present: bool, + pub team_present: bool, + pub user_present: bool, + pub membership_loaded: bool, + pub token_org: Option, + pub team_org: Option, + pub registry_loaded: bool, + pub registry: Option>, + pub covering: Option>, + pub grants: Vec>, + pub org_failed: bool, + pub io_failed: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { groups: Vec, stamp: bool }, + ReadRegistry, + ReadCovering, + ReadKey, + ReadTeam, + ReadMember { lookup: bool, include: bool }, + ReadProject, + ReadOrg { id: Option }, +} + +pub fn decide(facts: &Facts) -> Decision { + let empty = || Decision::Complete { + groups: vec![], + stamp: false, + }; + if facts.io_failed + || !facts.model_present + || !facts.token_present + || !facts.router_present + || !facts.database_present + { + return empty(); + } + if !facts.registry_loaded { + return Decision::ReadRegistry; + } + if facts.registry.as_ref().is_some_and(Vec::is_empty) { + return empty(); + } + let Some(covering) = &facts.covering else { + return Decision::ReadCovering; + }; + let budgeted: BTreeSet<_> = covering + .iter() + .filter(|group| { + facts + .registry + .as_ref() + .is_none_or(|registry| registry.contains(group)) + }) + .collect(); + if budgeted.is_empty() { + return empty(); + } + match facts.grants.len() { + 0 => return Decision::ReadKey, + 1 => return Decision::ReadTeam, + 2 => { + return Decision::ReadMember { + lookup: facts.team_present && facts.user_present && !facts.membership_loaded, + include: facts.team_present && facts.user_present, + }; + } + 3 => return Decision::ReadProject, + 4 => { + return Decision::ReadOrg { + id: facts + .token_org + .as_ref() + .filter(|id| !id.is_empty()) + .or(facts.team_org.as_ref()) + .cloned(), + }; + } + _ => (), + } + let granted: BTreeSet<_> = facts + .grants + .iter() + .take(if facts.org_failed { 4 } else { 5 }) + .flatten() + .collect(); + let groups: Vec<_> = budgeted + .intersection(&granted) + .map(|group| (*group).clone()) + .collect(); + Decision::Complete { + stamp: !groups.is_empty(), + groups, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::unknown(None, vec!["a", "b"])] + #[case::known(Some(vec!["a".into()]), vec!["a"])] + #[case::empty(Some(vec![]), vec![])] + fn intersects_budgeted_covering_and_granted( + #[case] registry: Option>, + #[case] expected: Vec<&str>, + ) { + let facts = Facts { + model_present: true, + token_present: true, + router_present: true, + database_present: true, + registry_loaded: true, + registry, + covering: Some(vec!["b".into(), "a".into(), "b".into()]), + grants: vec![ + vec!["a".into()], + vec![], + vec!["b".into()], + vec![], + vec!["c".into()], + ], + ..Default::default() + }; + assert_eq!( + decide(&facts), + Decision::Complete { + stamp: !expected.is_empty(), + groups: expected.into_iter().map(str::to_owned).collect() + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_budget_aliases.rs b/litellm-rust/crates/request-auth-core/src/model_budget_aliases.rs new file mode 100644 index 00000000000..dd1884f62b8 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_budget_aliases.rs @@ -0,0 +1,250 @@ +use serde::{Deserialize, Serialize}; + +use crate::model_budget_data::{Snapshot, check}; + +#[derive(Clone, Debug, Deserialize)] +pub struct Row { + pub key: String, + pub names: Vec, + pub prefix_known: Option, + pub provider_read: bool, + pub provider: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Facts { + pub model: String, + pub rows: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Read { + Entry, + Prefix, + Provider, +} + +#[derive(Clone, Debug, Serialize)] +pub struct Decision { + pub read: Option, + pub key: String, + pub models: Vec, +} + +fn reading(read: Read, key: &str) -> Decision { + Decision { + read: Some(read), + key: key.to_owned(), + models: vec![], + } +} + +fn row<'a>(facts: &'a Facts, key: &str) -> Result<&'a Row, Decision> { + facts + .rows + .iter() + .find(|row| row.key == key) + .ok_or_else(|| reading(Read::Entry, key)) +} + +fn first(row: &Row) -> Option<&str> { + row.names + .first() + .filter(|name| **name != row.key) + .map(String::as_str) +} + +fn provider(row: &Row) -> Result, Decision> { + if let Some((prefix, _)) = row.key.split_once('/') { + match row.prefix_known { + Some(true) => return Ok(Some(prefix)), + Some(false) => (), + None => return Err(reading(Read::Prefix, &row.key)), + } + } + if !row.provider_read { + return Err(reading(Read::Provider, &row.key)); + } + Ok(row.provider.as_deref()) +} + +fn target(facts: &Facts, key: &str, requested_provider: &str) -> Result, Decision> { + let row = row(facts, key)?; + let Some(name) = first(row) else { + return Ok(None); + }; + if provider(row)? != Some(requested_provider) { + return Ok(None); + } + Ok(Some( + name.strip_prefix(&format!("{requested_provider}/")) + .unwrap_or(name) + .to_owned(), + )) +} + +fn models(facts: &Facts) -> Result, Decision> { + let Some((provider, name)) = facts.model.split_once('/') else { + return Ok(match first(row(facts, &facts.model)?) { + Some(target) => vec![facts.model.clone(), target.to_owned()], + None => vec![facts.model.clone()], + }); + }; + let direct = target(facts, &facts.model, provider)?.filter(|name| !name.is_empty()); + let resolved = match direct { + Some(target) => Some(target), + None => target(facts, name, provider)?.filter(|name| !name.is_empty()), + }; + Ok(match resolved { + Some(target) if target != name => vec![facts.model.clone(), format!("{provider}/{target}")], + _ => vec![facts.model.clone()], + }) +} + +pub fn decide(facts: &Facts) -> Decision { + match models(facts) { + Ok(models) => Decision { + read: None, + key: facts.model.clone(), + models, + }, + Err(read) => read, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ScopeFacts { + pub models: Vec, + pub snapshots: Vec, + pub budget_present: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct ScopeDecision { + pub read_index: Option, + pub result_index: Option, +} + +pub fn scope(facts: &ScopeFacts) -> ScopeDecision { + if let Some(index) = facts.snapshots.iter().position(|value| check(value).denied) { + return ScopeDecision { + read_index: None, + result_index: Some(index), + }; + } + if !facts.budget_present { + return ScopeDecision { + read_index: None, + result_index: None, + }; + } + if facts.snapshots.len() < facts.models.len() { + return ScopeDecision { + read_index: Some(facts.snapshots.len()), + result_index: None, + }; + } + ScopeDecision { + read_index: None, + result_index: facts.snapshots.len().checked_sub(1), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + #[fixture] + fn alias() -> Row { + Row { + key: "provider/latest".into(), + names: vec!["provider/target".into()], + prefix_known: Some(true), + provider_read: false, + provider: None, + } + } + + #[rstest] + #[case::catalog_first(vec![], Some(Read::Entry), "provider/latest", vec![])] + #[case::prefix_before_metadata(vec![Row { prefix_known: None, ..alias() }], Some(Read::Prefix), "provider/latest", vec![])] + #[case::metadata_for_unknown_provider(vec![Row { prefix_known: Some(false), ..alias() }], Some(Read::Provider), "provider/latest", vec![])] + #[case::known_prefix_ignores_metadata(vec![alias()], None, "provider/latest", vec!["provider/latest", "provider/target"])] + #[case::absent_qualified_reads_bare(vec![Row { names: vec![], ..alias() }], Some(Read::Entry), "latest", vec![])] + #[case::matching_metadata(vec![Row { prefix_known: Some(false), provider_read: true, provider: Some("provider".into()), ..alias() }], None, "provider/latest", vec!["provider/latest", "provider/target"])] + #[case::wrong_provider_reads_bare(vec![Row { prefix_known: Some(false), provider_read: true, provider: Some("other".into()), ..alias() }], Some(Read::Entry), "latest", vec![])] + fn lazy_alias_reads( + #[case] rows: Vec, + #[case] read: Option, + #[case] key: &str, + #[case] models: Vec<&str>, + ) { + let result = decide(&Facts { + model: "provider/latest".into(), + rows, + }); + assert_eq!( + (result.read, result.key.as_str(), result.models), + (read, key, models.into_iter().map(str::to_owned).collect()) + ); + } + + #[rstest] + #[case::bare_alias("latest", vec!["target"], vec!["latest", "target"])] + #[case::self("latest", vec!["latest", "other"], vec!["latest"])] + #[case::empty_row("latest", vec![], vec!["latest"])] + #[case::empty_bare_target("latest", vec![""], vec!["latest", ""])] + fn bare_alias(#[case] model: &str, #[case] names: Vec<&str>, #[case] expected: Vec<&str>) { + let result = decide(&Facts { + model: model.into(), + rows: vec![Row { + key: model.into(), + names: names.into_iter().map(str::to_owned).collect(), + ..alias() + }], + }); + assert_eq!(result.read, None); + assert_eq!(result.models, expected); + } + + fn snapshot(limit: f64, spend: f64) -> Snapshot { + Snapshot { + max_budget: Some(limit), + local_spend: Some(spend), + legacy_spend: None, + shared_spend: None, + shared_started_at: None, + shared_window_seconds: None, + now: 0.0, + } + } + + #[rstest] + #[case::first(vec![], true, Some(0), None)] + #[case::target_after_allowed(vec![snapshot(1.0, 0.0)], true, Some(1), None)] + #[case::alias_denial(vec![snapshot(1.0, 1.0)], true, None, Some(0))] + #[case::target_denial(vec![snapshot(1.0, 0.0), snapshot(2.0, 3.0)], true, None, Some(1))] + #[case::all_allowed(vec![snapshot(1.0, 0.0), snapshot(2.0, 0.0)], true, None, Some(1))] + #[case::no_budgets(vec![], false, None, None)] + #[case::retain_denial(vec![snapshot(1.0, 1.0)], false, None, Some(0))] + #[case::nan_continues(vec![snapshot(f64::NAN, 1.0)], true, Some(1), None)] + fn scope_order( + #[case] snapshots: Vec, + #[case] budget_present: bool, + #[case] read_index: Option, + #[case] result_index: Option, + ) { + assert_eq!( + scope(&ScopeFacts { + models: vec!["alias".into(), "target".into()], + snapshots, + budget_present + }), + ScopeDecision { + read_index, + result_index + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_budget_data.rs b/litellm-rust/crates/request-auth-core/src/model_budget_data.rs new file mode 100644 index 00000000000..17e7fe3b0b6 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_budget_data.rs @@ -0,0 +1,331 @@ +use std::cmp::Ordering::{Equal, Greater, Less}; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Entry { + pub name: String, + pub max_budget: Option, + pub duration: Option, + pub duration_seconds: Option, +} + +impl Entry { + fn usable(&self) -> bool { + self.duration.is_some() && self.duration_seconds.is_some() + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ResolveFacts { + pub model: String, + pub bedrock_base: String, + pub bedrock_provider: String, + pub entries: Vec, + pub key_entries: Vec, + pub team: bool, + pub access_groups: Option>, + #[serde(default)] + pub excluded_entries: Vec, +} + +#[derive(Clone, Debug, Serialize)] +pub struct Resolution { + pub read_groups: bool, + pub entry: Option, + pub named_model: bool, +} + +fn candidates(facts: &ResolveFacts) -> Vec<&str> { + let base_names = if facts.bedrock_provider.starts_with("bedrock") { + vec![ + facts.bedrock_base.as_str(), + facts.bedrock_base.split_once('.').map_or("", |(_, s)| s), + ] + } else { + vec![] + }; + std::iter::once(facts.model.as_str()) + .chain(std::iter::once( + facts.model.rsplit('/').next().unwrap_or(&facts.model), + )) + .chain(base_names.into_iter().filter(|s| !s.is_empty())) + .collect() +} + +fn exact<'a>(entries: &'a [Entry], names: &[&str]) -> Option<&'a Entry> { + names.iter().find_map(|name| { + entries + .iter() + .find(|entry| entry.name == *name && entry.usable()) + }) +} + +pub fn resolve(facts: &ResolveFacts) -> Resolution { + if !facts.excluded_entries.is_empty() { + return resolve(&ResolveFacts { + entries: facts + .entries + .iter() + .filter(|entry| !facts.excluded_entries.contains(&entry.name)) + .cloned() + .collect(), + excluded_entries: vec![], + ..facts.clone() + }); + } + let names = candidates(facts); + let key_replaces = facts.team + && exact(&facts.key_entries, &names) + .is_some_and(|entry| entry.max_budget.is_some_and(|limit| limit >= 0.0)); + if !key_replaces && let Some(entry) = exact(&facts.entries, &names) { + return Resolution { + read_groups: false, + entry: Some(entry.clone()), + named_model: true, + }; + } + let group_entries: Vec<&Entry> = facts + .entries + .iter() + .filter(|entry| !entry.name.contains('*') && !names.contains(&entry.name.as_str())) + .collect(); + if !group_entries.is_empty() && facts.access_groups.is_none() { + return Resolution { + read_groups: true, + entry: None, + named_model: false, + }; + } + let group = group_entries + .into_iter() + .filter(|entry| { + entry.usable() + && facts + .access_groups + .as_ref() + .is_some_and(|groups| groups.contains(&entry.name)) + }) + .min_by(|a, b| a.name.cmp(&b.name)); + if let Some(entry) = group { + return Resolution { + read_groups: false, + entry: Some(entry.clone()), + named_model: false, + }; + } + let entry = facts + .entries + .iter() + .filter(|entry| { + entry.usable() + && !(key_replaces && names.contains(&entry.name.as_str())) + && entry.name.contains('*') + && names.iter().any(|name| glob(&entry.name, name)) + }) + .max_by_key(|entry| { + let wildcards = entry.name.chars().filter(|c| *c == '*').count(); + ( + entry.name.chars().count() - wildcards, + std::cmp::Reverse(wildcards), + &entry.name, + ) + }) + .cloned(); + Resolution { + read_groups: false, + entry, + named_model: false, + } +} + +pub fn team_exclusions(facts: &[ResolveFacts]) -> Vec { + facts + .iter() + .flat_map(|facts| { + let names = candidates(facts); + if exact(&facts.key_entries, &names) + .is_some_and(|entry| entry.max_budget.is_some_and(|limit| limit >= 0.0)) + { + names.into_iter().map(str::to_owned).collect::>() + } else { + vec![] + } + }) + .collect::>() + .into_iter() + .collect() +} + +fn glob(pattern: &str, name: &str) -> bool { + let parts: Vec<&str> = pattern.split('*').collect(); + if parts.len() == 1 { + return pattern == name; + } + let head = parts[0]; + let tail = parts[parts.len() - 1]; + if !name.starts_with(head) || !name.ends_with(tail) || head.len() + tail.len() > name.len() { + return false; + } + parts[1..parts.len() - 1] + .iter() + .try_fold(head.len(), |position, piece| { + name[position..name.len() - tail.len()] + .find(piece) + .map(|index| position + index + piece.len()) + }) + .is_some() +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Entity { + Key, + User, + Team, + TeamMember, + EndUser, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Snapshot { + pub max_budget: Option, + pub local_spend: Option, + pub legacy_spend: Option, + pub shared_spend: Option, + pub shared_started_at: Option, + pub shared_window_seconds: Option, + pub now: f64, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct QueryFacts { + pub entity: Entity, + pub snapshot: Snapshot, + pub duration_seconds: Option, + pub process_uptime: f64, + pub same_counter: bool, + pub redis_present: bool, + pub shared_read: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Read { + Local, + Legacy, + Shared, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct Decision { + pub read: Option, + pub denied: bool, + pub current_spend: f64, +} + +pub fn check(snapshot: &Snapshot) -> Decision { + let local = snapshot.local_spend.unwrap_or(0.0) + snapshot.legacy_spend.unwrap_or(0.0); + let current_spend = match ( + snapshot.shared_spend, + snapshot.shared_started_at, + snapshot.shared_window_seconds, + ) { + (Some(spend), Some(start), Some(seconds)) + if snapshot.now < start + seconds && spend > local => + { + spend + } + _ => local, + }; + Decision { + read: None, + denied: snapshot + .max_budget + .is_some_and(|limit| limit.partial_cmp(&0.0) != Some(Less) && current_spend >= limit), + current_spend, + } +} + +pub fn query(facts: &QueryFacts) -> Decision { + let done = check(&facts.snapshot); + let Some(limit) = facts.snapshot.max_budget else { + return done; + }; + if limit < 0.0 { + return done; + } + let read = if facts.snapshot.local_spend.is_none() { + Some(Read::Local) + } else if matches!(facts.entity, Entity::Key | Entity::EndUser) + && facts + .duration_seconds + .is_some_and(|seconds| facts.process_uptime < seconds) + && !facts.same_counter + && facts.snapshot.legacy_spend.is_none() + { + Some(Read::Legacy) + } else if !matches!( + (facts.snapshot.local_spend.unwrap_or(0.0) + facts.snapshot.legacy_spend.unwrap_or(0.0)) + .partial_cmp(&limit), + Some(Greater | Equal) + ) && !facts.redis_present + && facts.entity != Entity::EndUser + && !facts.shared_read + { + Some(Read::Shared) + } else { + None + }; + Decision { read, ..done } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::empty("*", "", true)] + #[case::unicode("é*界", "é世世界", true)] + #[case::overlap("ab*bc", "abc", false)] + #[case::middle("a*b*c", "axybzc", true)] + #[case::not_regex("a?*", "ab", false)] + fn wildcard(#[case] pattern: &str, #[case] name: &str, #[case] matches: bool) { + assert_eq!(glob(pattern, name), matches); + } + + #[rstest] + #[case::alias_zero("alias", Some(0.0), vec!["alias"])] + #[case::target_positive("target", Some(1.0), vec!["target"])] + #[case::target_infinite("target", Some(f64::INFINITY), vec!["target"])] + #[case::negative("target", Some(-1.0), vec![])] + #[case::nan("target", Some(f64::NAN), vec![])] + #[case::missing("alias", None, vec![])] + fn team_alias_overrides( + #[case] key: &str, + #[case] limit: Option, + #[case] expected: Vec<&str>, + ) { + let entry = Entry { + name: key.into(), + max_budget: limit, + duration: Some("1h".into()), + duration_seconds: Some(3600.0), + }; + let alias = ResolveFacts { + model: "alias".into(), + bedrock_base: "alias".into(), + bedrock_provider: String::new(), + entries: vec![], + key_entries: vec![entry], + team: true, + access_groups: None, + excluded_entries: vec![], + }; + let target = ResolveFacts { + model: "target".into(), + bedrock_base: "target".into(), + ..alias.clone() + }; + assert_eq!(team_exclusions(&[alias, target]), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_budgets.rs b/litellm-rust/crates/request-auth-core/src/model_budgets.rs new file mode 100644 index 00000000000..709a13c8a43 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_budgets.rs @@ -0,0 +1,312 @@ +use crate::{Error, protocol::*}; + +fn allowed() -> Result { + Step::allow(Grant::WithinModelBudgets) +} + +pub fn check_model_budgets() -> Result { + key_budgets(Box::new(allowed)) +} + +pub fn check_user_model_budgets() -> Result { + user_budgets(Box::new(allowed)) +} + +pub(crate) fn key_budgets(then: Then) -> Result { + Step::effect( + Effect::ReadModelBudgetScope, + move |scope: ModelBudgetScope| { + if !scope.key_has_model_budgets || scope.models.is_empty() { + return user_models(scope.user_has_model_budgets, scope.models, then); + } + key_models( + scope.key_budget_fallbacks, + scope.models, + scope.rewritable_model, + Box::new(move || user_budgets(then)), + ) + }, + ) +} + +fn user_budgets(then: Then) -> Result { + Step::effect( + Effect::ReadModelBudgetScope, + move |scope: ModelBudgetScope| { + user_models(scope.user_has_model_budgets, scope.models, then) + }, + ) +} + +fn key_models( + fallbacks: std::collections::BTreeMap>, + models: Vec, + rewritable: Option, + then: Then, +) -> Result { + let Some((model, rest)) = models.split_first() else { + return then(); + }; + let model = model.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::CheckKeyModelBudget { + model: model.clone(), + }, + move |reading: ModelBudgetRead| { + let exceeded = reading.exceeded(); + let Some(exceeded) = exceeded else { + return key_models(fallbacks, rest, rewritable, then); + }; + if Some(&model) != rewritable.as_ref() { + return Step::deny(exceeded); + } + let candidates = fallbacks.get(&model).cloned().unwrap_or_default(); + first_affordable( + candidates, + exceeded, + Box::new(move |fallback| key_models(fallbacks, rest, Some(fallback), then)), + ) + }, + ) +} + +type Found = Box Result + Send + Sync>; + +fn first_affordable( + candidates: Vec, + exceeded: ModelBudgetExceeded, + found: Found, +) -> Result { + let Some((candidate, rest)) = candidates.split_first() else { + return Step::deny(exceeded); + }; + let candidate = candidate.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::CheckKeyModelBudget { + model: candidate.clone(), + }, + move |reading: ModelBudgetRead| { + let over = reading.exceeded(); + if over.is_some() { + return first_affordable(rest, exceeded, found); + } + Step::effect( + Effect::AuthorizeBudgetFallback { + model: candidate.clone(), + }, + move |facts: crate::fallback_access::Facts| { + match crate::fallback_access::decide(facts)? { + crate::fallback_access::Decision::Allow => (), + crate::fallback_access::Decision::Deny { + denial: Some(denial @ crate::model_access::Decision::Depth { .. }), + } => return Step::deny(Denial::ModelAccessDenied { decision: denial }), + crate::fallback_access::Decision::Deny { .. } => { + return Step::deny(exceeded); + } + _ => return Err(Error::InvalidAnswer), + } + Step::effect( + Effect::RewriteRequestModel { + model: candidate.clone(), + }, + move |(): ()| found(candidate), + ) + }, + ) + }, + ) +} + +fn user_models(enabled: bool, models: Vec, then: Then) -> Result { + if !enabled { + return then(); + } + let Some((model, rest)) = models.split_first() else { + return then(); + }; + let rest = rest.to_vec(); + Step::effect( + Effect::CheckUserModelBudget { + model: model.clone(), + }, + move |reading: ModelBudgetRead| match reading.exceeded() { + Some(exceeded) => Step::deny(exceeded), + None => user_models(enabled, rest, then), + }, + ) +} + +pub fn check_team_member_model_budgets() -> Result { + Step::effect( + Effect::ReadTeamMemberModelBudgetScope, + |scope: TeamMemberModelBudgetScope| { + if !scope.has_model_budgets { + return allowed(); + } + member_models(scope.models, scope.rewritable_model, scope.can_fall_back) + }, + ) +} + +fn member_models( + models: Vec, + rewritable: Option, + can_fall_back: bool, +) -> Result { + let Some((model, rest)) = models.split_first() else { + return allowed(); + }; + let model = model.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::CheckTeamMemberModelBudget { + model: model.clone(), + }, + move |reading: ModelBudgetRead| { + let exceeded = reading.exceeded(); + let Some(exceeded) = exceeded else { + return member_models(rest, rewritable, can_fall_back); + }; + if Some(&model) != rewritable.as_ref() || !can_fall_back { + return Step::deny(exceeded); + } + Step::effect( + Effect::ReadTeamMemberFallbacks { + model: model.clone(), + }, + move |targets: Vec| { + first_usable( + targets + .into_iter() + .filter(|target| target.model != model) + .collect(), + exceeded, + Box::new(move |fallback| { + member_models(rest, Some(fallback), can_fall_back) + }), + ) + }, + ) + }, + ) +} + +fn first_usable( + targets: Vec, + exceeded: ModelBudgetExceeded, + found: Found, +) -> Result { + let Some((target, rest)) = targets.split_first() else { + return Step::deny(exceeded); + }; + let target = target.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::VetTeamMemberFallback { + target: target.clone(), + }, + move |facts: crate::fallback_candidate::Facts| { + let crate::fallback_candidate::Decision::Complete { usable, .. } = + crate::fallback_candidate::decide(&facts)? + else { + return Err(Error::InvalidAnswer); + }; + if !usable { + return first_usable(rest, exceeded, found); + } + Step::effect( + Effect::RerouteTeamMemberRequest { + model: target.model.clone(), + }, + move |(): ()| found(target.model), + ) + }, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn unauthorized_affordable_fallback_denies_without_trying_the_next_one() { + let exceeded = ModelBudgetExceeded { + message: "over budget".into(), + current_cost: 2.0, + max_budget: 1.0, + entity_type: None, + entity_id: None, + }; + let read = first_affordable( + vec!["first".into(), "second".into()], + exceeded, + Box::new(|_| allowed()), + ) + .unwrap(); + assert!( + matches!(read.node, Node::Effect(Effect::CheckKeyModelBudget { ref model }) if model == "first") + ); + let authorize = (read.resume.unwrap())(Answer::ModelBudgetData(ModelBudgetRead { + snapshot: crate::model_budget_data::Snapshot { + max_budget: None, + local_spend: None, + legacy_spend: None, + shared_spend: None, + shared_started_at: None, + shared_window_seconds: None, + now: 0.0, + }, + message: String::new(), + entity_type: None, + entity_id: None, + })) + .unwrap(); + assert!( + matches!(authorize.node, Node::Effect(Effect::AuthorizeBudgetFallback { ref model }) if model == "first") + ); + let result = (authorize.resume.unwrap())(Answer::FallbackAccess(Box::new( + crate::fallback_access::Facts { + refuse_agent_keys: true, + agent_id: Some("agent".into()), + key: None, + team: None, + embedded_team: None, + assigned_models: None, + assigned_scope: None, + member: None, + project: None, + team_present: false, + embedded_team_grants: vec![], + member_grants: vec![], + project_grants: vec![], + }, + ))) + .unwrap(); + assert!(matches!( + result.node, + Node::Outcome(Outcome::Deny { + denial: Denial::ModelBudgetExceeded { .. } + }) + )); + assert!(result.resume.is_none()); + } + + #[rstest] + fn user_entry_never_checks_key_budgets() { + let step = check_user_model_budgets().unwrap(); + let result = (step.resume.unwrap())(Answer::ModelScope(ModelBudgetScope { + models: vec!["model".into()], + rewritable_model: Some("model".into()), + key_has_model_budgets: true, + key_budget_fallbacks: Default::default(), + user_has_model_budgets: true, + })) + .unwrap(); + assert!( + matches!(result.node, Node::Effect(Effect::CheckUserModelBudget { ref model }) if model == "model") + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_groups.rs b/litellm-rust/crates/request-auth-core/src/model_groups.rs new file mode 100644 index 00000000000..7853d4b2c5f --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_groups.rs @@ -0,0 +1,73 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +pub struct Group { + pub assigned_team_ids: Vec, + pub assigned_key_ids: Vec, + pub resources: Vec, +} + +#[derive(Deserialize)] +#[serde(tag = "kind", content = "group")] +pub enum Lookup { + Found(Group), + Missing, + Failed, +} + +#[derive(Deserialize)] +pub struct Facts { + pub subject_present: bool, + pub database_present: bool, + pub cache_present: bool, + pub group_ids: Vec, + pub team_id: Option, + pub key_id: Option, + pub check_db_only: bool, + pub strict_errors: bool, + pub loaded: Vec, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Read { index: usize }, + Raise { index: usize }, + Complete { resources: Vec }, +} + +pub fn decide(facts: Facts) -> Decision { + if !facts.subject_present || !facts.database_present || !facts.cache_present { + return Decision::Complete { resources: vec![] }; + } + let mut resources = BTreeSet::new(); + for index in 0..facts.group_ids.len() { + let Some(loaded) = facts.loaded.get(index) else { + return Decision::Read { index }; + }; + match loaded { + Lookup::Missing | Lookup::Failed if facts.check_db_only || facts.strict_errors => { + return Decision::Raise { index }; + } + Lookup::Missing | Lookup::Failed => {} + Lookup::Found(group) => { + if facts + .team_id + .as_ref() + .is_some_and(|id| !id.is_empty() && group.assigned_team_ids.contains(id)) + || facts + .key_id + .as_ref() + .is_some_and(|id| !id.is_empty() && group.assigned_key_ids.contains(id)) + { + resources.extend(group.resources.iter().cloned()); + } + } + } + } + Decision::Complete { + resources: resources.into_iter().collect(), + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_helpers.rs b/litellm-rust/crates/request-auth-core/src/model_helpers.rs new file mode 100644 index 00000000000..9e4ac468a87 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_helpers.rs @@ -0,0 +1,148 @@ +use crate::{Error, model_access::Provider, patterns}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Group, + Helper, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub kind: Kind, + pub model: String, + pub grants: Option>, + pub team_id: Option, + pub team_aliases: Vec, + pub router_present: bool, + pub groups: Option>, + pub names: Option>, + pub provider_loaded: bool, + pub provider: Option, + pub provider_reads: Option>>, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool }, + ReadGroups, + ReadNames, + ReadProvider, +} + +pub fn decide(facts: &Facts) -> Result { + let done = |allowed| Ok(Decision::Complete { allowed }); + if facts.kind == Kind::Group + && facts + .grants + .as_ref() + .is_none_or(|grants| grants.contains(&facts.model)) + { + return done(true); + } + let grants = facts.grants.as_ref().ok_or(Error::InvalidAnswer)?; + if facts.router_present && facts.groups.is_none() { + return Ok(Decision::ReadGroups); + } + let groups = facts.groups.as_deref().unwrap_or_default(); + if facts.kind == Kind::Group { + return done(grants.iter().any(|grant| groups.contains(grant))); + } + let expanded = if facts.router_present + && facts.team_id.is_some() + && grants.iter().any(|grant| grant == "all-team-models") + { + let Some(names) = &facts.names else { + return Ok(Decision::ReadNames); + }; + let mut seen = BTreeSet::new(); + let values: Vec<_> = grants + .iter() + .filter(|grant| *grant != "all-team-models") + .chain(names) + .filter(|grant| seen.insert(*grant)) + .cloned() + .collect(); + if values.is_empty() { + grants.clone() + } else { + values + } + } else { + grants.clone() + }; + if expanded.iter().any(|grant| groups.contains(grant)) + || facts.team_aliases.contains(&facts.model) + { + return done(true); + } + let filtered: Vec<_> = expanded + .iter() + .filter(|grant| !groups.contains(grant)) + .collect(); + let wildcards: Vec<_> = filtered + .iter() + .filter(|grant| grant.contains('*')) + .collect(); + for pattern in &wildcards { + if patterns::model_pattern(&facts.model, pattern)? { + return done(true); + } + } + for (index, pattern) in wildcards.iter().enumerate() { + let provider = if let Some(providers) = &facts.provider_reads { + let Some(provider) = providers.get(index) else { + return Ok(Decision::ReadProvider); + }; + provider + } else { + if !facts.provider_loaded { + return Ok(Decision::ReadProvider); + } + &facts.provider + }; + if let Some(provider) = provider + && (provider.stripped_model != facts.model || !facts.model.contains('/')) + && patterns::model_pattern( + &format!("{}/{}", provider.provider, provider.stripped_model), + pattern, + )? + { + return done(true); + } + } + done( + (filtered.is_empty() && expanded.is_empty()) + || filtered.iter().any(|grant| { + matches!(grant.as_str(), "*" | "all-proxy-models") || **grant == facts.model + }), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::group(Kind::Group, Decision::Complete { allowed: true })] + #[case::helper(Kind::Helper, Decision::ReadGroups)] + fn direct_match_read_order(#[case] kind: Kind, #[case] expected: Decision) { + let facts = Facts { + kind, + model: "model".into(), + grants: Some(vec!["model".into()]), + team_id: None, + team_aliases: vec![], + router_present: true, + groups: None, + names: None, + provider_loaded: false, + provider: None, + provider_reads: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_projection.rs b/litellm-rust/crates/request-auth-core/src/model_projection.rs new file mode 100644 index 00000000000..77a5e286715 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_projection.rs @@ -0,0 +1,299 @@ +use crate::{ + Error, + model_access::{self, Decision as ModelDecision, Facts, ModelRecord, Provider, Registry}, + model_helpers, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum ReadKind { + Compaction, + Aliases, + AliasTarget, + Deployment, + Groups, + Names, + Provider, +} + +#[derive(Clone, Deserialize)] +pub struct Reading { + pub kind: ReadKind, + pub model: String, + pub values: Vec, + pub value: Option, + pub team: Option, + pub provider: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Evaluation { + pub grants: Vec, + pub reads: Vec, + pub aliases: Option, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { + decision: ModelDecision, + }, + Read { + read: ReadKind, + model: String, + team_id: Option, + }, +} + +enum Pending { + Read(ReadKind, String, Option), + Invalid(Error), +} +impl From for Pending { + fn from(error: Error) -> Self { + Self::Invalid(error) + } +} + +struct Reader<'a> { + reads: &'a [Reading], + index: usize, + team_id: Option, +} +impl<'a> Reader<'a> { + fn read(&mut self, kind: ReadKind, model: &str) -> Result<&'a Reading, Pending> { + let Some(value) = self.reads.get(self.index) else { + return Err(Pending::Read(kind, model.to_owned(), self.team_id.clone())); + }; + if value.kind != kind || value.model != model { + return Err(Pending::Invalid(Error::InvalidAnswer)); + } + self.index += 1; + Ok(value) + } +} + +fn helper( + model: &str, + facts: &Facts, + aliases: bool, + reader: &mut Reader<'_>, +) -> Result { + let mut helper = model_helpers::Facts { + kind: model_helpers::Kind::Helper, + model: model.to_owned(), + grants: Some(facts.grants.clone()), + team_id: facts.team_id.clone(), + team_aliases: if aliases { + facts.aliases.team.keys().cloned().collect() + } else { + vec![] + }, + router_present: facts + .registry + .helper_present + .unwrap_or(facts.registry.present), + groups: None, + names: None, + provider_loaded: false, + provider: None, + provider_reads: Some(vec![]), + }; + loop { + match model_helpers::decide(&helper)? { + model_helpers::Decision::Complete { allowed } => return Ok(allowed), + model_helpers::Decision::ReadGroups => { + helper.groups = Some(reader.read(ReadKind::Groups, model)?.values.clone()) + } + model_helpers::Decision::ReadNames => { + helper.names = Some(reader.read(ReadKind::Names, model)?.values.clone()) + } + model_helpers::Decision::ReadProvider => { + helper + .provider_reads + .as_mut() + .ok_or(Error::InvalidAnswer)? + .push(reader.read(ReadKind::Provider, model)?.provider.clone()); + } + } + } +} + +fn check_one(model: &str, facts: &Facts, reader: &mut Reader<'_>) -> Result { + let compaction = reader.read(ReadKind::Compaction, model)?.value.clone(); + let mut registry = Registry { + present: facts.registry.present, + ..Registry::default() + }; + if let Some(parent) = compaction { + registry.compaction_parents.insert(model.to_owned(), parent); + } + if !facts.aliases.global.contains_key(model) && registry.present { + let aliases = reader.read(ReadKind::Aliases, model)?; + if aliases.values.iter().any(|alias| alias == model) + && let Some(target) = &reader.read(ReadKind::AliasTarget, model)?.value + { + registry.aliases.insert(model.to_owned(), target.clone()); + } + } + let after_team = facts.aliases.team.get(model).map_or(model, String::as_str); + let after_key = facts + .aliases + .key + .get(after_team) + .map_or(after_team, String::as_str); + let after_global = facts + .aliases + .global + .get(after_key) + .map_or(after_key, String::as_str); + let dispatched = facts + .aliases + .key + .get(after_global) + .map_or(after_global, String::as_str); + let key_alias = after_key != after_team || dispatched != after_global; + if !key_alias + && registry.present + && facts + .team_id + .as_ref() + .is_some_and(|team| model.starts_with(&format!("model_name_{team}_"))) + { + let deployment = reader.read(ReadKind::Deployment, model)?; + registry.records.insert( + model.to_owned(), + ModelRecord { + deployment_team: deployment.team.clone(), + deployment_public_name: deployment.value.clone(), + ..ModelRecord::default() + }, + ); + } + let (candidates, owning) = + model_access::candidate_names(model, &facts.aliases, ®istry, facts.team_id.as_deref()); + for candidate in candidates { + if helper(&candidate, facts, true, reader)? { + return Ok(true); + } + } + for candidate in owning { + if helper(&candidate, facts, false, reader)? { + return Ok(true); + } + } + Ok(false) +} + +fn check(facts: &Facts, reader: &mut Reader<'_>) -> Result { + if facts.fallback_depth >= facts.max_depth { + return Ok(ModelDecision::Depth { model: None }); + } + for model in &facts.requested { + if facts.is_list && facts.fallback_depth + 1 >= facts.max_depth { + return Ok(ModelDecision::Depth { + model: Some(model.clone()), + }); + } + if !check_one(model, facts, reader)? { + return Ok(ModelDecision::Deny { + model: model.clone(), + grants: facts.grants.clone(), + object_type: facts.object_type, + }); + } + } + Ok(ModelDecision::Allow) +} + +fn customer(facts: &Facts, reader: &mut Reader<'_>) -> Result { + for model in &facts.requested { + let target = facts.aliases.team.get(model).unwrap_or(model); + let target = if facts.registry.present + && target.starts_with(&format!( + "model_name_{}_", + facts.team_id.as_deref().unwrap_or("None") + )) + && facts.aliases.team.contains_key(model) + && !reader.read(ReadKind::Names, model)?.values.contains(target) + { + model + } else { + target + }; + if target != model && facts.grants.contains(model) { + continue; + } + let scope = Facts { + requested: vec![target.clone()], + team_id: None, + is_list: false, + aliases: model_access::Aliases { + key: facts.aliases.key.clone(), + global: facts.aliases.global.clone(), + ..Default::default() + }, + ..facts.clone() + }; + let decision = check(&scope, reader)?; + if !decision.allowed() { + return Ok(decision); + } + } + Ok(ModelDecision::Allow) +} + +pub fn decide(facts: &Facts, reads: &[Reading]) -> Result { + let is_customer = facts.object_type == model_access::ObjectType::Customer; + let mut reader = Reader { + reads, + index: 0, + team_id: if is_customer { + None + } else { + facts.team_id.clone() + }, + }; + let decision = if is_customer { + customer(facts, &mut reader) + } else { + check(facts, &mut reader) + }; + match decision { + Ok(decision) if reader.index == reads.len() => Ok(Decision::Complete { decision }), + Ok(_) => Err(Error::InvalidAnswer), + Err(Pending::Read(read, model, team_id)) => Ok(Decision::Read { + read, + model, + team_id, + }), + Err(Pending::Invalid(error)) => Err(error), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + #[rstest] + #[case::empty(vec![], true)] + #[case::model(vec!["model".to_owned()], false)] + fn empty_list_never_reads_registry(#[case] requested: Vec, #[case] complete: bool) { + let facts = Facts { + requested, + grants: vec![], + team_id: None, + aliases: Default::default(), + registry: Default::default(), + object_type: model_access::ObjectType::Key, + fallback_depth: 0, + max_depth: 5, + is_list: true, + }; + assert_eq!( + matches!(decide(&facts, &[]).unwrap(), Decision::Complete { .. }), + complete + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/org_routes.rs b/litellm-rust/crates/request-auth-core/src/org_routes.rs new file mode 100644 index 00000000000..d361e46a25e --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/org_routes.rs @@ -0,0 +1,304 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Membership { + pub organization_id: Option, + pub role: Option, +} + +#[derive(Deserialize)] +pub struct Facts { + pub route: String, + pub user_present: bool, + pub user_role: Option, + pub memberships: Option>, + pub organization_id: Option, + pub org_admin_routes: Vec, + #[serde(default)] + pub organization_id_present: bool, + #[serde(default)] + pub identifier_pending: bool, + #[serde(default)] + pub hash_failure: Option, + #[serde(default)] + pub request_pending: bool, + #[serde(default)] + pub memberships_pending: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Denial { + CreateOrganization, + NoMemberships, + MissingOrganization, + MissingRole, + WrongRole, + TeamMissingOrganization, + TeamWrongRole, +} + +#[derive(Debug, Default, PartialEq, Eq, Serialize)] +pub struct Decision { + pub denial: Option, + pub organization_role: Option, + pub organization_ids: Vec, + pub read_identifier: bool, + pub hash_failure: Option, + pub read_request: bool, + pub read_memberships: bool, +} + +pub fn decide(facts: Facts) -> Decision { + if facts.user_present && facts.request_pending { + return Decision { + read_request: true, + ..Decision::default() + }; + } + if facts.user_present + && facts.user_role.as_deref() != Some("proxy_admin") + && facts.route != "/organization/new" + && (facts.org_admin_routes.contains(&facts.route) || facts.route == "/team/new") + && facts.memberships_pending + { + return Decision { + read_memberships: true, + ..Decision::default() + }; + } + let memberships = facts.memberships.as_deref().unwrap_or(&[]); + let ids = memberships + .iter() + .filter_map(|item| item.organization_id.clone()) + .collect(); + let role = memberships + .iter() + .rev() + .find(|item| { + item.organization_id.is_some() && item.organization_id == facts.organization_id + }) + .and_then(|item| item.role.clone()); + let present = facts.organization_id_present || facts.organization_id.is_some(); + let needs_lookup = facts.user_present + && facts.user_role.as_deref() != Some("proxy_admin") + && facts.route != "/organization/new" + && present + && ((facts.org_admin_routes.contains(&facts.route) && facts.memberships.is_some()) + || facts.route == "/team/new" && !memberships.is_empty()); + if needs_lookup && (facts.identifier_pending || facts.hash_failure.is_some()) { + return Decision { + denial: None, + organization_role: role, + organization_ids: ids, + read_identifier: facts.identifier_pending, + hash_failure: facts.hash_failure, + ..Decision::default() + }; + } + let denial = if !facts.user_present || facts.user_role.as_deref() == Some("proxy_admin") { + None + } else if facts.route == "/organization/new" { + Some(Denial::CreateOrganization) + } else if facts.org_admin_routes.contains(&facts.route) { + if facts.memberships.is_none() { + Some(Denial::NoMemberships) + } else if !present { + Some(Denial::MissingOrganization) + } else if role.is_none() { + Some(Denial::MissingRole) + } else if role.as_deref() != Some("org_admin") { + Some(Denial::WrongRole) + } else { + None + } + } else if facts.route == "/team/new" && !memberships.is_empty() { + if !present { + Some(Denial::TeamMissingOrganization) + } else if role.as_deref() != Some("org_admin") { + Some(Denial::TeamWrongRole) + } else { + None + } + } else { + None + }; + Decision { + denial, + organization_role: role, + organization_ids: ids, + read_identifier: false, + hash_failure: None, + ..Decision::default() + } +} + +pub fn user_is_admin(memberships: &[Membership], organizations: &[String]) -> bool { + !organizations.is_empty() + && organizations.iter().all(|id| { + memberships.iter().any(|item| { + item.organization_id.as_ref() == Some(id) + && item.role.as_deref() == Some("org_admin") + }) + }) +} + +#[derive(Deserialize)] +pub struct Identifier { + pub text: Option, + pub hash_failure: Option, +} + +#[derive(Deserialize)] +pub struct AdminFacts { + pub memberships: Option>, + pub request_read: bool, + pub singular_present: bool, + pub plural_count: usize, + pub candidate_index: usize, + pub candidate: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum AdminAction { + ReadRequest, + ReadSingular, + ReadPlural, + Allow, + Deny, + HashFailure, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct AdminDecision { + pub action: AdminAction, + pub candidate_index: usize, + pub plural_index: usize, + pub hash_failure: Option, +} + +pub fn admin(facts: &AdminFacts) -> AdminDecision { + let result = |action| AdminDecision { + action, + candidate_index: 0, + plural_index: 0, + hash_failure: None, + }; + let Some(memberships) = &facts.memberships else { + return result(AdminAction::Deny); + }; + if !facts.request_read { + return result(AdminAction::ReadRequest); + } + let count = usize::from(facts.singular_present) + facts.plural_count; + if count == 0 { + return result(AdminAction::Deny); + } + let next = if let Some(candidate) = &facts.candidate { + if let Some(message) = &candidate.hash_failure { + return AdminDecision { + hash_failure: Some(message.clone()), + ..result(AdminAction::HashFailure) + }; + } + if !candidate.text.as_ref().is_some_and(|id| { + memberships.iter().any(|item| { + item.organization_id.as_ref() == Some(id) + && item.role.as_deref() == Some("org_admin") + }) + }) { + return result(AdminAction::Deny); + } + facts.candidate_index + 1 + } else { + 0 + }; + if next == count { + return result(AdminAction::Allow); + } + AdminDecision { + action: if facts.singular_present && next == 0 { + AdminAction::ReadSingular + } else { + AdminAction::ReadPlural + }, + candidate_index: next, + plural_index: next.saturating_sub(usize::from(facts.singular_present)), + hash_failure: None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::no_memberships(false, false, None, None, AdminAction::Deny)] + #[case::read_request(true, false, None, None, AdminAction::ReadRequest)] + #[case::read_singular(true, true, None, None, AdminAction::ReadSingular)] + #[case::allowed(true, true, Some("org"), None, AdminAction::Allow)] + #[case::wrong_org(true, true, Some("other"), None, AdminAction::Deny)] + #[case::unhashable( + true, + true, + None, + Some("unhashable type: 'list'"), + AdminAction::HashFailure + )] + fn raw_admin_stages( + #[case] present: bool, + #[case] request_read: bool, + #[case] id: Option<&str>, + #[case] failure: Option<&str>, + #[case] expected: AdminAction, + ) { + let memberships = present.then(|| { + vec![Membership { + organization_id: Some("org".into()), + role: Some("org_admin".into()), + }] + }); + let candidate = (id.is_some() || failure.is_some()).then(|| Identifier { + text: id.map(str::to_owned), + hash_failure: failure.map(str::to_owned), + }); + assert_eq!( + admin(&AdminFacts { + memberships, + request_read, + singular_present: true, + plural_count: 0, + candidate_index: 0, + candidate + }) + .action, + expected + ); + } + + #[rstest] + #[case::none(None, Some(Denial::NoMemberships))] + #[case::empty(Some(vec![]), Some(Denial::MissingOrganization))] + fn membership_presence( + #[case] memberships: Option>, + #[case] expected: Option, + ) { + assert_eq!( + decide(Facts { + route: "/organization/update".into(), + user_present: true, + user_role: None, + memberships, + organization_id: None, + org_admin_routes: vec!["/organization/update".into()], + organization_id_present: false, + identifier_pending: false, + hash_failure: None, + request_pending: false, + memberships_pending: false, + }) + .denial, + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/owner_teams.rs b/litellm-rust/crates/request-auth-core/src/owner_teams.rs new file mode 100644 index 00000000000..d8816cf5407 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/owner_teams.rs @@ -0,0 +1,173 @@ +use crate::{Error, personal_budget}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Purpose { + Models, + Grants, +} + +#[derive(Clone, Deserialize)] +pub struct Load { + pub present: bool, + pub failed: bool, + pub blocked: bool, + pub org: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Budget { + pub index: usize, + pub facts: personal_budget::Facts, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub purpose: Purpose, + pub user_present: bool, + pub database_present: bool, + pub ids: Vec, + pub ui_team_id: String, + pub org: Option, + pub loaded: bool, + pub loads: Vec, + pub budgets: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadTeams { + ids: Vec, + }, + ReadBudgets { + indices: Vec, + }, + Complete { + selected: Vec, + withheld: Vec, + unreadable: bool, + }, +} + +pub fn decide(facts: &Facts) -> Result { + if facts.purpose == Purpose::Grants && (!facts.user_present || !facts.database_present) { + return Ok(Decision::Complete { + selected: vec![], + withheld: vec![], + unreadable: false, + }); + } + if !facts.loaded { + let mut seen = BTreeSet::new(); + let ids = facts + .ids + .iter() + .filter(|id| { + facts.purpose == Purpose::Models || (!id.is_empty() && **id != facts.ui_team_id) + }) + .filter(|id| seen.insert(*id)) + .cloned() + .collect(); + return Ok(Decision::ReadTeams { ids }); + } + let eligible: Vec<_> = facts + .loads + .iter() + .enumerate() + .filter(|(_, load)| load.present && !load.blocked) + .filter(|(_, load)| { + facts.purpose == Purpose::Models + || facts.org.is_none() + || load.org.is_none() + || load.org == facts.org + }) + .map(|(index, _)| index) + .collect(); + let unreadable = facts.loads.iter().any(|load| load.failed); + if facts.purpose == Purpose::Models { + return Ok(Decision::Complete { + selected: eligible, + withheld: vec![], + unreadable, + }); + } + let missing: Vec<_> = eligible + .iter() + .filter(|index| !facts.budgets.iter().any(|budget| budget.index == **index)) + .copied() + .collect(); + if !missing.is_empty() { + return Ok(Decision::ReadBudgets { indices: missing }); + } + let results: Result, _> = eligible + .iter() + .map(|index| { + let budget = facts + .budgets + .iter() + .find(|budget| budget.index == *index) + .ok_or(Error::InvalidAnswer)?; + match personal_budget::decide(&budget.facts)? { + personal_budget::Decision::Complete { allowed, .. } => Ok((*index, allowed)), + _ => Err(Error::InvalidAnswer), + } + }) + .collect(); + let results = results?; + let blocked = facts + .loads + .iter() + .enumerate() + .filter(|(_, load)| { + load.present + && load.blocked + && (facts.org.is_none() || load.org.is_none() || load.org == facts.org) + }) + .map(|(index, _)| index); + let withheld = blocked + .chain( + results + .iter() + .filter_map(|(index, allowed)| (!allowed).then_some(*index)), + ) + .collect(); + Ok(Decision::Complete { + selected: results + .into_iter() + .filter_map(|(index, allowed)| allowed.then_some(index)) + .collect(), + withheld, + unreadable, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::models(Purpose::Models, vec!["", "ui", "a"])] + #[case::grants(Purpose::Grants, vec!["a"])] + fn id_filter(#[case] purpose: Purpose, #[case] expected: Vec<&str>) { + let facts = Facts { + purpose, + user_present: true, + database_present: true, + ids: vec!["".into(), "ui".into(), "a".into(), "a".into()], + ui_team_id: "ui".into(), + org: None, + loaded: false, + loads: vec![], + budgets: vec![], + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::ReadTeams { + ids: expected.into_iter().map(str::to_owned).collect() + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/patterns.rs b/litellm-rust/crates/request-auth-core/src/patterns.rs new file mode 100644 index 00000000000..bb97e6841cd --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/patterns.rs @@ -0,0 +1,363 @@ +use fancy_regex::Regex; +use std::{ + collections::BTreeMap, + sync::{Arc, LazyLock, Mutex}, +}; + +use crate::Error; + +static EXPRESSIONS: LazyLock>>> = + LazyLock::new(|| Mutex::new(BTreeMap::new())); + +fn python_escape(escaped: char, in_class: bool) -> Option { + let word = r"\p{L}\p{N}_"; + let class = match escaped { + 'w' => format!("[{word}]"), + 'W' => format!("[^{word}]"), + 'b' if !in_class => format!("(?:(?<=[{word}])(?![{word}])|(? { + format!("(?:(?<=[{word}])(?=[{word}])|(? return None, + }; + Some(format!("(?-i:{class})")) +} + +fn python_class(pattern: &str) -> Option<(usize, String)> { + let negated = pattern.starts_with("[^"); + let start = if negated { 2 } else { 1 }; + let mut chars = pattern.char_indices().skip(start).peekable(); + let mut literals = String::new(); + let mut alternatives = Vec::new(); + let mut range = false; + if chars.peek().is_some_and(|(_, ch)| *ch == ']') { + literals.push_str(r"\]"); + chars.next(); + } + while let Some((index, ch)) = chars.next() { + match ch { + ']' => { + if alternatives.is_empty() { + return Some((index + 1, pattern[..index + 1].into())); + } + if !literals.is_empty() { + alternatives.push(format!("[{literals}]")); + } + let union = alternatives.join("|"); + return Some(( + index + 1, + if negated { + format!("(?:(?!(?:{union}))[\\s\\S])") + } else { + format!("(?:{union})") + }, + )); + } + '\\' => { + let (_, escaped) = chars.next()?; + if let Some(class) = python_escape(escaped, true) { + // A category escape cannot be a range endpoint in Python. + if range + || pattern[index + 2..].starts_with('-') + && !pattern[index + 2..].starts_with("-]") + { + return None; + } + alternatives.push(class); + } else { + literals.push(ch); + literals.push(escaped); + } + } + '[' => literals.push_str(r"\["), + '^' => literals.push_str(r"\^"), + _ => literals.push(ch), + } + range = ch == '-' && index > start; + } + None +} + +fn python_pattern(pattern: &str) -> String { + let mut output = String::with_capacity(pattern.len()); + let mut verbose = false; + let mut scopes = Vec::new(); + let mut chars = pattern.char_indices().peekable(); + while let Some((index, ch)) = chars.next() { + let rest = &pattern[index..]; + match ch { + '\\' => { + if let Some((_, escaped)) = chars.next() { + if let Some(replacement) = python_escape(escaped, false) { + output.push_str(&replacement); + } else { + output.push(ch); + output.push(escaped); + } + } else { + output.push(ch); + } + } + '[' => { + let Some((length, class)) = python_class(rest) else { + output.push_str(rest); + break; + }; + output.push_str(&class); + while chars.peek().is_some_and(|(next, _)| *next < index + length) { + chars.next(); + } + } + '#' if verbose => { + output.push(ch); + for (_, next) in chars.by_ref() { + output.push(next); + if next == '\n' { + break; + } + } + } + '(' if rest.starts_with("(?#") => { + output.push(ch); + let mut escaped = false; + for (_, next) in chars.by_ref() { + output.push(next); + if next == ')' && !escaped { + break; + } + escaped = next == '\\' && !escaped; + } + } + '(' => { + if let Some(header) = rest.strip_prefix("(?") { + let length = header + .find(|c: char| !"aiLmsux-".contains(c)) + .unwrap_or(header.len()); + if length > 0 && matches!(header.as_bytes().get(length), Some(b':' | b')')) { + let scoped = header.as_bytes()[length] == b':'; + if scoped { + scopes.push(verbose); + } + let mut disabled = false; + for flag in header[..length].chars() { + match flag { + 'x' => verbose = !disabled, + '-' => disabled = true, + _ => (), + } + } + output.push_str(&rest[..length + 3]); + while chars + .peek() + .is_some_and(|(next, _)| *next < index + length + 3) + { + chars.next(); + } + continue; + } + } + scopes.push(verbose); + output.push(ch); + } + ')' => { + verbose = scopes.pop().unwrap_or_default(); + output.push(ch); + } + _ => output.push(ch), + } + } + output +} + +fn expression(pattern: &str) -> Result, Error> { + if let Some(compiled) = EXPRESSIONS + .lock() + .unwrap_or_else(|error| error.into_inner()) + .get(pattern) + { + return Ok(Arc::clone(compiled)); + } + let compiled = Arc::new(Regex::new(&python_pattern(pattern)).map_err(Error::Pattern)?); + let mut expressions = EXPRESSIONS + .lock() + .unwrap_or_else(|error| error.into_inner()); + if expressions.len() >= 1024 { + expressions.clear(); + } + expressions.insert(pattern.to_owned(), Arc::clone(&compiled)); + Ok(compiled) +} + +pub fn regex_match(pattern: &str, value: &str) -> Result { + let expression = expression(pattern)?; + Ok(expression + .find(value) + .map_err(Error::Pattern)? + .is_some_and(|found| found.start() == 0)) +} + +pub fn model_pattern(model: &str, pattern: &str) -> Result { + if !pattern.contains('*') { + return Ok(false); + } + regex_match( + &format!("^{}(?=\\n?\\z)", pattern.replace('*', ".*")), + model, + ) +} + +pub fn route_pattern(route: &str, pattern: &str) -> Result { + let mut expression = String::from("^"); + let mut rest = pattern; + while let Some(start) = rest.find('{') { + expression.push_str(&rest[..start]); + let Some(end) = rest[start + 1..].find('}').map(|index| index + start + 1) else { + expression.push_str(&rest[start..]); + rest = ""; + break; + }; + let name = &rest[start + 1..end]; + if name.is_empty() { + expression.push_str("{}"); + } else if name.ends_with(":path") { + expression.push_str(if rest[end + 1..].contains(':') { + "[^:]+" + } else { + "[\\s\\S]+" + }); + } else { + expression.push_str("[^/]+"); + } + rest = &rest[end + 1..]; + } + expression.push_str(rest); + expression.push_str("(?=\\n?\\z)"); + regex_match(&expression, route) +} + +pub fn wildcard(route: &str, pattern: &str) -> bool { + pattern + .strip_suffix('*') + .map_or(route == pattern, |prefix| route.starts_with(prefix)) +} + +pub fn prefix(route: &str, pattern: &str) -> bool { + route == pattern + || route + .strip_prefix(pattern) + .is_some_and(|suffix| suffix.starts_with('/')) +} + +pub fn route_access(route: &str, patterns: &[String]) -> Result { + if patterns + .iter() + .any(|pattern| pattern == route || (pattern.ends_with('*') && wildcard(route, pattern))) + { + return Ok(true); + } + for pattern in patterns { + if route_pattern(route, pattern)? { + return Ok(true); + } + } + Ok(false) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + #[rstest] + #[case::combining(r"\w+", "\u{0301}a", false)] + #[case::joiner(r"\w+", "\u{200d}a", false)] + #[case::connector(r"\w+", "\u{203f}a", false)] + #[case::superscript(r"\w+", "²a", true)] + #[case::fraction(r"\w+", "¼", true)] + #[case::letter_number(r"\w+", "Ⅷ", true)] + #[case::underscore(r"\w+", "_", true)] + #[case::letter(r"\w+", "한", true)] + #[case::nonword_mark(r"\W+", "\u{0301}a", true)] + #[case::nonword_number(r"\W+", "²a", false)] + #[case::class_word(r"[\w]+", "²a", true)] + #[case::class_nonword(r"[\W]+", "\u{0301}a", true)] + #[case::negated_word(r"[^\w]+", "²a", false)] + #[case::negated_nonword(r"[^\W]+", "²a", true)] + #[case::negated_repeat(r"[^\w]+$", "!a", false)] + #[case::union(r"[\w\W]+", "\u{0301}a", true)] + #[case::negated_union(r"[^\w\W]+", "a", false)] + #[case::class_caret(r"[\w^]+", "^", true)] + #[case::class_dash(r"[\w-]+", "-", true)] + #[case::class_escaped_dash(r"[\-\w]+", "²", true)] + #[case::class_closing_bracket(r"[]\w]+", "]²", true)] + #[case::class_opening_bracket(r"[[\w]+", "[²", true)] + #[case::class_backspace(r"[\b\w]", "\u{8}", true)] + #[case::literal_escape(r"\\w", r"\w", true)] + #[case::literal_class_escape(r"[\\w]", "²", false)] + #[case::boundary_number(r"\b²", "²", true)] + #[case::boundary_mark(r"\b\u0301", "\u{0301}", false)] + #[case::nonboundary_mark(r"\B\u0301", "\u{0301}", true)] + #[case::nonboundary_number(r"\B²", "²", false)] + #[case::boundary_after_word(r"a\b\u0301", "a\u{0301}", true)] + #[case::boundary_after_nonword(r"!\b²", "!²", true)] + #[case::nonboundary_between_words(r"a\B²", "a²", true)] + #[case::nonboundary_between_nonwords(r"!\B\u0301", "!\u{0301}", true)] + #[case::end_boundary(r"²\b", "²", true)] + #[case::end_nonboundary(r"\u0301\B", "\u{0301}", true)] + #[case::empty_boundary(r"\b", "", false)] + #[case::empty_nonboundary(r"\B", "", false)] + #[case::unicode_flag(r"(?u:\w)", "²", true)] + #[case::case_insensitive_word(r"(?i:\w)", "\u{0345}", false)] + #[case::case_insensitive_class_word(r"(?i:[\w])", "\u{0345}", false)] + #[case::case_insensitive_class_literal(r"(?i:[\wΙ])", "\u{0345}", true)] + #[case::case_insensitive_boundary(r"(?i:\b\u0345)", "\u{0345}", false)] + #[case::case_insensitive_literal(r"(?i:a)\w", "A²", true)] + #[case::ascii_class(r"[a-zA-Z0-9_]", "²", false)] + #[case::comment(r"(?# [\b)\w", "²", true)] + #[case::verbose_comment("(?x:# [\\b)\n\\w)", "²", true)] + #[case::verbose_scope("(?x:# [\n)(?-x:#)\\w", "#²", true)] + #[case::verbose_restored("(?x:# [\n)#\\w", "#²", true)] + #[case::backreference(r"(\w)\1", "²²", true)] + #[case::lookbehind(r".(?<=[\w])a", "²a", true)] + fn python_word_semantics(#[case] pattern: &str, #[case] value: &str, #[case] expected: bool) { + assert_eq!(super::regex_match(pattern, value).unwrap(), expected); + } + + #[rstest] + #[case::ascii_flag(r"(?a:\w)")] + #[case::ascii_unicode_flag(r"(?-u:\w)")] + #[case::word_range_start(r"[\w-a]")] + #[case::word_range_end(r"[a-\w]")] + #[case::unclosed_class(r"[\w")] + #[case::unclosed_group(r"(\w")] + fn invalid_patterns_remain_errors(#[case] pattern: &str) { + assert!(super::regex_match(pattern, "a").is_err()); + } + + #[rstest] + #[case::segment("/key/k/reset", "/key/{id}/reset", true)] + #[case::segment_slash("/key/a/b/reset", "/key/{id}/reset", false)] + #[case::path_slash("/credentials/a/b/start", "/credentials/{id:path}/start", true)] + #[case::path_newline("/credentials/a\nb/start", "/credentials/{id:path}/start", true)] + #[case::google_suffix( + "/models/a:b:generateContent", + "/models/{id:path}:generateContent", + false + )] + #[case::final_newline("/models/a\n", "/models/{id}", true)] + fn route_templates(#[case] route: &str, #[case] pattern: &str, #[case] expected: bool) { + assert_eq!(super::route_pattern(route, pattern).unwrap(), expected); + } + + #[rstest] + #[case::provider("openai/a", "openai/*", true)] + #[case::exact_not_pattern("a", "a", false)] + #[case::regex_dot("aa-1", "a.-*", true)] + #[case::newline("a\nb", "*", false)] + #[case::final_newline("a\n", "a*", true)] + #[case::word_combining_mark("\u{0301}a", r"\w*", false)] + #[case::word_join_control("\u{200d}a", r"\w*", false)] + #[case::word_other_number("\u{00b2}a", r"\w*", true)] + fn model_patterns(#[case] model: &str, #[case] pattern: &str, #[case] expected: bool) { + assert_eq!(super::model_pattern(model, pattern).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_budget.rs b/litellm-rust/crates/request-auth-core/src/personal_budget.rs new file mode 100644 index 00000000000..7f2656f401d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_budget.rs @@ -0,0 +1,186 @@ +use serde::{Deserialize, Serialize}; + +use crate::{ + Error, + common_budgets::{self, BudgetFacts, MemberLimitFacts}, +}; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub team_id_present: bool, + pub check_db_only: bool, + pub token_org: Option, + pub team_org: Option, + pub source_ready: bool, + pub team: Option, + pub team_alerted: bool, + pub window_count: Option, + pub windows: Vec, + pub membership_loaded: bool, + pub member_limit: Option, + pub default_id: Option, + pub default_loaded: bool, + pub member: Option, + pub member_alerted: bool, + pub org_loaded: bool, + pub org_failed: bool, + pub org: Option, + pub org_alerted: bool, + pub io_failed: bool, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool, log_failure: bool }, + Source { org_id: Option }, + ReadTeam, + AlertTeam, + ReadWindows, + ReadWindow { index: usize }, + ReadMembership, + ReadMemberLimit, + ReadDefault { id: String }, + ReadMember { limit: Option }, + AlertMember, + ReadOrg { id: String }, + AlertOrg { id: String }, +} + +pub fn decide(facts: &Facts) -> Result { + let complete = |allowed, log_failure| { + Ok(Decision::Complete { + allowed, + log_failure, + }) + }; + let org_id = facts + .token_org + .as_ref() + .filter(|id| !id.is_empty()) + .or(facts.team_org.as_ref()); + if !facts.source_ready { + return Ok(Decision::Source { + org_id: org_id.cloned(), + }); + } + if facts.io_failed { + return complete(false, false); + } + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeam); + }; + let team = common_budgets::check(team); + if team.read_spend { + return Err(Error::InvalidAnswer); + } + if team.alert.is_some() && !facts.team_alerted { + return Ok(Decision::AlertTeam); + } + if team.denied { + return complete(false, false); + } + let Some(count) = facts.window_count else { + return Ok(Decision::ReadWindows); + }; + for index in 0..count { + let Some(window) = facts.windows.get(index) else { + return Ok(Decision::ReadWindow { index }); + }; + let budget = common_budgets::check(window); + if budget.read_spend { + return Err(Error::InvalidAnswer); + } + if budget.denied { + return complete(false, false); + } + } + if facts.check_db_only && !facts.user_present { + return complete(false, false); + } + if facts.user_present + && (facts.team_id_present || facts.check_db_only) + && !facts.membership_loaded + { + return Ok(Decision::ReadMembership); + } + if facts.user_present && facts.team_id_present { + let Some(limit) = &facts.member_limit else { + return Ok(Decision::ReadMemberLimit); + }; + if limit.member_limit.is_none() + && !facts.default_loaded + && let Some(id) = &facts.default_id + { + return Ok(Decision::ReadDefault { id: id.clone() }); + } + let limit = match common_budgets::member_limit(limit) { + Ok(limit) => limit, + Err(Error::IncompatibleDateTimes) => return complete(false, true), + Err(error) => return Err(error), + }; + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember { limit }); + }; + let member = common_budgets::check(member); + if member.read_spend { + return Err(Error::InvalidAnswer); + } + if member.alert.is_some() && !facts.member_alerted { + return Ok(Decision::AlertMember); + } + if member.denied { + return complete(false, false); + } + } + if let Some(id) = org_id { + if !facts.org_loaded { + return Ok(Decision::ReadOrg { id: id.clone() }); + } + if !facts.org_failed + && let Some(org) = &facts.org + { + let org = common_budgets::check(org); + if org.read_spend { + return Err(Error::InvalidAnswer); + } + if org.alert.is_some() && !facts.org_alerted { + return Ok(Decision::AlertOrg { id: id.clone() }); + } + if org.denied { + return complete(false, false); + } + } + } + complete(true, false) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::token(Some("token"), Some("team"), Some("token"))] + #[case::empty_token(Some(""), Some("team"), Some("team"))] + #[case::empty_team(None, Some(""), Some(""))] + #[case::absent(None, None, None)] + fn source_org( + #[case] token: Option<&str>, + #[case] team: Option<&str>, + #[case] selected: Option<&str>, + ) { + let facts = Facts { + token_org: token.map(str::to_owned), + team_org: team.map(str::to_owned), + ..Facts::default() + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::Source { + org_id: selected.map(str::to_owned) + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs b/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs new file mode 100644 index 00000000000..98c66742bbe --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs @@ -0,0 +1,81 @@ +use crate::{Error, personal_model_request::Loaded, personal_models, personal_owner}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub owner: personal_owner::Facts, + pub names: Option>, + pub models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadOwner, + ReadNames, + ReadModel { index: usize, teams: Vec }, + Complete { selected: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let teams = match personal_owner::decide(&facts.owner)? { + personal_owner::Decision::Complete { selected, .. } => selected, + _ => return Ok(Decision::ReadOwner), + }; + if facts + .owner + .teams + .as_ref() + .is_none_or(|teams| teams.loads.is_empty()) + { + return Ok(Decision::Complete { selected: vec![] }); + } + let Some(names) = &facts.names else { + return Ok(Decision::ReadNames); + }; + let mut selected = vec![]; + for index in 0..names.len() { + let Some(loaded) = facts.models.get(index) else { + return Ok(Decision::ReadModel { index, teams }); + }; + match personal_models::decide(personal_models::Request::Loaded { + present: loaded.present.clone(), + scopes: loaded.scopes.clone(), + })? { + personal_models::Decision::Complete { allowed } => { + if allowed { + selected.push(index); + } + } + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Complete { selected }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing_database(false, Decision::Complete { selected: vec![] })] + #[case::load(true, Decision::ReadOwner)] + fn owner_before_names(#[case] database_present: bool, #[case] expected: Decision) { + let facts = Facts { + owner: personal_owner::Facts { + user_present: true, + database_present, + cache_present: true, + logging_present: true, + owner_loaded: false, + owner_present: false, + owner_failed: false, + teams: None, + }, + names: None, + models: vec![], + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_model_request.rs b/litellm-rust/crates/request-auth-core/src/personal_model_request.rs new file mode 100644 index 00000000000..9db1e2a03cb --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_model_request.rs @@ -0,0 +1,86 @@ +use crate::{Error, owner_teams, personal_models}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Loaded { + pub present: Vec, + pub scopes: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub owner_present: bool, + pub teams: Option, + pub models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { model: String }, + ReadTeams, + ReadModel { index: usize, teams: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let teams = if facts.owner_present { + let Some(teams) = &facts.teams else { + return Ok(Decision::ReadTeams); + }; + match owner_teams::decide(teams)? { + owner_teams::Decision::Complete { selected, .. } => selected, + _ => return Err(Error::InvalidAnswer), + } + } else { + vec![] + }; + for (index, model) in facts.requested.iter().enumerate() { + if !facts.owner_present { + return Ok(Decision::Deny { + model: model.clone(), + }); + } + let Some(loaded) = facts.models.get(index) else { + return Ok(Decision::ReadModel { index, teams }); + }; + match personal_models::decide(personal_models::Request::Loaded { + present: loaded.present.clone(), + scopes: loaded.scopes.clone(), + })? { + personal_models::Decision::Complete { allowed: true } => (), + personal_models::Decision::Complete { allowed: false } => { + return Ok(Decision::Deny { + model: model.clone(), + }); + } + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Allow) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::empty(false, vec![], Decision::Allow)] + #[case::owner_lookup(true, vec![], Decision::ReadTeams)] + #[case::absent(false, vec!["model".into()], Decision::Deny { model: "model".into() })] + fn owner_and_empty_requests( + #[case] owner_present: bool, + #[case] requested: Vec, + #[case] expected: Decision, + ) { + let facts = Facts { + requested, + owner_present, + teams: None, + models: vec![], + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_models.rs b/litellm-rust/crates/request-auth-core/src/personal_models.rs new file mode 100644 index 00000000000..1fce2d6bd88 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_models.rs @@ -0,0 +1,120 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + model_access, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Scope { + pub user_present: bool, + pub team_id_present: bool, + pub team: Option, + pub member_grants: Option>, + pub member: Option, + pub io_failed: bool, +} + +#[derive(Deserialize)] +pub enum Request { + Scope(Box), + Loaded { + present: Vec, + scopes: Vec, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool }, + ReadTeam, + ReadGroups, + ReadMembership, + ReadMember, + ReadScopes { indices: Vec }, +} + +fn scope(facts: Scope) -> Result { + if facts.io_failed { + return Ok(Decision::Complete { allowed: false }); + } + let Some(team) = facts.team else { + return Ok(Decision::ReadTeam); + }; + let team = fallback_access::grouped(team, true)?; + if matches!( + team, + model_access::Decision::ReadGroups | model_access::Decision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadGroups); + } + if !team.allowed() { + return Ok(Decision::Complete { allowed: false }); + } + if !facts.user_present || !facts.team_id_present { + return Ok(Decision::Complete { allowed: true }); + } + let Some(grants) = facts.member_grants else { + return Ok(Decision::ReadMembership); + }; + if grants.is_empty() { + return Ok(Decision::Complete { allowed: true }); + } + let Some(member) = facts.member else { + return Ok(Decision::ReadMember); + }; + Ok(Decision::Complete { + allowed: model_access::check(&member)?.allowed(), + }) +} + +pub fn decide(request: Request) -> Result { + match request { + Request::Scope(facts) => scope(*facts), + Request::Loaded { present, scopes } => { + let indices: Vec<_> = present + .iter() + .enumerate() + .filter_map(|(index, present)| present.then_some(index)) + .collect(); + if scopes.is_empty() && !indices.is_empty() { + return Ok(Decision::ReadScopes { indices }); + } + if scopes.len() != indices.len() { + return Err(Error::InvalidAnswer); + } + let results: Result, _> = scopes + .into_iter() + .map(|facts| match scope(facts)? { + Decision::Complete { allowed } => Ok(allowed), + _ => Err(Error::InvalidAnswer), + }) + .collect(); + Ok(Decision::Complete { + allowed: results?.into_iter().any(|allowed| allowed), + }) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::none(vec![], Decision::Complete { allowed: false })] + #[case::missing(vec![false], Decision::Complete { allowed: false })] + #[case::read_all(vec![true, false, true], Decision::ReadScopes { indices: vec![0, 2] })] + fn all_team_reads(#[case] present: Vec, #[case] expected: Decision) { + assert_eq!( + decide(Request::Loaded { + present, + scopes: vec![] + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_owner.rs b/litellm-rust/crates/request-auth-core/src/personal_owner.rs new file mode 100644 index 00000000000..bc821d2d509 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_owner.rs @@ -0,0 +1,87 @@ +use crate::{Error, owner_teams}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub database_present: bool, + pub cache_present: bool, + pub logging_present: bool, + pub owner_loaded: bool, + pub owner_present: bool, + pub owner_failed: bool, + pub teams: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadOwner, + ReadTeams, + Complete { + selected: Vec, + withheld: Vec, + unreadable: bool, + }, +} + +pub fn decide(facts: &Facts) -> Result { + let empty = || { + Ok(Decision::Complete { + selected: vec![], + withheld: vec![], + unreadable: false, + }) + }; + if !facts.user_present + || !facts.database_present + || !facts.cache_present + || !facts.logging_present + { + return empty(); + } + if !facts.owner_loaded { + return Ok(Decision::ReadOwner); + } + if facts.owner_failed || !facts.owner_present { + return empty(); + } + let Some(teams) = &facts.teams else { + return Ok(Decision::ReadTeams); + }; + match owner_teams::decide(teams)? { + owner_teams::Decision::Complete { + selected, + withheld, + unreadable, + } => Ok(Decision::Complete { + selected, + withheld, + unreadable, + }), + _ => Err(Error::InvalidAnswer), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::no_database(false, Decision::Complete { selected: vec![], withheld: vec![], unreadable: false })] + #[case::lookup(true, Decision::ReadOwner)] + fn prerequisites(#[case] database_present: bool, #[case] expected: Decision) { + let facts = Facts { + user_present: true, + database_present, + cache_present: true, + logging_present: true, + owner_loaded: false, + owner_present: false, + owner_failed: false, + teams: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/protocol.rs b/litellm-rust/crates/request-auth-core/src/protocol.rs new file mode 100644 index 00000000000..96169a7be06 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/protocol.rs @@ -0,0 +1,463 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_secret_value::SecretValue; +use serde::{Deserialize, Serialize}; + +use crate::Error; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DateTime { + pub wall_microseconds: i64, + pub offset_microseconds: Option, + pub timezone_id: u64, + pub reference: usize, +} + +impl DateTime { + pub fn before(&self, other: &Self) -> Result { + if self.timezone_id == other.timezone_id { + return Ok(self.wall_microseconds < other.wall_microseconds); + } + match (self.offset_microseconds, other.offset_microseconds) { + (Some(left), Some(right)) => { + Ok(self.wall_microseconds - left < other.wall_microseconds - right) + } + (None, None) => Ok(self.wall_microseconds < other.wall_microseconds), + _ => Err(Error::IncompatibleDateTimes), + } + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct RequestFacts { + pub route: String, + pub token: Option, + pub is_llm_api_route: bool, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Settings { + pub master_key: Option, + pub allow_user_auth: bool, + pub has_db: bool, + pub master_key_only_routes: BTreeSet, + pub proxy_max_budget: f64, +} + +#[derive(Clone, veil::Redact, Serialize)] +#[serde(tag = "kind")] +pub enum Grant { + NoMasterKey { + #[redact(with = "[REDACTED]")] + token: Option, + }, + UserAuthRoute, + MasterKey, + CachedAdminKey, + VirtualKey { + budget_throttle_pct: Option, + }, + WithinModelBudgets, +} + +#[derive(Clone, veil::Redact, Serialize)] +#[serde(tag = "kind")] +pub enum Denial { + ModelAccessDenied { + decision: crate::model_access::Decision, + }, + MissingKey, + MalformedKey, + UserAuthDisabled, + ExpiredKey { + expires: DateTime, + now: DateTime, + #[redact(with = "[REDACTED]")] + key_label: String, + }, + MasterKeyOnlyRoute, + NoDatabase, + NotAVirtualKey, + KeyBlocked, + OwnerDeactivated { + user_id: String, + }, + TeamMemberBudgetExceeded { + user_id: String, + team_id: String, + spend: f64, + max_budget: f64, + }, + KeyBudgetExceeded { + spend: f64, + max_budget: f64, + }, + KeyWithoutToken, + ModelBudgetExceeded { + message: String, + current_cost: f64, + max_budget: f64, + entity_type: Option, + entity_id: Option, + }, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct ModelBudgetExceeded { + pub message: String, + pub current_cost: f64, + pub max_budget: f64, + pub entity_type: Option, + pub entity_id: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelBudgetRead { + pub snapshot: crate::model_budget_data::Snapshot, + pub message: String, + pub entity_type: Option, + pub entity_id: Option, +} + +impl ModelBudgetRead { + pub fn exceeded(self) -> Option { + let result = crate::model_budget_data::check(&self.snapshot); + if !result.denied { + return None; + } + Some(ModelBudgetExceeded { + message: self.message, + current_cost: result.current_spend, + max_budget: self.snapshot.max_budget?, + entity_type: self.entity_type, + entity_id: self.entity_id, + }) + } +} + +impl From for Denial { + fn from(value: ModelBudgetExceeded) -> Self { + Self::ModelBudgetExceeded { + message: value.message, + current_cost: value.current_cost, + max_budget: value.max_budget, + entity_type: value.entity_type, + entity_id: value.entity_id, + } + } +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum PresentedKey { + AdminKey { expires: Option }, + NonAdminKey, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct KeyFacts { + pub has_token: bool, + pub blocked: bool, + pub expires: Option, + pub spend: f64, + pub soft_budget: Option, + pub max_budget: Option, + pub team_member_spend: Option, + pub user_id: Option, + pub team_id: Option, + pub budget_throttle_pct: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct OwnerFacts { + pub deactivated: bool, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct TeamMemberBudget { + pub spend: f64, + pub max_budget: f64, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelBudgetScope { + pub models: Vec, + pub rewritable_model: Option, + pub key_has_model_budgets: bool, + pub key_budget_fallbacks: BTreeMap>, + pub user_has_model_budgets: bool, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TeamMemberModelBudgetScope { + pub models: Vec, + pub rewritable_model: Option, + pub has_model_budgets: bool, + pub can_fall_back: bool, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct FallbackTarget { + pub requested: String, + pub model: String, + pub served_model: String, +} + +#[derive(Clone, Debug, Serialize)] +#[serde(tag = "kind")] +pub enum Effect { + ResolveEndUser { + prefetch_key: bool, + }, + ResolvePresentedKey { + hashed_token: String, + }, + EvictPresentedKey { + hashed_token: String, + }, + LoadKey { + hashed_token: String, + }, + PrepareKey { + loaded_from_db: bool, + }, + InspectKey, + EnforceModelAccess, + PrefetchAuthObjects, + LoadKeyOwner { + user_id: String, + }, + CheckBudgetExemption, + ReadTeamMemberBudget { + user_id: String, + team_id: String, + }, + AlertTeamMemberBudget { + user_id: String, + team_id: String, + budget: TeamMemberBudget, + }, + SendKeyBudgetAlerts, + ReadKeySpend { + max_budget: f64, + }, + AlertKeyBudget { + spend: f64, + }, + AlertKeySoftBudget, + LoadKeyContext, + ReadProxySpend, + AlertProxyBudget { + spend: f64, + }, + ReadModelBudgetScope, + CheckKeyModelBudget { + model: String, + }, + AuthorizeBudgetFallback { + model: String, + }, + RewriteRequestModel { + model: String, + }, + CheckUserModelBudget { + model: String, + }, + ReadTeamMemberModelBudgetScope, + CheckTeamMemberModelBudget { + model: String, + }, + ReadTeamMemberFallbacks { + model: String, + }, + VetTeamMemberFallback { + target: FallbackTarget, + }, + RerouteTeamMemberRequest { + model: String, + }, +} + +#[derive(Debug, Serialize)] +#[serde(tag = "kind")] +pub enum Outcome { + Allow { grant: Grant }, + Deny { denial: Denial }, +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +pub enum Node { + Outcome(Outcome), + Effect(Effect), +} + +#[derive(Deserialize)] +#[serde(tag = "kind", content = "value")] +pub enum Answer { + ModelAccess(Box), + FallbackAccess(Box), + Candidate(Box), + Unit, + Bool(bool), + Exemption(crate::budget_pricing::ExemptionFacts), + Presented(Option), + Key(KeyFacts), + Owner(Option), + TeamBudget(Option), + Spend(f64), + ProxySpend(Option), + ModelScope(ModelBudgetScope), + Exceeded(Option), + ModelBudgetData(ModelBudgetRead), + TeamScope(TeamMemberModelBudgetScope), + Targets(Vec), +} + +pub trait FromAnswer: Sized { + fn from_answer(answer: Answer) -> Result; +} + +macro_rules! answer { + ($type:ty, $variant:ident) => { + impl FromAnswer for $type { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::$variant(value) => Ok(value), + _ => Err(Error::InvalidAnswer), + } + } + } + }; +} +answer!(bool, Bool); +answer!(crate::budget_pricing::ExemptionFacts, Exemption); +answer!(Option, Presented); +answer!(KeyFacts, Key); +answer!(Option, Owner); +answer!(Option, TeamBudget); +answer!(f64, Spend); +answer!(Option, ProxySpend); +answer!(ModelBudgetScope, ModelScope); +answer!(Option, Exceeded); +answer!(ModelBudgetRead, ModelBudgetData); +answer!(TeamMemberModelBudgetScope, TeamScope); +answer!(Vec, Targets); + +impl FromAnswer for crate::fallback_candidate::Facts { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::Candidate(value) => Ok(*value), + _ => Err(Error::InvalidAnswer), + } + } +} + +impl FromAnswer for crate::key_model_access::Facts { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::ModelAccess(value) => Ok(*value), + _ => Err(Error::InvalidAnswer), + } + } +} + +impl FromAnswer for crate::fallback_access::Facts { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::FallbackAccess(value) => Ok(*value), + _ => Err(Error::InvalidAnswer), + } + } +} + +impl FromAnswer for () { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::Unit => Ok(()), + _ => Err(Error::InvalidAnswer), + } + } +} + +pub type Continuation = Box Result + Send + Sync>; +pub type Then = Box Result + Send + Sync>; + +pub struct Step { + pub node: Node, + pub resume: Option, +} + +impl Step { + pub fn allow(grant: Grant) -> Result { + Ok(Self { + node: Node::Outcome(Outcome::Allow { grant }), + resume: None, + }) + } + pub fn deny(denial: impl Into) -> Result { + Ok(Self { + node: Node::Outcome(Outcome::Deny { + denial: denial.into(), + }), + resume: None, + }) + } + pub fn effect( + effect: Effect, + resume: impl FnOnce(T) -> Result + Send + Sync + 'static, + ) -> Result { + Ok(Self { + node: Node::Effect(effect), + resume: Some(Box::new(move |answer| resume(T::from_answer(answer)?))), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::same_zone_ignores_offset(10, Some(100), 1, 11, Some(200), 1, Ok(true))] + #[case::different_zones_use_offset(10, Some(100), 1, 11, Some(200), 2, Ok(false))] + #[case::equal_is_not_expired(10, None, 0, 10, None, 0, Ok(false))] + #[case::mixed_naive_aware(10, None, 0, 10, Some(0), 1, Err(()))] + #[allow(clippy::too_many_arguments)] + fn expiry_comparison( + #[case] wall: i64, + #[case] offset: Option, + #[case] zone: u64, + #[case] other_wall: i64, + #[case] other_offset: Option, + #[case] other_zone: u64, + #[case] expected: Result, + ) { + let expires = DateTime { + wall_microseconds: wall, + offset_microseconds: offset, + timezone_id: zone, + reference: 0, + }; + let now = DateTime { + wall_microseconds: other_wall, + offset_microseconds: other_offset, + timezone_id: other_zone, + reference: 1, + }; + assert_eq!(expires.before(&now).map_err(|_| ()), expected); + } + + #[rstest] + fn wrong_answer_never_advances_the_decision() { + let step = Step::effect(Effect::CheckBudgetExemption, |_: bool| { + Step::allow(Grant::MasterKey) + }) + .unwrap(); + assert!(matches!( + (step.resume.unwrap())(Answer::Unit), + Err(Error::InvalidAnswer) + )); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/reservation.rs b/litellm-rust/crates/request-auth-core/src/reservation.rs new file mode 100644 index 00000000000..185cfd19bfc --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/reservation.rs @@ -0,0 +1,312 @@ +use serde::{Deserialize, Serialize}; +use std::cmp::Ordering::Greater; + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Gate { + token_present: bool, + llm_route: Option, + unbilled: Option, + model_present: Option, + counter_count: Option, + cost_read: bool, + cost: Option, + }, + Sources { + key_present: bool, + team_present: bool, + team_id: Option, + user_present: bool, + user_id: Option, + apply_user_budget_to_team_keys: bool, + }, + Counter { + scope: Scope, + limit: Option, + }, + Window { + duration_present: bool, + limit: Option, + start_present: Option, + }, + Identity { + scope: IdentityScope, + primary: Option, + fallback: Option, + }, + MemberGate { + team_present: bool, + team_id: Option, + user_present: bool, + user_id: Option, + }, + MemberCap { + member_limit: Option, + default_read: bool, + }, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum IdentityScope { + EndUser, + Organization, + Project, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum Scope { + Key, + Team, + User, + EndUser, + Tag, + AccessGroup, + Member, + Organization, + Project, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadRoute, + ReadUnbilled, + ReadModel, + ReadCounters, + ReadEstimate, + Reserve, + Skip, + ReadStart, + WarnWindow, + AppendWindow, + ReadMember, + ReadDefault, + ReadCap, + Identity { id: Option }, + Sources { key: bool, team: bool, user: bool }, + Counter { include: bool }, +} + +pub fn decide(facts: Facts) -> Decision { + match facts { + Facts::Identity { + scope, + primary, + fallback, + } => Decision::Identity { + id: match scope { + IdentityScope::EndUser => primary.filter(|id| !id.is_empty()).or(fallback), + IdentityScope::Organization => primary.or(fallback), + IdentityScope::Project => primary, + }, + }, + Facts::MemberGate { + team_present, + team_id, + user_present, + user_id, + } => { + if team_present && team_id.is_some() && user_present && user_id.is_some() { + Decision::ReadMember + } else { + Decision::Skip + } + } + Facts::MemberCap { + member_limit, + default_read, + } => { + if member_limit.is_none() && !default_read { + Decision::ReadDefault + } else { + Decision::ReadCap + } + } + Facts::Gate { + token_present, + llm_route, + unbilled, + model_present, + counter_count, + cost_read, + cost, + } => { + if !token_present { + return Decision::Skip; + } + match llm_route { + None => return Decision::ReadRoute, + Some(false) => return Decision::Skip, + Some(true) => {} + } + match unbilled { + None => return Decision::ReadUnbilled, + Some(true) => return Decision::Skip, + Some(false) => {} + } + match model_present { + None => return Decision::ReadModel, + Some(false) => return Decision::Skip, + Some(true) => {} + } + match counter_count { + None => return Decision::ReadCounters, + Some(0) => return Decision::Skip, + Some(_) => {} + } + if !cost_read { + return Decision::ReadEstimate; + } + match cost { + None => Decision::Skip, + Some(value) if value <= 0.0 => Decision::Skip, + Some(_) => Decision::Reserve, + } + } + Facts::Sources { + key_present, + team_present, + team_id, + user_present, + user_id, + apply_user_budget_to_team_keys, + } => { + let team = team_present && team_id.is_some(); + Decision::Sources { + key: key_present, + team, + user: (!team || apply_user_budget_to_team_keys) + && user_present + && user_id.is_some(), + } + } + Facts::Counter { scope, limit } => Decision::Counter { + include: limit.is_some_and(|limit| match scope { + Scope::Key | Scope::Team | Scope::User => limit.partial_cmp(&0.0) == Some(Greater), + Scope::Project => limit.is_finite() && limit > 0.0, + _ => matches!(limit.partial_cmp(&0.0), None | Some(Greater)), + }), + }, + Facts::Window { + duration_present, + limit, + start_present, + } => { + if !duration_present || limit.is_none_or(|limit| limit <= 0.0) { + return Decision::Skip; + } + match start_present { + None => Decision::ReadStart, + Some(false) => Decision::WarnWindow, + Some(true) => Decision::AppendWindow, + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::customer_empty(IdentityScope::EndUser, Some(""), Some("fallback"), Some("fallback"))] + #[case::org_empty(IdentityScope::Organization, Some(""), Some("fallback"), Some(""))] + #[case::org_fallback(IdentityScope::Organization, None, Some("fallback"), Some("fallback"))] + #[case::project_no_fallback(IdentityScope::Project, None, Some("fallback"), None)] + fn identities( + #[case] scope: IdentityScope, + #[case] primary: Option<&str>, + #[case] fallback: Option<&str>, + #[case] expected: Option<&str>, + ) { + assert_eq!( + decide(Facts::Identity { + scope, + primary: primary.map(str::to_owned), + fallback: fallback.map(str::to_owned) + }), + Decision::Identity { + id: expected.map(str::to_owned) + } + ); + } + + #[rstest] + #[case::member_cap(Some(0.0), false, Decision::ReadCap)] + #[case::default_needed(None, false, Decision::ReadDefault)] + #[case::default_read(None, true, Decision::ReadCap)] + fn member_cap( + #[case] member_limit: Option, + #[case] default_read: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::MemberCap { + member_limit, + default_read + }), + expected + ); + } + + #[rstest] + #[case::all_present(true, Some(""), true, Some(""), Decision::ReadMember)] + #[case::missing_team(true, None, true, Some("user"), Decision::Skip)] + #[case::missing_user(true, Some("team"), false, Some("user"), Decision::Skip)] + fn member_gate( + #[case] team_present: bool, + #[case] team_id: Option<&str>, + #[case] user_present: bool, + #[case] user_id: Option<&str>, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::MemberGate { + team_present, + team_id: team_id.map(str::to_owned), + user_present, + user_id: user_id.map(str::to_owned) + }), + expected + ); + } + + #[rstest] + #[case::key_nan(Scope::Key, f64::NAN, false)] + #[case::customer_nan(Scope::EndUser, f64::NAN, true)] + #[case::project_infinite(Scope::Project, f64::INFINITY, false)] + #[case::team_infinite(Scope::Team, f64::INFINITY, true)] + #[case::zero(Scope::Organization, 0.0, false)] + #[case::positive(Scope::User, 10.0, true)] + fn limits(#[case] scope: Scope, #[case] limit: f64, #[case] include: bool) { + assert_eq!( + decide(Facts::Counter { + scope, + limit: Some(limit) + }), + Decision::Counter { include } + ); + } + + #[rstest] + #[case::no_duration(false, None, Decision::Skip)] + #[case::unread(true, None, Decision::ReadStart)] + #[case::invalid(true, Some(false), Decision::WarnWindow)] + #[case::valid(true, Some(true), Decision::AppendWindow)] + fn windows( + #[case] duration_present: bool, + #[case] start_present: Option, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::Window { + duration_present, + limit: Some(1.0), + start_present + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/resolved_models.rs b/litellm-rust/crates/request-auth-core/src/resolved_models.rs new file mode 100644 index 00000000000..d048a49745d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/resolved_models.rs @@ -0,0 +1,219 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + member_models, + model_access::{self, Decision as ModelDecision}, + personal_model_request, team_override, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct TeamRead { + pub found: bool, + pub failed: bool, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub model: String, + pub config_nonempty: bool, + pub key_models: Vec, + pub team_id: Option, + pub user_id: Option, + pub project_present: bool, + pub customer_present: bool, + pub database_present: bool, + pub personal_enabled: bool, + pub master: bool, + pub role: Option, + pub key: Option, + pub team_read: Option, + pub embedded_ready: bool, + pub team: Option, + pub member: Option, + pub owner_loaded: bool, + pub owner_present: bool, + pub owner_failed: bool, + pub personal: Option, + pub project_loaded: bool, + pub project_grants: Vec, + pub project: Option, + pub customer_loaded: bool, + pub customer_grants: Vec, + pub customer: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: ModelDecision, + member_models: Option>, + }, + ReadKey, + ReadKeyGroups, + ReadTeamObject, + ReadEmbeddedTeam, + ReadTeamScope, + ReadMember, + ReadOwner, + ReadPersonal, + ReadProjectObject, + ReadProjectScope, + ReadCustomerObject, + ReadCustomerScope, +} + +pub fn decide(facts: &Facts) -> Result { + let deny = |denial| { + Ok(Decision::Deny { + denial, + member_models: None, + }) + }; + if !facts.config_nonempty + && !facts + .key_models + .iter() + .any(|model| model == "all-team-models") + { + let Some(key) = &facts.key else { + return Ok(Decision::ReadKey); + }; + let key = fallback_access::grouped(key.clone(), false)?; + if matches!( + key, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadKeyGroups); + } + if !key.allowed() { + return deny(key); + } + } + if facts.team_id.is_some() { + let Some(read) = &facts.team_read else { + return Ok(Decision::ReadTeamObject); + }; + if read.failed && !facts.embedded_ready { + return Ok(Decision::ReadEmbeddedTeam); + } + if read.found || read.failed { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeamScope); + }; + match team_override::decide(team)? { + team_override::Decision::Allow => (), + team_override::Decision::Deny { denial } => return deny(denial), + _ => return Err(Error::InvalidAnswer), + } + if read.found && facts.user_id.is_some() { + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember); + }; + match member_models::decide(member)? { + member_models::Decision::Allow => (), + member_models::Decision::Deny { + denial, + member_models, + } => { + return Ok(Decision::Deny { + denial, + member_models, + }); + } + _ => return Err(Error::InvalidAnswer), + } + } + } + } + let model_access::Response::Predicate(personal) = + model_access::decide(model_access::Request::PersonalApplies { + enabled: facts.personal_enabled, + token_present: true, + team_id: facts.team_id.clone(), + user_id: facts.user_id.clone(), + master: facts.master, + role: facts.role.clone(), + api_key_present: true, + inheritance: false, + })? + else { + return Err(Error::InvalidAnswer); + }; + if personal { + if !facts.owner_loaded { + return Ok(Decision::ReadOwner); + } + if facts.owner_failed || !facts.owner_present { + return deny(ModelDecision::PersonalTeam { + model: facts.model.clone(), + }); + } + let Some(personal) = &facts.personal else { + return Ok(Decision::ReadPersonal); + }; + match personal_model_request::decide(personal)? { + personal_model_request::Decision::Allow => (), + personal_model_request::Decision::Deny { model } => { + return deny(ModelDecision::PersonalTeam { model }); + } + _ => return Err(Error::InvalidAnswer), + } + } + if facts.project_present { + if !facts.project_loaded { + return Ok(Decision::ReadProjectObject); + } + if !facts.project_grants.is_empty() { + let Some(project) = &facts.project else { + return Ok(Decision::ReadProjectScope); + }; + let decision = model_access::check(project)?; + if !decision.allowed() { + return deny(decision); + } + } + } + if facts.customer_present && facts.database_present { + if !facts.customer_loaded { + return Ok(Decision::ReadCustomerObject); + } + if !facts.customer_grants.is_empty() { + let Some(customer) = &facts.customer else { + return Ok(Decision::ReadCustomerScope); + }; + let model_access::Response::Decision(decision) = + model_access::decide(model_access::Request::Customer(customer.clone()))? + else { + return Err(Error::InvalidAnswer); + }; + if !decision.allowed() { + return deny(decision); + } + } + } + Ok(Decision::Allow) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::key(false, Decision::ReadKey)] + #[case::team(true, Decision::ReadTeamObject)] + fn key_before_team(#[case] config_nonempty: bool, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + config_nonempty, + team_id: Some("team".into()), + ..Facts::default() + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/route_access.rs b/litellm-rust/crates/request-auth-core/src/route_access.rs new file mode 100644 index 00000000000..6f090c91ad6 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_access.rs @@ -0,0 +1,567 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::{Error, patterns}; + +#[derive(Clone, Default, Deserialize)] +pub struct RouteCatalog { + pub groups: BTreeMap>, + pub viewer_blocked: Vec, + pub viewer_write_blocked: Vec, + pub self_credentials: Vec<(String, String)>, + pub managed_models: Vec, + pub managed_paths: Vec, + pub managed_realtime: Vec, + pub managed_live: Vec, + pub managed_mcp: Vec, + pub unbilled: Vec, +} + +impl RouteCatalog { + fn group(&self, name: &str) -> &[String] { + self.groups.get(name).map_or(&[], Vec::as_slice) + } + + fn access(&self, route: &str, name: &str) -> Result { + patterns::route_access(route, self.group(name)) + } + + fn contains(&self, route: &str, name: &str) -> bool { + self.group(name).iter().any(|entry| entry == route) + } +} + +#[derive(Clone, Default, Deserialize)] +pub struct RouteContext { + pub route: String, + pub method: Option, + pub catalog: RouteCatalog, + pub registered: bool, + pub registered_auth: Option, + pub configured_eu_jev: bool, +} + +#[derive(Clone, Default, Deserialize)] +pub struct PassThroughFacts { + pub key_allowed: Vec, + pub team_allowed: Vec, + pub denied: Vec, + pub forwarded: Vec, + pub jwt_team_routes: Vec, +} + +#[derive(Deserialize)] +pub struct VirtualRouteFacts { + pub context: RouteContext, + pub allowed_routes: Option>, + pub passthrough: PassThroughFacts, + pub password_reset_required: bool, +} + +#[derive(Deserialize)] +pub struct ViewerFacts { + pub context: RouteContext, + pub request_keys: Vec, +} + +#[derive(Deserialize)] +pub struct NonAdminFacts { + pub viewer: ViewerFacts, + pub role: Option, + pub custom_admin_routes: Vec, + pub premium: bool, + pub query_user_id: Option, + pub token_user_id: Option, + pub permissions: Vec, + pub organization_ids: Vec, + pub admin_organization_ids: Vec, + pub team_service_account: bool, + pub allowed_routes: Option>, + pub passthrough: PassThroughFacts, +} + +#[derive(Deserialize)] +pub struct MatchFacts { + pub route: String, + pub patterns: Vec, + pub mode: MatchMode, +} + +#[derive(Deserialize, strum::VariantArray)] +pub enum MatchMode { + Access, + Template, + Prefix, + Wildcard, + Denied, + Jwt, +} + +#[derive(Deserialize, strum::VariantArray)] +pub enum Classification { + Llm, + Unbilled, + Discovery, + SelfCredential, + Azure, + Managed, +} + +#[derive(Deserialize)] +pub enum Request { + Match(MatchFacts), + Classify { + context: RouteContext, + kind: Classification, + }, + Virtual(VirtualRouteFacts), + Viewer(ViewerFacts), + NonAdmin(NonAdminFacts), + PassThrough { + context: RouteContext, + facts: PassThroughFacts, + required: bool, + }, + DeniedPassThrough { + route: String, + facts: PassThroughFacts, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct Decision { + pub allowed: bool, + pub denial: Option, + pub detail: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Denial { + VirtualRoute, + PasswordReset, + PassThrough, + DeniedPassThrough, + CustomAdmin, + UserInfo, + AdminOnly, + ViewerInference, + ViewerWrite, + ViewerParameter, +} + +impl Decision { + fn allow() -> Self { + Self { + allowed: true, + denial: None, + detail: None, + } + } + fn predicate(value: bool) -> Self { + Self { + allowed: value, + denial: None, + detail: None, + } + } + fn deny(denial: Denial) -> Self { + Self { + allowed: false, + denial: Some(denial), + detail: None, + } + } + fn detail(denial: Denial, detail: String) -> Self { + Self { + allowed: false, + denial: Some(denial), + detail: Some(detail), + } + } +} + +fn denied_pattern(route: &str, pattern: &str) -> bool { + let normalized = pattern.trim_end_matches('/'); + normalized.is_empty() + || patterns::prefix(route, normalized) + || patterns::wildcard(route, pattern) +} + +fn jwt_grant(route: &str, entries: &[String]) -> bool { + entries.iter().any(|entry| { + !entry.trim_end_matches('*').trim_matches('/').is_empty() + && patterns::wildcard(route, entry) + }) +} + +fn passthrough_grant(route: &str, facts: &PassThroughFacts) -> bool { + let entries = if facts.key_allowed.is_empty() { + &facts.team_allowed + } else { + &facts.key_allowed + }; + entries.iter().any(|entry| patterns::prefix(route, entry)) +} + +fn denied_passthrough(facts: &PassThroughFacts) -> Option<&str> { + facts + .denied + .iter() + .find(|denied| { + facts + .forwarded + .iter() + .any(|route| denied_pattern(route, denied)) + }) + .map(String::as_str) +} + +fn require_passthrough(route: &str, facts: &PassThroughFacts) -> Decision { + if let Some(denied) = denied_passthrough(facts) { + return Decision::detail(Denial::DeniedPassThrough, denied.to_owned()); + } + if passthrough_grant(route, facts) || jwt_grant(route, &facts.jwt_team_routes) { + return Decision::allow(); + } + Decision::deny(Denial::PassThrough) +} + +fn azure(route: &str) -> Result { + Ok(patterns::regex_match( + "^/openai/deployments/[^/]+/[^/]+/chat/completions(?=\\n?\\z)", + route, + )? || patterns::regex_match("^/engines/[^/]+/chat/completions(?=\\n?\\z)", route)?) +} + +fn llm(context: &RouteContext) -> Result { + let route = &context.route; + let catalog = &context.catalog; + for group in ["openai_routes", "anthropic_routes", "google_routes"] { + if catalog.contains(route, group) { + return Ok(true); + } + } + for group in [ + "mcp_inference_routes", + "agent_inference_routes", + "litellm_native_routes", + ] { + if catalog.access(route, group)? { + return Ok(true); + } + } + for group in ["openai_routes", "google_routes", "anthropic_routes"] { + for pattern in catalog.group(group) { + if pattern.contains('{') && patterns::route_pattern(route, pattern)? { + return Ok(true); + } + if group == "openai_routes" + && pattern.ends_with('*') + && patterns::wildcard(route, pattern) + { + return Ok(true); + } + } + } + Ok(azure(route)? + || catalog + .group("mapped_pass_through_routes") + .iter() + .any(|pattern| patterns::prefix(route, pattern)) + || (route.trim_end_matches('/') == "/eu_jev/v1/systemone" && !context.configured_eu_jev)) +} + +fn discovery(context: &RouteContext) -> bool { + context.method.as_deref() == Some("GET") + && (context.route == "/v1/mcp/server" + || context + .route + .strip_prefix("/v1/mcp/server/") + .is_some_and(|tail| !tail.is_empty() && !tail.contains('/'))) +} + +fn self_credential(context: &RouteContext) -> Result { + for (method, pattern) in &context.catalog.self_credentials { + if context.method.as_ref() == Some(method) + && patterns::route_access(&context.route, std::slice::from_ref(pattern))? + { + return Ok(true); + } + } + Ok(false) +} + +fn managed(context: &RouteContext) -> Result { + let route = &context.route; + let method = context.method.as_deref(); + let catalog = &context.catalog; + if route == "/agents" || route == "/v1/agents" { + return Ok(matches!(method, None | Some("GET" | "HEAD"))); + } + if catalog.managed_realtime.contains(route) || catalog.managed_live.contains(route) { + return Ok(matches!(method, None | Some("GET"))); + } + if catalog.managed_models.contains(route) + || patterns::route_access(route, &catalog.managed_paths)? + { + return Ok(matches!(method, None | Some("POST"))); + } + Ok(patterns::route_access(route, &catalog.managed_mcp)? + || catalog.access(route, "agent_inference_routes")?) +} + +pub(crate) fn virtual_route(facts: &VirtualRouteFacts) -> Result { + let Some(allowed) = &facts.allowed_routes else { + return Ok(Decision::allow()); + }; + if allowed.is_empty() { + return Ok(Decision::allow()); + } + let context = &facts.context; + let route = &context.route; + if allowed + .iter() + .any(|pattern| patterns::prefix(route, pattern)) + { + return Ok(Decision::allow()); + } + let mut denied_passthrough = false; + for group in allowed { + if !context.catalog.groups.contains_key(group) { + continue; + } + let is_llm = matches!( + group.as_str(), + "llm_api_routes" | "llm_api_routes_without_mcp" + ); + if context.catalog.access(route, group)? { + if (is_llm || group == "openai_routes") && context.registered_auth == Some(true) { + if passthrough_grant(route, &facts.passthrough) { + return Ok(Decision::allow()); + } + denied_passthrough = true; + } else { + return Ok(Decision::allow()); + } + } + if !is_llm { + continue; + } + if route == "/auto_router/session" && context.method.as_deref() == Some("GET") { + return Ok(Decision::allow()); + } + if context.registered { + if context.registered_auth == Some(true) { + if passthrough_grant(route, &facts.passthrough) { + return Ok(Decision::allow()); + } + denied_passthrough = true; + } else { + return Ok(Decision::allow()); + } + } + if group == "llm_api_routes" + && (discovery(context) || context.catalog.access(route, "agent_management_routes")?) + { + return Ok(Decision::allow()); + } + } + if allowed + .iter() + .any(|pattern| patterns::wildcard(route, pattern)) + { + return Ok(Decision::allow()); + } + Ok(Decision::deny(if denied_passthrough { + Denial::PassThrough + } else if facts.password_reset_required { + Denial::PasswordReset + } else { + Denial::VirtualRoute + })) +} + +fn viewer(facts: &ViewerFacts) -> Result { + let context = &facts.context; + let route = &context.route; + let catalog = &context.catalog; + if llm(context)? { + return Ok(Decision::deny(Denial::ViewerInference)); + } + let invalid = facts + .request_keys + .iter() + .find(|key| key.as_str() != "user_email"); + if catalog.access(route, "management_routes")? { + if route == "/user/update" { + if let Some(key) = invalid { + return Ok(Decision::detail(Denial::ViewerParameter, key.clone())); + } + } else if patterns::route_access(route, &catalog.viewer_blocked)? + || (route.starts_with("/key/") + && (route.ends_with("/regenerate") || route.ends_with("/reset_spend"))) + { + return Ok(Decision::deny(Denial::ViewerWrite)); + } + } + if matches!( + context.method.as_deref().unwrap_or("GET"), + "GET" | "HEAD" | "OPTIONS" + ) { + return Ok(Decision::allow()); + } + if route == "/user/update" { + return Ok(invalid.map_or_else(Decision::allow, |key| { + Decision::detail(Denial::ViewerParameter, key.clone()) + })); + } + if route == "/user/password/change" || route == "/session/logout" || self_credential(context)? { + return Ok(Decision::allow()); + } + if patterns::route_access(route, &catalog.viewer_write_blocked)? + || (route.starts_with("/key/") && route.ends_with("/regenerate")) + { + return Ok(Decision::deny(Denial::ViewerWrite)); + } + if catalog.access(route, "admin_viewer_routes")? + || catalog.access(route, "global_spend_tracking_routes")? + { + return Ok(Decision::allow()); + } + Ok(Decision::deny(Denial::ViewerWrite)) +} + +fn non_admin(facts: &NonAdminFacts) -> Result { + let context = &facts.viewer.context; + let route = &context.route; + let catalog = &context.catalog; + let role = facts.role.as_deref(); + if facts.premium && facts.custom_admin_routes.contains(route) { + return Ok(Decision::deny(Denial::CustomAdmin)); + } + if context.registered_auth == Some(true) { + return Ok(require_passthrough(route, &facts.passthrough)); + } + if llm(context)? { + return Ok(Decision::allow()); + } + if catalog.contains(route, "info_routes") { + if route == "/user/info" + && facts + .query_user_id + .as_ref() + .is_some_and(|id| !id.is_empty()) + && facts.query_user_id != facts.token_user_id + && role != Some("proxy_admin_viewer") + { + return Ok(Decision::deny(Denial::UserInfo)); + } + return Ok(Decision::allow()); + } + if catalog.contains(route, "global_spend_tracking_routes") + && facts + .permissions + .iter() + .any(|entry| entry == "get_spend_routes") + { + return Ok(Decision::allow()); + } + if patterns::prefix(route, "/lens") { + return Ok(Decision::allow()); + } + if role == Some("proxy_admin_viewer") { + return viewer(&facts.viewer); + } + let org_admin = !facts.organization_ids.is_empty() + && facts + .organization_ids + .iter() + .all(|id| facts.admin_organization_ids.contains(id)); + if (role == Some("internal_user") + && (catalog.access(route, "internal_user_routes")? || self_credential(context)?)) + || (org_admin && catalog.access(route, "org_admin_allowed_routes")?) + || (role == Some("internal_user_viewer") + && catalog.access(route, "internal_user_view_only_routes")?) + || catalog.access(route, "self_managed_routes")? + || route.starts_with("/v1/mcp/") + || route.starts_with("/mcp-rest/") + || passthrough_grant(route, &facts.passthrough) + || (facts.team_service_account + && catalog.access(route, "team_service_account_key_routes")?) + { + return Ok(Decision::allow()); + } + if facts.allowed_routes.as_ref().is_some_and(|entries| { + entries + .iter() + .any(|entry| patterns::prefix(route, entry) || patterns::wildcard(route, entry)) + }) { + return Ok(Decision::allow()); + } + Ok(Decision::deny(Denial::AdminOnly)) +} + +pub fn decide(request: Request) -> Result { + match request { + Request::Virtual(facts) => virtual_route(&facts), + Request::Viewer(facts) => viewer(&facts), + Request::NonAdmin(facts) => non_admin(&facts), + Request::PassThrough { + context, + facts, + required, + } => Ok(if required { + require_passthrough(&context.route, &facts) + } else { + Decision::predicate(passthrough_grant(&context.route, &facts)) + }), + Request::DeniedPassThrough { route: _, facts } => Ok(denied_passthrough(&facts) + .map_or_else(Decision::allow, |entry| { + Decision::detail(Denial::DeniedPassThrough, entry.to_owned()) + })), + Request::Match(facts) => { + let matched = match facts.mode { + MatchMode::Access => patterns::route_access(&facts.route, &facts.patterns)?, + MatchMode::Jwt => jwt_grant(&facts.route, &facts.patterns), + MatchMode::Template => { + let mut found = false; + for pattern in &facts.patterns { + if patterns::route_pattern(&facts.route, pattern)? { + found = true; + break; + } + } + found + } + MatchMode::Prefix => facts + .patterns + .iter() + .any(|pattern| patterns::prefix(&facts.route, pattern)), + MatchMode::Wildcard => facts + .patterns + .iter() + .any(|pattern| patterns::wildcard(&facts.route, pattern)), + MatchMode::Denied => facts + .patterns + .iter() + .any(|pattern| denied_pattern(&facts.route, pattern)), + }; + Ok(Decision::predicate(matched)) + } + Request::Classify { context, kind } => Ok(Decision::predicate(match kind { + Classification::Llm => llm(&context)?, + Classification::Azure => azure(&context.route)?, + Classification::Discovery => discovery(&context), + Classification::SelfCredential => self_credential(&context)?, + Classification::Managed => managed(&context)?, + Classification::Unbilled => { + let tail = context.route.rsplit('/').next().unwrap_or(""); + let (resource, action) = tail.split_once(':').unwrap_or((tail, "")); + context.catalog.unbilled.contains(&context.route) + || matches!(resource, "count_tokens" | "count-tokens") + || action == "countTokens" + } + })), + } +} diff --git a/litellm-rust/crates/request-auth-core/src/route_enable.rs b/litellm-rust/crates/request-auth-core/src/route_enable.rs new file mode 100644 index 00000000000..4f4a42294d2 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_enable.rs @@ -0,0 +1,180 @@ +use crate::{ + Error, patterns, + route_access::{self, Classification, RouteContext}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum Feature { + Admin, + Llm, +} + +#[derive(Clone, Deserialize)] +pub struct Setting { + pub environment_present: bool, + pub premium: bool, + pub secret: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub context: RouteContext, + pub imported: bool, + pub failed: Option, + pub admin: Option, + pub llm: Option, + pub exemptions: Option>, + pub eu_jev_configured: Option, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum Failure { + Http, + Other, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Reason { + License, + Disabled, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + PropagateHttp, + ReadImport, + ReadExemptions, + ReadEuJev, + ReadSettings { feature: Feature }, + ReadSecret { feature: Feature }, + Deny { feature: Feature, reason: Reason }, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum LlmDecision { + Allow, + Deny, + ReadEuJev, +} + +pub fn llm_route(context: &RouteContext, configured: Option) -> Result { + let static_context = RouteContext { + configured_eu_jev: true, + ..context.clone() + }; + if route_access::decide(route_access::Request::Classify { + context: static_context, + kind: Classification::Llm, + })? + .allowed + { + return Ok(LlmDecision::Allow); + } + if context.route.trim_end_matches('/') != "/eu_jev/v1/systemone" { + return Ok(LlmDecision::Deny); + } + Ok(match configured { + None => LlmDecision::ReadEuJev, + Some(true) => LlmDecision::Deny, + Some(false) => LlmDecision::Allow, + }) +} + +fn setting(setting: Option<&Setting>, feature: Feature) -> Option { + let Some(setting) = setting else { + return Some(Decision::ReadSettings { feature }); + }; + if setting.environment_present && !setting.premium { + return Some(Decision::Deny { + feature, + reason: Reason::License, + }); + } + match setting.secret { + None => Some(Decision::ReadSecret { feature }), + Some(true) => Some(Decision::Deny { + feature, + reason: Reason::Disabled, + }), + Some(false) => None, + } +} + +pub fn decide(facts: &Facts) -> Result { + if let Some(failed) = facts.failed { + return Ok(match failed { + Failure::Http => Decision::PropagateHttp, + Failure::Other => Decision::Allow, + }); + } + if !facts.imported { + return Ok(Decision::ReadImport); + } + let management = facts + .context + .catalog + .groups + .get("management_routes") + .map(Vec::as_slice) + .unwrap_or_default(); + if patterns::route_access(&facts.context.route, management)? + && let Some(decision) = setting(facts.admin.as_ref(), Feature::Admin) + { + return Ok(decision); + } + match llm_route(&facts.context, facts.eu_jev_configured)? { + LlmDecision::Deny => return Ok(Decision::Allow), + LlmDecision::ReadEuJev => return Ok(Decision::ReadEuJev), + LlmDecision::Allow => (), + } + let Some(exemptions) = &facts.exemptions else { + return Ok(Decision::ReadExemptions); + }; + if exemptions.contains(&facts.context.route) { + return Ok(Decision::Allow); + } + Ok(setting(facts.llm.as_ref(), Feature::Llm).unwrap_or(Decision::Allow)) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(None, LlmDecision::ReadEuJev)] + #[case::builtin(Some(false), LlmDecision::Allow)] + #[case::configured(Some(true), LlmDecision::Deny)] + fn builtin_configuration(#[case] configured: Option, #[case] expected: LlmDecision) { + let context = RouteContext { + route: "/eu_jev/v1/systemone/".into(), + ..RouteContext::default() + }; + assert_eq!(llm_route(&context, configured).unwrap(), expected); + } + + #[rstest] + #[case::license(true, false, Decision::Deny { feature: Feature::Admin, reason: Reason::License })] + #[case::absent(false, false, Decision::ReadSecret { feature: Feature::Admin })] + #[case::premium(true, true, Decision::ReadSecret { feature: Feature::Admin })] + fn license_precedes_secret( + #[case] environment_present: bool, + #[case] premium: bool, + #[case] expected: Decision, + ) { + assert_eq!( + setting( + Some(&Setting { + environment_present, + premium, + secret: None + }), + Feature::Admin + ), + Some(expected) + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/route_roles.rs b/litellm-rust/crates/request-auth-core/src/route_roles.rs new file mode 100644 index 00000000000..059fa47075b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_roles.rs @@ -0,0 +1,325 @@ +use serde::{Deserialize, Serialize}; + +use crate::{Error, patterns}; + +#[derive( + Clone, Copy, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr, strum::VariantArray, +)] +enum Role { + #[strum(serialize = "proxy_admin")] + Admin, + #[strum(serialize = "proxy_admin_viewer")] + Viewer, + #[strum(serialize = "org_admin")] + Organization, + #[strum(serialize = "internal_user")] + Internal, + #[strum(serialize = "internal_user_viewer")] + InternalViewer, + #[strum(serialize = "team")] + Team, + #[strum(serialize = "customer")] + Customer, +} + +#[derive(Deserialize)] +pub struct RoleFacts { + pub user_present: bool, + pub role: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct RoleDecision { + pub effective: &'static str, + pub user_role: Option<&'static str>, + pub admin: bool, +} + +pub fn role(facts: &RoleFacts) -> RoleDecision { + let parsed = facts + .role + .as_deref() + .and_then(|value| value.parse::().ok()); + let effective = parsed.unwrap_or(Role::Internal).into(); + RoleDecision { + effective, + user_role: facts.user_present.then_some(effective), + admin: facts.user_present && parsed == Some(Role::Admin), + } +} + +#[derive(Deserialize)] +pub struct ApiFacts { + pub token_present: bool, + pub admin_read: bool, + pub user: RoleFacts, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum ApiDecision { + InvalidToken, + ReadAdmin, + NonAdmin, + Allow, +} + +pub fn api(facts: &ApiFacts) -> ApiDecision { + if !facts.token_present { + ApiDecision::InvalidToken + } else if !facts.admin_read { + ApiDecision::ReadAdmin + } else if role(&facts.user).admin { + ApiDecision::Allow + } else { + ApiDecision::NonAdmin + } +} + +#[derive(Deserialize)] +pub struct JwtFacts { + pub role: Option, + pub config_read: bool, + pub config_missing: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum JwtDecision { + ReadAdmin, + ReadTeam, + DefaultTeam, + Check, + Deny, +} + +pub fn jwt(facts: &JwtFacts) -> JwtDecision { + match facts.role.as_deref() { + Some("proxy_admin") if !facts.config_read => JwtDecision::ReadAdmin, + Some("team") if !facts.config_read => JwtDecision::ReadTeam, + Some("team") if facts.config_missing => JwtDecision::DefaultTeam, + Some("proxy_admin" | "team") => JwtDecision::Check, + _ => JwtDecision::Deny, + } +} + +#[derive(Deserialize)] +pub struct InsideFacts { + pub role: Option, + pub user_id: Option, + pub requested_user_id: Option, +} + +pub fn inside(facts: &InsideFacts) -> bool { + matches!( + facts.role.as_deref(), + Some("proxy_admin" | "proxy_admin_viewer") + ) || facts.user_id.is_some() && facts.user_id == facts.requested_user_id +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Entry { + pub text: Option, + pub python_type: String, + pub hash_failure: Option, + pub templates: Option>, +} + +#[derive(Deserialize)] +pub struct ListedFacts { + pub route: String, + pub entry: Option, + pub exhausted: bool, + pub template_index: usize, + pub compiled: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Action { + ReadEntry, + ReadTemplate, + Allow, + Deny, + Keep, + Expand, + Failure, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum FailureKind { + TypeError, + AttributeError, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct Failure { + pub kind: FailureKind, + pub message: String, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct EntryDecision { + pub action: Action, + pub template_index: usize, + pub failure: Option, +} + +impl EntryDecision { + fn action(action: Action) -> Self { + Self { + action, + template_index: 0, + failure: None, + } + } + + fn failure(kind: FailureKind, message: String) -> Self { + Self { + action: Action::Failure, + template_index: 0, + failure: Some(Failure { kind, message }), + } + } +} + +pub fn expand(entry: &Entry) -> EntryDecision { + if let Some(message) = &entry.hash_failure { + EntryDecision::failure(FailureKind::TypeError, message.clone()) + } else if entry.templates.is_some() { + EntryDecision::action(Action::Expand) + } else { + EntryDecision::action(Action::Keep) + } +} + +pub fn listed(facts: &ListedFacts) -> Result { + if facts.exhausted { + return Ok(EntryDecision::action(Action::Deny)); + } + let Some(entry) = &facts.entry else { + return Ok(EntryDecision::action(Action::ReadEntry)); + }; + if let Some(message) = &entry.hash_failure { + return Ok(EntryDecision::failure( + FailureKind::TypeError, + message.clone(), + )); + } + let Some(text) = &entry.text else { + return Ok(EntryDecision::failure( + FailureKind::AttributeError, + format!("'{}' object has no attribute 'endswith'", entry.python_type), + )); + }; + let Some(templates) = &entry.templates else { + return Ok(EntryDecision::action( + if patterns::wildcard(&facts.route, text) { + Action::Allow + } else { + Action::ReadEntry + }, + )); + }; + let next = if let Some(compiled) = &facts.compiled { + let expression = compiled + .strip_suffix('$') + .map_or_else(|| compiled.clone(), |prefix| format!("{prefix}(?=\\n?\\z)")); + if patterns::regex_match(&expression, &facts.route)? { + return Ok(EntryDecision::action(Action::Allow)); + } + facts.template_index + 1 + } else { + facts.template_index + }; + Ok(if next < templates.len() { + EntryDecision { + action: Action::ReadTemplate, + template_index: next, + failure: None, + } + } else { + EntryDecision::action(Action::ReadEntry) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(false, None, None, false)] + #[case::default(true, None, Some("internal_user"), false)] + #[case::unknown(true, Some("ADMIN"), Some("internal_user"), false)] + #[case::admin(true, Some("proxy_admin"), Some("proxy_admin"), true)] + #[case::viewer(true, Some("proxy_admin_viewer"), Some("proxy_admin_viewer"), false)] + fn roles( + #[case] present: bool, + #[case] raw: Option<&str>, + #[case] expected: Option<&str>, + #[case] admin: bool, + ) { + let result = role(&RoleFacts { + user_present: present, + role: raw.map(str::to_owned), + }); + assert_eq!(result.user_role, expected); + assert_eq!(result.admin, admin); + } + + #[rstest] + #[case::invalid(false, false, ApiDecision::InvalidToken)] + #[case::read(true, false, ApiDecision::ReadAdmin)] + #[case::nonadmin(true, true, ApiDecision::NonAdmin)] + fn api_order(#[case] token: bool, #[case] read: bool, #[case] expected: ApiDecision) { + assert_eq!( + api(&ApiFacts { + token_present: token, + admin_read: read, + user: RoleFacts { + user_present: false, + role: None + } + }), + expected + ); + } + + #[rstest] + #[case::read("team", false, false, JwtDecision::ReadTeam)] + #[case::default("team", true, true, JwtDecision::DefaultTeam)] + #[case::empty("team", true, false, JwtDecision::Check)] + #[case::other("internal_user", false, false, JwtDecision::Deny)] + fn jwt_selection( + #[case] role: &str, + #[case] read: bool, + #[case] missing: bool, + #[case] expected: JwtDecision, + ) { + assert_eq!( + jwt(&JwtFacts { + role: Some(role.into()), + config_read: read, + config_missing: missing + }), + expected + ); + } + + #[rstest] + #[case::self_empty(Some(""), Some(""), true)] + #[case::absent(None, None, false)] + #[case::different(Some("a"), Some("b"), false)] + fn inside_identity( + #[case] user: Option<&str>, + #[case] requested: Option<&str>, + #[case] expected: bool, + ) { + assert_eq!( + inside(&InsideFacts { + role: None, + user_id: user.map(str::to_owned), + requested_user_id: requested.map(str::to_owned) + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/route_scope.rs b/litellm-rust/crates/request-auth-core/src/route_scope.rs new file mode 100644 index 00000000000..2275dbc4881 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_scope.rs @@ -0,0 +1,97 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Proxy, + Admin, +} + +#[derive(Clone, Deserialize)] +pub struct Values { + pub text: Option, + pub names: Vec, + pub failure: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub kind: Kind, + pub route: String, + pub present: bool, + pub premium: bool, + pub warned: bool, + pub values: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny, + Warn, + ReadValues, + Failure { message: String }, +} + +pub fn decide(facts: &Facts) -> Decision { + if !facts.present { + return Decision::Allow; + } + if !facts.premium { + if !facts.warned { + return Decision::Warn; + } + if facts.kind == Kind::Admin { + return Decision::Allow; + } + } + let Some(values) = &facts.values else { + return Decision::ReadValues; + }; + if let Some(message) = &values.failure { + return Decision::Failure { + message: message.clone(), + }; + } + let member = values.text.as_ref().map_or_else( + || values.names.contains(&facts.route), + |text| text.contains(&facts.route), + ); + if member == (facts.kind == Kind::Proxy) { + Decision::Allow + } else { + Decision::Deny + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::proxy_warn(Kind::Proxy, false, false, Decision::Warn)] + #[case::admin_warn(Kind::Admin, false, false, Decision::Warn)] + #[case::proxy_after_warn(Kind::Proxy, false, true, Decision::ReadValues)] + #[case::admin_after_warn(Kind::Admin, false, true, Decision::Allow)] + #[case::licensed_proxy(Kind::Proxy, true, false, Decision::ReadValues)] + #[case::licensed_admin(Kind::Admin, true, false, Decision::ReadValues)] + fn licensing_retains_distinct_warning_behavior( + #[case] kind: Kind, + #[case] premium: bool, + #[case] warned: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(&Facts { + kind, + route: "/responses".into(), + present: true, + premium, + warned, + values: None + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/spend_counter_policy.rs b/litellm-rust/crates/request-auth-core/src/spend_counter_policy.rs new file mode 100644 index 00000000000..12213f6b87b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/spend_counter_policy.rs @@ -0,0 +1,299 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Spend { + current: f64, + verified: bool, + fallback: f64, + limit: Option, + window: bool, + fallback_authoritative: bool, + strict_before: Option, + authoritative_read: bool, + authoritative: Option, + strict_after: Option, + }, + ThrottleKey { + counter_key: String, + token_present: bool, + token: Option, + }, + Overflow { + current: f64, + limit: f64, + }, + OverBudget { + throttled: bool, + current: f64, + limit: f64, + reservation: f64, + strict: bool, + }, + ReservationFailure { + strict: bool, + }, + ReservationGroups { + currents: Vec, + limits: Vec, + reservation: f64, + }, + Adjustment { + current: Option, + adjustment: f64, + }, + ResizeConsistency { + inconsistent: bool, + reseed: bool, + }, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Value { spend: f64 }, + Repair { spend: f64 }, + ReadStrictBefore, + ReadStrictAfter, + ReadFloor, + Unverifiable, + ReadThrottle, + Skip, + Release, + Reject { current: f64, estimate: Option }, + Resize { cost: f64 }, + CheckOverflow, + Unavailable, + Compatible, + Inconsistent, + RejectResize, + ReservationGroups { groups: Vec> }, +} + +pub fn decide(facts: Facts) -> Decision { + match facts { + Facts::Adjustment { + current, + adjustment, + } => { + if current.is_some_and(|current| !(adjustment < 0.0 && current + adjustment < -1e-12)) { + Decision::Compatible + } else { + Decision::Inconsistent + } + } + Facts::ResizeConsistency { + inconsistent, + reseed, + } => { + if inconsistent && !reseed { + Decision::RejectResize + } else { + Decision::Compatible + } + } + Facts::Spend { + current, + verified, + fallback, + limit, + window, + fallback_authoritative, + strict_before, + authoritative_read, + authoritative, + strict_after, + } => { + if limit.is_none_or(|limit| current >= limit) { + return Decision::Value { spend: current }; + } + let obvious_floor = fallback > current || window; + if !obvious_floor && strict_before.is_none() { + return Decision::ReadStrictBefore; + } + if obvious_floor || strict_before == Some(true) { + if !authoritative_read { + return Decision::ReadFloor; + } + if let Some(spend) = authoritative { + if spend > current { + return Decision::Repair { spend }; + } + } else if fallback > current { + return Decision::Value { spend: fallback }; + } + } + if !(verified || fallback_authoritative || authoritative.is_some()) { + return match strict_after { + None => Decision::ReadStrictAfter, + Some(true) => Decision::Unverifiable, + Some(false) => Decision::Value { spend: current }, + }; + } + Decision::Value { spend: current } + } + Facts::ThrottleKey { + counter_key, + token_present, + token, + } => { + if token_present + && counter_key == format!("spend:key:{}", token.as_deref().unwrap_or("None")) + { + Decision::ReadThrottle + } else { + Decision::Skip + } + } + Facts::Overflow { current, limit } => { + if current > limit { + Decision::CheckOverflow + } else { + Decision::Skip + } + } + Facts::OverBudget { + throttled, + current, + limit, + reservation, + strict, + } => { + if throttled { + return Decision::Release; + } + let remaining = limit - (current - reservation); + if remaining <= 1e-12 { + return Decision::Reject { + current, + estimate: None, + }; + } + if strict && current - limit > 1e-12 { + return Decision::Reject { + current: current - reservation, + estimate: Some(reservation), + }; + } + Decision::Resize { cost: remaining } + } + Facts::ReservationFailure { strict } => { + if strict { + Decision::Unavailable + } else { + Decision::Skip + } + } + Facts::ReservationGroups { + currents, + limits, + reservation, + } => { + let indices: Vec = (0..currents.len()).collect(); + let groups = if currents.is_empty() { + vec![] + } else if currents + .iter() + .zip(&limits) + .all(|(current, limit)| current + reservation <= *limit) + { + vec![indices] + } else { + indices.into_iter().map(|index| vec![index]).collect() + }; + Decision::ReservationGroups { groups } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(None, 1.0, Decision::Inconsistent)] + #[case::positive(Some(1.0), 1.0, Decision::Compatible)] + #[case::negative(Some(1.0), -2.0, Decision::Inconsistent)] + #[case::zero(Some(1.0), -1.0, Decision::Compatible)] + fn counter_adjustment( + #[case] current: Option, + #[case] adjustment: f64, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::Adjustment { + current, + adjustment + }), + expected + ); + } + + #[rstest] + #[case::resize(true, false, Decision::RejectResize)] + #[case::settlement(true, true, Decision::Compatible)] + #[case::consistent(false, false, Decision::Compatible)] + fn consistency(#[case] inconsistent: bool, #[case] reseed: bool, #[case] expected: Decision) { + assert_eq!( + decide(Facts::ResizeConsistency { + inconsistent, + reseed + }), + expected + ); + } + + #[rstest] + #[case::throttled(true, 12.0, true, Decision::Release)] + #[case::empty(false, 12.0, false, Decision::Reject { current:12.0, estimate:None })] + #[case::partial(false, 11.0, false, Decision::Resize { cost:1.0 })] + #[case::strict(false, 11.0, true, Decision::Reject { current:9.0, estimate:Some(2.0) })] + fn reservation_outcome( + #[case] throttled: bool, + #[case] current: f64, + #[case] strict: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::OverBudget { + throttled, + current, + limit: 10.0, + reservation: 2.0, + strict + }), + expected + ); + } + + #[rstest] + #[case::unchecked(None, false, None, Decision::Value { spend:2.0 })] + #[case::strict_read(Some(3.0), false, None, Decision::ReadStrictBefore)] + #[case::window(Some(3.0), true, None, Decision::ReadFloor)] + #[case::strict_floor(Some(3.0), false, Some(true), Decision::ReadFloor)] + #[case::second_read(Some(3.0), false, Some(false), Decision::ReadStrictAfter)] + fn spend_read_order( + #[case] limit: Option, + #[case] window: bool, + #[case] strict_before: Option, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::Spend { + current: 2.0, + verified: false, + fallback: 0.0, + limit, + window, + fallback_authoritative: false, + strict_before, + authoritative_read: false, + authoritative: None, + strict_after: None + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/team_aliases.rs b/litellm-rust/crates/request-auth-core/src/team_aliases.rs new file mode 100644 index 00000000000..6abd1e477b8 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/team_aliases.rs @@ -0,0 +1,77 @@ +use crate::Error; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub aliases: BTreeMap, + pub team_id: Option, + pub router_present: bool, + pub names: Option>>, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadNames { count: usize }, + Complete { models: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let prefix = format!("model_name_{}_", facts.team_id.as_deref().unwrap_or("None")); + let needs_names = |target: &str| facts.router_present && target.starts_with(&prefix); + let count = facts + .requested + .iter() + .filter_map(|model| facts.aliases.get(model)) + .filter(|target| needs_names(target)) + .count(); + let names = match &facts.names { + None if count != 0 => return Ok(Decision::ReadNames { count }), + Some(names) if names.len() != count => return Err(Error::InvalidAnswer), + names => names.as_deref().unwrap_or_default(), + }; + let mut reads = names.iter(); + let models = facts + .requested + .iter() + .map(|model| { + let Some(target) = facts.aliases.get(model) else { + return model.clone(); + }; + if needs_names(target) && !reads.next().is_some_and(|names| names.contains(target)) { + model.clone() + } else { + target.clone() + } + }) + .collect(); + Ok(Decision::Complete { models }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(false, "alias")] + #[case::live(true, "model_name_team_target")] + fn deleted_target(#[case] live: bool, #[case] expected: &str) { + let target = "model_name_team_target".to_owned(); + let facts = Facts { + requested: vec!["alias".into()], + aliases: BTreeMap::from([("alias".into(), target.clone())]), + team_id: Some("team".into()), + router_present: true, + names: Some(vec![if live { vec![target] } else { vec![] }]), + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::Complete { + models: vec![expected.into()] + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/team_override.rs b/litellm-rust/crates/request-auth-core/src/team_override.rs new file mode 100644 index 00000000000..08e0cc20481 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/team_override.rs @@ -0,0 +1,70 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + model_access::{self, Decision as ModelDecision, Facts as ModelFacts, ObjectType}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub team: Option, + pub assigned_models: Option>, + pub assigned_scope: Option, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { denial: ModelDecision }, + ReadTeam, + ReadTeamGroups, + ReadAssignedGroups, + ReadAssignedScope, +} + +pub fn decide(facts: &Facts) -> Result { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeam); + }; + let denial = fallback_access::grouped(team.clone(), true)?; + match denial { + ModelDecision::Allow => return Ok(Decision::Allow), + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } => { + return Ok(Decision::ReadTeamGroups); + } + ModelDecision::Deny { + object_type: ObjectType::Team, + .. + } => (), + _ => return Ok(Decision::Deny { denial }), + } + let Some(assigned) = &facts.assigned_models else { + return Ok(Decision::ReadAssignedGroups); + }; + if assigned.is_empty() { + return Ok(Decision::Deny { denial }); + } + let Some(scope) = &facts.assigned_scope else { + return Ok(Decision::ReadAssignedScope); + }; + match model_access::check(scope)? { + ModelDecision::Allow => Ok(Decision::Allow), + ModelDecision::Deny { .. } => Ok(Decision::Deny { denial }), + error => Ok(Decision::Deny { denial: error }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn team_is_read_before_override_groups() { + assert!(matches!( + decide(&Facts::default()).unwrap(), + Decision::ReadTeam + )); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/unpriced_models.rs b/litellm-rust/crates/request-auth-core/src/unpriced_models.rs new file mode 100644 index 00000000000..bc6073555c3 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/unpriced_models.rs @@ -0,0 +1,144 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Metric { + pub name: String, + pub scalar: Option, + pub nested: Vec>, +} + +#[derive(Clone, Deserialize)] +pub struct Entry { + pub tiered_present: bool, + pub metrics: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Deployment { + pub parameter_keys: Vec, + pub has_id: bool, + pub info_loaded: bool, + pub info: Option, + pub cost_keys: Option>, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub model: Option, + pub router_present: bool, + pub group_found: Option, + pub deployments: Option>, + pub cost_deployments: Option>, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { unpriced: bool }, + ReadGroup, + ReadDeployments, + ReadCostDeployments, + ReadInfo { index: usize }, + ReadCost { index: usize }, +} + +pub fn entry_has_priced_metric(entry: &Entry) -> bool { + entry.tiered_present + || entry.metrics.iter().any(|metric| { + metric.name.contains("cost_per") + && (metric.scalar.is_some_and(|price| price > 0.0) + || metric + .nested + .iter() + .any(|price| price.is_some_and(|value| value > 0.0))) + }) +} + +pub fn decide(facts: &Facts) -> Decision { + if !facts.router_present || facts.model.as_deref().is_none_or(str::is_empty) { + return Decision::Complete { unpriced: false }; + } + match facts.group_found { + None => return Decision::ReadGroup, + Some(false) => return Decision::Complete { unpriced: false }, + Some(true) => {} + } + let Some(deployments) = &facts.deployments else { + return Decision::ReadDeployments; + }; + for (index, deployment) in deployments.iter().enumerate() { + if deployment + .parameter_keys + .iter() + .any(|key| key.contains("cost_per") || key == "tiered_pricing") + { + return Decision::Complete { unpriced: false }; + } + if !deployment.has_id { + continue; + } + if !deployment.info_loaded { + return Decision::ReadInfo { index }; + } + if deployment + .info + .as_ref() + .is_some_and(entry_has_priced_metric) + { + return Decision::Complete { unpriced: false }; + } + } + let Some(cost_deployments) = &facts.cost_deployments else { + return Decision::ReadCostDeployments; + }; + for (index, deployment) in cost_deployments.iter().enumerate() { + if !deployment.has_id { + continue; + } + let Some(keys) = &deployment.cost_keys else { + return Decision::ReadCost { index }; + }; + if keys.iter().any(|key| { + matches!( + key.as_str(), + "input_cost_per_token" | "output_cost_per_token" + ) + }) { + return Decision::Complete { unpriced: false }; + } + } + Decision::Complete { unpriced: true } +} + +#[cfg(test)] +mod tests { + use super::{Entry, Metric, entry_has_priced_metric}; + use rstest::rstest; + + #[rstest] + #[case::positive("input_cost_per_character", Some(0.1), vec![], false, true)] + #[case::zero("input_cost_per_token", Some(0.0), vec![], false, false)] + #[case::negative("cost_per_query", Some(-1.0), vec![], false, false)] + #[case::nested("cost_per_image", None, vec![None, Some(0.2)], false, true)] + #[case::non_price("context_window", Some(1000.0), vec![], false, false)] + #[case::tiered("context_window", None, vec![], true, true)] + fn metrics( + #[case] name: &str, + #[case] scalar: Option, + #[case] nested: Vec>, + #[case] tiered_present: bool, + #[case] expected: bool, + ) { + assert_eq!( + entry_has_priced_metric(&Entry { + tiered_present, + metrics: vec![Metric { + name: name.to_owned(), + scalar, + nested + }] + }), + expected + ); + } +} diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md new file mode 100644 index 00000000000..0c78570ae64 --- /dev/null +++ b/litellm/_auth/AGENTS.md @@ -0,0 +1,18 @@ +# litellm/_auth + +## Why this exists + +- Lots of Python code like `litellm/proxy/auth/user_api_key_auth.py` is too complex +- Soon, Rust will implement the same auth stack, and expose common functionality as Python API through pyo3, to avoid maintaining 2 copy of code. +- So the goal of `_auth` here is to define that interface first, not writing really good python code. Should be behavior preserving refactor though + +## The swap point + +- `decision.py` is the interface. `AuthCore.start(RequestFacts, Settings, now) -> Step` is the only thing a core implements, and `python_core.PythonCore` is the Python one. A Rust core will implement the same protocol through pyo3, and the builder picks one without changing anything else +- A core does no IO and reads no globals. The same facts, settings, time and effect answers always give the same outcome +- A `Step` is either an `Outcome` or an effect. Each effect carries its typed `resume` continuation, and `driver.drive` answers effects through an `EffectHandler` until an outcome comes back. The proxy's handler (`_BuilderEffects` in `litellm/proxy/auth/user_api_key_auth.py`) owns all cache and DB IO, so both cores share it +- Projection stays in the proxy. `_authenticated` in the builder module turns an `Allow` grant into `UserAPIKeyAuth`, and `litellm/proxy/auth/auth_denial.py` maps every `Denial` to the existing error contract +- The key checks run in the core, in the order the proxy has always applied them: blocked key, model access, SCIM-deactivated owner, budget exemption, team member budget, expiry, budget alerts, the key's own max budget (or its throttle), the soft budget alert, per-model budgets, the team and project context, the proxy budget alert, then the key's token. A budget-exempt request skips everything from the team member budget to per-model budgets except expiry. `start_resolved_key` runs the same checks for keys the proxy resolved itself, such as internal jobs +- `model_budgets.py` decides the key and user per-model budgets: which models to check, when an over-budget key may fall back (only for the top-level request model, to the first `budget_fallbacks` entry it can afford and may call), and that the user check sees the rerouted model. The key chain runs it inline, custom auth runs it through `AuthCore.check_model_budgets`, and JWT runs the user half through `AuthCore.check_user_model_budgets`. Those two entry points end in `Allow(WithinModelBudgets())` and are driven by `driver.drive_model_budgets`, which only needs a `ModelBudgetHandler` +- The team member per-model budget runs through `AuthCore.check_team_member_model_budgets`, which has its own step type (`TeamMemberStep`) and driver (`drive_team_member_model_budgets`). The core decides when an over-budget member may fall back and walks the router's fallback targets in order, skipping the requested model. The proxy handler (`_TeamMemberModelBudgetEffects`) reads those targets and vets each one (pricing, customer access, every per-model budget, model access and access group budgets) +- The team and end-user per-model budgets still run in the proxy after the core, as plain loops with no fallback diff --git a/litellm/_auth/__init__.py b/litellm/_auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/_auth/access.py b/litellm/_auth/access.py new file mode 100644 index 00000000000..3e9ad03d9a5 --- /dev/null +++ b/litellm/_auth/access.py @@ -0,0 +1,248 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Blocked: + blocked: bool + + +@dataclass(frozen=True, slots=True) +class UserParamRequest: + enabled: bool + + +@dataclass(frozen=True, slots=True) +class ClientTagRequest: + enabled: bool + llm_route: bool | None = None + + +@dataclass(frozen=True, slots=True) +class GuardrailRequest: + requested_keys: tuple[str, ...] + + +TraceSource: TypeAlias = Literal["Gate", "Header", "Body", "Traceparent", "Selected", "Other"] + + +@dataclass(frozen=True, slots=True) +class TraceRequest: + required: bool + + +@dataclass(frozen=True, slots=True) +class Search: + name: str + allowed: tuple[str, ...] + strict: bool + + +@dataclass(frozen=True, slots=True) +class TraceGate: + required: bool + method: str + llm_route: bool + passthrough: bool + telemetry: bool + + +@dataclass(frozen=True, slots=True) +class TraceStep: + source: TraceSource + value: str | None + present: bool + headers_only: bool + + +@dataclass(frozen=True, slots=True) +class UserParam: + enabled: bool + method: str | None + llm_route: bool + mcp_route: bool + has_user: bool + + +@dataclass(frozen=True, slots=True) +class ClientTags: + enabled: bool + llm_route: bool + has_tags: bool + + +@dataclass(frozen=True, slots=True) +class Guardrails: + requested_keys: tuple[str, ...] + modify_guardrails: bool | None + + +@dataclass(frozen=True, slots=True) +class AgentTrace: + required: bool + header: str | None + + +@dataclass(frozen=True, slots=True) +class ConcurrentFailures: + failures: tuple[bool, ...] + terminal: tuple[bool, ...] + + +@dataclass(frozen=True, slots=True) +class TeamTrace: + required: bool + method: str + llm_route: bool + passthrough: bool + telemetry: bool + headers_only: bool + header: str | None + body_trace: str | None + body_trace_present: bool + traceparent_id: str | None + selected_trace: str | None + selected_trace_present: bool + other_trace: str | None + + +@dataclass(frozen=True, slots=True) +class ToolRequest: + token_present: bool + call_types: tuple[str, ...] | None = None + capable_types: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class ToolAllowlist: + key_nonempty_list: bool + team_nonempty_list: bool | None = None + requested_count: int | None = None + + +@dataclass(frozen=True, slots=True) +class Tools: + requested: tuple[str, ...] + key_allowed: tuple[str, ...] | None + team_allowed: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class GrantLayers: + deny_by_default: bool + identity_present: bool + master_key: bool + dashboard: bool + virtual_key: bool + session_token: bool + team_present: bool + + +@dataclass(frozen=True, slots=True) +class VectorRequest: + client_present: bool + strict: bool + ids_read: bool = False + unreadable: bool = False + grants: tuple[tuple[str, ...] | None, ...] | None = None + requested_count: int | None = None + + +@dataclass(frozen=True, slots=True) +class VectorStores: + requested: tuple[str, ...] + allowed: tuple[str, ...] | None + strict: bool + + +@dataclass(frozen=True, slots=True) +class TeamLookup: + dashboard: bool + missing: bool + recorded_models: tuple[str, ...] + allow_unavailable: bool | None + + +AccessFacts: TypeAlias = ( + ConcurrentFailures + | UserParamRequest + | ClientTagRequest + | GuardrailRequest + | TraceRequest + | Search + | Blocked + | TraceGate + | TraceStep + | UserParam + | ClientTags + | Guardrails + | AgentTrace + | TeamTrace + | ToolRequest + | ToolAllowlist + | Tools + | GrantLayers + | VectorRequest + | VectorStores + | TeamLookup +) +AccessDenial: TypeAlias = Literal[ + "Search", + "Blocked", + "UserParam", + "ClientTags", + "Guardrails", + "AgentTrace", + "TeamTrace", + "Tools", + "VectorStores", + "MissingVectorGrant", + "MalformedVectorIds", + "TeamLookup", +] +GrantLayer: TypeAlias = Literal["key", "team", "user"] + + +@dataclass(frozen=True, slots=True) +class AccessDecision: + read: ( + Literal[ + "UserParam", + "ClientTagRoute", + "ClientTags", + "GuardrailPermission", + "UnavailablePolicy", + "ToolRoute", + "ToolAllowlist", + "TeamToolList", + "ToolNames", + "KeyToolGrant", + "TeamToolGrant", + "StrictVectorIds", + "LegacyVectorIds", + "LegacyVectorGrants", + "RagVectorIds", + "VectorGrants", + ] + | None + ) = None + denial: AccessDenial | None = None + rejected: tuple[str, ...] = () + layers: tuple[GrantLayer, ...] = () + rejected_index: int | None = None + next_trace: TraceSource | None = None + failure_index: int | None = None + + +_DECISION: Final = TypeAdapter(AccessDecision) + + +@native_boundary +def decide(facts: AccessFacts) -> AccessDecision: + return _DECISION.validate_python(_native.request_auth_access(encode_facts(facts))) diff --git a/litellm/_auth/admission.py b/litellm/_auth/admission.py new file mode 100644 index 00000000000..691cb6a37ff --- /dev/null +++ b/litellm/_auth/admission.py @@ -0,0 +1,77 @@ +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class CommonGate: + force_virtual_key_checks: bool + public_route: bool + additional_public_route: bool + unauthenticated_passthrough: bool + master_key_present: bool + jwt_enabled: bool + oauth_enabled: bool + oauth_proxy_enabled: bool + custom_auth: bool + custom_auth_common_checks: bool + + +@dataclass(frozen=True, slots=True) +class Reservation: + skip_budget_checks: bool + disabled: bool + apply_user_budget_to_team_keys: bool + fail_closed: bool + billing_agent: bool + + +@dataclass(frozen=True, slots=True) +class Ownership: + virtual_key: bool + custom_auth: bool + jwt: bool + oauth: bool + + +@dataclass(frozen=True, slots=True) +class CommonPlan: + skip: Literal["PublicRoute", "UnauthenticatedPassthrough", "NoAuthDevelopment", "CustomAuth"] | None + + +@dataclass(frozen=True, slots=True) +class ReservationPlan: + reserve: bool + apply_user_budget_to_team_keys: bool + fail_closed: bool + + +@dataclass(frozen=True, slots=True) +class OwnershipPlan: + owner: Literal["Python"] + candidate: Literal["VirtualKey", "CustomAuth", "Validator", "Other"] + + +_COMMON: Final = TypeAdapter(CommonPlan) +_RESERVATION: Final = TypeAdapter(ReservationPlan) +_OWNERSHIP: Final = TypeAdapter(OwnershipPlan) + + +@native_boundary +def common_plan(facts: CommonGate) -> CommonPlan: + return _COMMON.validate_python(_native.request_auth_admission(encode_facts(facts))) + + +@native_boundary +def reservation_plan(facts: Reservation) -> ReservationPlan: + return _RESERVATION.validate_python(_native.request_auth_admission(encode_facts(facts))) + + +@native_boundary +def ownership_plan(facts: Ownership) -> OwnershipPlan: + return _OWNERSHIP.validate_python(_native.request_auth_admission(encode_facts(facts))) diff --git a/litellm/_auth/agent_access.py b/litellm/_auth/agent_access.py new file mode 100644 index 00000000000..1655dc3445a --- /dev/null +++ b/litellm/_auth/agent_access.py @@ -0,0 +1,68 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.managed import Agent +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class AgentAccessFacts: + registry_managed: bool + database_present: bool + target_loaded: bool = False + target: Agent | None = None + subject_present: bool = False + key_hash_present: bool = False + subject_managed: bool = False + session: bool = False + custom: bool = False + authority_loaded: bool = False + explicit_grants: tuple[str, ...] | None = None + legacy_loaded: bool = False + legacy_grants: tuple[str, ...] | None = None + target_ids: tuple[str, ...] | None = None + + +AgentAccessDecision: TypeAlias = Literal[ + "ReadTarget", "ReadAuthority", "ReadExplicitGrants", "ReadLegacyGrants", "ReadTargetIds", "Allow", "Deny" +] +_DECISION: Final[TypeAdapter[AgentAccessDecision]] = TypeAdapter(AgentAccessDecision) + + +@native_boundary +def decide(facts: AgentAccessFacts) -> AgentAccessDecision: + return _DECISION.validate_python(_native.request_auth_agent_access(encode_facts(facts))) + + +@dataclass(frozen=True, slots=True) +class InvocationFacts: + registered: Agent | None + has_store: bool + target_loaded: bool = False + target: Agent | None = None + actor_managed: bool = False + actor_id_present: bool = False + permissions: AgentAccessFacts | None = None + assigned: bool = False + fee_loaded: bool = False + fee: float | None = None + + +@dataclass(frozen=True, slots=True) +class InvocationPlan: + action: Literal["Complete", "ReadTarget", "ReadPermissions", "Assign", "ReadFee", "Deny"] + denial: Literal["DatabaseRequired", "Removed", "CallerNotPermitted", "InvalidFee"] | None + use_target: bool + assign_billing: bool + + +_INVOCATION: Final = TypeAdapter(InvocationPlan) + + +@native_boundary +def invocation(facts: InvocationFacts) -> InvocationPlan: + return _INVOCATION.validate_python(_native.request_auth_agent_invocation(encode_facts(facts))) diff --git a/litellm/_auth/agent_ceilings.py b/litellm/_auth/agent_ceilings.py new file mode 100644 index 00000000000..544bf05d556 --- /dev/null +++ b/litellm/_auth/agent_ceilings.py @@ -0,0 +1,60 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class AgentCeilingGroup: + status: Literal["Found", "Missing", "Unavailable"] + models: tuple[str, ...] = () + mcp_server_ids: tuple[str, ...] = () + agent_ids: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class AgentCeilingFacts: + ids: tuple[str, ...] | None = None + groups: tuple[AgentCeilingGroup, ...] | None = None + check_db_only: bool = False + + +@dataclass(frozen=True, slots=True) +class AgentCeilingDecision: + kind: Literal["ReadIds", "ReadGroups", "Complete"] + authoritative: bool = False + present: bool = False + ids: tuple[str, ...] = () + models: frozenset[str] = frozenset() + mcp_server_ids: frozenset[str] = frozenset() + agent_ids: frozenset[str] = frozenset() + indices: tuple[int, ...] = () + + +@dataclass(frozen=True, slots=True) +class AgentGroupLoadFacts: + database_present: bool + check_db_only: bool = False + status: Literal["Success", "NotFound", "HttpError"] | None = None + + +@dataclass(frozen=True, slots=True) +class AgentGroupLoadDecision: + kind: Literal["Read", "Return", "DatabaseUnavailable", "HttpUnavailable", "Raise"] + + +_DECISION: Final = TypeAdapter(AgentCeilingDecision) +_LOAD: Final = TypeAdapter(AgentGroupLoadDecision) + + +@native_boundary +def agent_ceiling(facts: AgentCeilingFacts) -> AgentCeilingDecision: + return _DECISION.validate_python(_native.request_auth_agent_ceiling(asdict(facts))) + + +@native_boundary +def agent_group_load(facts: AgentGroupLoadFacts) -> AgentGroupLoadDecision: + return _LOAD.validate_python(_native.request_auth_agent_group_load(asdict(facts))) diff --git a/litellm/_auth/agent_grants.py b/litellm/_auth/agent_grants.py new file mode 100644 index 00000000000..a0af67f5273 --- /dev/null +++ b/litellm/_auth/agent_grants.py @@ -0,0 +1,128 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Source: + team_id: str | None + grants: tuple[str, ...] | None + unavailable: bool = False + + +@dataclass(frozen=True, slots=True) +class Resolve: + managed: bool + strict: bool + role: str | None + key_team_loaded: bool = False + key_team: tuple[str, ...] | None = None + default_deny: bool | None = None + caller_loaded: bool = False + own: tuple[str, ...] | None = None + ceiling_loaded: bool = False + ceiling: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class Failure: + strict: bool + closed: bool + + +@dataclass(frozen=True, slots=True) +class OwnerTeam: + strict: bool + permission_id_present: bool + permission_present: bool + + +@dataclass(frozen=True, slots=True) +class OwnerUnion: + teams: tuple[Source, ...] + withheld: tuple[Source, ...] + unreadable: bool + + +@dataclass(frozen=True, slots=True) +class BillingTeam: + sources: tuple[Source, ...] + agent_id: str + + +@dataclass(frozen=True, slots=True) +class Human: + sources: tuple[Source, ...] + team_id: str | None + + +@dataclass(frozen=True, slots=True) +class Delegated: + sources: tuple[Source, ...] + agent_id: str + team_id: str | None + explicit: bool + allowed: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class TeamLookup: + subject: bool + team_id: str | None + database: bool + strict: bool + personal: bool | None = None + loaded: bool = False + row_present: bool = False + + +@dataclass(frozen=True, slots=True) +class ConfiguredAgent: + agent_id: str + groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class GroupMatches: + agents: tuple[ConfiguredAgent, ...] + groups: tuple[str, ...] + + +Action: TypeAlias = Literal[ + "Complete", + "ReadManaged", + "ReadKeyTeam", + "ReadDefault", + "ReadCaller", + "ReadCeiling", + "ReadPersonalEligibility", + "ReadOwnerTeams", + "ReadTeam", + "ReadDeclared", + "Unavailable", + "SelectTeam", +] + + +@dataclass(frozen=True, slots=True) +class Decision: + action: Action + grants: tuple[str, ...] | None + team_id: str | None + + +_DECISION: Final = TypeAdapter(Decision) + +Facts: TypeAlias = ( + Resolve | Failure | OwnerTeam | OwnerUnion | BillingTeam | Human | Delegated | TeamLookup | GroupMatches +) + + +@native_boundary +def decide(facts: Facts) -> Decision: + return _DECISION.validate_python(_native.request_auth_agent_grants(encode_facts(facts))) diff --git a/litellm/_auth/auth_context.py b/litellm/_auth/auth_context.py new file mode 100644 index 00000000000..a0754b1ed11 --- /dev/null +++ b/litellm/_auth/auth_context.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Fetch: + team_id: str | None + ui_team_id: str + user_id: str | None + project_id: str | None + end_user_id: str | None + + +@dataclass(frozen=True, slots=True) +class Team: + team_id: str | None + ui_team_id: str + failed: bool + missing: bool + recorded_models: tuple[str, ...] + allow_unavailable: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Admin: + role: str | None + user_id: str | None + admin_id: str + user_row: bool + + +@dataclass(frozen=True, slots=True) +class FetchPlan: + team: bool + user: bool + project: bool + end_user: bool + + +@dataclass(frozen=True, slots=True) +class TeamPlan: + source: Literal["None", "Token", "Loaded", "ReadUnavailablePolicy", "Failure"] + + +@dataclass(frozen=True, slots=True) +class AdminPlan: + force_admin: bool + preserve_budget: bool + + +_FETCH: Final = TypeAdapter(FetchPlan) +_TEAM: Final = TypeAdapter(TeamPlan) +_ADMIN: Final = TypeAdapter(AdminPlan) + + +@native_boundary +def fetch(facts: Fetch) -> FetchPlan: + return _FETCH.validate_python(_native.request_auth_context(encode_facts(facts))) + + +@native_boundary +def team(facts: Team) -> TeamPlan: + return _TEAM.validate_python(_native.request_auth_context(encode_facts(facts))) + + +@native_boundary +def admin(facts: Admin) -> AdminPlan: + return _ADMIN.validate_python(_native.request_auth_context(encode_facts(facts))) diff --git a/litellm/_auth/budget_alerts.py b/litellm/_auth/budget_alerts.py new file mode 100644 index 00000000000..d5546243790 --- /dev/null +++ b/litellm/_auth/budget_alerts.py @@ -0,0 +1,69 @@ +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: str + message: str + + +@dataclass(frozen=True, slots=True) +class ThresholdFacts: + key: str + digit: bool + digits: str | None = None + ratio: float | None = None + integer_failure: Failure | None = None + ratio_failure: Failure | None = None + text: str | None = None + items: tuple[str | None, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ConfigFacts: + global_entries: tuple[ThresholdFacts, ...] = () + local_entries: tuple[ThresholdFacts, ...] = () + global_failure: Failure | None = None + local_failure: Failure | None = None + order: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class Facts: + member: bool + spend: float | None + limit: float | None + default_fraction: float + config: ConfigFacts | None = None + + +@dataclass(frozen=True, slots=True) +class Threshold: + key: str + recipients: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Decision: + read_config: bool = False + read_order: bool = False + global_keys: tuple[str, ...] = () + local_keys: tuple[str, ...] = () + emit: bool = False + thresholds: tuple[Threshold, ...] | None = None + failure: Failure | None = None + + +_FACTS: Final = TypeAdapter(Facts) +_DECISION: Final = TypeAdapter(Decision) + + +@native_boundary +def plan(facts: Facts) -> Decision: + return _DECISION.validate_python(_native.request_auth_budget_alert_plan(_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/budget_context.py b/litellm/_auth/budget_context.py new file mode 100644 index 00000000000..ae052b3605f --- /dev/null +++ b/litellm/_auth/budget_context.py @@ -0,0 +1,174 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelCheckFacts: + scope: Literal["Team", "EndUser", "CustomEndUser"] + entity_id: str | None + budget_count: int + model_count: int + + +@dataclass(frozen=True, slots=True) +class BudgetFields: + max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + tpd_limit: int | None = None + model_budget_count: int | None = None + + +@dataclass(frozen=True, slots=True) +class EndUserFacts: + mode: Literal["Models", "Token", "Default"] + object_present: bool + end_user_id: str | None + prisma_present: bool + key_budget_id: str | None + key_read: bool = False + key_budget: BudgetFields | None = None + global_config_read: bool = False + global_budget_id: str | None = None + global_read: bool = False + global_budget: BudgetFields | None = None + + +EndUserSource: TypeAlias = Literal["Own", "ReadKey", "Key", "ReadGlobalConfig", "ReadGlobal", "Global", "None"] + + +@dataclass(frozen=True, slots=True) +class TokenDefaultsFacts: + current: BudgetFields + keep_token_limits: bool + default_present: bool + + +TokenField: TypeAlias = Literal["MaxBudget", "Tpm", "Rpm", "Tpd", "Models"] + + +@dataclass(frozen=True, slots=True) +class MembershipSnapshot: + budget_id: str | None + model_budget_count: int | None + + +@dataclass(frozen=True, slots=True) +class MemberFacts: + user_id: str | None + team_id: str | None + prisma_present: bool + default_id: str | None + skip_budget_checks: bool + membership_read: bool = False + membership: MembershipSnapshot | None = None + default_read: bool = False + default_budget: BudgetFields | None = None + lookup_failure: Literal["Membership", "Default"] | None = None + + +MemberSource: TypeAlias = Literal[ + "Done", "ReadMembership", "ReadDefault", "Member", "Default", "Clear", "RaiseLookup", "SkipLookup" +] + + +@dataclass(frozen=True, slots=True) +class IncreaseFacts: + number: float | None + truthy: bool + type_name: str + + +@dataclass(frozen=True, slots=True) +class TemporaryFacts: + max_budget: float | None + metadata_read: bool = False + has_increase: bool = False + has_expiry: bool = False + expires_at: int | None = None + now: int | None = None + increase: IncreaseFacts | None = None + + +@dataclass(frozen=True, slots=True) +class TemporaryDecision: + read: Literal["Metadata", "Expiry", "Now", "Increase"] | None = None + copy: bool = False + max_budget: float | None = None + invalid_increase_type: str | None = None + + +@dataclass(frozen=True, slots=True) +class NumericFacts: + number: float | None + boolean: bool + + +@dataclass(frozen=True, slots=True) +class ThrottleFacts: + opt_in: bool + tpm_limit: int | None + rpm_limit: int | None + percentage: NumericFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ThrottleDecision: + read_percentage: bool = False + percentage: float | None = None + + +_MODELS: Final = TypeAdapter(ModelCheckFacts) +_INDICES: Final = TypeAdapter(tuple[int, ...]) +_END: Final = TypeAdapter(EndUserFacts) +_END_SOURCE: Final = TypeAdapter[EndUserSource](EndUserSource) +_TOKEN: Final = TypeAdapter(TokenDefaultsFacts) +_FIELDS: Final = TypeAdapter(tuple[TokenField, ...]) +_MEMBER: Final = TypeAdapter(MemberFacts) +_MEMBER_SOURCE: Final = TypeAdapter[MemberSource](MemberSource) +_TEMP: Final = TypeAdapter(TemporaryFacts) +_TEMP_DECISION: Final = TypeAdapter(TemporaryDecision) +_THROTTLE: Final = TypeAdapter(ThrottleFacts) +_THROTTLE_DECISION: Final = TypeAdapter(ThrottleDecision) +_NUMBER: Final = TypeAdapter(NumericFacts) +_PERCENTAGE: Final = TypeAdapter[float | None](float | None) + + +@native_boundary +def model_checks(facts: ModelCheckFacts) -> tuple[int, ...]: + return _INDICES.validate_python(_native.request_auth_context_models(_MODELS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def end_user_source(facts: EndUserFacts) -> EndUserSource: + return _END_SOURCE.validate_python(_native.request_auth_context_end_user(_END.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def token_assignments(facts: TokenDefaultsFacts) -> tuple[TokenField, ...]: + return _FIELDS.validate_python(_native.request_auth_context_token(_TOKEN.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def member_source(facts: MemberFacts) -> MemberSource: + return _MEMBER_SOURCE.validate_python(_native.request_auth_context_member(_MEMBER.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def temporary(facts: TemporaryFacts) -> TemporaryDecision: + return _TEMP_DECISION.validate_python(_native.request_auth_context_temporary(_TEMP.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def percentage(facts: NumericFacts) -> float | None: + return _PERCENTAGE.validate_python(_native.request_auth_context_percentage(_NUMBER.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def throttle(facts: ThrottleFacts) -> ThrottleDecision: + return _THROTTLE_DECISION.validate_python(_native.request_auth_context_throttle(_THROTTLE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/budget_lookup.py b/litellm/_auth/budget_lookup.py new file mode 100644 index 00000000000..a980a1dee26 --- /dev/null +++ b/litellm/_auth/budget_lookup.py @@ -0,0 +1,90 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class OrganizationFacts: + token_present: bool + client_present: bool + team_present: bool + token_id: bool | None = None + team_id: bool | None = None + row_present: bool | None = None + lookup_failed: bool = False + + +OrganizationAction: TypeAlias = Literal["Skip", "ReadTokenId", "ReadTeamId", "ReadOrganization", "UseOrganization"] + + +@dataclass(frozen=True, slots=True) +class MemberFacts: + team_present: bool + token_present: bool + membership_loaded: bool + team_id: bool | None = None + user_id: bool | None = None + limits_read: bool = False + maximum_present: bool = False + default_id_string: bool = False + default_read: bool = False + + +MemberAction: TypeAlias = Literal[ + "Skip", "ReadTeamId", "ReadUserId", "ReadMembership", "ReadLimits", "ReadDefault", "UseLimits" +] +_ORGANIZATION: Final = TypeAdapter[OrganizationAction](OrganizationAction) +_MEMBER: Final = TypeAdapter[MemberAction](MemberAction) +_ORGANIZATION_FACTS: Final = TypeAdapter(OrganizationFacts) +_MEMBER_FACTS: Final = TypeAdapter(MemberFacts) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_EMAILS: Final = TypeAdapter(tuple[str, ...]) +_METADATA: Final = TypeAdapter(Mapping[object, object]) + + +def soft_email_value(raw: object) -> object: + return _METADATA.validate_python(raw).get("soft_budget_alerting_emails") if isinstance(raw, dict) else None + + +@dataclass(frozen=True, slots=True) +class SoftFacts: + limit: float | None + spend: float | None + token_present: bool + + +_SOFT: Final = TypeAdapter(SoftFacts) + + +@native_boundary +def read_soft_emails(facts: SoftFacts) -> bool: + return _native.request_auth_budget_lookup_soft(_SOFT.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def organization(facts: OrganizationFacts) -> OrganizationAction: + return _ORGANIZATION.validate_python( + _native.request_auth_budget_lookup_organization(_ORGANIZATION_FACTS.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + ) + + +@native_boundary +def member(facts: MemberFacts) -> MemberAction: + return _MEMBER.validate_python(_native.request_auth_budget_lookup_member(_MEMBER_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def emails(raw: object) -> tuple[str, ...]: + facts: Final = ( + {"Text": raw} + if isinstance(raw, str) + else {"List": tuple(value if isinstance(value, str) else None for value in _ITEMS.validate_python(raw))} + if isinstance(raw, list) + else "Other" + ) + return _EMAILS.validate_python(_native.request_auth_budget_lookup_emails(facts)) diff --git a/litellm/_auth/budget_pricing.py b/litellm/_auth/budget_pricing.py new file mode 100644 index 00000000000..fda075c1730 --- /dev/null +++ b/litellm/_auth/budget_pricing.py @@ -0,0 +1,90 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class DeploymentPrice: + has_input_price: bool = False + has_output_price: bool = False + input_price: float | None = None + output_price: float | None = None + query_price: float | None = None + ptu_count: float | None = None + cost_per_ptu_per_hour: float | None = None + + +@dataclass(frozen=True, slots=True) +class ModelPriceFacts: + cached: bool | None = None + lookup_failed: bool = False + group_found: bool = False + input_price: float | None = None + output_price: float | None = None + deployments: tuple[DeploymentPrice, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class PriceDecision: + zero_cost: bool + cache_value: bool | None + read_deployments: bool = False + + +@dataclass(frozen=True, slots=True) +class RouteBudgetFacts: + side_effect_route: bool = False + discovery_route: bool = False + llm_route: bool = False + registered_pass_through: bool | None = None + route: str = "" + method: str | None = None + body_present: bool = False + + +RouteBudgetDecision: TypeAlias = Literal["ReadPassThrough", "ReadPricing", "Skip"] +_PRICE: Final = TypeAdapter(ModelPriceFacts) +_PRICE_DECISION: Final = TypeAdapter(PriceDecision) +_ROUTE: Final = TypeAdapter(RouteBudgetFacts) +_ROUTE_DECISION: Final = TypeAdapter[RouteBudgetDecision](RouteBudgetDecision) + + +@native_boundary +def price(facts: ModelPriceFacts) -> PriceDecision: + return _PRICE_DECISION.validate_python(_native.request_auth_budget_price(_PRICE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def route(facts: RouteBudgetFacts) -> RouteBudgetDecision: + return _ROUTE_DECISION.validate_python(_native.request_auth_budget_route(_ROUTE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@dataclass(frozen=True, slots=True) +class ExemptionFacts: + route: RouteBudgetFacts + model_input_present: bool + router_present: bool + model_count: int + prices: tuple[ModelPriceFacts, ...] = () + cache_read: bool = False + + +@dataclass(frozen=True, slots=True) +class ExemptionDecision: + read_pass_through: bool + read_price_cache: bool + read_model: int | None + exempt: bool + + +_EXEMPTION: Final = TypeAdapter(ExemptionFacts) +_EXEMPTION_DECISION: Final = TypeAdapter(ExemptionDecision) + + +@native_boundary +def exemption(facts: ExemptionFacts) -> ExemptionDecision: + return _EXEMPTION_DECISION.validate_python(_native.request_auth_budget_exemption(_EXEMPTION.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/claude_normalization.py b/litellm/_auth/claude_normalization.py new file mode 100644 index 00000000000..f3fd934df04 --- /dev/null +++ b/litellm/_auth/claude_normalization.py @@ -0,0 +1,74 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ClaudeMembershipFacts: + name: str + router_present: bool + alias_count: int + global_alias_exists: bool | None = None + alias_entries: tuple[bool | None, ...] = () + registry_loaded: bool = False + registry_alias_exists: bool | None = None + deployment: bool | None = None + names: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ClaudeMembershipDecision: + kind: Literal["Complete", "ReadGlobal", "ReadAlias", "ReadRegistry", "ReadAliases", "ReadDeployment", "ReadNames"] + present: bool = False + index: int = 0 + + +@dataclass(frozen=True, slots=True) +class ClaudeNormalizationFacts: + route: str + request_present: bool + router_present: bool | None = None + normalized: bool | None = None + model_loaded: bool = False + requested: str | None = None + body_loaded: bool = False + body_model: str | None = None + settings_loaded: bool = False + membership: tuple[ClaudeMembershipFacts, ...] = () + marked: bool = False + + +@dataclass(frozen=True, slots=True) +class ClaudeNormalizationDecision: + kind: Literal[ + "Complete", + "ReadRouter", + "ReadNormalized", + "ReadModel", + "ReadBodyModel", + "ReadSettings", + "ReadMembership", + "MarkNormalized", + "Rewrite", + ] + name: str | None = None + index: int = 0 + target: str | None = None + + +_MEMBERSHIP: Final = TypeAdapter(ClaudeMembershipDecision) +_NORMALIZATION: Final = TypeAdapter(ClaudeNormalizationDecision) + + +@native_boundary +def claude_membership(facts: ClaudeMembershipFacts) -> ClaudeMembershipDecision: + return _MEMBERSHIP.validate_python(_native.request_auth_claude_membership(asdict(facts))) + + +@native_boundary +def claude_normalization(facts: ClaudeNormalizationFacts) -> ClaudeNormalizationDecision: + return _NORMALIZATION.validate_python(_native.request_auth_claude_normalization(asdict(facts))) diff --git a/litellm/_auth/codex_models.py b/litellm/_auth/codex_models.py new file mode 100644 index 00000000000..f4b2e13613d --- /dev/null +++ b/litellm/_auth/codex_models.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_helpers import ModelHelperFacts +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + +CodexModelKind: TypeAlias = Literal["Routes", "NativeRoutes", "Chatgpt", "Grants"] +CandidateField: TypeAlias = Literal["Provider", "Model", "Name", "PublicName"] + + +@dataclass(frozen=True, slots=True) +class CodexCandidate: + provider_loaded: bool = False + provider: str | None = None + model: str | None = None + name: str | None = None + public_loaded: bool = False + public_name: str | None = None + + +@dataclass(frozen=True, slots=True) +class CodexGrants: + key: tuple[str, ...] + team: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class CodexModelFacts: + kind: CodexModelKind + model: str | None + router_present: bool + team_id: str | None = None + bundled: tuple[str, ...] | None = None + priced: bool | None = None + candidates: tuple[CodexCandidate, ...] | None = None + alias_loaded: bool = False + alias: str | None = None + grants: CodexGrants | None = None + prefixed_grants: CodexGrants | None = None + bare: tuple[ModelHelperFacts, ...] = () + prefixed: tuple[ModelHelperFacts, ...] = () + route: str | None = None + + +@dataclass(frozen=True, slots=True) +class CodexModelDecision: + kind: Literal[ + "Complete", + "ReadBundled", + "ReadCost", + "ReadCandidates", + "ReadAlias", + "ReadCandidate", + "ReadGrants", + "ReadHelper", + ] + allowed: bool = False + target: str | None = None + model: str | None = None + index: int = 0 + field: CandidateField | None = None + prefixed: bool = False + + +_DECISION: Final = TypeAdapter(CodexModelDecision) + + +@native_boundary +def codex_model(facts: CodexModelFacts) -> CodexModelDecision: + return _DECISION.validate_python(_native.request_auth_codex_model(encode_facts(facts))) diff --git a/litellm/_auth/codex_normalization.py b/litellm/_auth/codex_normalization.py new file mode 100644 index 00000000000..adf3ecedf6d --- /dev/null +++ b/litellm/_auth/codex_normalization.py @@ -0,0 +1,64 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.codex_models import CodexModelFacts +from litellm._auth.codex_turn import CodexTurnFacts +from litellm._auth.failure import native_boundary +from litellm._auth.routes import RouteContext +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class CodexNormalizationFacts: + kind: Literal["Alias", "Prefix"] + route: str + request_present: bool + router_present: bool | None = None + context: RouteContext | None = None + eu_jev: bool | None = None + resolved: bool | None = None + marked: bool = False + model_loaded: bool = False + model: str | None = None + explicit: bool | None = None + provider: bool | None = None + routing: CodexModelFacts | None = None + body_present: bool | None = None + settings_loaded: bool = False + target: str | None = None + provider_after: bool | None = None + turn: CodexTurnFacts | None = None + prefix: CodexModelFacts | None = None + + +@dataclass(frozen=True, slots=True) +class CodexNormalizationDecision: + kind: Literal[ + "Complete", + "ReadRouter", + "ReadContext", + "ReadEuJev", + "ReadResolved", + "MarkResolved", + "ReadModel", + "ReadExplicit", + "ReadProvider", + "ReadRouting", + "ReadBody", + "ReadSettings", + "ReadProviderAfter", + "ReadTurn", + "ReadPrefix", + "Rewrite", + ] + target: str | None = None + + +_DECISION: Final = TypeAdapter(CodexNormalizationDecision) + + +@native_boundary +def codex_normalization(facts: CodexNormalizationFacts) -> CodexNormalizationDecision: + return _DECISION.validate_python(_native.request_auth_codex_normalization(asdict(facts))) diff --git a/litellm/_auth/codex_turn.py b/litellm/_auth/codex_turn.py new file mode 100644 index 00000000000..dac1bbc1dda --- /dev/null +++ b/litellm/_auth/codex_turn.py @@ -0,0 +1,100 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.codex_models import CodexModelFacts +from litellm._auth.failure import native_boundary +from litellm._auth.route_enable import RouteEnableFacts +from litellm._auth.route_scope import RouteScopeFacts +from litellm._auth.routes import VirtualRouteFacts +from litellm.rust_bridge import _native + +JwtRead: TypeAlias = Literal["Returned", "Exception"] + + +@dataclass(frozen=True, slots=True) +class CodexKeyFacts: + proxy: RouteScopeFacts | None = None + enabled: RouteEnableFacts | None = None + virtual_route: VirtualRouteFacts | None = None + role_loaded: bool = False + role: str | None = None + admin: RouteScopeFacts | None = None + + +@dataclass(frozen=True, slots=True) +class CodexAccessFacts: + key: CodexKeyFacts + jwt_token_present: bool + jwt: JwtRead | None = None + + +@dataclass(frozen=True, slots=True) +class CodexCredentials: + bearer_present: bool + gateway_present: bool + equal: bool + + +@dataclass(frozen=True, slots=True) +class CodexTurnFacts: + model: str | None + native: bool = False + credentials: CodexCredentials | None = None + access: CodexAccessFacts | None = None + routing: CodexModelFacts | None = None + catalog_loaded: bool = False + slugs: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class CodexTurnDecision: + kind: Literal[ + "Complete", + "Failure", + "ReadProxy", + "ReadEnabled", + "ReadVirtual", + "ReadRole", + "ReadAdmin", + "ReadCredentials", + "ReadAccess", + "ReadRouting", + "ReadJwt", + "ReadCatalog", + ] + allowed: bool = False + message: str | None = None + + +_DECISION: Final = TypeAdapter(CodexTurnDecision) + + +@native_boundary +def codex_key(facts: CodexKeyFacts) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Key": asdict(facts)})) + + +@native_boundary +def codex_access(facts: CodexAccessFacts) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Access": asdict(facts)})) + + +@native_boundary +def codex_turn(facts: CodexTurnFacts) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Turn": asdict(facts)})) + + +@native_boundary +def codex_jwt(read: JwtRead) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Jwt": read})) + + +@native_boundary +def codex_jwt_token(*, claims_present: bool, token_present: bool) -> CodexTurnDecision: + return _DECISION.validate_python( + _native.request_auth_codex_turn( + {"JwtToken": {"claims_present": claims_present, "token_present": token_present}} + ) + ) diff --git a/litellm/_auth/common_budgets.py b/litellm/_auth/common_budgets.py new file mode 100644 index 00000000000..54e6a020a20 --- /dev/null +++ b/litellm/_auth/common_budgets.py @@ -0,0 +1,115 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +BudgetScope: TypeAlias = Literal[ + "Team", + "Key", + "TeamWindow", + "KeyWindow", + "Project", + "Organization", + "Tag", + "AccessGroup", + "User", + "TeamMember", + "EndUser", + "Proxy", + "KeySoft", + "TeamSoft", + "ProjectSoft", +] +BudgetAlert: TypeAlias = Literal[ + "token_budget", + "user_budget", + "team_budget", + "project_budget", + "organization_budget", + "team_member_threshold", + "soft_budget", +] + + +@dataclass(frozen=True, slots=True) +class BudgetFacts: + scope: BudgetScope + limit: float | None + spend: float | None = None + token_present: bool = False + team_key: bool = False + apply_user_budget_to_team_keys: bool = False + info_route: bool = False + llm_route: bool = True + model_list_route: bool = False + skip: bool = False + alert_recipients: tuple[str, ...] = () + throttle_percentage: float | None = None + throttle_opt_in: bool = False + tpm_limit: int | None = None + rpm_limit: int | None = None + + +@dataclass(frozen=True, slots=True) +class BudgetDecision: + read_spend: bool = False + denied: bool = False + alert: BudgetAlert | None = None + throttle_percentage: float | None = None + + +@dataclass(frozen=True, slots=True) +class AlertFacts: + spend: float | None + max_budget: float | None + configured: bool = False + minimum_percent: float | None = None + default_ratio: float = 0.8 + team_member: bool = False + + +@dataclass(frozen=True, slots=True) +class MemberLimitFacts: + member_limit: float | None + default_limit: float | None + temporary_increase: float | None + expires_at: float | None + now: float + now_naive: bool = False + + +@dataclass(frozen=True, slots=True) +class WindowFacts: + duration_seconds: float | None + reset_at: float | None + now: float + + +_FACTS: Final = TypeAdapter(BudgetFacts) +_DECISION: Final = TypeAdapter(BudgetDecision) +_ALERT: Final = TypeAdapter(AlertFacts) +_MEMBER: Final = TypeAdapter(MemberLimitFacts) +_WINDOW: Final = TypeAdapter(WindowFacts) + + +@native_boundary +def check(facts: BudgetFacts) -> BudgetDecision: + return _DECISION.validate_python(_native.request_auth_common_budget(_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def alert(facts: AlertFacts) -> bool: + return _native.request_auth_budget_alert(_ALERT.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def member_limit(facts: MemberLimitFacts) -> float | None: + return _native.request_auth_member_limit(_MEMBER.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def window_start(facts: WindowFacts) -> float | None: + return _native.request_auth_budget_window_start(_WINDOW.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary diff --git a/litellm/_auth/common_flow.py b/litellm/_auth/common_flow.py new file mode 100644 index 00000000000..208f68cd253 --- /dev/null +++ b/litellm/_auth/common_flow.py @@ -0,0 +1,148 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class BudgetSkip: + skip: bool + exempt: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Pricing: + enabled: bool + llm_route: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Unpriced: + count: int + + +@dataclass(frozen=True, slots=True) +class MemberFallback: + router: bool + multipart: bool | None = None + resource_bound: bool | None = None + + +@dataclass(frozen=True, slots=True) +class ModelBudgets: + scope: Literal["Key", "User", "Member"] + entity_present: bool | None = None + budget_present: bool | None = None + mapping: bool | None = None + team_present: bool | None = None + + +@dataclass(frozen=True, slots=True) +class FallbackAllowed: + router: bool + disabled: bool | None = None + + +ReadPlan: TypeAlias = Literal[ + "ReadExemption", + "SkipBudgets", + "CheckBudgets", + "ReadRoute", + "ReadUnpriced", + "RejectUnpriced", + "ReadContentType", + "ReadResourceBinding", + "ReadFallbackSettings", + "CheckModelBudgets", + "ReadModelBudgetPresence", + "ReadModelBudgetShape", + "ReadModelBudgetToken", + "ReadModelBudgetUser", + "ReadModelBudgetTeam", + "AllowFallback", + "ReadFallbackDisabled", + "Continue", +] +_READ: Final = TypeAdapter(ReadPlan) + + +@native_boundary +def read(facts: BudgetSkip | Pricing | Unpriced | MemberFallback | ModelBudgets | FallbackAllowed) -> ReadPlan: + return _READ.validate_python(_native.request_auth_common_read(encode_facts(facts))) + + +@dataclass(frozen=True, slots=True) +class Facts: + model_present: bool + skip_budgets: bool + team_present: bool + user_present: bool + token_present: bool + user_id: str | None + agent_id: str | None + customer_present: bool + customer_models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Plan: + membership: bool + team_models: bool + member_models: bool + user_models: bool + customer_models: bool + agent_trace: bool + + +@dataclass(frozen=True, slots=True) +class BudgetFacts: + token_present: bool + group_count: int + customer_present: bool + customer_budget_present: bool + + +BudgetScope: TypeAlias = Literal[ + "Team", "TeamWindows", "KeyWindows", "TeamSoft", "Organization", "Tags", "Groups", "User", "Member", "Customer" +] + + +@dataclass(frozen=True, slots=True) +class CustomFacts: + end_user_id: str | None + end_user_present: bool + common_checks: bool + model_present: bool = False + router_present: bool = False + + +@dataclass(frozen=True, slots=True) +class CustomPlan: + read_end_user: bool + check_end_user: bool + check_model_access: bool + read_pricing: bool + + +_PLAN: Final = TypeAdapter(Plan) +_BUDGETS: Final = TypeAdapter(tuple[BudgetScope, ...]) +_CUSTOM: Final = TypeAdapter(CustomPlan) + + +@native_boundary +def plan(facts: Facts) -> Plan: + return _PLAN.validate_python(_native.request_auth_common_flow(asdict(facts))) + + +@native_boundary +def budgets(facts: BudgetFacts) -> tuple[BudgetScope, ...]: + return _BUDGETS.validate_python(_native.request_auth_common_budget_scopes(asdict(facts))) + + +@native_boundary +def custom(facts: CustomFacts) -> CustomPlan: + return _CUSTOM.validate_python(_native.request_auth_custom_flow(asdict(facts))) diff --git a/litellm/_auth/credential.py b/litellm/_auth/credential.py new file mode 100644 index 00000000000..105eb4afe73 --- /dev/null +++ b/litellm/_auth/credential.py @@ -0,0 +1,141 @@ +import re +from dataclasses import dataclass +from enum import Enum +from typing import Final + +_AWS4_PREFIX: Final = "AWS4-HMAC-SHA256" +_AWS4_BEARER_CREDENTIAL: Final = re.compile(r"Credential=Bearer\s+([^/\s,]+)") +_AWS4_CREDENTIAL: Final = re.compile(r"Credential=([^/\s,]+)") +_TOKEN_SCHEMES: Final = frozenset({"bearer", "basic"}) +_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX: Final = "openai-insecure-api-key." + + +class KeySource(Enum): + CUSTOM_HEADER = "custom_header" + LITELLM_HEADER = "x-litellm-api-key" + AUTHORIZATION = "authorization" + API_KEY = "api-key" + X_API_KEY = "x-api-key" + GOOGLE = "x-goog-api-key" + APIM = "ocp-apim-subscription-key" + QUERY = "key" + PASS_THROUGH = "pass_through_endpoint" + + +@dataclass(frozen=True, slots=True) +class PresentedKeys: + custom_header: str | None = None + litellm_header: str | None = None + authorization: str | None = None + api_key: str | None = None + x_api_key: str | None = None + google: str | None = None + apim: str | None = None + query: str | None = None + pass_through: str | None = None + + +@dataclass(frozen=True, slots=True) +class SelectedKey: + source: KeySource + presented: str + token: str + + +class WebSocketDenial(Enum): + NO_API_KEY = "no_api_key" + INVALID_AUTHORIZATION = "invalid_authorization" + + +@dataclass(frozen=True, slots=True) +class WebSocketKeys: + authorization: str | None = None + api_key: str | None = None + subprotocols: str | None = None + + +def _aws4_access_key(value: str) -> str | None: + match: Final = _AWS4_BEARER_CREDENTIAL.search(value) or _AWS4_CREDENTIAL.search(value) + return match.group(1) if match else None + + +def _split_scheme(value: str) -> tuple[str, str] | None: + scheme, separator, credential = value.partition(" ") + return (scheme.lower(), credential) if separator else None + + +def scheme_credential(value: str) -> str | None: + if value.startswith(_AWS4_PREFIX): + return _aws4_access_key(value) + split: Final = _split_scheme(value) + if split is None: + return None + scheme, credential = split + return credential if scheme in _TOKEN_SCHEMES else None + + +def litellm_header_token(value: str) -> str: + token: Final = scheme_credential(value) + return value if token is None else token + + +def _raw_key(keys: PresentedKeys) -> SelectedKey | None: + raw_sources: Final = ( + (KeySource.API_KEY, keys.api_key), + (KeySource.X_API_KEY, keys.x_api_key), + (KeySource.GOOGLE, keys.google), + (KeySource.APIM, keys.apim), + (KeySource.QUERY, keys.query or None), + (KeySource.PASS_THROUGH, keys.pass_through), + ) + return next( + ( + SelectedKey(source=source, presented=value, token=value) + for source, value in raw_sources + if value is not None + ), + None, + ) + + +def select_api_key(keys: PresentedKeys) -> SelectedKey | None: + if keys.custom_header is not None: + return SelectedKey( + source=KeySource.CUSTOM_HEADER, + presented=keys.custom_header, + token=scheme_credential(keys.custom_header) or "", + ) + if keys.litellm_header is not None: + return SelectedKey( + source=KeySource.LITELLM_HEADER, + presented=keys.litellm_header, + token=litellm_header_token(keys.litellm_header), + ) + if keys.authorization: + return SelectedKey( + source=KeySource.AUTHORIZATION, + presented=keys.authorization, + token=scheme_credential(keys.authorization) or "", + ) + return _raw_key(keys) + + +def _subprotocol_api_key(subprotocols: str | None) -> str | None: + protocols: Final = tuple(protocol.strip() for protocol in (subprotocols or "").split(",")) + return next( + ( + protocol.removeprefix(_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX) + for protocol in protocols + if protocol.startswith(_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX) + ), + None, + ) + + +def select_websocket_api_key(keys: WebSocketKeys) -> str | WebSocketDenial: + if keys.authorization: + split: Final = _split_scheme(keys.authorization) + if split is None or split[0] != "bearer": + return WebSocketDenial.INVALID_AUTHORIZATION + return split[1].strip() + return keys.api_key or _subprotocol_api_key(keys.subprotocols) or WebSocketDenial.NO_API_KEY diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py new file mode 100644 index 00000000000..8703d126772 --- /dev/null +++ b/litellm/_auth/decision.py @@ -0,0 +1,460 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import Protocol, TypeAlias + +from litellm._auth.budget_pricing import ExemptionFacts +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts +from litellm._auth.model_access import ModelAccessDecision +from litellm._auth.model_budget_data import ModelBudgetSnapshot + + +@dataclass(frozen=True, slots=True) +class RequestFacts: + route: str + token: str | None + is_llm_api_route: bool + + +@dataclass(frozen=True, slots=True) +class Settings: + master_key: str | None + allow_user_auth: bool + has_db: bool + master_key_only_routes: frozenset[str] + proxy_max_budget: float + + +@dataclass(frozen=True, slots=True) +class NoMasterKey: + token: str | None + + +@dataclass(frozen=True, slots=True) +class UserAuthRoute: + pass + + +@dataclass(frozen=True, slots=True) +class MasterKey: + pass + + +@dataclass(frozen=True, slots=True) +class CachedAdminKey: + pass + + +@dataclass(frozen=True, slots=True) +class VirtualKey: + budget_throttle_pct: float | None + + +@dataclass(frozen=True, slots=True) +class WithinModelBudgets: + pass + + +Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey | WithinModelBudgets + + +@dataclass(frozen=True, slots=True) +class MissingKey: + pass + + +@dataclass(frozen=True, slots=True) +class MalformedKey: + pass + + +@dataclass(frozen=True, slots=True) +class UserAuthDisabled: + pass + + +@dataclass(frozen=True, slots=True) +class ExpiredKey: + expires: datetime + now: datetime + key_label: str + + +@dataclass(frozen=True, slots=True) +class MasterKeyOnlyRoute: + pass + + +@dataclass(frozen=True, slots=True) +class NoDatabase: + pass + + +@dataclass(frozen=True, slots=True) +class NotAVirtualKey: + pass + + +@dataclass(frozen=True, slots=True) +class KeyBlocked: + pass + + +@dataclass(frozen=True, slots=True) +class OwnerDeactivated: + user_id: str + + +@dataclass(frozen=True, slots=True) +class TeamMemberBudgetExceeded: + user_id: str + team_id: str + spend: float + max_budget: float + + +@dataclass(frozen=True, slots=True) +class KeyBudgetExceeded: + spend: float + max_budget: float + + +@dataclass(frozen=True, slots=True) +class KeyWithoutToken: + pass + + +@dataclass(frozen=True, slots=True) +class ModelBudgetExceeded: + message: str + current_cost: float + max_budget: float + entity_type: str | None + entity_id: str | None + + +@dataclass(frozen=True, slots=True) +class ModelBudgetRead: + snapshot: ModelBudgetSnapshot + message: str = "" + entity_type: str | None = None + entity_id: str | None = None + + +@dataclass(frozen=True, slots=True) +class ModelAccessDenied: + decision: ModelAccessDecision + + +Denial: TypeAlias = ( + ModelAccessDenied + | MissingKey + | MalformedKey + | UserAuthDisabled + | ExpiredKey + | MasterKeyOnlyRoute + | NoDatabase + | NotAVirtualKey + | KeyBlocked + | OwnerDeactivated + | TeamMemberBudgetExceeded + | KeyBudgetExceeded + | KeyWithoutToken + | ModelBudgetExceeded +) + + +@dataclass(frozen=True, slots=True) +class Allow: + grant: Grant + + +@dataclass(frozen=True, slots=True) +class Deny: + denial: Denial + + +Outcome: TypeAlias = Allow | Deny + + +@dataclass(frozen=True, slots=True) +class AdminKey: + expires: datetime | None + + +@dataclass(frozen=True, slots=True) +class NonAdminKey: + pass + + +PresentedKey: TypeAlias = AdminKey | NonAdminKey + + +@dataclass(frozen=True, slots=True) +class KeyFacts: + has_token: bool + blocked: bool + expires: datetime | None + spend: float + soft_budget: float | None + max_budget: float | None + team_member_spend: float | None + user_id: str | None + team_id: str | None + budget_throttle_pct: float | None + + +@dataclass(frozen=True, slots=True) +class OwnerFacts: + deactivated: bool + + +@dataclass(frozen=True, slots=True) +class TeamMemberBudget: + spend: float + max_budget: float + + +@dataclass(frozen=True, slots=True) +class ModelBudgetScope: + models: tuple[str, ...] + rewritable_model: str | None + key_has_model_budgets: bool + key_budget_fallbacks: Mapping[str, tuple[str, ...]] + user_has_model_budgets: bool + + +@dataclass(frozen=True, slots=True) +class TeamMemberModelBudgetScope: + models: tuple[str, ...] + rewritable_model: str | None + has_model_budgets: bool + can_fall_back: bool + + +@dataclass(frozen=True, slots=True) +class FallbackTarget: + requested: str + model: str + served_model: str + + +@dataclass(frozen=True, slots=True) +class ResolveEndUser: + prefetch_key: bool + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ResolvePresentedKey: + hashed_token: str + resume: "Callable[[PresentedKey | None], Step]" + + +@dataclass(frozen=True, slots=True) +class EvictPresentedKey: + hashed_token: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKey: + hashed_token: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class PrepareKey: + loaded_from_db: bool + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class InspectKey: + resume: "Callable[[KeyFacts], Step]" + + +@dataclass(frozen=True, slots=True) +class EnforceModelAccess: + resume: "Callable[[KeyAndFallbackModelAccessFacts], Step]" + + +@dataclass(frozen=True, slots=True) +class PrefetchAuthObjects: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKeyOwner: + user_id: str + resume: "Callable[[OwnerFacts | None], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckBudgetExemption: + resume: "Callable[[ExemptionFacts], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberBudget: + user_id: str + team_id: str + resume: "Callable[[TeamMemberBudget | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertTeamMemberBudget: + user_id: str + team_id: str + budget: TeamMemberBudget + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class SendKeyBudgetAlerts: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadKeySpend: + max_budget: float + resume: "Callable[[float], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertKeyBudget: + spend: float + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertKeySoftBudget: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKeyContext: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadProxySpend: + resume: "Callable[[float | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertProxyBudget: + spend: float + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadModelBudgetScope: + resume: "Callable[[ModelBudgetScope], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckKeyModelBudget: + model: str + resume: "Callable[[ModelBudgetRead], Step]" + + +@dataclass(frozen=True, slots=True) +class AuthorizeBudgetFallback: + model: str + resume: "Callable[[FallbackAuthorizationFacts], Step]" + + +@dataclass(frozen=True, slots=True) +class RewriteRequestModel: + model: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckUserModelBudget: + model: str + resume: "Callable[[ModelBudgetRead], Step]" + + +IdentityEffect: TypeAlias = ( + ResolveEndUser + | ResolvePresentedKey + | EvictPresentedKey + | LoadKey + | PrepareKey + | InspectKey + | EnforceModelAccess + | PrefetchAuthObjects + | LoadKeyOwner + | CheckBudgetExemption + | LoadKeyContext +) +BudgetEffect: TypeAlias = ( + ReadTeamMemberBudget + | AlertTeamMemberBudget + | SendKeyBudgetAlerts + | ReadKeySpend + | AlertKeyBudget + | AlertKeySoftBudget + | ReadProxySpend + | AlertProxyBudget +) +ModelBudgetEffect: TypeAlias = ( + ReadModelBudgetScope | CheckKeyModelBudget | AuthorizeBudgetFallback | RewriteRequestModel | CheckUserModelBudget +) +Effect: TypeAlias = IdentityEffect | BudgetEffect | ModelBudgetEffect +Step: TypeAlias = Outcome | Effect + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberModelBudgetScope: + resume: "Callable[[TeamMemberModelBudgetScope], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class CheckTeamMemberModelBudget: + model: str + resume: "Callable[[ModelBudgetRead], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberFallbacks: + model: str + resume: "Callable[[tuple[FallbackTarget, ...]], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class VetTeamMemberFallback: + target: FallbackTarget + resume: "Callable[[FallbackCandidateFacts], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class RerouteTeamMemberRequest: + model: str + resume: "Callable[[], TeamMemberStep]" + + +TeamMemberEffect: TypeAlias = ( + ReadTeamMemberModelBudgetScope + | CheckTeamMemberModelBudget + | ReadTeamMemberFallbacks + | VetTeamMemberFallback + | RerouteTeamMemberRequest +) +TeamMemberStep: TypeAlias = Outcome | TeamMemberEffect + + +class AuthCore(Protocol): + def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... + + def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... + + def check_model_budgets(self) -> Step: ... + + def check_user_model_budgets(self) -> Step: ... + + def check_team_member_model_budgets(self) -> TeamMemberStep: ... diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py new file mode 100644 index 00000000000..00cfc9e2333 --- /dev/null +++ b/litellm/_auth/driver.py @@ -0,0 +1,274 @@ +from typing import Protocol, assert_never + +from litellm._auth.budget_pricing import ExemptionFacts +from litellm._auth.decision import ( + AlertKeyBudget, + AlertKeySoftBudget, + AlertProxyBudget, + AlertTeamMemberBudget, + Allow, + AuthorizeBudgetFallback, + BudgetEffect, + CheckBudgetExemption, + CheckKeyModelBudget, + CheckTeamMemberModelBudget, + CheckUserModelBudget, + Deny, + Effect, + EnforceModelAccess, + EvictPresentedKey, + FallbackTarget, + IdentityEffect, + InspectKey, + KeyFacts, + LoadKey, + LoadKeyContext, + LoadKeyOwner, + ModelBudgetEffect, + ModelBudgetRead, + ModelBudgetScope, + Outcome, + OwnerFacts, + PrefetchAuthObjects, + PrepareKey, + PresentedKey, + ReadKeySpend, + ReadModelBudgetScope, + ReadProxySpend, + ReadTeamMemberBudget, + ReadTeamMemberFallbacks, + ReadTeamMemberModelBudgetScope, + RerouteTeamMemberRequest, + ResolveEndUser, + ResolvePresentedKey, + RewriteRequestModel, + SendKeyBudgetAlerts, + Step, + TeamMemberBudget, + TeamMemberEffect, + TeamMemberModelBudgetScope, + TeamMemberStep, + VetTeamMemberFallback, +) +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts + + +class ModelBudgetHandler(Protocol): + async def read_model_budget_scope(self) -> ModelBudgetScope: ... + + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: ... + + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: ... + + async def rewrite_request_model(self, model: str) -> None: ... + + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: ... + + +class TeamMemberModelBudgetHandler(Protocol): + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: ... + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetRead: ... + + async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: ... + + async def vet_team_member_fallback(self, target: FallbackTarget) -> FallbackCandidateFacts: ... + + async def reroute_team_member_request(self, model: str) -> None: ... + + +class EffectHandler(ModelBudgetHandler, Protocol): + async def resolve_end_user(self, prefetch_key: bool) -> None: ... + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: ... + + async def evict_presented_key(self, hashed_token: str) -> None: ... + + async def load_key(self, hashed_token: str) -> None: ... + + async def prepare_key(self, loaded_from_db: bool) -> None: ... + + async def inspect_key(self) -> KeyFacts: ... + + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: ... + + async def prefetch_auth_objects(self) -> None: ... + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ... + + async def check_budget_exemption(self) -> ExemptionFacts: ... + + async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: ... + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: ... + + async def send_key_budget_alerts(self) -> None: ... + + async def read_key_spend(self, max_budget: float) -> float: ... + + async def alert_key_budget(self, spend: float) -> None: ... + + async def alert_key_soft_budget(self) -> None: ... + + async def load_key_context(self) -> None: ... + + async def read_proxy_spend(self) -> float | None: ... + + async def alert_proxy_budget(self, spend: float) -> None: ... + + +async def _answer(step: Effect, handler: EffectHandler) -> Step: + match step: + case ( + ResolveEndUser() + | ResolvePresentedKey() + | EvictPresentedKey() + | LoadKey() + | PrepareKey() + | InspectKey() + | EnforceModelAccess() + | PrefetchAuthObjects() + | LoadKeyOwner() + | CheckBudgetExemption() + | LoadKeyContext() + ): + return await _answer_identity(step, handler) + case ( + ReadModelBudgetScope() + | CheckKeyModelBudget() + | AuthorizeBudgetFallback() + | RewriteRequestModel() + | CheckUserModelBudget() + ): + return await _answer_model_budget(step, handler) + case _: + return await _answer_budget(step, handler) + + +async def _answer_identity(step: IdentityEffect, handler: EffectHandler) -> Step: + match step: + case ResolveEndUser(prefetch_key=prefetch_key, resume=resume): + await handler.resolve_end_user(prefetch_key) + return resume() + case ResolvePresentedKey(hashed_token=hashed_token, resume=resume): + return resume(await handler.resolve_presented_key(hashed_token)) + case EvictPresentedKey(hashed_token=hashed_token, resume=resume): + await handler.evict_presented_key(hashed_token) + return resume() + case LoadKey(hashed_token=hashed_token, resume=resume): + await handler.load_key(hashed_token) + return resume() + case PrepareKey(loaded_from_db=loaded_from_db, resume=resume): + await handler.prepare_key(loaded_from_db) + return resume() + case InspectKey(resume=resume): + return resume(await handler.inspect_key()) + case EnforceModelAccess(resume=resume): + return resume(await handler.enforce_model_access()) + case PrefetchAuthObjects(resume=resume): + await handler.prefetch_auth_objects() + return resume() + case LoadKeyOwner(user_id=user_id, resume=resume): + return resume(await handler.load_key_owner(user_id)) + case CheckBudgetExemption(resume=resume): + return resume(await handler.check_budget_exemption()) + case LoadKeyContext(resume=resume): + await handler.load_key_context() + return resume() + case _: + assert_never(step) + + +async def _answer_budget(step: BudgetEffect, handler: EffectHandler) -> Step: + match step: + case ReadTeamMemberBudget(user_id=user_id, team_id=team_id, resume=resume): + return resume(await handler.read_team_member_budget(user_id, team_id)) + case AlertTeamMemberBudget(user_id=user_id, team_id=team_id, budget=budget, resume=resume): + await handler.alert_team_member_budget(user_id, team_id, budget) + return resume() + case SendKeyBudgetAlerts(resume=resume): + await handler.send_key_budget_alerts() + return resume() + case ReadKeySpend(max_budget=max_budget, resume=resume): + return resume(await handler.read_key_spend(max_budget)) + case AlertKeyBudget(spend=spend, resume=resume): + await handler.alert_key_budget(spend) + return resume() + case AlertKeySoftBudget(resume=resume): + await handler.alert_key_soft_budget() + return resume() + case ReadProxySpend(resume=resume): + return resume(await handler.read_proxy_spend()) + case AlertProxyBudget(spend=spend, resume=resume): + await handler.alert_proxy_budget(spend) + return resume() + case _: + assert_never(step) + + +async def _answer_model_budget(step: ModelBudgetEffect, handler: ModelBudgetHandler) -> Step: + match step: + case ReadModelBudgetScope(resume=resume): + return resume(await handler.read_model_budget_scope()) + case CheckKeyModelBudget(model=model, resume=resume): + return resume(await handler.check_key_model_budget(model)) + case AuthorizeBudgetFallback(model=model, resume=resume): + return resume(await handler.authorize_budget_fallback(model)) + case RewriteRequestModel(model=model, resume=resume): + await handler.rewrite_request_model(model) + return resume() + case CheckUserModelBudget(model=model, resume=resume): + return resume(await handler.check_user_model_budget(model)) + case _: + assert_never(step) + + +async def drive_model_budgets(step: Step, handler: ModelBudgetHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case ( + ReadModelBudgetScope() + | CheckKeyModelBudget() + | AuthorizeBudgetFallback() + | RewriteRequestModel() + | CheckUserModelBudget() + ): + return await drive_model_budgets(await _answer_model_budget(step, handler), handler) + case _: + raise RuntimeError(f"{type(step).__name__} is not a model budget effect") + + +async def drive(step: Step, handler: EffectHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case _: + return await drive(await _answer(step, handler), handler) + + +async def _answer_team_member(step: TeamMemberEffect, handler: TeamMemberModelBudgetHandler) -> TeamMemberStep: + match step: + case ReadTeamMemberModelBudgetScope(resume=resume): + return resume(await handler.read_team_member_model_budget_scope()) + case CheckTeamMemberModelBudget(model=model, resume=resume): + return resume(await handler.check_team_member_model_budget(model)) + case ReadTeamMemberFallbacks(model=model, resume=resume): + return resume(await handler.read_team_member_fallbacks(model)) + case VetTeamMemberFallback(target=target, resume=resume): + return resume(await handler.vet_team_member_fallback(target)) + case RerouteTeamMemberRequest(model=model, resume=resume): + await handler.reroute_team_member_request(model) + return resume() + case _: + assert_never(step) + + +async def drive_team_member_model_budgets(step: TeamMemberStep, handler: TeamMemberModelBudgetHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case _: + return await drive_team_member_model_budgets(await _answer_team_member(step, handler), handler) diff --git a/litellm/_auth/end_user_context.py b/litellm/_auth/end_user_context.py new file mode 100644 index 00000000000..457cf5e6af6 --- /dev/null +++ b/litellm/_auth/end_user_context.py @@ -0,0 +1,171 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class IdFacts: + value: str | None + + +@dataclass(frozen=True, slots=True) +class DefaultFacts: + prisma_present: bool + explicit_id: str | None + global_id: str | None + global_read: bool = False + + +@dataclass(frozen=True, slots=True) +class DefaultDecision: + read_global: bool + budget_id: str | None + + +@dataclass(frozen=True, slots=True) +class ApplyFacts: + own_budget_id: str | None + own_budget_present: bool + key_id: str | None + global_read: bool = False + global_id: str | None = None + default_read: bool = False + default_present: bool = False + + +ApplyDecision: TypeAlias = Literal["Keep", "ReadGlobal", "ReadDefault", "Copy"] + + +@dataclass(frozen=True, slots=True) +class RegistryFacts: + end_user_id: str + global_id: str | None + key_id: str | None + validate_id: bool + token_max_present: bool + registry_read: bool = False + registry: tuple[str, ...] | None = None + + +RegistryDecision: TypeAlias = Literal["Read", "Restricted", "Unrestricted"] +Field: TypeAlias = Literal["EndUserId", "Tpm", "Rpm", "Tpd", "MaxBudget", "Models", "Region"] + + +@dataclass(frozen=True, slots=True) +class Fields: + tpm: bool = False + rpm: bool = False + tpd: bool = False + max_budget: bool = False + models: bool = False + region: bool = False + + +@dataclass(frozen=True, slots=True) +class AssignmentFacts: + mode: Literal["Params", "Token"] + fields: Fields + + +@dataclass(frozen=True, slots=True) +class LookupFacts: + key_id: str | None + own_budget_present: bool = False + end_user_read: bool = False + end_user_present: bool = False + global_read: bool = False + global_id: str | None = None + default_read: bool = False + default_present: bool = False + token_max_present: bool = False + applied: bool = False + read_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class LookupDecision: + action: Literal["ReadEndUser", "ReadGlobal", "ReadDefault", "ApplyOwn", "ApplyDefault", "Done"] + set_max_budget: bool + project_budget: bool + + +@dataclass(frozen=True, slots=True) +class ValidationFacts: + mode: Literal["Resolve", "Exists"] + end_user_id: str | None + enabled: bool = False + prisma_present: bool = False + global_id: str | None = None + key_id: str | None = None + cache_read: bool = False + cached: str | None = None + end_user_read: bool = False + end_user_present: bool = False + user_read: bool = False + user_present: bool = False + cache_written: bool = False + + +@dataclass(frozen=True, slots=True) +class ValidationDecision: + read: Literal["Cache", "EndUser", "User", "WriteCache"] | None + end_user_id: str | None + valid: bool + cache_value: str | None + cache_ttl: int + user_email: str | None + + +_ID: Final = TypeAdapter(IdFacts) +_STRING: Final = TypeAdapter[str | None](str | None) +_DEFAULT: Final = TypeAdapter(DefaultFacts) +_DEFAULT_RESULT: Final = TypeAdapter(DefaultDecision) +_APPLY: Final = TypeAdapter(ApplyFacts) +_APPLY_RESULT: Final = TypeAdapter[ApplyDecision](ApplyDecision) +_REGISTRY: Final = TypeAdapter(RegistryFacts) +_REGISTRY_RESULT: Final = TypeAdapter[RegistryDecision](RegistryDecision) +_ASSIGNMENTS: Final = TypeAdapter(AssignmentFacts) +_FIELDS: Final = TypeAdapter(tuple[Field, ...]) +_LOOKUP: Final = TypeAdapter(LookupFacts) +_LOOKUP_RESULT: Final = TypeAdapter(LookupDecision) +_VALIDATION: Final = TypeAdapter(ValidationFacts) +_VALIDATION_RESULT: Final = TypeAdapter(ValidationDecision) + + +@native_boundary +def normalize_id(facts: IdFacts) -> str | None: + return _STRING.validate_python(_native.request_auth_end_user_normalize_id(_ID.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def default_id(facts: DefaultFacts) -> DefaultDecision: + return _DEFAULT_RESULT.validate_python(_native.request_auth_end_user_default_id(_DEFAULT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def apply_default(facts: ApplyFacts) -> ApplyDecision: + return _APPLY_RESULT.validate_python(_native.request_auth_end_user_apply(_APPLY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def registry(facts: RegistryFacts) -> RegistryDecision: + return _REGISTRY_RESULT.validate_python(_native.request_auth_end_user_registry(_REGISTRY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def assignments(facts: AssignmentFacts) -> tuple[Field, ...]: + return _FIELDS.validate_python(_native.request_auth_end_user_assignments(_ASSIGNMENTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def lookup(facts: LookupFacts) -> LookupDecision: + return _LOOKUP_RESULT.validate_python(_native.request_auth_end_user_lookup(_LOOKUP.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def validation(facts: ValidationFacts) -> ValidationDecision: + return _VALIDATION_RESULT.validate_python(_native.request_auth_end_user_validation(_VALIDATION.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts diff --git a/litellm/_auth/failure.py b/litellm/_auth/failure.py new file mode 100644 index 00000000000..66b13582096 --- /dev/null +++ b/litellm/_auth/failure.py @@ -0,0 +1,39 @@ +from collections.abc import Callable +from functools import wraps +from typing import ParamSpec, TypeVar + +_P = ParamSpec("_P") # rebind-ok: ParamSpec declarations cannot be annotated Final +_R = TypeVar("_R") # rebind-ok: TypeVar declarations cannot be annotated Final + + +class NativeAuthFailure(Exception): + pass + + +class NativeAuthTypeError(NativeAuthFailure, TypeError): + pass + + +class NativeAuthValueError(NativeAuthFailure, ValueError): + pass + + +class NativeAuthRuntimeError(NativeAuthFailure, RuntimeError): + pass + + +def native_boundary(function: Callable[_P, _R]) -> Callable[_P, _R]: + @wraps(function) + def call(*args: _P.args, **kwargs: _P.kwargs) -> _R: # kwargs-ok: ParamSpec preserves the wrapped signature + try: + return function(*args, **kwargs) + except NativeAuthFailure: + raise + except TypeError as error: + raise NativeAuthTypeError(str(error)) from error + except ValueError as error: + raise NativeAuthValueError(str(error)) from error + except Exception as error: # noqa: BLE001 # native faults must bypass legacy IO fail-open catches + raise NativeAuthRuntimeError(str(error)) from error + + return call diff --git a/litellm/_auth/fallback_access.py b/litellm/_auth/fallback_access.py new file mode 100644 index 00000000000..057f3a47745 --- /dev/null +++ b/litellm/_auth/fallback_access.py @@ -0,0 +1,58 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class GroupedModelScope: + facts: ModelAccessFacts + group_ids: tuple[str, ...] = () + group_models: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class FallbackAuthorizationFacts: + key: GroupedModelScope | None = None + refuse_agent_keys: bool = False + agent_id: str | None = None + team: GroupedModelScope | None = None + embedded_team: ModelAccessFacts | None = None + assigned_models: tuple[str, ...] | None = None + member: ModelAccessFacts | None = None + project: ModelAccessFacts | None = None + team_present: bool = False + embedded_team_grants: tuple[str, ...] = () + member_grants: tuple[str, ...] = () + project_grants: tuple[str, ...] = () + assigned_scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class FallbackAuthorizationDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadKey", + "ReadTeam", + "ReadEmbeddedTeam", + "ReadMember", + "ReadProject", + "ReadKeyGroups", + "ReadTeamGroups", + "ReadAssignedGroups", + "ReadAssignedScope", + ] + denial: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(FallbackAuthorizationDecision) + + +@native_boundary +def fallback_authorization(facts: FallbackAuthorizationFacts) -> FallbackAuthorizationDecision: + return _DECISION.validate_python(_native.request_auth_fallback_access(asdict(facts))) diff --git a/litellm/_auth/fallback_candidate.py b/litellm/_auth/fallback_candidate.py new file mode 100644 index 00000000000..39c50cdfd63 --- /dev/null +++ b/litellm/_auth/fallback_candidate.py @@ -0,0 +1,88 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.common_budgets import BudgetFacts +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.model_budget_data import ModelBudgetEntity, ModelBudgetSnapshot +from litellm._auth.routes import RouteContext +from litellm._auth.unpriced_models import UnpricedModelFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class CandidateBudgetScope: + entity: ModelBudgetEntity + identity_present: bool + configured: bool + + +@dataclass(frozen=True, slots=True) +class CandidateBudgetRead: + model_index: int + entity: ModelBudgetEntity + snapshot: ModelBudgetSnapshot + + +@dataclass(frozen=True, slots=True) +class FallbackCandidateFacts: + router_present: bool + user_present: bool + team_present: bool + membership: Literal["Unread", "Loaded", "Failed"] = "Unread" + block_unpriced: bool = False + route: RouteContext | None = None + unpriced: UnpricedModelFacts | None = None + customer_present: bool = False + customer_grants: tuple[str, ...] | None = None + customer: ModelAccessFacts | None = None + price: ModelPriceFacts | None = None + budget_models: tuple[str, ...] | None = None + budget_scopes: tuple[CandidateBudgetScope, ...] | None = None + budgets: tuple[CandidateBudgetRead, ...] = () + authorization: FallbackAuthorizationFacts | None = None + attribution_reset: bool = False + attribution: ModelGroupAttributionFacts | None = None + attribution_stamped: bool = False + group_budgets_read: bool = False + group_budgets: tuple[BudgetFacts, ...] = () + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class FallbackCandidateDecision: + kind: Literal[ + "Complete", + "ReadMembership", + "ReadRoute", + "ReadUnpriced", + "ReadCustomerGrants", + "ReadCustomer", + "ReadPrice", + "ReadBudgetModels", + "ReadBudgetScopes", + "ReadBudget", + "ReadAuthorization", + "ResetAttribution", + "ReadAttribution", + "StampAttribution", + "ReadGroupBudgets", + ] + usable: bool = False + log_failure: bool = False + model_index: int = 0 + entity: ModelBudgetEntity = "Key" + groups: tuple[str, ...] = () + + +_DECISION: Final = TypeAdapter(FallbackCandidateDecision) + + +@native_boundary +def fallback_candidate(facts: FallbackCandidateFacts) -> FallbackCandidateDecision: + return _DECISION.validate_python(_native.request_auth_fallback_candidate(asdict(facts))) diff --git a/litellm/_auth/key.py b/litellm/_auth/key.py new file mode 100644 index 00000000000..af63778e3cf --- /dev/null +++ b/litellm/_auth/key.py @@ -0,0 +1,12 @@ +import hashlib +import secrets + + +def hash_token(token: str) -> str: + return hashlib.sha256(token.encode()).hexdigest() + + +def is_master_key(token: str | None, master_key: str | None) -> bool: + if token is None or master_key is None: + return False + return secrets.compare_digest(token.encode(), master_key.encode()) diff --git a/litellm/_auth/key_model_access.py b/litellm/_auth/key_model_access.py new file mode 100644 index 00000000000..00666dd2a57 --- /dev/null +++ b/litellm/_auth/key_model_access.py @@ -0,0 +1,45 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessDecision +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class KeyAndFallbackModelAccessFacts: + config_nonempty: bool + key_models: tuple[str, ...] + primary_loaded: bool = False + primary: GroupedModelScope | None = None + fallback_names: tuple[str, ...] | None = None + fallback_scopes: tuple[GroupedModelScope, ...] = () + routing: tuple[Literal["Routed", "Failed"], ...] = () + + +@dataclass(frozen=True, slots=True) +class KeyAndFallbackModelAccessDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadPrimary", + "ReadPrimaryGroups", + "ReadFallbackNames", + "ReadFallbackScope", + "ReadFallbackGroups", + "RouteFallback", + "RoutingFailure", + ] + index: int = 0 + denial: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(KeyAndFallbackModelAccessDecision) + + +@native_boundary +def key_and_fallback_model_access(facts: KeyAndFallbackModelAccessFacts) -> KeyAndFallbackModelAccessDecision: + return _DECISION.validate_python(_native.request_auth_key_model_access(asdict(facts))) diff --git a/litellm/_auth/managed.py b/litellm/_auth/managed.py new file mode 100644 index 00000000000..19c9c64c960 --- /dev/null +++ b/litellm/_auth/managed.py @@ -0,0 +1,129 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class BoundIdentity: + active: bool + revision: str + + +@dataclass(frozen=True, slots=True) +class Agent: + agent_id: str + enabled: bool + identity_managed: bool + identity: BoundIdentity | None + execution_mode: str + + +@dataclass(frozen=True, slots=True) +class Context: + agent_id: str + binding_revision: str | None + mode: str + user_id: str | None + + +@dataclass(frozen=True, slots=True) +class Actor: + agent: Agent + context: Context | None + + +@dataclass(frozen=True, slots=True) +class ActorLookup: + agent_id: str | None + has_database: bool + registry_loaded: bool = False + registry_managed: bool = False + agent_loaded: bool = False + agent: Agent | None = None + retired: bool | None = None + context: Context | None = None + jwt_claims: bool = False + delegation_verified: bool = False + human_grants: tuple[str, ...] | None = None + policy_assigned: bool = False + + +@dataclass(frozen=True, slots=True) +class UnionGrants: + grants: tuple[tuple[str, ...] | None, ...] + + +@dataclass(frozen=True, slots=True) +class IntersectGrants: + left: tuple[str, ...] | None + right: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class DeclaredGrants: + direct: tuple[str, ...] + declared_groups: tuple[str, ...] + unified_groups: tuple[str, ...] + group_agents: tuple[str, ...] | None = None + unified_agents: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ManagedAgentGrants: + permission_present: bool + mode: str | None + user_present: bool + own: tuple[str, ...] | None = None + ceilings: tuple[tuple[str, ...], ...] | None = None + caller_loaded: bool = False + caller: tuple[str, ...] | None = None + human: tuple[str, ...] | None = None + + +ManagedDenial: TypeAlias = Literal[ + "Disabled", + "MissingContext", + "ChangedIdentity", + "ExecutionMode", + "MissingHuman", + "DatabaseRequired", + "Removed", + "VerifiedIdentityRequired", + "HumanNotPermitted", +] + + +@dataclass(frozen=True, slots=True) +class ManagedDecision: + action: Literal[ + "Allow", + "Deny", + "ReadRegistry", + "ReadAgent", + "ReadRetired", + "AssignPolicy", + "ReadHumanGrants", + "Grants", + "ReadDeclaredGroups", + "ReadUnifiedGroups", + "ReadOwn", + "ReadCeilings", + "ReadCaller", + ] + denial: ManagedDenial | None + grants: tuple[str, ...] | None + + +_DECISION: Final = TypeAdapter(ManagedDecision) + + +@native_boundary +def decide( + facts: Actor | ActorLookup | UnionGrants | IntersectGrants | DeclaredGrants | ManagedAgentGrants, +) -> ManagedDecision: + return _DECISION.validate_python(_native.request_auth_managed(encode_facts(facts))) diff --git a/litellm/_auth/managed_models.py b/litellm/_auth/managed_models.py new file mode 100644 index 00000000000..2a41347d329 --- /dev/null +++ b/litellm/_auth/managed_models.py @@ -0,0 +1,95 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ManagedModelFacts: + purpose: Literal["All", "Ceilings", "Caller"] + model_present: bool + token_present: bool + managed_loaded: bool = False + managed_present: bool = False + managed_models_valid: bool = False + managed_models_count: int = 0 + managed: ModelAccessFacts | None = None + ceilings: tuple[tuple[str, ...], ...] | None = None + aliases_loaded: bool = False + scopes: tuple[ModelAccessFacts, ...] = () + caller_loaded: bool = False + caller_present: bool = False + team_loaded: bool = False + team_present: bool = False + team: GroupedModelScope | None = None + member: MemberModelFacts | None = None + user_loaded: bool = False + user_present: bool = False + user_grants: tuple[str, ...] = () + user: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ManagedModelDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadManaged", + "ReadManagedScope", + "ReadCeilings", + "ReadAliases", + "ReadCeilingScope", + "ReadCaller", + "ReadTeam", + "ReadTeamScope", + "ReadTeamGroups", + "ReadMember", + "ReadUser", + "ReadUserScope", + ] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + index: int = 0 + + +@dataclass(frozen=True, slots=True) +class AgentCeilingSourceFacts: + agent_present: bool + managed_loaded: bool = False + managed_present: bool = False + loaded: bool = False + count: int = 0 + + +@dataclass(frozen=True, slots=True) +class AgentCeilingSourceDecision: + kind: Literal["ReadManaged", "ReadManagedCeilings", "ReadUnmanagedCeiling", "Complete"] + indices: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(ManagedModelDecision) +_CEILINGS: Final = TypeAdapter(AgentCeilingSourceDecision) +_GRANTS: Final = TypeAdapter(frozenset[str]) + + +@native_boundary +def managed_models(facts: ManagedModelFacts) -> ManagedModelDecision: + return _DECISION.validate_python(_native.request_auth_managed_models(asdict(facts))) + + +@native_boundary +def agent_ceiling_source(facts: AgentCeilingSourceFacts) -> AgentCeilingSourceDecision: + return _CEILINGS.validate_python(_native.request_auth_agent_ceiling_source(asdict(facts))) + + +@native_boundary +def managed_model_grants( + models: tuple[str | None, ...] | None, ceilings: tuple[tuple[str, ...], ...] +) -> frozenset[str]: + return _GRANTS.validate_python(_native.request_auth_managed_model_grants({"models": models, "ceilings": ceilings})) diff --git a/litellm/_auth/member_models.py b/litellm/_auth/member_models.py new file mode 100644 index 00000000000..a2e845b69f4 --- /dev/null +++ b/litellm/_auth/member_models.py @@ -0,0 +1,32 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class MemberModelFacts: + user_present: bool + team_present: bool + membership_loaded: bool = False + grants: tuple[str, ...] | None = None + scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class MemberModelDecision: + kind: Literal["Allow", "Deny", "ReadMembership", "ReadGrants", "ReadScope"] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + + +_DECISION: Final = TypeAdapter(MemberModelDecision) + + +@native_boundary +def member_models(facts: MemberModelFacts) -> MemberModelDecision: + return _DECISION.validate_python(_native.request_auth_member_models(asdict(facts))) diff --git a/litellm/_auth/model_access.py b/litellm/_auth/model_access.py new file mode 100644 index 00000000000..e3c29ad4c64 --- /dev/null +++ b/litellm/_auth/model_access.py @@ -0,0 +1,261 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import asdict, dataclass, field +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +ModelObject: TypeAlias = Literal["user", "customer", "team", "key", "org", "project", "agent"] +ModelReadKind: TypeAlias = Literal["Compaction", "Aliases", "AliasTarget", "Deployment", "Groups", "Names", "Provider"] + + +@dataclass(frozen=True, slots=True) +class ModelProvider: + stripped_model: str + provider: str + + +@dataclass(frozen=True, slots=True) +class ModelRecord: + access_groups: tuple[str, ...] = () + provider: ModelProvider | None = None + deployment_team: str | None = None + deployment_public_name: str | None = None + + +@dataclass(frozen=True, slots=True) +class ModelReading: + kind: ModelReadKind + model: str + values: tuple[str, ...] = () + value: str | None = None + team: str | None = None + provider: ModelProvider | None = None + + +@dataclass(frozen=True, slots=True) +class ModelEvaluation: + grants: tuple[str, ...] + reads: tuple[ModelReading, ...] + aliases: ModelAliases | None = None + + +@dataclass(frozen=True, slots=True) +class ModelRegistry: + present: bool = False + names: tuple[str, ...] = () + aliases: Mapping[str, str] = field(default_factory=dict[str, str]) + records: Mapping[str, ModelRecord] = field(default_factory=dict[str, ModelRecord]) + compaction_parents: Mapping[str, str] = field(default_factory=dict[str, str]) + helper_present: bool | None = None + evaluations: tuple[ModelEvaluation, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ModelAliases: + team: Mapping[str, str] = field(default_factory=dict[str, str]) + key: Mapping[str, str] = field(default_factory=dict[str, str]) + global_: Mapping[str, str] = field(default_factory=dict[str, str]) + team_order: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ModelAccessFacts: + requested: tuple[str, ...] + grants: tuple[str, ...] = () + team_id: str | None = None + aliases: ModelAliases = field(default_factory=ModelAliases) + registry: ModelRegistry = field(default_factory=ModelRegistry) + object_type: ModelObject = "user" + fallback_depth: int = 0 + max_depth: int = 5 + is_list: bool = False + + +@dataclass(frozen=True, slots=True) +class ModelAccessDecision: + kind: Literal[ + "Allow", + "Deny", + "Depth", + "NoDefaultModels", + "NoAgentModels", + "EmptyAgentCeiling", + "PersonalTeam", + "ReadGroups", + "ReadScope", + "ReadGroupScope", + ] + model: str | None = None + grants: tuple[str, ...] = () + object_type: ModelObject = "user" + + +@dataclass(frozen=True, slots=True) +class AssignedModelGroup: + assigned_team_ids: tuple[str, ...] + assigned_key_ids: tuple[str, ...] + models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class PersonalModelTeam: + blocked: bool + available: bool + team: ModelAccessFacts + member: ModelAccessFacts | None = None + + +_DECISION: Final = TypeAdapter(ModelAccessDecision) +_MODELS: Final = TypeAdapter(tuple[str, ...]) +_BOOL: Final = TypeAdapter(bool) + + +@native_boundary +def model_access(facts: ModelAccessFacts, *, kind: Literal["Object", "Customer"] = "Object") -> ModelAccessDecision: + return _DECISION.validate_python(_native.request_auth_model_access({kind: asdict(facts)})) + + +@native_boundary +def grouped_model_access( + facts: ModelAccessFacts, + group_ids: tuple[str, ...], + group_models: tuple[str, ...] | None = None, + *, + merge: bool = False, +) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access( + { + "Grouped": { + "facts": asdict(facts), + "group_ids": group_ids, + "group_models": group_models, + "merge": merge, + } + } + ) + ) + + +@native_boundary +def user_model_access(facts: ModelAccessFacts, *, present: bool) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access({"User": {"facts": asdict(facts), "present": present}}) + ) + + +@native_boundary +def user_model_precheck(*, present: bool, grants: tuple[str, ...]) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access({"UserPrecheck": {"present": present, "grants": grants}}) + ) + + +@native_boundary +def agent_model_access( + facts: ModelAccessFacts, managed_grants: tuple[str, ...] | None, ceilings: tuple[tuple[str, ...], ...] +) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access( + { + "Agent": { + "facts": asdict(facts), + "managed_grants": managed_grants, + "ceilings": ceilings, + } + } + ) + ) + + +@native_boundary +def personal_model_access(requested: tuple[str, ...], teams: tuple[PersonalModelTeam, ...]) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access( + { + "Personal": { + "requested": requested, + "teams": tuple(asdict(team) for team in teams), + } + } + ) + ) + + +@native_boundary +def model_candidates(facts: ModelAccessFacts) -> tuple[str, ...]: + return _MODELS.validate_python(_native.request_auth_model_access({"Candidates": asdict(facts)})) + + +@native_boundary +def model_access_helper(model: str, facts: ModelAccessFacts) -> bool: + return _BOOL.validate_python( + _native.request_auth_model_access({"Helper": {"model": model, "facts": asdict(facts)}}) + ) + + +@native_boundary +def key_model_grants(grants: tuple[str, ...], team_id: str | None, team_grants: tuple[str, ...]) -> tuple[str, ...]: + return _MODELS.validate_python( + _native.request_auth_model_access( + { + "KeyGrants": { + "grants": grants, + "team_id": team_id, + "team_grants": team_grants, + } + } + ) + ) + + +@native_boundary +def assigned_model_grants( + groups: tuple[AssignedModelGroup, ...], team_id: str | None, key_id: str | None +) -> tuple[str, ...]: + return _MODELS.validate_python( + _native.request_auth_model_access( + { + "AssignedGroups": { + "groups": tuple(asdict(group) for group in groups), + "team_id": team_id, + "key_id": key_id, + } + } + ) + ) + + +@native_boundary +def personal_model_access_applies( + *, + enabled: bool, + token_present: bool, + team_id: str | None, + user_id: str | None, + master: bool, + role: str | None, + api_key_present: bool, + inheritance: bool = False, +) -> bool: + return _BOOL.validate_python( + _native.request_auth_model_access( + { + "PersonalApplies": { + "enabled": enabled, + "token_present": token_present, + "team_id": team_id, + "user_id": user_id, + "master": master, + "role": role, + "api_key_present": api_key_present, + "inheritance": inheritance, + } + } + ) + ) diff --git a/litellm/_auth/model_attribution.py b/litellm/_auth/model_attribution.py new file mode 100644 index 00000000000..aa2dd98fdd4 --- /dev/null +++ b/litellm/_auth/model_attribution.py @@ -0,0 +1,46 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelGroupAttributionFacts: + model_present: bool = False + token_present: bool = False + router_present: bool = False + database_present: bool = False + team_present: bool = False + user_present: bool = False + membership_loaded: bool = False + token_org: str | None = None + team_org: str | None = None + registry_loaded: bool = False + registry: tuple[str, ...] | None = None + covering: tuple[str, ...] | None = None + grants: tuple[tuple[str, ...], ...] = () + org_failed: bool = False + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class ModelGroupAttributionDecision: + kind: Literal[ + "Complete", "ReadRegistry", "ReadCovering", "ReadKey", "ReadTeam", "ReadMember", "ReadProject", "ReadOrg" + ] + groups: tuple[str, ...] = () + stamp: bool = False + lookup: bool = False + include: bool = False + id: str | None = None + + +_DECISION: Final = TypeAdapter(ModelGroupAttributionDecision) + + +@native_boundary +def model_group_attribution(facts: ModelGroupAttributionFacts) -> ModelGroupAttributionDecision: + return _DECISION.validate_python(_native.request_auth_model_attribution(asdict(facts))) diff --git a/litellm/_auth/model_budget_aliases.py b/litellm/_auth/model_budget_aliases.py new file mode 100644 index 00000000000..b9fedf9f091 --- /dev/null +++ b/litellm/_auth/model_budget_aliases.py @@ -0,0 +1,59 @@ +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_budget_data import ModelBudgetSnapshot +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Row: + key: str + names: tuple[str, ...] + prefix_known: bool | None = None + provider_read: bool = False + provider: str | None = None + + +@dataclass(frozen=True, slots=True) +class Facts: + model: str + rows: tuple[Row, ...] = () + + +@dataclass(frozen=True, slots=True) +class Decision: + read: Literal["Entry", "Prefix", "Provider"] | None + key: str + models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class ScopeFacts: + models: tuple[str, ...] + snapshots: tuple[ModelBudgetSnapshot, ...] = () + budget_present: bool = True + + +@dataclass(frozen=True, slots=True) +class ScopeDecision: + read_index: int | None + result_index: int | None + + +_FACTS: Final = TypeAdapter(Facts) +_DECISION: Final = TypeAdapter(Decision) +_SCOPE: Final = TypeAdapter(ScopeFacts) +_SCOPE_DECISION: Final = TypeAdapter(ScopeDecision) + + +@native_boundary +def decide(facts: Facts) -> Decision: + return _DECISION.validate_python(_native.request_auth_model_budget_aliases(_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes typed native facts + + +@native_boundary +def scope(facts: ScopeFacts) -> ScopeDecision: + return _SCOPE_DECISION.validate_python(_native.request_auth_model_budget_scope(_SCOPE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes typed native facts diff --git a/litellm/_auth/model_budget_data.py b/litellm/_auth/model_budget_data.py new file mode 100644 index 00000000000..59e39ce8da6 --- /dev/null +++ b/litellm/_auth/model_budget_data.py @@ -0,0 +1,94 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +ModelBudgetEntity: TypeAlias = Literal["Key", "User", "Team", "TeamMember", "EndUser"] + + +@dataclass(frozen=True, slots=True) +class ModelBudgetEntry: + name: str + max_budget: float | None + duration: str | None + duration_seconds: float | None + + +@dataclass(frozen=True, slots=True) +class ModelBudgetResolveFacts: + model: str + bedrock_base: str + bedrock_provider: str + entries: tuple[ModelBudgetEntry, ...] + key_entries: tuple[ModelBudgetEntry, ...] = () + team: bool = False + access_groups: tuple[str, ...] | None = None + excluded_entries: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class ModelBudgetResolution: + read_groups: bool + entry: ModelBudgetEntry | None + named_model: bool + + +@dataclass(frozen=True, slots=True) +class ModelBudgetSnapshot: + max_budget: float | None = None + local_spend: float | None = None + legacy_spend: float | None = None + shared_spend: float | None = None + shared_started_at: float | None = None + shared_window_seconds: float | None = None + now: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class ModelBudgetQueryFacts: + entity: ModelBudgetEntity + snapshot: ModelBudgetSnapshot + duration_seconds: float | None = None + process_uptime: float = 0.0 + same_counter: bool = False + redis_present: bool = False + shared_read: bool = False + + +@dataclass(frozen=True, slots=True) +class ModelBudgetDecision: + read: Literal["Local", "Legacy", "Shared"] | None + denied: bool + current_spend: float + + +_RESOLVE: Final = TypeAdapter(ModelBudgetResolveFacts) +_RESOLUTION: Final = TypeAdapter(ModelBudgetResolution) +_QUERY: Final = TypeAdapter(ModelBudgetQueryFacts) +_SNAPSHOT: Final = TypeAdapter(ModelBudgetSnapshot) +_DECISION: Final = TypeAdapter(ModelBudgetDecision) +_TEAM_FACTS: Final = TypeAdapter(tuple[ModelBudgetResolveFacts, ...]) +_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@native_boundary +def team_exclusions(facts: tuple[ModelBudgetResolveFacts, ...]) -> tuple[str, ...]: + return _NAMES.validate_python(_native.request_auth_model_budget_team_exclusions(_TEAM_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes typed native facts + + +@native_boundary +def resolve(facts: ModelBudgetResolveFacts) -> ModelBudgetResolution: + return _RESOLUTION.validate_python(_native.request_auth_model_budget_resolve(_RESOLVE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def query(facts: ModelBudgetQueryFacts) -> ModelBudgetDecision: + return _DECISION.validate_python(_native.request_auth_model_budget_query(_QUERY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def check(facts: ModelBudgetSnapshot) -> ModelBudgetDecision: + return _DECISION.validate_python(_native.request_auth_model_budget_check(_SNAPSHOT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary diff --git a/litellm/_auth/model_groups.py b/litellm/_auth/model_groups.py new file mode 100644 index 00000000000..43a92d972a2 --- /dev/null +++ b/litellm/_auth/model_groups.py @@ -0,0 +1,48 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class GroupResources: + assigned_team_ids: tuple[str, ...] + assigned_key_ids: tuple[str, ...] + resources: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class GroupLookup: + kind: Literal["Found", "Missing", "Failed"] + group: GroupResources | None = None + + +@dataclass(frozen=True, slots=True) +class ResourceGroupFacts: + subject_present: bool + database_present: bool + cache_present: bool + group_ids: tuple[str, ...] + team_id: str | None + key_id: str | None + check_db_only: bool = False + strict_errors: bool = False + loaded: tuple[GroupLookup, ...] = () + + +@dataclass(frozen=True, slots=True) +class ResourceGroupDecision: + kind: Literal["Read", "Raise", "Complete"] + index: int = 0 + resources: tuple[str, ...] = () + + +_DECISION: Final = TypeAdapter(ResourceGroupDecision) + + +@native_boundary +def resource_group_decision(facts: ResourceGroupFacts) -> ResourceGroupDecision: + return _DECISION.validate_python(_native.request_auth_model_groups(asdict(facts))) diff --git a/litellm/_auth/model_helpers.py b/litellm/_auth/model_helpers.py new file mode 100644 index 00000000000..8de363b28ac --- /dev/null +++ b/litellm/_auth/model_helpers.py @@ -0,0 +1,37 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelProvider +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelHelperFacts: + kind: Literal["Group", "Helper"] + model: str + grants: tuple[str, ...] | None + team_id: str | None = None + team_aliases: tuple[str, ...] = () + router_present: bool = False + groups: tuple[str, ...] | None = None + names: tuple[str, ...] | None = None + provider_loaded: bool = False + provider: ModelProvider | None = None + provider_reads: tuple[ModelProvider | None, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ModelHelperDecision: + kind: Literal["Complete", "ReadGroups", "ReadNames", "ReadProvider"] + allowed: bool = False + + +_DECISION: Final = TypeAdapter(ModelHelperDecision) + + +@native_boundary +def model_helper(facts: ModelHelperFacts) -> ModelHelperDecision: + return _DECISION.validate_python(_native.request_auth_model_helpers(asdict(facts))) diff --git a/litellm/_auth/model_projection.py b/litellm/_auth/model_projection.py new file mode 100644 index 00000000000..832b1cf13b9 --- /dev/null +++ b/litellm/_auth/model_projection.py @@ -0,0 +1,27 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts, ModelReading, ModelReadKind +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelProjectionDecision: + kind: Literal["Complete", "Read"] + read: ModelReadKind | None = None + model: str | None = None + team_id: str | None = None + decision: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(ModelProjectionDecision) + + +@native_boundary +def model_projection(facts: ModelAccessFacts, reads: tuple[ModelReading, ...]) -> ModelProjectionDecision: + return _DECISION.validate_python( + _native.request_auth_model_projection(asdict(facts), tuple(asdict(read) for read in reads)) + ) diff --git a/litellm/_auth/org_routes.py b/litellm/_auth/org_routes.py new file mode 100644 index 00000000000..66563e6a5c1 --- /dev/null +++ b/litellm/_auth/org_routes.py @@ -0,0 +1,95 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class OrganizationMembership: + organization_id: str | None + role: str | None + + +@dataclass(frozen=True, slots=True) +class OrganizationRouteFacts: + route: str + user_present: bool + user_role: str | None + memberships: tuple[OrganizationMembership, ...] | None + organization_id: str | None + org_admin_routes: tuple[str, ...] + organization_id_present: bool = False + identifier_pending: bool = False + hash_failure: str | None = None + request_pending: bool = False + memberships_pending: bool = False + + +OrgRouteDenial: TypeAlias = Literal[ + "CreateOrganization", + "NoMemberships", + "MissingOrganization", + "MissingRole", + "WrongRole", + "TeamMissingOrganization", + "TeamWrongRole", +] + + +@dataclass(frozen=True, slots=True) +class OrganizationRouteDecision: + denial: OrgRouteDenial | None + organization_role: str | None + organization_ids: tuple[str, ...] + read_identifier: bool = False + hash_failure: str | None = None + read_request: bool = False + read_memberships: bool = False + + +_DECISION: Final = TypeAdapter(OrganizationRouteDecision) + + +@native_boundary +def organization_route(facts: OrganizationRouteFacts) -> OrganizationRouteDecision: + return _DECISION.validate_python(_native.request_auth_org_route(asdict(facts))) + + +@native_boundary +def user_is_org_admin(memberships: tuple[OrganizationMembership, ...], organizations: tuple[str, ...]) -> bool: + return _native.request_auth_org_admin(tuple(asdict(item) for item in memberships), organizations) + + +@dataclass(frozen=True, slots=True) +class Identifier: + text: str | None + hash_failure: str | None = None + + +@dataclass(frozen=True, slots=True) +class AdminFacts: + memberships: tuple[OrganizationMembership, ...] | None + request_read: bool = False + singular_present: bool = False + plural_count: int = 0 + candidate_index: int = 0 + candidate: Identifier | None = None + + +@dataclass(frozen=True, slots=True) +class AdminDecision: + action: Literal["ReadRequest", "ReadSingular", "ReadPlural", "Allow", "Deny", "HashFailure"] + candidate_index: int + plural_index: int + hash_failure: str | None + + +_ADMIN: Final = TypeAdapter(AdminDecision) + + +@native_boundary +def admin(facts: AdminFacts) -> AdminDecision: + return _ADMIN.validate_python(_native.request_auth_org_admin_raw(asdict(facts))) diff --git a/litellm/_auth/owner_teams.py b/litellm/_auth/owner_teams.py new file mode 100644 index 00000000000..be1919eb626 --- /dev/null +++ b/litellm/_auth/owner_teams.py @@ -0,0 +1,53 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.personal_budget import PersonalTeamBudgetFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class OwnerTeamLoad: + present: bool + failed: bool = False + blocked: bool = False + org: str | None = None + + +@dataclass(frozen=True, slots=True) +class OwnerTeamBudget: + index: int + facts: PersonalTeamBudgetFacts + + +@dataclass(frozen=True, slots=True) +class OwnerTeamFacts: + purpose: Literal["Models", "Grants"] + ids: tuple[str, ...] + ui_team_id: str + user_present: bool = True + database_present: bool = True + org: str | None = None + loaded: bool = False + loads: tuple[OwnerTeamLoad, ...] = () + budgets: tuple[OwnerTeamBudget, ...] = () + + +@dataclass(frozen=True, slots=True) +class OwnerTeamDecision: + kind: Literal["ReadTeams", "ReadBudgets", "Complete"] + ids: tuple[str, ...] = () + indices: tuple[int, ...] = () + selected: tuple[int, ...] = () + withheld: tuple[int, ...] = () + unreadable: bool = False + + +_DECISION: Final = TypeAdapter(OwnerTeamDecision) + + +@native_boundary +def owner_teams(facts: OwnerTeamFacts) -> OwnerTeamDecision: + return _DECISION.validate_python(_native.request_auth_owner_teams(asdict(facts))) diff --git a/litellm/_auth/personal_budget.py b/litellm/_auth/personal_budget.py new file mode 100644 index 00000000000..b82fe1d1937 --- /dev/null +++ b/litellm/_auth/personal_budget.py @@ -0,0 +1,66 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.common_budgets import BudgetFacts, MemberLimitFacts +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalTeamBudgetFacts: + user_present: bool = False + team_id_present: bool = True + check_db_only: bool = False + token_org: str | None = None + team_org: str | None = None + source_ready: bool = False + team: BudgetFacts | None = None + team_alerted: bool = False + window_count: int | None = None + windows: tuple[BudgetFacts, ...] = () + membership_loaded: bool = False + member_limit: MemberLimitFacts | None = None + default_id: str | None = None + default_loaded: bool = False + member: BudgetFacts | None = None + member_alerted: bool = False + org_loaded: bool = False + org_failed: bool = False + org: BudgetFacts | None = None + org_alerted: bool = False + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class PersonalTeamBudgetDecision: + kind: Literal[ + "Complete", + "Source", + "ReadTeam", + "AlertTeam", + "ReadWindows", + "ReadWindow", + "ReadMembership", + "ReadMemberLimit", + "ReadDefault", + "ReadMember", + "AlertMember", + "ReadOrg", + "AlertOrg", + ] + allowed: bool = False + log_failure: bool = False + org_id: str | None = None + id: str = "" + index: int = 0 + limit: float | None = None + + +_DECISION: Final = TypeAdapter(PersonalTeamBudgetDecision) + + +@native_boundary +def personal_team_budget(facts: PersonalTeamBudgetFacts) -> PersonalTeamBudgetDecision: + return _DECISION.validate_python(_native.request_auth_personal_budget(asdict(facts))) diff --git a/litellm/_auth/personal_model_listing.py b/litellm/_auth/personal_model_listing.py new file mode 100644 index 00000000000..e90cc43c388 --- /dev/null +++ b/litellm/_auth/personal_model_listing.py @@ -0,0 +1,32 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm._auth.personal_owner import PersonalOwnerFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelListingFacts: + owner: PersonalOwnerFacts + names: tuple[str, ...] | None = None + models: tuple[LoadedTeamModelFacts, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelListingDecision: + kind: Literal["ReadOwner", "ReadNames", "ReadModel", "Complete"] + index: int = 0 + teams: tuple[int, ...] = () + selected: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelListingDecision) + + +@native_boundary +def personal_model_listing(facts: PersonalModelListingFacts) -> PersonalModelListingDecision: + return _DECISION.validate_python(_native.request_auth_personal_model_listing(asdict(facts))) diff --git a/litellm/_auth/personal_model_request.py b/litellm/_auth/personal_model_request.py new file mode 100644 index 00000000000..aa885980426 --- /dev/null +++ b/litellm/_auth/personal_model_request.py @@ -0,0 +1,33 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.owner_teams import OwnerTeamFacts +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelRequestFacts: + requested: tuple[str, ...] + owner_present: bool + teams: OwnerTeamFacts | None = None + models: tuple[LoadedTeamModelFacts, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelRequestDecision: + kind: Literal["Allow", "Deny", "ReadTeams", "ReadModel"] + model: str = "" + index: int = 0 + teams: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelRequestDecision) + + +@native_boundary +def personal_model_request(facts: PersonalModelRequestFacts) -> PersonalModelRequestDecision: + return _DECISION.validate_python(_native.request_auth_personal_model_request(asdict(facts))) diff --git a/litellm/_auth/personal_models.py b/litellm/_auth/personal_models.py new file mode 100644 index 00000000000..88a13391eec --- /dev/null +++ b/litellm/_auth/personal_models.py @@ -0,0 +1,45 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelScope: + user_present: bool + team_id_present: bool + team: GroupedModelScope | None = None + member_grants: tuple[str, ...] | None = None + member: ModelAccessFacts | None = None + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class LoadedTeamModelFacts: + present: tuple[bool, ...] + scopes: tuple[PersonalModelScope, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelDecision: + kind: Literal["Complete", "ReadTeam", "ReadGroups", "ReadMembership", "ReadMember", "ReadScopes"] + allowed: bool = False + indices: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelDecision) + + +@native_boundary +def personal_model_scope(facts: PersonalModelScope) -> PersonalModelDecision: + return _DECISION.validate_python(_native.request_auth_personal_models({"Scope": asdict(facts)})) + + +@native_boundary +def loaded_team_models(facts: LoadedTeamModelFacts) -> PersonalModelDecision: + return _DECISION.validate_python(_native.request_auth_personal_models({"Loaded": asdict(facts)})) diff --git a/litellm/_auth/personal_owner.py b/litellm/_auth/personal_owner.py new file mode 100644 index 00000000000..7ec4c9a16c8 --- /dev/null +++ b/litellm/_auth/personal_owner.py @@ -0,0 +1,36 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.owner_teams import OwnerTeamFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerFacts: + user_present: bool + database_present: bool + cache_present: bool + logging_present: bool + owner_loaded: bool = False + owner_present: bool = False + owner_failed: bool = False + teams: OwnerTeamFacts | None = None + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerDecision: + kind: Literal["ReadOwner", "ReadTeams", "Complete"] + selected: tuple[int, ...] = () + withheld: tuple[int, ...] = () + unreadable: bool = False + + +_DECISION: Final = TypeAdapter(PersonalOwnerDecision) + + +@native_boundary +def personal_owner(facts: PersonalOwnerFacts) -> PersonalOwnerDecision: + return _DECISION.validate_python(_native.request_auth_personal_owner(asdict(facts))) diff --git a/litellm/_auth/reservation.py b/litellm/_auth/reservation.py new file mode 100644 index 00000000000..3fe0ac82678 --- /dev/null +++ b/litellm/_auth/reservation.py @@ -0,0 +1,131 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Gate: + token_present: bool + llm_route: bool | None = None + unbilled: bool | None = None + model_present: bool | None = None + counter_count: int | None = None + cost_read: bool = False + cost: float | None = None + + +@dataclass(frozen=True, slots=True) +class Sources: + key_present: bool + team_present: bool + team_id: str | None + user_present: bool + user_id: str | None + apply_user_budget_to_team_keys: bool + + +Scope: TypeAlias = Literal["Key", "Team", "User", "EndUser", "Tag", "AccessGroup", "Member", "Organization", "Project"] + + +@dataclass(frozen=True, slots=True) +class Counter: + scope: Scope + limit: float | None + + +@dataclass(frozen=True, slots=True) +class Window: + duration_present: bool + limit: float | None + start_present: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Identity: + scope: Literal["EndUser", "Organization", "Project"] + primary: str | None + fallback: str | None = None + + +@dataclass(frozen=True, slots=True) +class MemberGate: + team_present: bool + team_id: str | None + user_present: bool + user_id: str | None + + +@dataclass(frozen=True, slots=True) +class MemberCap: + member_limit: float | None + default_read: bool = False + + +@dataclass(frozen=True, slots=True) +class IdentityPlan: + id: str | None + kind: Literal["Identity"] = "Identity" + + +@dataclass(frozen=True, slots=True) +class Action: + kind: Literal[ + "ReadRoute", + "ReadUnbilled", + "ReadModel", + "ReadCounters", + "ReadEstimate", + "Reserve", + "Skip", + "ReadStart", + "WarnWindow", + "AppendWindow", + "ReadMember", + "ReadDefault", + "ReadCap", + ] + + +@dataclass(frozen=True, slots=True) +class SourcePlan: + key: bool + team: bool + user: bool + kind: Literal["Sources"] = "Sources" + + +@dataclass(frozen=True, slots=True) +class CounterPlan: + include: bool + kind: Literal["Counter"] = "Counter" + + +_ACTION: Final = TypeAdapter(Action) +_SOURCES: Final = TypeAdapter(SourcePlan) +_COUNTER: Final = TypeAdapter(CounterPlan) +_IDENTITY: Final = TypeAdapter(IdentityPlan) + + +@native_boundary +def action(facts: Gate | Window | MemberGate | MemberCap) -> Action: + return _ACTION.validate_python(_native.request_auth_reservation(encode_facts(facts))) + + +@native_boundary +def sources(facts: Sources) -> SourcePlan: + return _SOURCES.validate_python(_native.request_auth_reservation(encode_facts(facts))) + + +@native_boundary +def counter(scope: Scope, limit: float | None) -> bool: + return _COUNTER.validate_python(_native.request_auth_reservation(encode_facts(Counter(scope, limit)))).include + + +@native_boundary +def identity(facts: Identity) -> str | None: + return _IDENTITY.validate_python(_native.request_auth_reservation(encode_facts(facts))).id diff --git a/litellm/_auth/resolved_models.py b/litellm/_auth/resolved_models.py new file mode 100644 index 00000000000..5e784b035d8 --- /dev/null +++ b/litellm/_auth/resolved_models.py @@ -0,0 +1,78 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm._auth.personal_model_request import PersonalModelRequestFacts +from litellm._auth.team_override import TeamOverrideFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ResolvedTeamRead: + found: bool = False + failed: bool = False + + +@dataclass(frozen=True, slots=True) +class ResolvedModelFacts: + model: str = "" + config_nonempty: bool = False + key_models: tuple[str, ...] = () + team_id: str | None = None + user_id: str | None = None + project_present: bool = False + customer_present: bool = False + database_present: bool = False + personal_enabled: bool = False + master: bool = False + role: str | None = None + key: GroupedModelScope | None = None + team_read: ResolvedTeamRead | None = None + embedded_ready: bool = False + team: TeamOverrideFacts | None = None + member: MemberModelFacts | None = None + owner_loaded: bool = False + owner_present: bool = False + owner_failed: bool = False + personal: PersonalModelRequestFacts | None = None + project_loaded: bool = False + project_grants: tuple[str, ...] = () + project: ModelAccessFacts | None = None + customer_loaded: bool = False + customer_grants: tuple[str, ...] = () + customer: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ResolvedModelDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadKey", + "ReadKeyGroups", + "ReadTeamObject", + "ReadEmbeddedTeam", + "ReadTeamScope", + "ReadMember", + "ReadOwner", + "ReadPersonal", + "ReadProjectObject", + "ReadProjectScope", + "ReadCustomerObject", + "ReadCustomerScope", + ] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + + +_DECISION: Final = TypeAdapter(ResolvedModelDecision) + + +@native_boundary +def resolved_models(facts: ResolvedModelFacts) -> ResolvedModelDecision: + return _DECISION.validate_python(_native.request_auth_resolved_models(asdict(facts))) diff --git a/litellm/_auth/route_enable.py b/litellm/_auth/route_enable.py new file mode 100644 index 00000000000..828bc618cfc --- /dev/null +++ b/litellm/_auth/route_enable.py @@ -0,0 +1,52 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.routes import RouteContext +from litellm.rust_bridge import _native + +RouteFeature: TypeAlias = Literal["Admin", "Llm"] +LlmRouteDecision: TypeAlias = Literal["Allow", "Deny", "ReadEuJev"] + + +@dataclass(frozen=True, slots=True) +class RouteSetting: + environment_present: bool + premium: bool + secret: bool | None = None + + +@dataclass(frozen=True, slots=True) +class RouteEnableFacts: + context: RouteContext + imported: bool = False + failed: Literal["Http", "Other"] | None = None + admin: RouteSetting | None = None + llm: RouteSetting | None = None + exemptions: tuple[str, ...] | None = None + eu_jev_configured: bool | None = None + + +@dataclass(frozen=True, slots=True) +class RouteEnableDecision: + kind: Literal[ + "Allow", "PropagateHttp", "ReadImport", "ReadExemptions", "ReadEuJev", "ReadSettings", "ReadSecret", "Deny" + ] + feature: RouteFeature | None = None + reason: Literal["License", "Disabled"] | None = None + + +_DECISION: Final = TypeAdapter(RouteEnableDecision) +_LLM: Final = TypeAdapter[LlmRouteDecision](LlmRouteDecision) + + +@native_boundary +def route_enable(facts: RouteEnableFacts) -> RouteEnableDecision: + return _DECISION.validate_python(_native.request_auth_route_enable(asdict(facts))) + + +@native_boundary +def llm_route(context: RouteContext, configured: bool | None = None) -> LlmRouteDecision: + return _LLM.validate_python(_native.request_auth_llm_route(asdict(context), configured)) diff --git a/litellm/_auth/route_roles.py b/litellm/_auth/route_roles.py new file mode 100644 index 00000000000..4ad919ca8e0 --- /dev/null +++ b/litellm/_auth/route_roles.py @@ -0,0 +1,120 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class RoleFacts: + user_present: bool + role: str | None + + +@dataclass(frozen=True, slots=True) +class RoleDecision: + effective: str + user_role: str | None + admin: bool + + +@dataclass(frozen=True, slots=True) +class ApiFacts: + token_present: bool + admin_read: bool + user: RoleFacts + + +ApiDecision: TypeAlias = Literal["InvalidToken", "ReadAdmin", "NonAdmin", "Allow"] + + +@dataclass(frozen=True, slots=True) +class JwtFacts: + role: str | None + config_read: bool = False + config_missing: bool = False + + +JwtDecision: TypeAlias = Literal["ReadAdmin", "ReadTeam", "DefaultTeam", "Check", "Deny"] + + +@dataclass(frozen=True, slots=True) +class InsideFacts: + role: str | None + user_id: str | None + requested_user_id: str | None + + +@dataclass(frozen=True, slots=True) +class Entry: + text: str | None + python_type: str + hash_failure: str | None = None + templates: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ListedFacts: + route: str + entry: Entry | None = None + exhausted: bool = False + template_index: int = 0 + compiled: str | None = None + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: Literal["TypeError", "AttributeError"] + message: str + + +@dataclass(frozen=True, slots=True) +class EntryDecision: + action: Literal["ReadEntry", "ReadTemplate", "Allow", "Deny", "Keep", "Expand", "Failure"] + template_index: int + failure: Failure | None + + +_ROLE: Final = TypeAdapter(RoleFacts) +_ROLE_RESULT: Final = TypeAdapter(RoleDecision) +_API: Final = TypeAdapter(ApiFacts) +_API_RESULT: Final = TypeAdapter[ApiDecision](ApiDecision) +_JWT: Final = TypeAdapter(JwtFacts) +_JWT_RESULT: Final = TypeAdapter[JwtDecision](JwtDecision) +_INSIDE: Final = TypeAdapter(InsideFacts) +_BOOL: Final = TypeAdapter(bool) +_ENTRY: Final = TypeAdapter(Entry) +_LISTED: Final = TypeAdapter(ListedFacts) +_ENTRY_RESULT: Final = TypeAdapter(EntryDecision) + + +@native_boundary +def role(facts: RoleFacts) -> RoleDecision: + return _ROLE_RESULT.validate_python(_native.request_auth_route_role(_ROLE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def api(facts: ApiFacts) -> ApiDecision: + return _API_RESULT.validate_python(_native.request_auth_route_role_api(_API.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def jwt(facts: JwtFacts) -> JwtDecision: + return _JWT_RESULT.validate_python(_native.request_auth_route_role_jwt(_JWT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def inside(facts: InsideFacts) -> bool: + return _BOOL.validate_python(_native.request_auth_route_role_inside(_INSIDE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def listed(facts: ListedFacts) -> EntryDecision: + return _ENTRY_RESULT.validate_python(_native.request_auth_route_role_listed(_LISTED.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def expand(facts: Entry) -> EntryDecision: + return _ENTRY_RESULT.validate_python(_native.request_auth_route_role_expand(_ENTRY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/route_scope.py b/litellm/_auth/route_scope.py new file mode 100644 index 00000000000..a2f0616f942 --- /dev/null +++ b/litellm/_auth/route_scope.py @@ -0,0 +1,40 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +RouteScopeKind: TypeAlias = Literal["Proxy", "Admin"] + + +@dataclass(frozen=True, slots=True) +class RouteValues: + text: str | None = None + names: tuple[str, ...] = () + failure: str | None = None + + +@dataclass(frozen=True, slots=True) +class RouteScopeFacts: + kind: RouteScopeKind + route: str + present: bool + premium: bool + warned: bool = False + values: RouteValues | None = None + + +@dataclass(frozen=True, slots=True) +class RouteScopeDecision: + kind: Literal["Allow", "Deny", "Warn", "ReadValues", "Failure"] + message: str | None = None + + +_DECISION: Final = TypeAdapter(RouteScopeDecision) + + +@native_boundary +def route_scope(facts: RouteScopeFacts) -> RouteScopeDecision: + return _DECISION.validate_python(_native.request_auth_route_scope(asdict(facts))) diff --git a/litellm/_auth/routes.py b/litellm/_auth/routes.py new file mode 100644 index 00000000000..70e28eb404d --- /dev/null +++ b/litellm/_auth/routes.py @@ -0,0 +1,146 @@ +from collections.abc import Mapping +from dataclasses import asdict, dataclass, field +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class RouteCatalog: + groups: Mapping[str, tuple[str, ...]] = field(default_factory=dict[str, tuple[str, ...]]) + viewer_blocked: tuple[str, ...] = () + viewer_write_blocked: tuple[str, ...] = () + self_credentials: tuple[tuple[str, str], ...] = () + managed_models: tuple[str, ...] = () + managed_paths: tuple[str, ...] = () + managed_realtime: tuple[str, ...] = () + managed_live: tuple[str, ...] = () + managed_mcp: tuple[str, ...] = () + unbilled: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class RouteContext: + route: str + method: str | None = None + catalog: RouteCatalog = field(default_factory=RouteCatalog) + registered: bool = False + registered_auth: bool | None = None + configured_eu_jev: bool = False + + +@dataclass(frozen=True, slots=True) +class PassThroughFacts: + key_allowed: tuple[str, ...] = () + team_allowed: tuple[str, ...] = () + denied: tuple[str, ...] = () + forwarded: tuple[str, ...] = () + jwt_team_routes: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class VirtualRouteFacts: + context: RouteContext + allowed_routes: tuple[str, ...] | None + passthrough: PassThroughFacts = field(default_factory=PassThroughFacts) + password_reset_required: bool = False + + +@dataclass(frozen=True, slots=True) +class ViewerFacts: + context: RouteContext + request_keys: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class NonAdminFacts: + viewer: ViewerFacts + role: str | None = None + custom_admin_routes: tuple[str, ...] = () + premium: bool = False + query_user_id: str | None = None + token_user_id: str | None = None + permissions: tuple[str, ...] = () + organization_ids: tuple[str, ...] = () + admin_organization_ids: tuple[str, ...] = () + team_service_account: bool = False + allowed_routes: tuple[str, ...] | None = None + passthrough: PassThroughFacts = field(default_factory=PassThroughFacts) + + +RouteDenial: TypeAlias = Literal[ + "VirtualRoute", + "PasswordReset", + "PassThrough", + "DeniedPassThrough", + "CustomAdmin", + "UserInfo", + "AdminOnly", + "ViewerInference", + "ViewerWrite", + "ViewerParameter", +] + + +@dataclass(frozen=True, slots=True) +class RouteDecision: + allowed: bool + denial: RouteDenial | None = None + detail: str | None = None + + +_DECISION: Final = TypeAdapter(RouteDecision) + + +@native_boundary +def match_route( + route: str, + patterns: tuple[str, ...], + mode: Literal["Access", "Template", "Prefix", "Wildcard", "Denied", "Jwt"] = "Access", +) -> bool: + return _DECISION.validate_python( + _native.request_auth_route_access({"Match": {"route": route, "patterns": patterns, "mode": mode}}) + ).allowed + + +@native_boundary +def classify_route( + context: RouteContext, kind: Literal["Llm", "Unbilled", "Discovery", "SelfCredential", "Azure", "Managed"] +) -> bool: + return _DECISION.validate_python( + _native.request_auth_route_access({"Classify": {"context": asdict(context), "kind": kind}}) + ).allowed + + +@native_boundary +def virtual_route(facts: VirtualRouteFacts) -> RouteDecision: + return _DECISION.validate_python(_native.request_auth_route_access({"Virtual": asdict(facts)})) + + +@native_boundary +def non_admin_route(facts: NonAdminFacts) -> RouteDecision: + return _DECISION.validate_python(_native.request_auth_route_access({"NonAdmin": asdict(facts)})) + + +@native_boundary +def viewer_route(facts: ViewerFacts) -> RouteDecision: + return _DECISION.validate_python(_native.request_auth_route_access({"Viewer": asdict(facts)})) + + +@native_boundary +def passthrough_route(context: RouteContext, facts: PassThroughFacts, *, required: bool = False) -> RouteDecision: + return _DECISION.validate_python( + _native.request_auth_route_access( + {"PassThrough": {"context": asdict(context), "facts": asdict(facts), "required": required}} + ) + ) + + +@native_boundary +def denied_passthrough_route(route: str, facts: PassThroughFacts) -> RouteDecision: + return _DECISION.validate_python( + _native.request_auth_route_access({"DeniedPassThrough": {"route": route, "facts": asdict(facts)}}) + ) diff --git a/litellm/_auth/runtime.py b/litellm/_auth/runtime.py new file mode 100644 index 00000000000..58b83bc7296 --- /dev/null +++ b/litellm/_auth/runtime.py @@ -0,0 +1,23 @@ +from litellm._auth import rust_core +from litellm._auth.credential import PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys +from litellm._auth.decision import AuthCore + + +def get_auth_core() -> AuthCore: + return rust_core.RustCore() + + +def scheme_credential(value: str) -> str | None: + return rust_core.scheme_credential(value) + + +def litellm_header_token(value: str) -> str: + return rust_core.litellm_header_token(value) + + +def select_api_key(keys: PresentedKeys) -> SelectedKey | None: + return rust_core.select_api_key(keys) + + +def select_websocket_api_key(keys: WebSocketKeys) -> str | WebSocketDenial: + return rust_core.select_websocket_api_key(keys) diff --git a/litellm/_auth/rust_core.py b/litellm/_auth/rust_core.py new file mode 100644 index 00000000000..87b2ea3ad66 --- /dev/null +++ b/litellm/_auth/rust_core.py @@ -0,0 +1,270 @@ +from collections.abc import Callable, Mapping +from dataclasses import fields, is_dataclass +from datetime import datetime +from enum import Enum +from typing import ( + Final, + TypeAlias, + cast, # noqa: TID251 # dataclass and enum reflection erases field types +) + +from pydantic import TypeAdapter + +from litellm._auth import decision +from litellm._auth.credential import KeySource, PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +Wire: TypeAlias = None | bool | int | float | str | list["Wire"] | dict[str, "Wire"] +CoreStep: TypeAlias = decision.Step | decision.TeamMemberStep + +OBJECT_MAPPING: Final = TypeAdapter(dict[str, object]) +OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) +AUTH_STEP: Final = TypeAdapter[decision.Step](decision.Step) +MEMBER_STEP: Final = TypeAdapter[decision.TeamMemberStep](decision.TeamMemberStep) +CORE_STEP: Final = TypeAdapter[CoreStep](CoreStep) +TEAM_BUDGET: Final = TypeAdapter(decision.TeamMemberBudget) +FALLBACK_TARGET: Final = TypeAdapter(decision.FallbackTarget) +TEAM_ALERT: Final = TypeAdapter(decision.AlertTeamMemberBudget) +MEMBER_FALLBACK: Final = TypeAdapter(decision.VetTeamMemberFallback) +MODEL_DENIAL: Final = TypeAdapter(decision.ModelAccessDenied) + +_PUBLIC_TYPES: Final[Mapping[str, Callable[..., object]]] = { + cls.__name__: cls + for cls in ( + decision.NoMasterKey, + decision.UserAuthRoute, + decision.MasterKey, + decision.CachedAdminKey, + decision.VirtualKey, + decision.WithinModelBudgets, + decision.MissingKey, + decision.MalformedKey, + decision.UserAuthDisabled, + decision.ExpiredKey, + decision.MasterKeyOnlyRoute, + decision.NoDatabase, + decision.NotAVirtualKey, + decision.KeyBlocked, + decision.OwnerDeactivated, + decision.TeamMemberBudgetExceeded, + decision.KeyBudgetExceeded, + decision.KeyWithoutToken, + decision.ModelBudgetExceeded, + decision.ModelAccessDenied, + decision.Allow, + decision.Deny, + decision.ResolveEndUser, + decision.ResolvePresentedKey, + decision.EvictPresentedKey, + decision.LoadKey, + decision.PrepareKey, + decision.InspectKey, + decision.EnforceModelAccess, + decision.PrefetchAuthObjects, + decision.LoadKeyOwner, + decision.CheckBudgetExemption, + decision.ReadTeamMemberBudget, + decision.AlertTeamMemberBudget, + decision.SendKeyBudgetAlerts, + decision.ReadKeySpend, + decision.AlertKeyBudget, + decision.AlertKeySoftBudget, + decision.LoadKeyContext, + decision.ReadProxySpend, + decision.AlertProxyBudget, + decision.ReadModelBudgetScope, + decision.CheckKeyModelBudget, + decision.AuthorizeBudgetFallback, + decision.RewriteRequestModel, + decision.CheckUserModelBudget, + decision.ReadTeamMemberModelBudgetScope, + decision.CheckTeamMemberModelBudget, + decision.ReadTeamMemberFallbacks, + decision.VetTeamMemberFallback, + decision.RerouteTeamMemberRequest, + ) +} +_ANSWER_TYPES: Final = { + "EnforceModelAccess": "ModelAccess", + "ResolvePresentedKey": "Presented", + "InspectKey": "Key", + "LoadKeyOwner": "Owner", + "CheckBudgetExemption": "Exemption", + "ReadTeamMemberBudget": "TeamBudget", + "ReadKeySpend": "Spend", + "ReadProxySpend": "ProxySpend", + "ReadModelBudgetScope": "ModelScope", + "CheckKeyModelBudget": "ModelBudgetData", + "AuthorizeBudgetFallback": "FallbackAccess", + "CheckUserModelBudget": "ModelBudgetData", + "ReadTeamMemberModelBudgetScope": "TeamScope", + "CheckTeamMemberModelBudget": "ModelBudgetData", + "ReadTeamMemberFallbacks": "Targets", + "VetTeamMemberFallback": "Candidate", +} + + +class _Projection: + def __init__(self) -> None: + self.datetimes: tuple[datetime, ...] = () + + @native_boundary + def encode(self, value: object) -> Wire: + if isinstance(value, datetime): + offset: Final = value.utcoffset() + reference: Final = len(self.datetimes) + self.datetimes = (*self.datetimes, value) + return { + "wall_microseconds": ( + (value.toordinal() * 86400 + value.hour * 3600 + value.minute * 60 + value.second) * 1000000 + + value.microsecond + ), + "offset_microseconds": None + if offset is None + else (offset.days * 86400 + offset.seconds) * 1000000 + offset.microseconds, + "timezone_id": id(value.tzinfo), + "reference": reference, + } + if isinstance(value, Enum): + return self.encode(cast(object, value.value)) # cast-ok: Enum.value is dynamically typed + if value is None or isinstance(value, (str, bool, int, float)): + return value + if is_dataclass(value) and not isinstance(value, type): + return { + "kind": type(value).__name__, + **{field.name: self.encode(field_value(value, field.name)) for field in fields(value)}, + } + if isinstance(value, Mapping): + return {key: self.encode(item) for key, item in OBJECT_MAPPING.validate_python(value).items()} + if isinstance(value, (tuple, list, frozenset)): + return [self.encode(item) for item in OBJECT_TUPLE.validate_python(value)] + raise TypeError("unsupported auth protocol value") + + def decode(self, value: object) -> object: + if not isinstance(value, dict): + return value + data: Final = OBJECT_MAPPING.validate_python(value) + if "wall_microseconds" in data: + reference: Final = data["reference"] + if not isinstance(reference, int): + raise TypeError("invalid auth datetime reference") + return self.datetimes[reference] + kind: Final = data.get("kind") + if kind == "ModelAccessDenied": + return MODEL_DENIAL.validate_python({"decision": data["decision"]}) + if kind is None: + return {key: self.decode(item) for key, item in data.items()} + if not isinstance(kind, str): + raise TypeError("invalid auth protocol tag") + constructor: Final = public_type(kind) + return constructor(**{key: self.decode(item) for key, item in data.items() if key != "kind"}) + + @native_boundary + def step(self, result: tuple[Mapping[str, object], _native.NativeAuthContinuation | None]) -> CoreStep: + data, continuation = result + kind: Final = data["kind"] + if not isinstance(kind, str): + raise TypeError("invalid auth protocol tag") + if continuation is None: + return CORE_STEP.validate_python(self.decode(data)) + + @native_boundary + def resume(*answers: object) -> CoreStep: + answer: Final = { + "kind": _ANSWER_TYPES.get(kind, "Unit"), + "value": self.encode(answers[0]) if answers else None, + } + return self.step(continuation.resume(answer)) + + payload: Final = {key: self.decode(value) for key, value in data.items() if key != "kind"} + if kind == "AlertTeamMemberBudget": + return TEAM_ALERT.validate_python( + {**payload, "budget": TEAM_BUDGET.validate_python(payload["budget"]), "resume": resume} + ) + if kind == "VetTeamMemberFallback": + return MEMBER_FALLBACK.validate_python( + {**payload, "target": FALLBACK_TARGET.validate_python(payload["target"]), "resume": resume} + ) + constructor: Final = public_type(kind) + return CORE_STEP.validate_python(constructor(**payload, resume=resume)) + + +def field_value(value: object, name: str) -> object: + return cast(object, getattr(value, name)) # cast-ok: dataclass fields are discovered dynamically + + +@native_boundary +def encode_facts(value: object) -> Wire: + return _Projection().encode(value) + + +@native_boundary +def expired(expires: datetime, now: datetime) -> bool: + return _native.request_auth_expired(encode_facts(expires), encode_facts(now)) + + +def public_type(kind: str) -> Callable[..., object]: + return _PUBLIC_TYPES[kind] + + +class RustCore: + @native_boundary + def start(self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime) -> decision.Step: + projection: Final = _Projection() + return AUTH_STEP.validate_python( + projection.step( + _native.request_auth_start( + projection.encode(facts), projection.encode(settings), projection.encode(now), False + ) + ), + ) + + @native_boundary + def start_resolved_key( + self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime + ) -> decision.Step: + projection: Final = _Projection() + return AUTH_STEP.validate_python( + projection.step( + _native.request_auth_start( + projection.encode(facts), projection.encode(settings), projection.encode(now), True + ) + ), + ) + + @native_boundary + def check_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python(_Projection().step(_native.request_auth_budgets("Key"))) + + @native_boundary + def check_user_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python(_Projection().step(_native.request_auth_budgets("User"))) + + @native_boundary + def check_team_member_model_budgets(self) -> decision.TeamMemberStep: + return MEMBER_STEP.validate_python(_Projection().step(_native.request_auth_budgets("TeamMember"))) + + +@native_boundary +def select_api_key(keys: PresentedKeys) -> SelectedKey | None: + selected: Final = _native.request_auth_select_api_key(_Projection().encode(keys)) + if selected is None: + return None + return SelectedKey(source=KeySource(selected["source"]), presented=selected["presented"], token=selected["token"]) + + +@native_boundary +def select_websocket_api_key(keys: WebSocketKeys) -> str | WebSocketDenial: + selected: Final = _native.request_auth_select_websocket_api_key(_Projection().encode(keys)) + return WebSocketDenial(selected["value"]) if selected["kind"] == "Denial" else selected["value"] + + +@native_boundary +def scheme_credential(value: str) -> str | None: + return _native.request_auth_scheme_credential(value) + + +@native_boundary +def litellm_header_token(value: str) -> str: + return _native.request_auth_litellm_header_token(value) diff --git a/litellm/_auth/spend_counter_policy.py b/litellm/_auth/spend_counter_policy.py new file mode 100644 index 00000000000..924cc58ca21 --- /dev/null +++ b/litellm/_auth/spend_counter_policy.py @@ -0,0 +1,129 @@ +from dataclasses import dataclass +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import Field, TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Spend: + current: float + verified: bool + fallback: float + limit: float | None + window: bool + fallback_authoritative: bool + strict_before: bool | None = None + authoritative_read: bool = False + authoritative: float | None = None + strict_after: bool | None = None + + +@dataclass(frozen=True, slots=True) +class ThrottleKey: + counter_key: str + token_present: bool + token: str | None + + +@dataclass(frozen=True, slots=True) +class Overflow: + current: float + limit: float + + +@dataclass(frozen=True, slots=True) +class OverBudget: + throttled: bool + current: float + limit: float + reservation: float + strict: bool + + +@dataclass(frozen=True, slots=True) +class ReservationFailure: + strict: bool + + +@dataclass(frozen=True, slots=True) +class ReservationGroups: + currents: tuple[float, ...] + limits: tuple[float, ...] + reservation: float + + +@dataclass(frozen=True, slots=True) +class SpendValue: + kind: Literal["Value", "Repair"] + spend: float + + +@dataclass(frozen=True, slots=True) +class Action: + kind: Literal[ + "ReadStrictBefore", + "ReadStrictAfter", + "ReadFloor", + "Unverifiable", + "ReadThrottle", + "Skip", + "Release", + "CheckOverflow", + "Unavailable", + "Compatible", + "Inconsistent", + "RejectResize", + ] + + +@dataclass(frozen=True, slots=True) +class Rejected: + kind: Literal["Reject"] + current: float + estimate: float | None + + +@dataclass(frozen=True, slots=True) +class Resize: + kind: Literal["Resize"] + cost: float + + +@dataclass(frozen=True, slots=True) +class Groups: + kind: Literal["ReservationGroups"] + groups: tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True, slots=True) +class Adjustment: + current: float | None + adjustment: float + + +@dataclass(frozen=True, slots=True) +class ResizeConsistency: + inconsistent: bool + reseed: bool + + +Decision: TypeAlias = Annotated[SpendValue | Action | Rejected | Resize | Groups, Field(discriminator="kind")] +_DECISION: Final = TypeAdapter[Decision](Decision) + + +@native_boundary +def decide( + facts: Spend + | ThrottleKey + | Overflow + | OverBudget + | ReservationFailure + | ReservationGroups + | Adjustment + | ResizeConsistency, +) -> Decision: + return _DECISION.validate_python(_native.request_auth_spend_counter_policy(encode_facts(facts))) diff --git a/litellm/_auth/team_aliases.py b/litellm/_auth/team_aliases.py new file mode 100644 index 00000000000..c72ff652111 --- /dev/null +++ b/litellm/_auth/team_aliases.py @@ -0,0 +1,34 @@ +from collections.abc import Mapping +from dataclasses import asdict, dataclass, replace +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class TeamAliasFacts: + requested: tuple[str, ...] + aliases: Mapping[str, str] + team_id: str | None + router_present: bool + names: tuple[tuple[str, ...], ...] | None = None + + +@dataclass(frozen=True, slots=True) +class TeamAliasDecision: + kind: Literal["ReadNames", "Complete"] + count: int = 0 + models: tuple[str, ...] = () + + +_DECISION: Final = TypeAdapter(TeamAliasDecision) + + +@native_boundary +def team_aliases(facts: TeamAliasFacts) -> TeamAliasDecision: + return _DECISION.validate_python( + _native.request_auth_team_aliases(asdict(replace(facts, aliases=dict(facts.aliases)))) + ) diff --git a/litellm/_auth/team_override.py b/litellm/_auth/team_override.py new file mode 100644 index 00000000000..504f5cc83b7 --- /dev/null +++ b/litellm/_auth/team_override.py @@ -0,0 +1,30 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class TeamOverrideFacts: + team: GroupedModelScope | None = None + assigned_models: tuple[str, ...] | None = None + assigned_scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class TeamOverrideDecision: + kind: Literal["Allow", "Deny", "ReadTeam", "ReadTeamGroups", "ReadAssignedGroups", "ReadAssignedScope"] + denial: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(TeamOverrideDecision) + + +@native_boundary +def team_override(facts: TeamOverrideFacts) -> TeamOverrideDecision: + return _DECISION.validate_python(_native.request_auth_team_override(asdict(facts))) diff --git a/litellm/_auth/unpriced_models.py b/litellm/_auth/unpriced_models.py new file mode 100644 index 00000000000..2944bc8294b --- /dev/null +++ b/litellm/_auth/unpriced_models.py @@ -0,0 +1,53 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PriceMetric: + name: str + scalar: float | None = None + nested: tuple[float | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class PriceEntry: + tiered_present: bool + metrics: tuple[PriceMetric, ...] + + +@dataclass(frozen=True, slots=True) +class DeploymentPricing: + parameter_keys: tuple[str, ...] + has_id: bool + info_loaded: bool = False + info: PriceEntry | None = None + cost_keys: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class UnpricedModelFacts: + model: str | None + router_present: bool + group_found: bool | None = None + deployments: tuple[DeploymentPricing, ...] | None = None + cost_deployments: tuple[DeploymentPricing, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class UnpricedModelDecision: + kind: Literal["Complete", "ReadGroup", "ReadDeployments", "ReadInfo", "ReadCostDeployments", "ReadCost"] + unpriced: bool = False + index: int = 0 + + +_DECISION: Final = TypeAdapter(UnpricedModelDecision) + + +@native_boundary +def unpriced_model(facts: UnpricedModelFacts) -> UnpricedModelDecision: + return _DECISION.validate_python(_native.request_auth_unpriced_model(asdict(facts))) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index 557d1c7bc0a..34bb6781df4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -1,11 +1,7 @@ -import asyncio from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Final, TypeAlias -from fastapi import HTTPException - -from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_AccessGroupTable if TYPE_CHECKING: @@ -31,13 +27,6 @@ class AgentAccessGroupCeiling: CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] -class _MissingGroup: - __slots__ = () - - -_MISSING: Final = _MissingGroup() - - async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through @@ -46,29 +35,9 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup: - from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + from litellm.proxy.auth.agent_ceiling_adapter import load_agent_access_group - if prisma_client is None: - verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) - return None - try: - return await get_access_object( - access_group_id=access_group_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - except AccessGroupNotFoundError: - raise - except HTTPException as e: - if check_db_only: - raise - verbose_proxy_logger.warning( - "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail - ) - return None + return await load_agent_access_group(access_group_id, check_db_only=check_db_only) async def resolve_agent_access_group_ceiling( @@ -82,31 +51,10 @@ async def resolve_agent_access_group_ceiling( provably does not exist is dropped from the ceiling, it grants and bounds nothing, while a group that exists but could not be read stays named and loads as ``None``. The ceiling itself stays once any group is named, so an agent whose every group is missing is still capped to nothing.""" - named_ids: Final = await load_access_group_ids(agent_id) - if not named_ids: - return None + from litellm.proxy.auth.agent_ceiling_adapter import project_agent_ceiling, read_agent_ceiling_facts - from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError - - async def load(group_id: str) -> LoadedAccessGroup | _MissingGroup: - try: - return await ( - _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id) - ) - except AccessGroupNotFoundError: - return _MISSING - - results: Final = await asyncio.gather(*(load(group_id) for group_id in named_ids)) - access_group_ids: Final = tuple( - group_id for group_id, result in zip(named_ids, results, strict=True) if not isinstance(result, _MissingGroup) - ) - groups: Final = tuple(result for result in results if isinstance(result, LiteLLM_AccessGroupTable)) - return AgentAccessGroupCeiling( - access_group_ids=access_group_ids, - models=frozenset(model for group in groups for model in group.access_model_names), - mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), - agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), - access_groups=groups, + return project_agent_ceiling( + await read_agent_ceiling_facts(agent_id, load_access_group_ids, load_access_group, check_db_only=check_db_only) ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 9be9d4c7121..e628fce03d4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -7,14 +7,15 @@ Follows the same pattern as MCP permission handling. import asyncio from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass -from itertools import chain +from dataclasses import dataclass, replace from types import MappingProxyType -from typing import Final, Protocol, TypeAlias +from typing import Final, Protocol, TypeAlias, assert_never from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError +from litellm._auth import agent_access, agent_grants, managed +from litellm._auth.failure import NativeAuthFailure from litellm._logging import verbose_logger from litellm.proxy._types import ( UI_TEAM_ID, @@ -30,6 +31,8 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.proxy.mcp_server.ui_session_utils import build_effective_auth_contexts +from litellm.proxy.utils import PrismaClient +from litellm.repositories.native_auth_repository import AgentIdRow, native_auth_repository from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse @@ -58,6 +61,8 @@ class AgentAccessLookup(Protocol): def __call__(self, user_api_key_auth: UserAPIKeyAuth | None, /, *, strict: bool) -> Awaitable[AgentAccess]: ... +_AGENT_GROUPS: Final = TypeAdapter(tuple[str, ...]) + _DEFAULT_DENY_FLAG: Final = TypeAdapter(bool) @@ -70,13 +75,6 @@ def agent_access_default_deny() -> bool: return True -def _is_proxy_admin(user_api_key_auth: UserAPIKeyAuth | None) -> bool: - return user_api_key_auth is not None and user_api_key_auth.user_role in ( - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN.value, - ) - - def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry @@ -89,24 +87,20 @@ def _restricted_ids(access: AgentAccess) -> frozenset[str] | None: return _to_stable_ids(access.agent_ids) -def _union_agent_access(accesses: Sequence[AgentAccess]) -> AgentAccess: - if accesses and all(isinstance(access, UnrestrictedAgentAccess) for access in accesses): - return UnrestrictedAgentAccess() - return RestrictedAgentAccess( - frozenset( - chain.from_iterable(access.agent_ids for access in accesses if isinstance(access, RestrictedAgentAccess)) - ) - ) +def _project_agent_grants(decision: managed.ManagedDecision | agent_grants.Decision) -> AgentAccess: + return UnrestrictedAgentAccess() if decision.grants is None else RestrictedAgentAccess(frozenset(decision.grants)) def _intersect_agent_access(key_access: AgentAccess, team_access: AgentAccess) -> AgentAccess: key_ids: Final = _restricted_ids(key_access) team_ids: Final = _restricted_ids(team_access) - if key_ids is None: - return UnrestrictedAgentAccess() if team_ids is None else RestrictedAgentAccess(team_ids) - if team_ids is None: - return RestrictedAgentAccess(key_ids) - return RestrictedAgentAccess(key_ids & team_ids) + return _project_agent_grants( + managed.decide( + managed.IntersectGrants( + tuple(key_ids) if key_ids is not None else None, tuple(team_ids) if team_ids is not None else None + ) + ) + ) class AgentRequestHandler: @@ -135,27 +129,18 @@ class AgentRequestHandler: ) -> AgentAccess: """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" - if managed_agent_policy(user_api_key_auth) is not None: - return await _managed_actor_agent_access(user_api_key_auth) - key_team_access: Final = await AgentRequestHandler.resolve_key_team_agent_access( + return await _read_resolved_agent_grants( user_api_key_auth, - strict=strict, - lookup_key_access=lookup_key_access, - lookup_team_access=lookup_team_access, + resolve_ceiling, + default_deny, + lookup_key_access, + lookup_team_access, + agent_grants.Resolve( + managed=managed_agent_policy(user_api_key_auth) is not None, + strict=strict, + role=user_api_key_auth.user_role if user_api_key_auth is not None else None, + ), ) - deny_undefined: Final = strict or (not _is_proxy_admin(user_api_key_auth) and default_deny()) - if deny_undefined and isinstance(key_team_access, UnrestrictedAgentAccess): - return RestrictedAgentAccess(frozenset()) - caller_access: Final = await AgentRequestHandler.agent_caller_access(user_api_key_auth, strict=strict) - own_access: Final = _intersect_agent_access(key_team_access, caller_access) - agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling( - user_api_key_auth, resolve_ceiling, strict=strict - ) - if agent_ceiling is None: - return own_access - if isinstance(own_access, UnrestrictedAgentAccess): - return RestrictedAgentAccess(agent_ceiling) - return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling) @staticmethod async def agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None, *, strict: bool = False) -> AgentAccess: @@ -181,11 +166,14 @@ class AgentRequestHandler: ) key_access: Final = await key_lookup(user_api_key_auth, strict=strict) team_access: Final = await team_lookup(user_api_key_auth, strict=strict) + except NativeAuthFailure: + raise except Exception as e: - if strict: + failure: Final = agent_grants.decide(agent_grants.Failure(strict, closed=False)) + if failure.action == "Unavailable": raise HTTPException(503, "Agent invocation policy is unavailable") from e verbose_logger.warning("Failed to get allowed agents: %s", e) - return UnrestrictedAgentAccess() + return _project_agent_grants(failure) return _intersect_agent_access(key_access, team_access) @staticmethod @@ -212,71 +200,8 @@ class AgentRequestHandler: user_api_key_auth: UserAPIKeyAuth | None = None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> bool: - """ - Check if a specific agent is allowed for the given user/key. - - Args: - agent_id: The agent ID to check - user_api_key_auth: User authentication info - - Returns: - bool: True if agent is allowed, False otherwise - """ - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore - from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure - from litellm.proxy.proxy_server import prisma_client - from litellm.types.proxy.agent_identity import AgentIdentityFailure - - registered: Final = global_agent_registry.get_agent_by_id(agent_id) - registry_managed: Final = isinstance(registered, AgentResponse) and registered.identity_managed - if registry_managed or prisma_client is not None: - target: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id) - if isinstance(target, AgentIdentityFailure): - raise_identity_failure(target) - elif target is None and registry_managed: - return False - elif isinstance(target, AgentResponse) and target.identity_managed: - if ( - not target.enabled - or target.identity is None - or not target.identity.active - or user_api_key_auth is None - ): - return False - from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler - - key_hash: Final = user_api_key_auth.api_key or user_api_key_auth.token - authority: Final = ( - await MCPRequestHandler._reload_admitted_key(key_hash, check_db_only=True) # pyright: ignore[reportPrivateUsage] # the authoritative key reload has no public seam - if key_hash - and managed_agent_policy(user_api_key_auth) is None - and not user_api_key_auth.is_session_token - and not user_api_key_auth.authenticated_by_custom_auth - else user_api_key_auth - ) - fresh_auth: Final = authority.model_copy( - update=MappingProxyType( - {"requires_fresh_policy": True, "agent_caller": user_api_key_auth.agent_caller} - ) - ) - explicit: Final = await _granted_agent_ids( - fresh_auth, - _strict_agent_access, - build_effective_auth_contexts, - ) - return target.agent_id in explicit - - match await AgentRequestHandler.resolve_agent_access( - user_api_key_auth, - resolve_ceiling, - strict=bool(user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy), - ): - case UnrestrictedAgentAccess(): - return True - case RestrictedAgentAccess(allowed_agent_ids): - stable_id: Final = global_agent_registry.stable_agent_id(agent_id) - return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agent_ids) + facts: Final = await read_agent_access_facts(agent_id, user_api_key_auth, resolve_ceiling) + return agent_access.decide(facts) == "Allow" @staticmethod def _get_key_object_permission( @@ -361,29 +286,11 @@ class AgentRequestHandler: # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups) key_access_group_ids: Final = tuple(user_api_key_auth.access_group_ids or ()) - if not direct_agents and not declared_access_groups and not key_access_group_ids: - return UnrestrictedAgentAccess() - - access_group_agents: Final = ( - tuple( - await AgentRequestHandler._get_agents_from_access_groups( - declared_access_groups, check_db_only=strict - ) - ) - if declared_access_groups - else () + return await _read_declared_grants( + managed.DeclaredGrants(direct_agents, declared_access_groups, key_access_group_ids), strict ) - unified_agents: Final = ( - tuple( - await AgentRequestHandler._get_unified_access_group_agents( - key_access_group_ids, check_db_only=strict - ) - ) - if key_access_group_ids - else () - ) - - return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) + except NativeAuthFailure: + raise except Exception as e: if strict: raise HTTPException(503, "Agent invocation policy is unavailable") from e @@ -408,6 +315,8 @@ class AgentRequestHandler: proxy_logging_obj=proxy_logging_obj, check_db_only=strict, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # fail closed: an owner whose teams cannot be read grants nothing if strict: raise HTTPException(503, "Agent invocation policy is unavailable") from e @@ -419,22 +328,16 @@ class AgentRequestHandler: withheld_results: Final = await asyncio.gather( *(AgentRequestHandler._owner_team_agent_grants(team, strict=strict) for team in owner_teams.withheld) ) - if any(isinstance(access, AgentPolicyUnavailable) for access in team_results): + decision: Final = agent_grants.decide( + agent_grants.OwnerUnion( + tuple(_agent_grant_source(None, access) for access in team_results), + tuple(_agent_grant_source(None, access) for access in withheld_results), + owner_teams.unreadable, + ) + ) + if decision.action == "Unavailable": raise HTTPException(503, "Agent invocation policy is unavailable") - team_accesses: Final = tuple( - access for access in team_results if not isinstance(access, AgentPolicyUnavailable) - ) - withheld_restrictions: Final = tuple( - RestrictedAgentAccess(frozenset()) - for access in withheld_results - if isinstance(access, (AgentPolicyUnavailable, RestrictedAgentAccess)) - ) - accesses: Final = ( - (*team_accesses, *withheld_restrictions, RestrictedAgentAccess(frozenset())) - if owner_teams.unreadable - else (*team_accesses, *withheld_restrictions) - ) - return _union_agent_access(accesses) + return _project_agent_grants(decision) @staticmethod async def personal_key_agent_billing_auth( @@ -469,23 +372,27 @@ class AgentRequestHandler: ) stable_agent_id: Final = global_agent_registry.stable_agent_id(agent_id) - restricted_sources: Final = tuple( - (team, access) - for team, access in zip(teams, accesses, strict=True) - if isinstance(access, RestrictedAgentAccess) and stable_agent_id in _to_stable_ids(access.agent_ids) + selection: Final = agent_grants.decide( + agent_grants.BillingTeam( + tuple( + agent_grants.Source( + team.team_id, + tuple(_to_stable_ids(access.agent_ids)) + if isinstance(access, RestrictedAgentAccess) + else None, + isinstance(access, AgentPolicyUnavailable), + ) + for team, access in zip(teams, accesses, strict=True) + ), + stable_agent_id, + ) ) - unrestricted_sources: Final = ( - tuple(zip(teams, accesses, strict=True)) - if not restricted_sources - and teams - and all(isinstance(access, UnrestrictedAgentAccess) for access in accesses) - else () - ) - sources: Final = restricted_sources or unrestricted_sources - if not sources: + if selection.team_id is None: verbose_logger.warning("No owner team grants agent %s; billing against the personal key", agent_id) return auth - team: Final = min((source[0] for source in sources), key=lambda source: source.team_id) + team: Final = next(team for team in teams if team.team_id == selection.team_id) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # attribution must never fail an authorized call verbose_logger.warning("Failed to attribute agent %s spend to an owner team: %s", agent_id, e) return auth @@ -522,6 +429,8 @@ class AgentRequestHandler: else None ) team_auth_update: Final = await team_grant_billing_auth_update(auth, team, membership) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # a chosen team's controls must survive billing lookup failures verbose_logger.warning( "Failed to load billing controls for agent %s from owner team %s: %s", @@ -538,18 +447,27 @@ class AgentRequestHandler: *, strict: bool, ) -> AgentAccess | AgentPolicyUnavailable: - if team.object_permission_id is not None and team.object_permission is None: - if strict: - return AgentPolicyUnavailable(team_id=team.team_id) - verbose_logger.warning("Owner team %s has an unreadable object_permission, granting nothing", team.team_id) - return RestrictedAgentAccess(frozenset()) + decision: Final = agent_grants.decide( + agent_grants.OwnerTeam( + strict, + team.object_permission_id is not None, + team.object_permission is not None, + ) + ) + if decision.action == "Unavailable": + return AgentPolicyUnavailable(team_id=team.team_id) + if decision.action == "Complete": + return _project_agent_grants(decision) try: return await AgentRequestHandler._team_agent_grants(team, strict=strict) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # per-team isolation: a team whose grants cannot be read grants nothing - if strict: + failure: Final = agent_grants.decide(agent_grants.Failure(strict, closed=True)) + if failure.action == "Unavailable": return AgentPolicyUnavailable(team_id=team.team_id) verbose_logger.warning("Failed to get allowed agents for owner team %s: %s", team.team_id, e) - return RestrictedAgentAccess(frozenset()) + return _project_agent_grants(failure) @staticmethod async def _team_agent_grants(team: LiteLLM_TeamTable, *, strict: bool) -> AgentAccess: @@ -564,25 +482,9 @@ class AgentRequestHandler: # 2. Also include agents from team's access_group_ids (unified access groups) team_access_group_ids: Final = tuple(team.access_group_ids or ()) - if not direct_agents and not declared_access_groups and not team_access_group_ids: - return UnrestrictedAgentAccess() - - access_group_agents: Final = ( - tuple( - await AgentRequestHandler._get_agents_from_access_groups(declared_access_groups, check_db_only=strict) - ) - if declared_access_groups - else () + return await _read_declared_grants( + managed.DeclaredGrants(direct_agents, declared_access_groups, team_access_group_ids), strict ) - unified_agents: Final = ( - tuple( - await AgentRequestHandler._get_unified_access_group_agents(team_access_group_ids, check_db_only=strict) - ) - if team_access_group_ids - else () - ) - - return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) @staticmethod async def _get_allowed_agents_for_team( @@ -599,14 +501,27 @@ class AgentRequestHandler: Fetches the team object once and reuses it for both permission sources. Declared-but-empty grants stay restricted; see `get_allowed_agents_for_key`. """ - if user_api_key_auth is None: - return UnrestrictedAgentAccess() - - if user_api_key_auth.team_id is None: + facts: Final = agent_grants.TeamLookup( + subject=user_api_key_auth is not None, + team_id=user_api_key_auth.team_id if user_api_key_auth is not None else None, + database=True, + strict=strict, + ) + initial: Final = agent_grants.decide(facts) + if initial.action == "Complete": + return _project_agent_grants(initial) + assert user_api_key_auth is not None + if initial.action == "ReadPersonalEligibility": from litellm.proxy.auth.auth_checks import personal_key_inherits_team_grants - if not personal_key_inherits_team_grants(user_api_key_auth): - return UnrestrictedAgentAccess() + personal: Final = agent_grants.decide( + replace( + facts, + personal=personal_key_inherits_team_grants(user_api_key_auth), + ) + ) + if personal.action == "Complete": + return _project_agent_grants(personal) return await AgentRequestHandler._personal_key_owner_team_agent_access(user_api_key_auth, strict=strict) try: @@ -617,11 +532,13 @@ class AgentRequestHandler: user_api_key_cache, ) - if not prisma_client: - return RestrictedAgentAccess(frozenset()) if strict else UnrestrictedAgentAccess() + database: Final = agent_grants.decide(replace(facts, database=bool(prisma_client))) + if database.action == "Complete": + return _project_agent_grants(database) # Fetch the team object once for both permission sources - team_obj: Final = await get_team_object( + assert user_api_key_auth.team_id is not None + team_obj: Final[LiteLLM_TeamTable | None] = await get_team_object( team_id=user_api_key_auth.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -630,18 +547,23 @@ class AgentRequestHandler: check_db_only=strict, ) - if team_obj is None: - return RestrictedAgentAccess(frozenset()) if strict else UnrestrictedAgentAccess() + loaded: Final = agent_grants.decide(replace(facts, loaded=True, row_present=team_obj is not None)) # pyright: ignore[reportUnnecessaryComparison] # legacy adapters may return no row + if loaded.action == "Complete": + return _project_agent_grants(loaded) + assert team_obj is not None return await AgentRequestHandler._team_agent_grants(team_obj, strict=strict) + except NativeAuthFailure: + raise except Exception as e: - if strict: + failure: Final = agent_grants.decide(agent_grants.Failure(strict, closed=False)) + if failure.action == "Unavailable": raise HTTPException(503, "Agent invocation policy is unavailable") from e # litellm-dashboard is the default UI team and will never have agents; # skip noisy warnings for it. if user_api_key_auth.team_id != UI_TEAM_ID: verbose_logger.warning("Failed to get allowed agents for team: %s", e) - return UnrestrictedAgentAccess() + return _project_agent_grants(failure) @staticmethod def _get_config_agent_ids_for_access_groups( @@ -650,17 +572,22 @@ class AgentRequestHandler: """ Helper to get agent_ids from config-loaded agents that match any of the given access groups. """ - server_ids: Final[set[str]] = set() - for agent in config_agents: - agent_access_groups = getattr(agent, "agent_access_groups", None) - if agent_access_groups: - if any(group in agent_access_groups for group in access_groups): - server_ids.add(agent.agent_id) - return server_ids + decision: Final = agent_grants.decide( + agent_grants.GroupMatches( + tuple( + agent_grants.ConfiguredAgent( + agent.agent_id, _AGENT_GROUPS.validate_python(getattr(agent, "agent_access_groups", None) or ()) + ) + for agent in config_agents + ), + tuple(access_groups), + ) + ) + return set(decision.grants or ()) @staticmethod async def _get_db_agent_ids_for_access_groups( - prisma_client, access_groups: Sequence[str], *, check_db_only: bool = False + prisma_client: PrismaClient | None, access_groups: Sequence[str], *, check_db_only: bool = False ) -> set[str]: """ Helper to get agent_ids from DB agents that match any of the given access groups. @@ -671,8 +598,13 @@ class AgentRequestHandler: if not access_groups or prisma_client is None: return set() - agents: Final = await AgentsRepository(prisma_client, use_writer=check_db_only).table.find_many( - where={"agent_access_groups": {"hasSome": access_groups}} + native: Final = native_auth_repository(prisma_client) + agents: Final = ( + await native.many("agent_access_group", access_groups, AgentIdRow, writer=check_db_only) + if native is not None + else await AgentsRepository(prisma_client, use_writer=check_db_only).table.find_many( + where={"agent_access_groups": {"hasSome": access_groups}} + ) ) return {agent.agent_id for agent in agents} @@ -685,7 +617,9 @@ class AgentRequestHandler: """ from litellm.proxy.auth.auth_checks import get_agent_ids_from_access_groups - return await get_agent_ids_from_access_groups(access_group_ids=access_group_ids, check_db_only=check_db_only) + return await get_agent_ids_from_access_groups( + access_group_ids=list(access_group_ids), check_db_only=check_db_only + ) @staticmethod async def _get_agents_from_access_groups( @@ -765,6 +699,8 @@ class AgentRequestHandler: return [] return key_object_permission.agent_access_groups or [] + except NativeAuthFailure: + raise except Exception as e: verbose_logger.warning("Failed to get agent access groups for key: %s", e) return [] @@ -799,7 +735,7 @@ class AgentRequestHandler: parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - if team_obj is None: + if team_obj is None: # pyright: ignore[reportUnnecessaryComparison] # legacy adapters may return no row verbose_logger.debug("team_obj is None") return [] @@ -808,6 +744,8 @@ class AgentRequestHandler: return [] return object_permissions.agent_access_groups or [] + except NativeAuthFailure: + raise except Exception as e: verbose_logger.warning("Failed to get agent access groups for team: %s", e) return [] @@ -836,7 +774,7 @@ async def _granted_agent_ids( accesses: Final = await asyncio.gather( *(resolve_access(auth_context) for auth_context in await effective_contexts(user_api_key_auth)) ) - return frozenset().union(*(_granted_ids(access) for access in accesses)) + return frozenset[str]().union(*(_granted_ids(access) for access in accesses)) async def accessible_agents( @@ -881,25 +819,76 @@ async def _strict_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: async def _managed_actor_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: agent: Final = managed_agent_policy(auth) - if agent is None or not agent.object_permission: - return RestrictedAgentAccess(frozenset()) - permission: Final = LiteLLM_ObjectPermissionTable.model_validate(agent.object_permission or MappingProxyType({})) - own_auth: Final = UserAPIKeyAuth(object_permission=permission) - own: Final = _granted_ids(await AgentRequestHandler.get_allowed_agents_for_key(own_auth, strict=True)) - - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - - ceilings: Final = await resolve_managed_agent_ceilings(agent) - grouped: Final = frozenset(target for target in own if all(target in ceiling.agent_ids for ceiling in ceilings)) - caller: Final = await AgentRequestHandler.agent_caller_access(auth, strict=True) - capped: Final = grouped if isinstance(caller, UnrestrictedAgentAccess) else grouped & caller.agent_ids context: Final = auth.managed_agent_context - if context is None or context.mode == "autonomous": - return RestrictedAgentAccess(capped) - if context.user_id is None: - return RestrictedAgentAccess(frozenset()) - human_ids: Final = await verified_human_agent_grants(context.user_id, auth.team_id) - return RestrictedAgentAccess(capped.intersection(human_ids)) + return await _read_managed_agent_grants( + auth, + managed.ManagedAgentGrants( + permission_present=bool(agent is not None and agent.object_permission), + mode=context.mode if context is not None else None, + user_present=context is not None and context.user_id is not None, + ), + ) + + +async def _read_managed_agent_grants(auth: UserAPIKeyAuth, facts: managed.ManagedAgentGrants) -> AgentAccess: + decision: Final = managed.decide(facts) + agent: Final = managed_agent_policy(auth) + match decision.action: + case "Grants": + return _project_agent_grants(decision) + case "ReadOwn": + permission: Final = LiteLLM_ObjectPermissionTable.model_validate( + agent.object_permission if agent is not None else MappingProxyType({}) + ) + own_auth: Final = UserAPIKeyAuth(object_permission=permission) + own: Final = await AgentRequestHandler.get_allowed_agents_for_key(own_auth, strict=True) + return await _read_managed_agent_grants(auth, replace(facts, own=tuple(_granted_ids(own)))) + case "ReadCeilings": + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings + + if agent is None: + raise RuntimeError("native requested absent managed agent") + ceilings: Final = await resolve_managed_agent_ceilings(agent) + return await _read_managed_agent_grants( + auth, replace(facts, ceilings=tuple(tuple(ceiling.agent_ids) for ceiling in ceilings)) + ) + case "ReadCaller": + caller: Final = await AgentRequestHandler.agent_caller_access(auth, strict=True) + return await _read_managed_agent_grants( + auth, + replace( + facts, + caller_loaded=True, + caller=tuple(caller.agent_ids) if isinstance(caller, RestrictedAgentAccess) else None, + ), + ) + case "ReadHumanGrants": + context: Final = auth.managed_agent_context + human: Final = await verified_human_agent_grants( + context.user_id if context is not None else None, auth.team_id + ) + return await _read_managed_agent_grants(auth, replace(facts, human=tuple(human))) + case _: + raise RuntimeError("invalid native managed grants transition") + + +async def _read_declared_grants(facts: managed.DeclaredGrants, strict: bool) -> AgentAccess: + decision: Final = managed.decide(facts) + match decision.action: + case "Grants": + return _project_agent_grants(decision) + case "ReadDeclaredGroups": + groups: Final = await AgentRequestHandler._get_agents_from_access_groups( # pyright: ignore[reportPrivateUsage] # module-owned staged reader + facts.declared_groups, check_db_only=strict + ) + return await _read_declared_grants(replace(facts, group_agents=tuple(groups)), strict) + case "ReadUnifiedGroups": + unified: Final = await AgentRequestHandler._get_unified_access_group_agents( # pyright: ignore[reportPrivateUsage] # module-owned staged reader + facts.unified_groups, check_db_only=strict + ) + return await _read_declared_grants(replace(facts, unified_agents=tuple(unified)), strict) + case _: + raise RuntimeError("invalid native declared grants transition") async def _verified_human_agent_sources( @@ -919,7 +908,13 @@ async def verified_human_agent_grants(user_id: str | None, team_id: str | None = sources: Final = await _verified_human_agent_sources( user_id, allowed_team_ids=frozenset((team_id,)) if team_id else frozenset() ) - return frozenset().union(*(grants for source, grants in sources if source is None or source == team_id)) + decision: Final = agent_grants.decide( + agent_grants.Human( + tuple(agent_grants.Source(source, tuple(grants)) for source, grants in sources), + team_id, + ) + ) + return frozenset(decision.grants or ()) async def resolve_delegated_agent_team( @@ -931,15 +926,202 @@ async def resolve_delegated_agent_team( allowed_team_ids: frozenset[str] | None = None, ) -> str | None: sources: Final = await _verified_human_agent_sources(user_id) - if any(source is None and agent_id in grants for source, grants in sources): - return team_id - granting_teams: Final = frozenset( - source - for source, grants in sources - if source is not None and agent_id in grants and (allowed_team_ids is None or source in allowed_team_ids) + decision: Final = agent_grants.decide( + agent_grants.Delegated( + tuple(agent_grants.Source(source, tuple(grants)) for source, grants in sources), + agent_id, + team_id, + explicit_team, + tuple(allowed_team_ids) if allowed_team_ids is not None else None, + ) ) - if team_id in granting_teams: - return team_id - if not explicit_team and granting_teams: - return min(granting_teams) - raise HTTPException(403, "Select a team that grants access to this agent using x-litellm-team-id") + if decision.action == "SelectTeam": + raise HTTPException(403, "Select a team that grants access to this agent using x-litellm-team-id") + return decision.team_id + + +def _agent_grant_source(team_id: str | None, access: AgentAccess | AgentPolicyUnavailable) -> agent_grants.Source: + return agent_grants.Source( + team_id, + tuple(access.agent_ids) if isinstance(access, RestrictedAgentAccess) else None, + isinstance(access, AgentPolicyUnavailable), + ) + + +async def _read_resolved_agent_grants( + auth: UserAPIKeyAuth | None, + ceiling: CeilingResolver, + default_deny: Callable[[], bool], + key_lookup: AgentAccessLookup | None, + team_lookup: AgentAccessLookup | None, + facts: agent_grants.Resolve, +) -> AgentAccess: + decision: Final = agent_grants.decide(facts) + match decision.action: + case "Complete": + return _project_agent_grants(decision) + case "ReadManaged": + assert auth is not None + return await _managed_actor_agent_access(auth) + case "ReadKeyTeam": + access: Final = await AgentRequestHandler.resolve_key_team_agent_access( + auth, + strict=facts.strict, + lookup_key_access=key_lookup, + lookup_team_access=team_lookup, + ) + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + key_team_loaded=True, + key_team=_agent_grant_source(None, access).grants, + ), + ) + case "ReadDefault": + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + default_deny=default_deny(), + ), + ) + case "ReadCaller": + caller: Final = await AgentRequestHandler.agent_caller_access(auth, strict=facts.strict) + own: Final = managed.decide( + managed.IntersectGrants( + facts.key_team, + _agent_grant_source(None, caller).grants, + ) + ) + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + caller_loaded=True, + own=own.grants, + ), + ) + case "ReadCeiling": + ids: Final = await AgentRequestHandler._agent_access_group_ceiling(auth, ceiling, strict=facts.strict) # pyright: ignore[reportPrivateUsage] # module-owned staged reader + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + ceiling_loaded=True, + ceiling=tuple(ids) if ids is not None else None, + ), + ) + case _: + raise NativeAuthFailure(f"unexpected agent grant directive {decision.action}") + + +async def read_agent_access_facts( + agent_id: str, + auth: UserAPIKeyAuth | None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> agent_access.AgentAccessFacts: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.proxy_server import prisma_client + + registered: Final = global_agent_registry.get_agent_by_id(agent_id) + facts: Final = agent_access.AgentAccessFacts( + registry_managed=isinstance(registered, AgentResponse) and registered.identity_managed, + database_present=prisma_client is not None, + subject_present=auth is not None, + key_hash_present=bool(auth.api_key or auth.token) if auth is not None else False, + subject_managed=managed_agent_policy(auth) is not None, + session=auth.is_session_token if auth is not None else False, + custom=auth.authenticated_by_custom_auth if auth is not None else False, + ) + return await _read_agent_access(agent_id, auth, resolve_ceiling, facts) + + +async def _read_agent_access( + agent_id: str, + auth: UserAPIKeyAuth | None, + resolve_ceiling: CeilingResolver, + facts: agent_access.AgentAccessFacts, + authority: UserAPIKeyAuth | None = None, +) -> agent_access.AgentAccessFacts: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.auth.managed_authorization import project_agent + from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore + from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure + from litellm.proxy.proxy_server import prisma_client + from litellm.types.proxy.agent_identity import AgentIdentityFailure + + action: Final = agent_access.decide(facts) + match action: + case "Allow" | "Deny": + return facts + case "ReadTarget": + target: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id) + if isinstance(target, AgentIdentityFailure): + raise_identity_failure(target) + return await _read_agent_access( + agent_id, + auth, + resolve_ceiling, + replace(facts, target_loaded=True, target=project_agent(target) if target is not None else None), + authority, + ) + case "ReadAuthority": + from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + assert auth is not None + key_hash: Final = auth.api_key or auth.token + assert key_hash is not None + reloaded: Final = await MCPRequestHandler._reload_admitted_key(key_hash, check_db_only=True) # pyright: ignore[reportPrivateUsage] # authoritative key reload has no public seam + return await _read_agent_access( + agent_id, auth, resolve_ceiling, replace(facts, authority_loaded=True), reloaded + ) + case "ReadExplicitGrants": + assert auth is not None + source: Final = authority if authority is not None else auth + fresh: Final = source.model_copy( + update=MappingProxyType({"requires_fresh_policy": True, "agent_caller": auth.agent_caller}) + ) + explicit: Final = await _granted_agent_ids(fresh, _strict_agent_access, build_effective_auth_contexts) + return await _read_agent_access( + agent_id, auth, resolve_ceiling, replace(facts, explicit_grants=tuple(explicit)), authority + ) + case "ReadLegacyGrants": + grants: Final = await AgentRequestHandler.resolve_agent_access( + auth, resolve_ceiling, strict=bool(auth is not None and auth.requires_fresh_policy) + ) + return await _read_agent_access( + agent_id, + auth, + resolve_ceiling, + replace( + facts, + legacy_loaded=True, + legacy_grants=tuple(grants.agent_ids) if isinstance(grants, RestrictedAgentAccess) else None, + ), + authority, + ) + case "ReadTargetIds": + stable: Final = global_agent_registry.stable_agent_id(agent_id) + ids: Final = global_agent_registry.ids_for_agent(stable) + return await _read_agent_access( + agent_id, auth, resolve_ceiling, replace(facts, target_ids=tuple(ids)), authority + ) + case _: + assert_never(action) diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index 57e99edecba..b0aeeeb0164 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -1,10 +1,12 @@ from collections.abc import Mapping +from dataclasses import replace from itertools import product from types import MappingProxyType -from typing import Annotated, Final, Literal +from typing import Final, Literal, assert_never -from pydantic import Field, TypeAdapter, ValidationError +from pydantic import TypeAdapter, ValidationError +from litellm._auth import agent_access, managed from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure @@ -94,17 +96,18 @@ _MODEL_ROUTE_KINDS: Final[ def managed_agent_route_allowed(route: str, method: str | None) -> bool: - from litellm.proxy.auth.route_checks import RouteChecks + from litellm.proxy.auth import route_policy_adapter - if route in ("/agents", "/v1/agents"): - return method in (None, "GET", "HEAD") - if route in _MANAGED_REALTIME_ROUTES or route in _MANAGED_LIVE_ROUTES: - return method in (None, "GET") - if route in _MANAGED_MODEL_ROUTES or RouteChecks.check_route_access(route, _MANAGED_MODEL_PATHS): - return method in (None, "POST") - return RouteChecks.check_route_access(route, _MANAGED_MCP_ROUTES) or RouteChecks.check_route_access( - route, LiteLLMRoutes.agent_inference_routes.value - ) + return route_policy_adapter.classify(route, "Managed", method) + + +def _model_route_kind( + route: str, +) -> Literal["completion", "image_generation", "image_edit", "moderation", "speech", "body", "path"]: + for suffix, kind in _MODEL_ROUTE_KINDS.items(): + if route.endswith(suffix): + return kind + return "completion" def managed_inference_request( @@ -134,7 +137,7 @@ def managed_inference_request( return dict(body) from litellm.proxy.common_utils.http_parsing_utils import resolve_inference_model - kind: Final = next((kind for suffix, kind in _MODEL_ROUTE_KINDS.items() if route.endswith(suffix)), "completion") + kind: Final = _model_route_kind(route) endpoint_model: Final = path_model or ( query_model if route.endswith(("/completions", "/embeddings", "/images/generations", "/images/edits")) else None ) @@ -159,71 +162,131 @@ def managed_agent_policy(auth: "UserAPIKeyAuth | None") -> AgentResponse | None: async def admit_managed_actor(auth: UserAPIKeyAuth, store: AgentIdentityStore | None) -> None: delegation_verified: Final = auth._managed_delegation_verified # pyright: ignore[reportPrivateUsage] # the one-shot delegation marker is a PrivateAttr by design auth._managed_delegation_verified = False # pyright: ignore[reportPrivateUsage] # consumed here so a replayed token cannot reuse it - if auth.agent_id is None: + facts: Final = managed.ActorLookup( + agent_id=auth.agent_id, + has_database=store is not None, + context=_managed_context(auth.managed_agent_context), + jwt_claims=bool(auth.jwt_claims), + delegation_verified=delegation_verified, + ) + await _drive_actor(auth, store, facts) + + +async def _drive_actor( + auth: UserAPIKeyAuth, + store: AgentIdentityStore | None, + facts: managed.ActorLookup, + agent: AgentResponse | None = None, +) -> None: + decision: Final = managed.decide(facts) + if decision.denial is not None: + raise_identity_failure(_managed_failure(decision.denial)) + if decision.action == "Allow": return - if store is None: + if decision.action == "ReadRegistry": from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + assert auth.agent_id is not None registered: Final = global_agent_registry.get_agent_by_id(auth.agent_id) - if auth.managed_agent_context is not None or ( - registered is not None and (registered.identity_managed or registered.identity is not None) - ): - raise_identity_failure( - AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database") - ) + await _drive_actor( + auth, + store, + replace( + facts, + registry_loaded=True, + registry_managed=registered is not None + and (registered.identity_managed or registered.identity is not None), + ), + ) return - agent: Final = await store.agent(auth.agent_id) - if isinstance(agent, AgentIdentityFailure): - raise_identity_failure(agent) - if agent is None: + if decision.action == "ReadAgent": + assert store is not None and auth.agent_id is not None + loaded: Final = await store.agent(auth.agent_id) + if isinstance(loaded, AgentIdentityFailure): + raise_identity_failure(loaded) + await _drive_actor( + auth, + store, + replace(facts, agent_loaded=True, agent=project_agent(loaded) if loaded is not None else None), + loaded, + ) + return + if decision.action == "ReadRetired": + assert store is not None and auth.agent_id is not None retired: Final = await store.retired_agent(auth.agent_id) if isinstance(retired, AgentIdentityFailure): raise_identity_failure(retired) - if auth.managed_agent_context is not None or retired: - raise_identity_failure(AgentIdentityFailure(message="Agent no longer exists")) + await _drive_actor(auth, store, replace(facts, retired=retired), agent) return - if not agent.identity_managed: + if decision.action == "AssignPolicy": + assert agent is not None + auth.managed_agent_policy = agent + auth.billing_agent_policy = agent + auth.requires_fresh_policy = True + await _drive_actor(auth, store, replace(facts, policy_assigned=True), agent) return - if auth.jwt_claims and auth.managed_agent_context is None: - raise_identity_failure(AgentIdentityFailure(message="A managed agent requires a matching verified identity")) - failure: Final = actor_admission_failure(agent, auth.managed_agent_context) - if failure is not None: - raise_identity_failure(failure) - auth.managed_agent_policy = agent - auth.billing_agent_policy = agent - auth.requires_fresh_policy = True - if ( - auth.managed_agent_context is not None - and auth.managed_agent_context.mode == "delegated" - and not delegation_verified - ): + if decision.action == "ReadHumanGrants": from litellm.proxy.agent_endpoints.auth.agent_permission_handler import verified_human_agent_grants + assert auth.managed_agent_context is not None grants: Final = await verified_human_agent_grants(auth.managed_agent_context.user_id, auth.team_id) - if agent.agent_id not in grants: - raise_identity_failure( - AgentIdentityFailure(message="The delegated user is not permitted to invoke this agent") - ) + await _drive_actor(auth, store, replace(facts, human_grants=tuple(grants)), agent) + return + raise RuntimeError("Invalid managed actor decision") + + +def project_agent(agent: AgentResponse) -> managed.Agent: + return managed.Agent( + agent_id=agent.agent_id, + enabled=agent.enabled, + identity_managed=agent.identity_managed, + identity=managed.BoundIdentity(active=agent.identity.active, revision=agent.identity.revision) + if agent.identity is not None + else None, + execution_mode=agent.execution_mode, + ) + + +def _managed_context(context: ManagedAgentContext | None) -> managed.Context | None: + return ( + managed.Context( + agent_id=context.agent_id, + binding_revision=context.binding_revision, + mode=context.mode, + user_id=context.user_id, + ) + if context is not None + else None + ) + + +def _managed_failure(denial: managed.ManagedDenial) -> AgentIdentityFailure: + if denial == "DatabaseRequired": + return AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database") + messages: Final[Mapping[managed.ManagedDenial, str]] = MappingProxyType( + { + "Disabled": "Agent execution is disabled", + "MissingContext": "This agent requires its bound identity provider token", + "ChangedIdentity": "Agent identity changed during authentication; retry", + "ExecutionMode": "Agent is not enabled for this execution mode", + "MissingHuman": "A verified human subject is required", + "Removed": "Agent no longer exists", + "VerifiedIdentityRequired": "A managed agent requires a matching verified identity", + "HumanNotPermitted": "The delegated user is not permitted to invoke this agent", + } + ) + return AgentIdentityFailure(message=messages[denial]) def actor_admission_failure( agent: AgentResponse, context: ManagedAgentContext | None, ) -> AgentIdentityFailure | None: - if not agent.enabled or agent.identity is None or not agent.identity.active: - return AgentIdentityFailure(message="Agent execution is disabled") - if context is None: - return AgentIdentityFailure(message="This agent requires its bound identity provider token") - if context.agent_id != agent.agent_id or context.binding_revision != agent.identity.revision: - return AgentIdentityFailure(message="Agent identity changed during authentication; retry") - if agent.execution_mode not in (context.mode, "both"): - return AgentIdentityFailure(message="Agent is not enabled for this execution mode") - if context.mode == "delegated" and not context.user_id: - return AgentIdentityFailure(message="A verified human subject is required") - return None + decision: Final = managed.decide(managed.Actor(project_agent(agent), _managed_context(context))) + return _managed_failure(decision.denial) if decision.denial is not None else None -_INVOCATION_COST: Final = TypeAdapter(Annotated[float, Field(ge=0, allow_inf_nan=False)]) +_INVOCATION_COST: Final = TypeAdapter(float) def invocation_target(route: str, body: Mapping[str, object]) -> str | None: @@ -238,36 +301,93 @@ def invocation_target(route: str, body: Mapping[str, object]) -> str | None: async def prepare_agent_invocation( auth: UserAPIKeyAuth, target_name: str, store: AgentIdentityStore | None, *, billable: bool = True ) -> None: - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentRequestHandler from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through registered: Final = await get_agent_with_read_through(target_name) - if registered is None: - return - registered_managed: Final = registered.identity_managed or registered.identity is not None - if store is None and registered_managed: - raise_identity_failure( - AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database") - ) - target: Final = await store.agent(registered.agent_id) if store is not None else None - if isinstance(target, AgentIdentityFailure): - raise_identity_failure(target) - if target is None and registered_managed: - raise_identity_failure(AgentIdentityFailure(message="Invoked agent no longer exists")) - effective: Final = target if target is not None else registered - if not effective.identity_managed and auth.managed_agent_policy is None: - return - if not await AgentRequestHandler.is_agent_allowed(effective.agent_id, auth): - raise_identity_failure(AgentIdentityFailure(message="The caller is not permitted to invoke this agent")) - auth.invoked_agent_id = effective.agent_id - auth.invoked_agent_policy = effective - if auth.agent_id is None and effective.identity_managed: - auth.billing_agent_policy = effective - raw_fee: Final = (effective.litellm_params or MappingProxyType({})).get("cost_per_query", 0.0) if billable else 0.0 + facts: Final = agent_access.InvocationFacts( + registered=project_agent(registered) if registered is not None else None, + has_store=store is not None, + actor_managed=auth.managed_agent_policy is not None, + actor_id_present=auth.agent_id is not None, + ) + await _drive_invocation(auth, store, registered, None, billable, facts) + + +async def _drive_invocation( + auth: UserAPIKeyAuth, + store: AgentIdentityStore | None, + registered: AgentResponse | None, + target: AgentResponse | None, + billable: bool, + facts: agent_access.InvocationFacts, +) -> None: + plan: Final = agent_access.invocation(facts) + effective: Final = target if plan.use_target else registered + match plan.action: + case "Complete": + if facts.fee_loaded: + assert facts.fee is not None + auth.agent_invocation_cost = facts.fee + return + case "Deny": + failures: Final = MappingProxyType( + { + "DatabaseRequired": AgentIdentityFailure( + code="policy_unavailable", message="Managed agent policy requires a database" + ), + "Removed": AgentIdentityFailure(message="Invoked agent no longer exists"), + "CallerNotPermitted": AgentIdentityFailure( + message="The caller is not permitted to invoke this agent" + ), + "InvalidFee": AgentIdentityFailure( + code="policy_unavailable", message="Agent invocation price is invalid" + ), + } + ) + assert plan.denial is not None + raise_identity_failure(failures[plan.denial]) + case "ReadTarget": + assert registered is not None + loaded: Final = await store.agent(registered.agent_id) if store is not None else None + if isinstance(loaded, AgentIdentityFailure): + raise_identity_failure(loaded) + await _drive_invocation( + auth, + store, + registered, + loaded, + billable, + replace(facts, target_loaded=True, target=project_agent(loaded) if loaded is not None else None), + ) + case "ReadPermissions": + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import read_agent_access_facts + + assert effective is not None + permissions: Final = await read_agent_access_facts(effective.agent_id, auth) + await _drive_invocation(auth, store, registered, target, billable, replace(facts, permissions=permissions)) + case "Assign": + assert effective is not None + auth.invoked_agent_id = effective.agent_id + auth.invoked_agent_policy = effective + if plan.assign_billing: + auth.billing_agent_policy = effective + await _drive_invocation(auth, store, registered, target, billable, replace(facts, assigned=True)) + case "ReadFee": + assert effective is not None + raw: Final = ( + (effective.litellm_params or MappingProxyType[str, object]({})).get("cost_per_query", 0.0) + if billable + else 0.0 + ) + await _drive_invocation( + auth, store, registered, target, billable, replace(facts, fee_loaded=True, fee=_invocation_fee(raw)) + ) + case _: + assert_never(plan.action) + + +def _invocation_fee(raw: object) -> float | None: try: - fee: Final = _INVOCATION_COST.validate_python(raw_fee) + return _INVOCATION_COST.validate_python(raw) except ValidationError: - raise_identity_failure( - AgentIdentityFailure(code="policy_unavailable", message="Agent invocation price is invalid") - ) - auth.agent_invocation_cost = fee + return None diff --git a/litellm/proxy/agent_endpoints/identity_store.py b/litellm/proxy/agent_endpoints/identity_store.py index c3e2064cbd5..9767ca94224 100644 --- a/litellm/proxy/agent_endpoints/identity_store.py +++ b/litellm/proxy/agent_endpoints/identity_store.py @@ -3,9 +3,11 @@ from collections.abc import Mapping from datetime import datetime, timezone from typing import TYPE_CHECKING, Final +from litellm._auth.failure import NativeAuthFailure from litellm._internal_context import with_service_target from litellm.proxy.agent_endpoints.managed_identity import classify_agent_subject from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl +from litellm.repositories.native_auth_repository import NativeAuthRepository, native_auth_repository from litellm.repositories.table_repositories import ( AgentIdentityRepository, AgentsRepository, @@ -48,6 +50,7 @@ class AgentIdentityStore: RetiredAgentIdentityRepository(client, use_writer=True), RetiredAgentRepository(client, use_writer=True), cache=cache, + native=native_auth_repository(client), ) def __init__( @@ -59,6 +62,7 @@ class AgentIdentityStore: retired_agents: RetiredAgentRepository | None = None, *, cache: UserApiKeyCache | None = None, + native: NativeAuthRepository | None = None, ) -> None: self.agents = agents self.identities = identities @@ -66,9 +70,12 @@ class AgentIdentityStore: self.retired = retired self.retired_agents = retired_agents self.cache = cache + self.native = native async def agent(self, agent_id: str) -> AgentResponse | AgentIdentityFailure | None: try: + if self.native is not None: + return await self.native.one("agent", [agent_id], AgentResponse, writer=True) where: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id} include: Final[LiteLLM_AgentsTableInclude] = { "identity": True, @@ -78,6 +85,8 @@ class AgentIdentityStore: if row is None: return None return AgentResponse.model_validate(row.model_dump()) + except NativeAuthFailure: + raise except Exception: return AgentIdentityFailure(code="policy_unavailable", message="Agent policy could not be loaded") @@ -185,6 +194,8 @@ class AgentIdentityStore: return AgentIdentityFailure(code="policy_unavailable", message="Subject classification is unavailable") async def retired_agent(self, agent_id: str) -> bool | AgentIdentityFailure: + if self.native is not None: + return bool(await self.native.read("retired_agent", [agent_id], writer=True)) if self.retired_agents is None: return AgentIdentityFailure(code="policy_unavailable", message="Agent history is unavailable") try: diff --git a/litellm/proxy/auth/agent_ceiling_adapter.py b/litellm/proxy/auth/agent_ceiling_adapter.py new file mode 100644 index 00000000000..aee14fc9c6b --- /dev/null +++ b/litellm/proxy/auth/agent_ceiling_adapter.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Final + +from fastapi import HTTPException + +from litellm._auth.agent_ceilings import ( + AgentCeilingFacts, + AgentCeilingGroup, + AgentGroupLoadFacts, + agent_ceiling, + agent_group_load, +) +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable + +if TYPE_CHECKING: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AccessGroupIdsLoader, + AccessGroupLoader, + AgentAccessGroupCeiling, + ) + + +@dataclass(frozen=True, slots=True) +class AgentCeilingReading: + facts: AgentCeilingFacts + rows: tuple[LiteLLM_AccessGroupTable | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + agent_id: str + ids: AccessGroupIdsLoader + group: AccessGroupLoader + + +async def read_agent_ceiling_facts( + agent_id: str, + load_access_group_ids: AccessGroupIdsLoader, + load_access_group: AccessGroupLoader, + *, + check_db_only: bool = False, +) -> AgentCeilingReading: + return await _read( + AgentCeilingReading(AgentCeilingFacts(check_db_only=check_db_only)), + _Source(agent_id, load_access_group_ids, load_access_group), + ) + + +def project_agent_ceiling(reading: AgentCeilingReading) -> AgentAccessGroupCeiling | None: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + decision: Final = agent_ceiling(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete agent ceiling facts") + if not decision.present: + return None + return AgentAccessGroupCeiling( + decision.ids, + decision.models, + decision.mcp_server_ids, + decision.agent_ids, + tuple(_row(reading, index) for index in decision.indices), + ) + + +def _row(reading: AgentCeilingReading, index: int) -> LiteLLM_AccessGroupTable: + row: Final = reading.rows[index] + if row is None: + raise NativeAuthRuntimeError("agent ceiling selected a missing row") + return row + + +async def _group(group_id: str, loader: AccessGroupLoader) -> tuple[AgentCeilingGroup, LiteLLM_AccessGroupTable | None]: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + try: + row: Final = await loader(group_id) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError: + return AgentCeilingGroup("Missing"), None + if not isinstance(row, LiteLLM_AccessGroupTable): + return AgentCeilingGroup("Unavailable"), None + return AgentCeilingGroup( + "Found", tuple(row.access_model_names), tuple(row.access_mcp_server_ids), tuple(row.access_agent_ids) + ), row + + +async def _authoritative(group_id: str) -> LiteLLM_AccessGroupTable | None: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + _load_access_group, # pyright: ignore[reportPrivateUsage] # authoritative raw reader preserves the public injected-loader contract + ) + + return await _load_access_group(group_id, check_db_only=True) + + +async def _read(reading: AgentCeilingReading, source: _Source) -> AgentCeilingReading: + decision: Final = agent_ceiling(reading.facts) + match decision.kind: + case "Complete": + return reading + case "ReadIds": + ids: Final = await source.ids(source.agent_id) + return await _read(replace(reading, facts=replace(reading.facts, ids=ids)), source) + case "ReadGroups": + if reading.facts.ids is None: + raise NativeAuthRuntimeError("agent group reads without named IDs") + loader: Final = _authoritative if decision.authoritative else source.group + rows: Final = tuple(await asyncio.gather(*(_group(group_id, loader) for group_id in reading.facts.ids))) + return AgentCeilingReading( + replace(reading.facts, groups=tuple(row[0] for row in rows)), tuple(row[1] for row in rows) + ) + + +async def load_agent_access_group( + access_group_id: str, *, check_db_only: bool = False +) -> LiteLLM_AccessGroupTable | None: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + initial: Final = AgentGroupLoadFacts(prisma_client is not None, check_db_only) + initial_decision: Final = agent_group_load(initial) + if initial_decision.kind == "DatabaseUnavailable": + verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) + return None + if initial_decision.kind != "Read": + raise NativeAuthRuntimeError("invalid initial agent group read decision") + try: + return await get_access_object( + access_group_id=access_group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError: + if agent_group_load(replace(initial, status="NotFound")).kind != "Raise": + raise NativeAuthRuntimeError("invalid missing agent group decision") from None + raise + except HTTPException as error: + decision: Final = agent_group_load(replace(initial, status="HttpError")) + if decision.kind == "Raise": + raise + if decision.kind != "HttpUnavailable": + raise NativeAuthRuntimeError("invalid failed agent group decision") from error + verbose_proxy_logger.warning( + "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, error.detail + ) + return None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4514882a45f..ae8324f88ff 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -10,22 +10,24 @@ Run checks for: """ import asyncio -import math import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime from functools import partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, Protocol, TypeAlias +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, Protocol, TypeAlias, assert_never from fastapi import HTTPException, Request, status from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack import litellm +from litellm._auth import access as access_policy +from litellm._auth import common_flow +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict @@ -34,8 +36,6 @@ from litellm.constants import ( CLI_SESSION_KEY_PREFIX, DEFAULT_ACCESS_GROUP_CACHE_TTL, DEFAULT_IN_MEMORY_TTL, - DEFAULT_MAX_RECURSE_DEPTH, - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, INVALID_API_KEY_ERROR_MESSAGE, LITELLM_PROXY_MASTER_KEY_ALIAS, @@ -43,9 +43,7 @@ from litellm.constants import ( REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.dd_tracing import tracer -from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.tool_fields import DropPolicy @@ -77,7 +75,6 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, - ModelAccessDeniedProxyException, NewTeamRequest, ProxyErrorTypes, ProxyException, @@ -91,19 +88,9 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( resolve_agent_access_group_ceiling, ) from litellm.proxy.agent_endpoints.auth.agent_caller import ( - agent_caller_auth, load_agent_caller_team, load_agent_caller_user, ) -from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy -from litellm.proxy.auth.budget_throttle import ( - budget_throttle_percentage, - should_throttle_budget_exceeded, -) -from litellm.proxy.auth.model_access_denied import ( - customer_model_access_denied_client_message, - model_access_denied_client_message, -) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( broadcast_evictions, @@ -147,24 +134,18 @@ from litellm.proxy.guardrails.tool_name_extraction import ( extract_request_tool_names, ) from litellm.proxy.route_llm_request import route_request -from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start -from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.redis.stores.spend_counters import ( SPEND_COUNTERS_TARGET, - end_user_counter, - key_counter, - key_window_counter, - model_access_group_counter, - org_counter, - project_counter, - tag_counter, - team_counter, team_member_counter, - team_window_counter, - user_counter, ) from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.native_auth_repository import ( + EndUserRegistryRow, + ModelAccessGroupRow, + TagRegistryRow, + native_auth_repository, +) from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import DatabaseClient, RowT_co @@ -401,22 +382,6 @@ def _model_access_group_budget_table( return repo.table -class _MemberModelScope(Protocol): - @property - def allowed_models(self) -> Sequence[str] | None: ... - - -class _TeamMembershipModelScope(Protocol): - @property - def litellm_budget_table(self) -> _MemberModelScope | None: ... - - -def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[str]: - """The member's own model scope, read through a narrowed view of the membership row.""" - budget_table: Final = membership.litellm_budget_table - return () if budget_table is None else (budget_table.allowed_models or ()) - - class _RawCacheRead(Protocol): async def async_get_cache(self, *, key: str, local_only: bool = ...) -> object: ... @@ -435,18 +400,6 @@ typed_general_settings: Final = _typed_request_body _TRACE_ID_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -def _metadata_mapping_has_trace_id(metadata: Mapping[str, object]) -> bool: - trace_id: Final[object] = metadata.get("trace_id") - return isinstance(trace_id, str) and bool(trace_id) - - -def _metadata_has_trace_id(metadata: object) -> bool: - if not isinstance(metadata, dict): - return False - metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(metadata) - return _metadata_mapping_has_trace_id(metadata_dict) - - def _team_metadata_requires_trace_id(team_object: LiteLLM_TeamTable) -> bool: if not isinstance( team_object.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # narrow the untyped field @@ -460,31 +413,45 @@ def _team_metadata_requires_trace_id(team_object: LiteLLM_TeamTable) -> bool: def _request_has_trace_id(request_body: Mapping[str, object], request: Request, *, headers_only: bool = False) -> bool: + return _read_trace_id(request_body, request, "Header", headers_only=headers_only) + + +def _read_trace_id( + request_body: Mapping[str, object], + request: Request, + source: access_policy.TraceSource, + *, + headers_only: bool, +) -> bool: from litellm.proxy.litellm_pre_call_utils import ( _trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser get_chain_id_from_headers, metadata_variable_name_for_request, ) - if get_chain_id_from_headers(dict(request.headers)): - return True - if headers_only: - return False - litellm_trace_id: Final[object] = request_body.get("litellm_trace_id") - if isinstance(litellm_trace_id, str) and litellm_trace_id: - return True - if "litellm_trace_id" not in request_body: - traceparent: Final[str | None] = request.headers.get("traceparent") - if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent): - return True - metadata_variable_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body) - selected_metadata: Final[object] = request_body.get(metadata_variable_name) - if isinstance(selected_metadata, dict): - selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata) - if "trace_id" in selected_metadata_dict: - return _metadata_mapping_has_trace_id(selected_metadata_dict) - other_metadata_variable_name: Final = "litellm_metadata" if metadata_variable_name == "metadata" else "metadata" - return _metadata_has_trace_id(request_body.get(other_metadata_variable_name)) + if source == "Header": + value, present = get_chain_id_from_headers(dict(request.headers)), True + elif source == "Body": + value, present = request_body.get("litellm_trace_id"), "litellm_trace_id" in request_body + elif source == "Traceparent": + traceparent: Final = request.headers.get("traceparent") + value, present = _trace_id_from_traceparent(traceparent) if isinstance(traceparent, str) else None, True + else: + selected_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body) + name: Final = ( + selected_name + if source == "Selected" + else ("litellm_metadata" if selected_name == "metadata" else "metadata") + ) + metadata: Final = request_body.get(name) + mapping: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(metadata) if isinstance(metadata, dict) else {} + value, present = mapping.get("trace_id"), "trace_id" in mapping + result: Final = access_policy.decide( + access_policy.TraceStep(source, value if isinstance(value, str) else None, present, headers_only) + ) + if result.next_trace is not None: + return _read_trace_id(request_body, request, result.next_trace, headers_only=headers_only) + return result.denial is None class _JsonLoadsObj(Protocol): @@ -541,138 +508,10 @@ def _log_budget_lookup_failure(entity: str, error: Exception) -> None: ) -def _get_router_zero_cost_cache(llm_router: RegistryReader) -> dict[str, bool] | None: - """ - Return the router's per-instance zero-cost cache, or ``None`` for objects - that don't expose one (e.g. ``MagicMock`` stand-ins in unit tests). - - The cache lives on the ``Router`` instance so it: - * is invalidated by ``Router._invalidate_model_group_info_cache`` on - any model add/remove/upsert (including in-place pricing changes via - ``/model/update``, which go through ``upsert_deployment``); - * dies with the router itself — no risk of CPython reusing the - previous router's ``id()`` and serving its cached entries. - """ - cache: Final = llm_router.registry().zero_cost() - return cache if isinstance(cache, dict) else None - - def is_model_cost_zero(model: str | list[str] | None, llm_router: RegistryReader | None) -> bool: - """ - Check if a model has zero cost (no configured pricing). + from litellm.proxy.auth import budget_pricing_adapter - Uses the router registry's group info to get pricing information. - - Args: - model: The model name or list of model names - llm_router: The LiteLLM router instance - - Returns: - bool: True if all costs for the model are zero, False otherwise - """ - if model is None or llm_router is None: - return False - - # Handle list of models - model_list: Final = [model] if isinstance(model, str) else model - - zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router) - - for model_name in model_list: - if zero_cost_cache is not None: - cached = zero_cost_cache.get(model_name) - if cached is not None: - if cached is False: - return False - continue - try: - model_group_info = llm_router.registry().group(model_name, include_hidden=True) - - if model_group_info is None: - # Model not found or no pricing info available - # Conservative approach: assume it has cost - verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - # Check costs for this model - # Only allow bypass if BOTH costs are explicitly set to 0 (not None) - input_cost = model_group_info.input_cost_per_token - output_cost = model_group_info.output_cost_per_token - - # If costs are not explicitly configured (None), assume it has cost - if input_cost is None or output_cost is None: - verbose_proxy_logger.debug( - "Model %s has undefined cost (input: %s, output: %s), assuming it has cost", - model_name, - input_cost, - output_cost, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - # If either cost is non-zero, return False - if input_cost > 0 or output_cost > 0: - verbose_proxy_logger.debug( - "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - # Costs are 0 — verify this is from explicit configuration, - # not from defaulted sparse auto-registration entries. - # See: https://github.com/BerriAI/litellm/issues/24770 - safe_name = str(model_name).replace("\n", "").replace("\r", "") - served_deployments = _deployments_served_for(model_name, llm_router) - if not _is_cost_explicitly_configured(served_deployments): - verbose_proxy_logger.debug( - "Model %s has zero cost but no explicit cost " - "configuration in model_cost entry — treating as unknown " - "cost (enforce budget)", - safe_name, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - if _has_explicit_positive_cost(served_deployments): - verbose_proxy_logger.debug( - "Model %s routes to a deployment with an explicit positive per-token price (enforce budget)", - safe_name, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - if _has_ptu_flat_cost(served_deployments): - verbose_proxy_logger.debug( - "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " - "rate is not a free model (enforce budget)", - safe_name, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - verbose_proxy_logger.debug( - "Model %s has zero cost explicitly configured (input: %s, output: %s)", - safe_name, - input_cost, - output_cost, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = True - - except Exception as e: - # If we can't determine the cost, assume it has cost (conservative approach) - verbose_proxy_logger.debug("Error checking cost for model %s: %s, assuming it has cost", model_name, e) - return False - - # All models checked have zero cost - return True + return budget_pricing_adapter.is_model_cost_zero(model, llm_router) _is_model_cost_zero: Final = is_model_cost_zero @@ -682,7 +521,7 @@ _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) _TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) -def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: +def deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: """ The deployments a request to ``model`` routes to, for the gates that waive budget checks. @@ -706,128 +545,16 @@ def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[ ) -def _has_ptu_flat_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: - """Whether any of the deployments bills reserved PTU capacity as a flat cost. - - Such a deployment carries an explicit zero per-token price so the flat cost is not charged - twice, which otherwise reads here as a free model and waives every budget check for it. - """ - for deployment in deployments: - model_info = deployment.get("model_info") or _NO_MODEL_INFO - if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None: - return True - return False - - -def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) -> bool: - """ - Check if any of the deployments has cost fields explicitly set in its - litellm.model_cost entry. - - When Router._create_deployment() registers a model not in the global - cost map, it creates a sparse entry like {"id": ""} with no cost - fields. _get_model_info_helper() then defaults missing costs to 0. - This function detects that scenario by checking the raw model_cost entry. - - It also reaches a deployment that prices itself through its ``model_info`` block, whose entry - lands in the cost map under the deployment id. - """ - for deployment in deployments: - model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") - if model_id is None: - continue - raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) - if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: - return True - return False - - -def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: - """Whether any of the deployments carries an explicit positive per-token or per-query price. - - The group's price is read from the deployments named after the routed group and from - that group's own alias target, while an alias chain is served from the deployments named - after its routed group or from a wildcard route; a priced route among those serves the - request at its own rate whatever the group's price reads. - """ - for deployment in deployments: - model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") - if model_id is None: - continue - raw_entry = as_str_mapping(litellm.model_cost.get(model_id)) or _EMPTY_COST_ENTRY - if ( - _is_positive_cost(raw_entry.get("input_cost_per_token")) - or _is_positive_cost(raw_entry.get("output_cost_per_token")) - or _is_positive_cost(_deployment_query_price(deployment["litellm_params"], raw_entry)) - ): - return True - return False - - -def _deployment_query_price(litellm_params: Mapping[str, object], cost_entry: Mapping[str, object]) -> float | None: - provider: Final = litellm_params.get("custom_llm_provider") or str(litellm_params.get("model")).partition("/")[0] - return decisions_request_price({**cost_entry, "litellm_provider": provider}, decisions_pricing_fields(cost_entry)) +_deployments_served_for: Final = deployments_served_for _EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) -def _is_positive_cost(value: object) -> bool: - return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 - - -def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool: - if entry.get("tiered_pricing") is not None: - return True - for key, value in entry.items(): - if "cost_per" not in key: - continue - if _is_positive_cost(value): - return True - if isinstance(value, dict) and any(_is_positive_cost(nested) for nested in value.values()): - return True - return False - - -def _entry_declares_price(entry: Mapping[str, object]) -> bool: - return any("cost_per" in key or key == "tiered_pricing" for key in entry) - - -def _model_group_has_pricing(model: str, llm_router: "Router") -> bool: - """ - A model group counts as priced when a deployment overrides any *cost_per* field or - tiered_pricing in its litellm_params, even at zero, or when its resolved model info carries - tiered_pricing or a positive price on any billed metric (tokens, characters, seconds, pages, - images, queries, ...), so models billed by a non-token metric are not treated as unpriced. - """ - registry: Final = llm_router.registry() - for deployment in registry.deployments(model) or (): - litellm_params = deployment.get("litellm_params") or _EMPTY_COST_ENTRY - if _entry_declares_price(litellm_params): - return True - - model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") - if model_id is None: - continue - - model_info = registry.deployment_model_info(model_id, litellm_params.get("model") or "") - if model_info is not None and _entry_has_priced_metric(model_info): - return True - - return False - - def model_has_no_cost_mapping(model: str | None, llm_router: Router | None) -> bool: - if not model or llm_router is None: - return False + from litellm.proxy.auth.unpriced_model_adapter import model_has_no_cost_mapping as native_model_has_no_cost_mapping - if llm_router.registry().group(model) is None: - return False - - if _model_group_has_pricing(model=model, llm_router=llm_router): - return False - - return not _is_cost_explicitly_configured(llm_router.registry().deployments(model) or ()) + return native_model_has_no_cost_mapping(model, llm_router) def _unpriced_models_in_request(model: str | list[str] | None, llm_router: Router | None) -> tuple[str, ...]: @@ -863,7 +590,7 @@ async def _run_project_checks( return # 1.1. If project is blocked - if project_object.blocked is True: + if access_policy.decide(access_policy.Blocked(project_object.blocked is True)).denial is not None: raise Exception( f"Project={project_object.project_id} is blocked. Update via `/project/update` if you're an admin." ) @@ -894,26 +621,43 @@ async def _run_project_checks( def _enforce_user_param_check(general_settings: dict, request: Request, request_body: dict, route: str) -> None: - if not general_settings.get("enforce_user_param", False): + if ( + access_policy.decide( + access_policy.UserParamRequest(bool(general_settings.get("enforce_user_param", False))) + ).read + is None + ): return - http_method: Final = request.method if hasattr(request, "method") else None - is_post_method: Final = http_method and http_method.upper() == "POST" is_openai_route: Final = RouteChecks.is_llm_api_route(route=route) is_mcp_route: Final = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value ) - if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body: + decision: Final = access_policy.decide( + access_policy.UserParam( + enabled=bool(general_settings.get("enforce_user_param", False)), + method=http_method, + llm_route=is_openai_route, + mcp_route=is_mcp_route, + has_user="user" in request_body, + ) + ) + if decision.denial is not None: raise Exception(f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}") def _reject_clientside_metadata_tags_check( general_settings: dict, request_body: dict, route: str, request: Request | None = None ) -> None: - if not general_settings.get("reject_clientside_metadata_tags", False): + initial: Final = access_policy.ClientTagRequest( + bool(general_settings.get("reject_clientside_metadata_tags", False)) + ) + if access_policy.decide(initial).read is None: + return + routed: Final = replace(initial, llm_route=RouteChecks.is_llm_api_route(route=route)) + if access_policy.decide(routed).read is None: return - original_client_tags: Final = ( getattr(request.state, "litellm_baggage_client_metadata_tags", None) if request is not None else None ) @@ -926,7 +670,14 @@ def _reject_clientside_metadata_tags_check( and "tags" in request_body["metadata"] ) ) - if RouteChecks.is_llm_api_route(route=route) and has_client_tags: + decision: Final = access_policy.decide( + access_policy.ClientTags( + enabled=initial.enabled, + llm_route=bool(routed.llm_route), + has_tags=has_client_tags, + ) + ) + if decision.denial is not None: raise ProxyException( message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via key or team metadata, or by configured request_tag_mappings.", type=ProxyErrorTypes.bad_request_error, @@ -936,19 +687,9 @@ def _reject_clientside_metadata_tags_check( def _global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_checks: bool, route: str) -> None: - if ( - litellm.max_budget > 0 - and not skip_budget_checks - and global_proxy_spend is not None - and RouteChecks.is_llm_api_route(route=route) - and route not in MODEL_LIST_ROUTES - ): - if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget: - raise litellm.BudgetExceededError( - current_cost=global_proxy_spend, - max_budget=litellm.max_budget, - entity_type=Litellm_EntityType.PROXY.value, - ) + from litellm.proxy.auth import common_budget_adapter + + return common_budget_adapter.global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route) _GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = ( @@ -964,17 +705,22 @@ def requests_guardrail_modification(request_body: Mapping[str, object]) -> bool: default-on guardrails run (``disable_global_guardrails``, ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``), at the top level or in either metadata key""" return ( - _requests_guardrail_modification_in(request_body.get("metadata")) - or _requests_guardrail_modification_in(request_body.get("litellm_metadata")) - or _requests_guardrail_modification_in(request_body) + access_policy.decide(access_policy.GuardrailRequest(_requested_guardrail_keys(request_body))).read is not None ) -def _requests_guardrail_modification_in(container: object) -> bool: - coerced: Final = _guardrail_container_as_dict(container) - if coerced is None: - return False - return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) +def _requested_guardrail_keys(request_body: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + chain.from_iterable( + tuple(container) + for container in ( + _guardrail_container_as_dict(request_body.get("metadata")), + _guardrail_container_as_dict(request_body.get("litellm_metadata")), + request_body, + ) + if container is not None + ) + ) def _guardrail_container_as_dict(container: object) -> Mapping[str, object] | None: @@ -1008,10 +754,19 @@ def guardrail_modification_check(request_body: Mapping[str, object], team_object require_guardrail_modification_permission(team_object) -def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: - from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails +def _team_modify_guardrails(team_object: LiteLLM_TeamTable | None) -> bool | None: + metadata: Final = _typed_request_body(team_object.metadata or {}) if team_object is not None else {} + guardrails: Final = metadata.get("guardrails") + values: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(guardrails) if isinstance(guardrails, dict) else {} + value: Final = values.get("modify_guardrails") + return value if isinstance(value, bool) else None - if not can_modify_guardrails(team_object): + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: + if ( + access_policy.decide(access_policy.Guardrails(("guardrails",), _team_modify_guardrails(team_object))).denial + is not None + ): raise HTTPException( status_code=403, detail={"error": "Your team does not have permission to modify guardrails."}, @@ -1033,26 +788,46 @@ async def check_tools_allowlist( get_call_types_for_route, ) - if valid_token is None: + if access_policy.decide(access_policy.ToolRequest(valid_token is not None)).read is None: return + assert valid_token is not None call_types: Final = get_call_types_for_route(route) - if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types): + route_plan: Final = access_policy.decide( + access_policy.ToolRequest(True, tuple(ct.value for ct in call_types or ()), tuple(TOOL_CAPABLE_CALL_TYPES)) + ) + if route_plan.read is None: return key_meta: Final = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} team_meta: Final = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} key_allowed: Final = key_meta.get("allowed_tools") team_allowed: Final = team_meta.get("allowed_tools") - effective: Final = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed - if not isinstance(effective, list) or len(effective) == 0: + key_facts: Final = access_policy.ToolAllowlist(isinstance(key_allowed, list) and len(key_allowed) > 0) + key_plan: Final = access_policy.decide(key_facts) + facts: Final = ( + replace(key_facts, team_nonempty_list=isinstance(team_allowed, list) and len(team_allowed) > 0) + if key_plan.read == "TeamToolList" + else key_facts + ) + if access_policy.decide(facts).read is None: return tool_names: Final = extract_request_tool_names(route, request_body) - if not tool_names: + grant_plan: Final = access_policy.decide(replace(facts, requested_count=len(tool_names))) + if grant_plan.read is None: return - allowed_set: Final = {str(t) for t in effective} - disallowed: Final = [n for n in tool_names if n not in allowed_set] - if disallowed: + if grant_plan.read not in ("KeyToolGrant", "TeamToolGrant"): + raise NativeAuthRuntimeError("Unexpected tool-grant decision") + effective: Final = key_allowed if grant_plan.read == "KeyToolGrant" else team_allowed + assert isinstance(effective, list) + decision: Final = access_policy.decide( + access_policy.Tools( + requested=tuple(tool_names), + key_allowed=tuple(str(t) for t in effective), + team_allowed=None, + ) + ) + if decision.denial is not None: raise ProxyException( - message=f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.", + message=f"Tool(s) {list(decision.rejected)} are not in the allowed tools list for this key/team.", type=ProxyErrorTypes.tool_access_denied, param="tools", code=status.HTTP_403_FORBIDDEN, @@ -1107,17 +882,9 @@ def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] | def route_skips_budget_checks(route: str) -> bool: - if route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: - return False - if route in MODEL_DISCOVERY_ROUTES: - return True - if RouteChecks.is_llm_api_route(route=route): - return False - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack - InitPassThroughEndpointHelpers, - ) + from litellm.proxy.auth import budget_pricing_adapter - return not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route) + return budget_pricing_adapter.route_skips_budget_checks(route) def request_skips_budget_checks( @@ -1127,11 +894,9 @@ def request_skips_budget_checks( *, request_body: Mapping[str, object] | None = None, ) -> bool: - if route_skips_budget_checks(route=route): - return True - if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body): - return True - return is_model_cost_zero(model=model, llm_router=llm_router) + from litellm.proxy.auth import budget_pricing_adapter + + return budget_pricing_adapter.request_skips_budget_checks(route, model, llm_router, request_body=request_body) async def common_checks( @@ -1181,35 +946,60 @@ async def common_checks( team_id=valid_token.team_id if valid_token is not None else None, ) - skip_all_budget_checks: Final = skip_budget_checks or request_skips_budget_checks( - route=route, - model=_model, - llm_router=llm_router, - request_body=_typed_request_body(request_body), + budget_skip: Final = common_flow.BudgetSkip(skip_budget_checks) + budget_gate: Final = ( + common_flow.read( + replace( + budget_skip, + exempt=request_skips_budget_checks( + route=route, model=_model, llm_router=llm_router, request_body=_typed_request_body(request_body) + ), + ) + ) + if common_flow.read(budget_skip) == "ReadExemption" + else common_flow.read(budget_skip) ) + skip_all_budget_checks: Final = budget_gate == "SkipBudgets" - membership_user_id: Final = ( - valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None + flow_facts: Final = common_flow.Facts( + model_present=bool(_model), + skip_budgets=skip_all_budget_checks, + team_present=team_object is not None, + user_present=user_object is not None, + token_present=valid_token is not None, + user_id=valid_token.user_id if valid_token is not None else None, + agent_id=None, + customer_present=False, + customer_models=(), ) - team_membership_loaded: Final = team_object is not None and membership_user_id is not None - loaded_team_membership: Final = ( - await get_team_membership( + flow: Final = common_flow.plan(flow_facts) + membership_user_id: Final = valid_token.user_id if valid_token is not None else None + team_membership_loaded: Final = flow.membership + + async def read_membership() -> LiteLLM_TeamMembership | None: + if not flow.membership: + return None + assert membership_user_id is not None and team_object is not None + return await get_team_membership( user_id=membership_user_id, team_id=team_object.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - if team_object is not None and membership_user_id is not None - else None - ) - unpriced_models: Final = ( - _unpriced_models_in_request(model=_model, llm_router=llm_router) - if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route) - else () + loaded_team_membership: Final = await read_membership() + + pricing: Final = common_flow.Pricing(litellm.block_requests_for_models_without_pricing) + pricing_read: Final = ( + common_flow.read(replace(pricing, llm_route=RouteChecks.is_llm_api_route(route=route))) + if common_flow.read(pricing) == "ReadRoute" + else common_flow.read(pricing) ) - if unpriced_models: + unpriced_models: Final = ( + _unpriced_models_in_request(model=_model, llm_router=llm_router) if pricing_read == "ReadUnpriced" else () + ) + if common_flow.read(common_flow.Unpriced(len(unpriced_models))) == "RejectUnpriced": raise ProxyException( message=_unpriced_models_block_message(unpriced_models), type=ProxyErrorTypes.model_cost_map_missing, @@ -1218,34 +1008,30 @@ async def common_checks( ) # 1. If team is blocked - if team_object is not None and team_object.blocked is True: + if ( + team_object is not None + and access_policy.decide(access_policy.Blocked(team_object.blocked is True)).denial is not None + ): raise Exception(f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin.") # 2. If team can call model (or key's access_group_ids grant it) - if _model and team_object: + if flow.team_models: + assert _model is not None and team_object is not None with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"): - try: - await can_team_access_model( - model=_model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=(valid_token.team_model_aliases if valid_token else None), - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - except ProxyException as team_denial: - if team_denial.type != ProxyErrorTypes.team_model_access_denied: - raise - if not await key_access_group_grants_model( - model=_model, - valid_token=valid_token, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ): - raise + from litellm.proxy.auth.team_override_adapter import can_team_call_model_with_override + + await can_team_call_model_with_override( + model=_model, + valid_token=valid_token, + team_object=team_object, + llm_router=llm_router, + team_model_aliases=(valid_token.team_model_aliases if valid_token else None), + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) # 2.2. If team member has per-member model scope, enforce it - if _model and team_object and valid_token and valid_token.user_id: + if flow.member_models: + assert _model is not None and team_object is not None and valid_token is not None with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_model_access"): await check_team_member_model_access( model=_model, @@ -1261,17 +1047,20 @@ async def common_checks( ) # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent - if valid_token is not None and valid_token.agent_id: + if common_flow.plan( + replace(flow_facts, agent_id=valid_token.agent_id if valid_token is not None else None) + ).agent_trace: + assert valid_token is not None and valid_token.agent_id is not None from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers agent: Final = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) if agent is not None: require_trace_id: Final = (agent.litellm_params or {}).get("require_trace_id_on_calls_by_agent") - if require_trace_id: + if access_policy.decide(access_policy.TraceRequest(bool(require_trace_id))).next_trace is not None: headers_dict: Final = dict(request.headers) trace_id: Final = get_chain_id_from_headers(headers_dict) - if not trace_id: + if access_policy.decide(access_policy.AgentTrace(bool(require_trace_id), trace_id)).denial is not None: raise ProxyException( message="Requests made with this agent's key must include the x-litellm-trace-id header.", type=ProxyErrorTypes.bad_request_error, @@ -1294,32 +1083,28 @@ async def common_checks( ) or pass_through_route ) - team_trace_id_gate_applies: Final = ( - team_object is not None - and _team_metadata_requires_trace_id(team_object) - and request.method not in ("GET", "HEAD", "OPTIONS") - and (RouteChecks.is_llm_api_route(route=route) or pass_through_route) - and not RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes.trace_telemetry_routes.value, - ) + trace_request: Final = access_policy.decide( + access_policy.TraceRequest(team_object is not None and _team_metadata_requires_trace_id(team_object)) ) - team_metadata: Final[Mapping[str, object] | None] = ( - _TRACE_ID_METADATA_ADAPTER.validate_python( - team_object.metadata # pyright: ignore[reportUnknownMemberType] # validate the narrowed untyped field + team_trace: Final = ( + access_policy.decide( + access_policy.TraceGate( + required=True, + method=request.method, + llm_route=RouteChecks.is_llm_api_route(route=route), + passthrough=pass_through_route, + telemetry=RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.trace_telemetry_routes.value + ), + ) ) - if team_trace_id_gate_applies and team_object is not None - else None + if trace_request.next_trace is not None + else trace_request ) if ( - team_trace_id_gate_applies + team_trace.next_trace is not None and team_object is not None - and team_metadata is not None - and not _request_has_trace_id( - request_body=_typed_request_body(request_body), - request=request, - headers_only=headers_only, - ) + and not _request_has_trace_id(_typed_request_body(request_body), request, headers_only=headers_only) ): raise ProxyException( message=( @@ -1342,7 +1127,8 @@ async def common_checks( ) ## 2.1 If user can call model (if personal key) - if _model and team_object is None and user_object is not None: + if flow.user_models: + assert _model is not None and user_object is not None with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"): await can_user_call_model( model=_model, @@ -1351,7 +1137,15 @@ async def common_checks( key_model_aliases=key_model_aliases_for_auth_check(valid_token), ) - if end_user_object is not None and end_user_object.models: + customer_plan: Final = common_flow.plan( + replace( + flow_facts, + customer_present=end_user_object is not None, + customer_models=tuple(end_user_object.models or ()) if end_user_object is not None else (), + ) + ) + if customer_plan.customer_models: + assert end_user_object is not None with tracer.trace("litellm.proxy.auth.common_checks.can_customer_call_model"): if _model: can_customer_access_model( @@ -1427,106 +1221,96 @@ async def common_checks( LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body) async def _user_max_budget_check() -> None: - # 4.1 personal budget - if user_object is None or user_object.max_budget is None: - return - is_team_key: Final = team_object is not None and team_object.team_id is not None - if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True: - return + from litellm.proxy.auth import common_budget_adapter - from litellm.proxy.proxy_server import get_current_spend - - user_budget: Final = user_object.max_budget - user_spend: Final = await get_current_spend( - counter_key=user_counter(user_object.user_id), - fallback_spend=user_object.spend or 0.0, - max_budget=user_budget, + return await common_budget_adapter.user_max_budget_check( + user_object, team_object, general_settings, proxy_logging_obj ) - call_info: Final = CallInfo( - spend=user_spend, - max_budget=user_budget, - user_id=user_object.user_id, - user_email=user_object.user_email, - event_group=Litellm_EntityType.USER, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="user_budget", - user_info=call_info, - ) - ) - if math.isfinite(user_budget) and user_spend >= user_budget: - raise litellm.BudgetExceededError( - current_cost=user_spend, - max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", - entity_type=Litellm_EntityType.USER.value, - entity_id=user_object.user_id, - ) # Each scope reads a distinct counter key with no cross-scope ordering # dependency, so the per-scope Redis-first reads run concurrently instead # of one sequential await per scope. return_exceptions lets every scope # settle, then the first error in scope-priority order propagates exactly # as the sequential path raised. - budget_check_coros: Final = tuple( - coro - for coro in ( - team_max_budget_check( - team_object=team_object, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ), - _team_multi_budget_check(team_object=team_object), - _virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None, - _team_soft_budget_check( - team_object=team_object, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ), - organization_max_budget_check( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - _tag_max_budget_check( - request_body=request_body, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ), - _model_access_group_max_budget_check( - matched_model_access_groups=matched_model_access_groups, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - if matched_model_access_groups - else None, - _user_max_budget_check(), - _check_team_member_budget( - team_object=team_object, - user_object=user_object, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=loaded_team_membership, - team_membership_loaded=team_membership_loaded, - ), - check_end_user_budget(end_user_obj=end_user_object, route=route) - if end_user_object is not None and end_user_object.litellm_budget_table is not None - else None, + async def _check_common_budget(scope: common_flow.BudgetScope) -> None: + match scope: + case "Team": + await team_max_budget_check( + team_object=team_object, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token + ) + case "TeamWindows": + await _team_multi_budget_check(team_object=team_object) + case "KeyWindows": + assert valid_token is not None + await _virtual_key_multi_budget_check(valid_token=valid_token) + case "TeamSoft": + await _team_soft_budget_check( + team_object=team_object, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token + ) + case "Organization": + await organization_max_budget_check( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + case "Tags": + await _tag_max_budget_check( + request_body=request_body, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ) + case "Groups": + await _model_access_group_max_budget_check( + matched_model_access_groups=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + case "User": + await _user_max_budget_check() + case "Member": + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, + ) + case "Customer": + assert end_user_object is not None + await check_end_user_budget(end_user_obj=end_user_object, route=route) + case _: + assert_never(scope) + + budget_scopes: Final = common_flow.budgets( + common_flow.BudgetFacts( + token_present=valid_token is not None, + group_count=len(matched_model_access_groups), + customer_present=end_user_object is not None, + customer_budget_present=end_user_object is not None + and end_user_object.litellm_budget_table is not None, ) - if coro is not None ) + budget_check_coros: Final = tuple(_check_common_budget(scope) for scope in budget_scopes) with tracer.trace("litellm.proxy.auth.common_checks.budget_checks"): budget_results: Final = await asyncio.gather(*budget_check_coros, return_exceptions=True) - budget_error: Final = next((r for r in budget_results if isinstance(r, BaseException)), None) - if budget_error is not None: + failure: Final = access_policy.decide( + access_policy.ConcurrentFailures( + tuple(isinstance(result, BaseException) for result in budget_results), + tuple(isinstance(result, NativeAuthFailure) for result in budget_results), + ) + ).failure_index + if failure is not None: + budget_error: Final = budget_results[failure] + if not isinstance(budget_error, BaseException): + raise NativeAuthFailure("native selected an absent budget failure") raise budget_error _enforce_user_param_check(general_settings, request, request_body, route) @@ -1587,18 +1371,17 @@ async def common_checks( def effective_user_role(user_role: str | None) -> LitellmUserRoles: - try: - return LitellmUserRoles(user_role) - except ValueError: - return LitellmUserRoles.INTERNAL_USER + from litellm.proxy.auth import route_role_adapter + + return route_role_adapter.effective_user_role(user_role) def get_user_role( user_obj: LiteLLM_UserTable | None, ) -> LitellmUserRoles | None: - if user_obj is None: - return None - return effective_user_role(user_obj.user_role) + from litellm.proxy.auth import route_role_adapter + + return route_role_adapter.get_user_role(user_obj) _get_user_role: Final = get_user_role @@ -1611,59 +1394,24 @@ def _is_api_route_allowed( valid_token: UserAPIKeyAuth | None, user_obj: LiteLLM_UserTable | None = None, ) -> bool: - """ - - Route b/w api token check and normal token check - """ - _user_role: Final = get_user_role(user_obj=user_obj) + from litellm.proxy.auth import route_role_adapter - if valid_token is None: - raise Exception("Invalid proxy server token passed. valid_token=None.") - - if not is_user_proxy_admin(user_obj=user_obj): # if non-admin - RouteChecks.non_proxy_admin_allowed_routes_check( - user_obj=user_obj, - _user_role=_user_role, - route=route, - request=request, - request_data=request_data, - valid_token=valid_token, - ) - return True + return route_role_adapter.is_api_route_allowed(route, request, request_data, valid_token, user_obj) def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: - if user_obj is None: - return False + from litellm.proxy.auth import route_role_adapter - if user_obj.user_role is not None and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return True - - return False + return route_role_adapter.is_user_proxy_admin(user_obj) _is_user_proxy_admin: Final = is_user_proxy_admin -def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: - """ - Return if a user is allowed to access route. Helper function for `allowed_routes_check`. +def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint + from litellm.proxy.auth import route_role_adapter - Parameters: - - user_route: str - the route the user is trying to call - - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name - (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*"). - """ - from starlette.routing import compile_path - - for allowed_route in allowed_routes: - if allowed_route in LiteLLMRoutes.__members__: - for template in LiteLLMRoutes[allowed_route].value: - regex, _, _ = compile_path(template) - if regex.match(user_route): - return True - elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route): - return True - return False + return route_role_adapter.allowed_routes_list_check(user_route, allowed_routes) def allowed_routes_check( @@ -1671,62 +1419,24 @@ def allowed_routes_check( user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth, ) -> bool: - """ - Check if user -> not admin - allowed to access these routes - """ + from litellm.proxy.auth import route_role_adapter - if user_role == LitellmUserRoles.PROXY_ADMIN: - is_allowed = _allowed_routes_check( - user_route=user_route, - allowed_routes=litellm_proxy_roles.admin_allowed_routes, - ) - return is_allowed - - elif user_role == LitellmUserRoles.TEAM: - if litellm_proxy_roles.team_allowed_routes is None: - """ - By default allow a team to call openai + info routes - """ - is_allowed = _allowed_routes_check(user_route=user_route, allowed_routes=["openai_routes", "info_routes"]) - return is_allowed - elif litellm_proxy_roles.team_allowed_routes is not None: - is_allowed = _allowed_routes_check( - user_route=user_route, - allowed_routes=litellm_proxy_roles.team_allowed_routes, - ) - return is_allowed - return False + return route_role_adapter.allowed_routes_check(user_role, user_route, litellm_proxy_roles) def allowed_route_check_inside_route( user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None, ) -> bool: - ret_val = True - if ( - user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ): - ret_val = False - if requested_user_id is not None and user_api_key_dict.user_id is not None: - if user_api_key_dict.user_id == requested_user_id: - ret_val = True - return ret_val + from litellm.proxy.auth import route_role_adapter + + return route_role_adapter.allowed_route_check_inside_route(user_api_key_dict, requested_user_id) def get_actual_routes(allowed_routes: list) -> list: - actual_routes: Final[list] = [] - for route_name in allowed_routes: - try: - route_value = LiteLLMRoutes[route_name].value - if isinstance(route_value, set): - actual_routes.extend(list(route_value)) - else: - actual_routes.extend(route_value) + from litellm.proxy.auth import route_role_adapter - except KeyError: - actual_routes.append(route_name) - return actual_routes + return route_role_adapter.get_actual_routes(allowed_routes) KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" @@ -1735,11 +1445,9 @@ _OBJECTS_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[obj def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: - """The default budget a key assigns to end users that carry no budget of their own.""" - if key_metadata is None: - return None - budget_id: Final = key_metadata.get(KEY_END_USER_BUDGET_ID_METADATA_FIELD) - return budget_id if isinstance(budget_id, str) and budget_id != "" else None + from litellm.proxy.auth import end_user_context_adapter + + return end_user_context_adapter.get_key_end_user_budget_id(key_metadata) @with_service_target(AUTH_OBJECTS_TARGET) @@ -1749,63 +1457,11 @@ async def get_default_end_user_budget( parent_otel_span: Span | None = None, budget_id: str | None = None, ) -> LiteLLM_BudgetTable | None: - """ - Fetches the default end user budget from the database. + from litellm.proxy.auth import end_user_context_adapter - ``budget_id`` selects the budget row; when omitted the proxy-wide - ``litellm.max_end_user_budget_id`` is used. This budget is applied to end - users who don't have an explicit budget_id set. Results are cached for performance. - - Args: - prisma_client: Database client instance - user_api_key_cache: Cache for storing/retrieving budget data - parent_otel_span: Optional OpenTelemetry span for tracing - budget_id: Budget row to load instead of the proxy-wide default - - Returns: - LiteLLM_BudgetTable if configured and found, None otherwise - """ - default_budget_id: Final = budget_id if budget_id is not None else litellm.max_end_user_budget_id - if prisma_client is None or default_budget_id is None: - return None - - cache_key: Final = f"default_end_user_budget:{default_budget_id}" - - # Check cache first - cached_budget: Final = await user_api_key_cache.async_get_cache( - key=cache_key, - model_type=LiteLLM_BudgetTable, + return await end_user_context_adapter.get_default_end_user_budget( + prisma_client, user_api_key_cache, parent_otel_span, budget_id ) - if cached_budget is not None: - return cached_budget - - # Fetch from database - try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( - where={"budget_id": default_budget_id} - ) - - if budget_record is None: - verbose_proxy_logger.warning( - "Default end user budget not found in database: %s", - default_budget_id.replace("\r", "").replace("\n", ""), - ) - return None - - _budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) - # Cache the budget for 60 seconds - await user_api_key_cache.async_set_cache( - key=cache_key, - value=_budget_obj, - model_type=LiteLLM_BudgetTable, - ttl=get_management_object_ttl(user_api_key_cache), - ) - - return _budget_obj - - except Exception as e: - verbose_proxy_logger.error("Error fetching default end user budget: %s", e) - return None @log_db_metrics @@ -1831,6 +1487,9 @@ async def get_team_member_default_budget( Returns: LiteLLM_BudgetTable if found, None otherwise """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: return None @@ -1844,9 +1503,16 @@ async def get_team_member_default_budget( return cached_budget try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( - where={"budget_id": budget_id} + native: Final = native_auth_repository(prisma_client) + budget_record: Final = ( + await native.one("budget", [budget_id], LiteLLM_BudgetTable) + if native is not None + else await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( + where={"budget_id": budget_id} + ) ) + except NativeAuthFailure: + raise except Exception: verbose_proxy_logger.exception("Error fetching team-default member budget %s", budget_id) if raise_on_lookup_error: @@ -1857,8 +1523,11 @@ async def get_team_member_default_budget( verbose_proxy_logger.warning("Team-default member budget not found in database: %s", budget_id) return None - budget: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) + budget: Final = LiteLLM_BudgetTable.model_validate( + budget_record.model_dump() if isinstance(budget_record, BaseModel) else budget_record.dict() + ) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=budget, model_type=LiteLLM_BudgetTable, @@ -1873,13 +1542,9 @@ def resolve_effective_team_member_budget( team_default_budget: LiteLLM_BudgetTable | None, now: datetime, ) -> float | None: - if member_budget is not None and member_budget.max_budget is not None: - return member_budget.effective_max_budget(now) - if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: - return None - return team_default_budget.max_budget + ( - member_budget.active_temp_budget_increase(now) if member_budget is not None else 0.0 - ) + from litellm.proxy.auth import common_budget_adapter + + return common_budget_adapter.resolve_effective_team_member_budget(member_budget, team_default_budget, now) async def resolve_default_end_user_budget( @@ -1888,30 +1553,10 @@ async def resolve_default_end_user_budget( key_end_user_budget_id: str | None, parent_otel_span: Span | None = None, ) -> LiteLLM_BudgetTable | None: - """ - The default budget for an end user with no budget of its own. + from litellm.proxy.auth import budget_context_adapter - The key's ``end_user_budget_id`` takes precedence over the proxy-wide - ``litellm.max_end_user_budget_id``; the proxy-wide default is the fallback when the key - names no budget or its budget row is missing. - """ - if key_end_user_budget_id is not None: - key_budget: Final = await get_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - budget_id=key_end_user_budget_id, - ) - if key_budget is not None: - return key_budget - - if litellm.max_end_user_budget_id is None: - return None - - return await get_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, + return await budget_context_adapter.resolve_default_end_user_budget( + prisma_client, user_api_key_cache, key_end_user_budget_id, parent_otel_span ) @@ -1922,83 +1567,20 @@ async def _apply_default_budget_to_end_user( parent_otel_span: Span | None = None, key_end_user_budget_id: str | None = None, ) -> LiteLLM_EndUserTable: - """ - Returns the end user with the resolved default budget when it has no budget of its own. + from litellm.proxy.auth import end_user_context_adapter - A row whose own ``budget_id`` resolved to a budget is returned unchanged. Otherwise the - default is resolved on every call and set on a copy: the cached row carries at most the - proxy-wide default (readers such as the Prometheus customer gauges rely on that), never a - key's, so requests through keys with different defaults never observe each other's budget. - - Args: - end_user_obj: The end user object to potentially apply default budget to - prisma_client: Database client instance - user_api_key_cache: Cache for storing/retrieving data - parent_otel_span: Optional OpenTelemetry span for tracing - key_end_user_budget_id: The requesting key's ``end_user_budget_id``, if any - """ - if end_user_obj.budget_id is not None and end_user_obj.litellm_budget_table is not None: - return end_user_obj - - if key_end_user_budget_id is None and litellm.max_end_user_budget_id is None: - return end_user_obj - - default_budget: Final = await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, + return await end_user_context_adapter.apply_default_budget_to_end_user( + end_user_obj, prisma_client, user_api_key_cache, parent_otel_span, key_end_user_budget_id ) - if default_budget is None: - return end_user_obj - - verbose_proxy_logger.debug( - "Applied default budget %s to end user %s", default_budget.budget_id, end_user_obj.user_id - ) - return end_user_obj.model_copy(update=MappingProxyType({"litellm_budget_table": default_budget})) - async def check_end_user_budget( end_user_obj: LiteLLM_EndUserTable, route: str, ) -> None: - """ - Check if end user is within their budget limit. + from litellm.proxy.auth import common_budget_adapter - Args: - end_user_obj: The end user object to check - route: The request route - - Raises: - litellm.BudgetExceededError: If end user has exceeded their budget - """ - if RouteChecks.is_info_route(route): - return - - if end_user_obj.litellm_budget_table is None: - return - - end_user_budget: Final = end_user_obj.litellm_budget_table.max_budget - if end_user_budget is None: - return - - from litellm.proxy.proxy_server import get_current_spend - - end_user_spend: Final = await get_current_spend( - counter_key=end_user_counter(end_user_obj.user_id), - fallback_spend=end_user_obj.spend or 0.0, - max_budget=end_user_budget, - fallback_authoritative=True, - ) - if end_user_spend > end_user_budget: - raise litellm.BudgetExceededError( - current_cost=end_user_spend, - max_budget=end_user_budget, - message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", - entity_type=Litellm_EntityType.END_USER.value, - entity_id=end_user_obj.user_id, - ) + return await common_budget_adapter.check_end_user_budget(end_user_obj, route) _check_end_user_budget: Final = check_end_user_budget @@ -2059,10 +1641,13 @@ async def _cache_registry_answer( value: tuple[str, ...] | str, ttl: float, user_api_key_cache: UserApiKeyCache, + fill_generation: int | None = None, ) -> None: """Best-effort: a cache backend failure must not turn a registry load into a failed request.""" try: - await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl) + await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl, fill_generation=fill_generation) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort cache write: auth must survive a cache backend error verbose_proxy_logger.warning("Failed to cache registry %s: %s", cache_key, e) @@ -2073,10 +1658,14 @@ async def _fetch_and_cache_registry( max_size: int, fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], user_api_key_cache: UserApiKeyCache, + native: bool = False, ) -> frozenset[str] | None: """The registry as the database has it, cached whole, or ``None`` when it is unusable.""" + fill_generation: Final = user_api_key_cache.current_generation() if native else None try: registry_ids: Final = await fetch_ids() + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # fail-safe: any registry load error must degrade to per-id lookups, never break auth verbose_proxy_logger.warning( "Registry %s could not be loaded from the database, so per-id lookups will run and the " @@ -2086,6 +1675,7 @@ async def _fetch_and_cache_registry( e, ) await _cache_registry_answer( + fill_generation=fill_generation, cache_key=cache_key, value=overflow_sentinel, ttl=REGISTRY_ERROR_NEGATIVE_CACHE_TTL, @@ -2095,6 +1685,7 @@ async def _fetch_and_cache_registry( if len(registry_ids) > max_size: await _cache_registry_answer( + fill_generation=fill_generation, cache_key=cache_key, value=overflow_sentinel, ttl=get_management_object_ttl(user_api_key_cache), @@ -2103,6 +1694,7 @@ async def _fetch_and_cache_registry( return None await _cache_registry_answer( + fill_generation=fill_generation, cache_key=cache_key, value=registry_ids, ttl=get_management_object_ttl(user_api_key_cache), @@ -2118,6 +1710,7 @@ async def _load_bounded_registry( load_lock: asyncio.Lock, fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], user_api_key_cache: UserApiKeyCache, + native: bool = False, ) -> frozenset[str] | None: """ A bounded id set under one cache key, so an id outside it costs no DB read. @@ -2137,6 +1730,7 @@ async def _load_bounded_registry( return cached_after_wait return await _fetch_and_cache_registry( + native=native, cache_key=cache_key, overflow_sentinel=overflow_sentinel, max_size=max_size, @@ -2145,20 +1739,28 @@ async def _load_bounded_registry( ) -async def _load_end_user_restricted_registry( +async def load_end_user_restricted_registry( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, ) -> frozenset[str] | None: """The set of end-user ids whose ``LiteLLM_EndUserTable`` row carries a restriction.""" async def fetch_ids() -> tuple[str, ...]: - restricted_rows: Final = await _end_user_table(EndUserRepository(prisma_client)).find_many( - where=_restricted_end_user_where(), - take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1, + native: Final = native_auth_repository(prisma_client) + restricted_rows: Final = ( + await native.many( + "end_user_restricted_registry", [str(END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1)], EndUserRegistryRow + ) + if native is not None + else await _end_user_table(EndUserRepository(prisma_client)).find_many( + where=_restricted_end_user_where(), + take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1, + ) ) return tuple(row.user_id for row in restricted_rows) return await _load_bounded_registry( + native=native_auth_repository(prisma_client) is not None, cache_key=end_user_restricted_registry_cache_key(), overflow_sentinel=END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, max_size=END_USER_RESTRICTED_REGISTRY_MAX_SIZE, @@ -2168,6 +1770,9 @@ async def _load_end_user_restricted_registry( ) +_load_end_user_restricted_registry: Final = load_end_user_restricted_registry + + async def _end_user_is_known_unrestricted( end_user_id: str, prisma_client: PrismaClient, @@ -2175,30 +1780,11 @@ async def _end_user_is_known_unrestricted( token_end_user_max_budget: float | None, key_end_user_budget_id: str | None = None, ) -> bool: - """ - True when the cached registry proves the id restricts nothing, so its row need not be read. + from litellm.proxy.auth import end_user_context_adapter - Every field ``get_end_user_object`` callers consume (budget, spend under that budget, region, - default model, models, object permission, blocked) is part of the registry predicate, so an id - outside it is indistinguishable from one with no row at all. The skip is off whenever mere existence of - the row is meaningful: ``max_end_user_budget_id`` or the key's ``end_user_budget_id`` grafts a - default budget onto any row that exists, ``validate_end_user_id_in_db`` rejects ids that resolve - to no row, and a token-supplied ``end_user_max_budget`` (a ``user_custom_auth`` callable can set - one against an otherwise unrestricted row) is enforced against the row's recorded spend. - """ - if ( - litellm.max_end_user_budget_id is not None - or key_end_user_budget_id is not None - or litellm.validate_end_user_id_in_db - or token_end_user_max_budget is not None - ): - return False - - registry: Final = await _load_end_user_restricted_registry( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + return await end_user_context_adapter.end_user_is_known_unrestricted( + end_user_id, prisma_client, user_api_key_cache, token_end_user_max_budget, key_end_user_budget_id ) - return registry is not None and end_user_id not in registry @log_db_metrics @@ -2234,6 +1820,9 @@ async def get_end_user_object( Returns: LiteLLM_EndUserTable if found, None otherwise """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No db connected") @@ -2267,22 +1856,30 @@ async def get_end_user_object( # Fetch from database try: - response: Final = await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique( - where={"user_id": end_user_id}, - include={"litellm_budget_table": True, "object_permission": True}, + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("end_user", [end_user_id], LiteLLM_EndUserTable) + if native is not None + else await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique( + where={"user_id": end_user_id}, + include={"litellm_budget_table": True, "object_permission": True}, + ) ) if response is None: raise Exception end_user_row: Final = await _apply_default_budget_to_end_user( - end_user_obj=LiteLLM_EndUserTable.model_validate(response.dict()), + end_user_obj=LiteLLM_EndUserTable.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ), prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=_key, value=end_user_row, model_type=LiteLLM_EndUserTable, @@ -2300,6 +1897,8 @@ async def get_end_user_object( key_end_user_budget_id=key_end_user_budget_id, ) + except NativeAuthFailure: + raise except Exception: return None @@ -2318,61 +1917,20 @@ async def resolve_and_validate_end_user_id( route: str = "", key_end_user_budget_id: str | None = None, ) -> str | None: - """Optionally drop end-user ids that don't resolve to a known DB row. + from litellm.proxy.auth import end_user_context_adapter - Default: pass-through. LiteLLM's documented pattern is that the `user` - field is an arbitrary caller-supplied identifier, so validation is - opt-in behind ``litellm.validate_end_user_id_in_db`` to preserve - backwards compatibility. - - When the flag is set: accept the id when it matches any of - - LiteLLM_EndUserTable.user_id - - LiteLLM_UserTable.user_id - - LiteLLM_UserTable.user_email (case-insensitive) - - If the id doesn't match but a default end-user budget is configured - (``litellm.max_end_user_budget_id`` or the key's ``end_user_budget_id``), - we still preserve the id so that budget is applied downstream; otherwise - we return None. - - DB lookups reuse ``get_end_user_object`` / ``get_user_object`` so they - share the same cache as the rest of the auth path instead of adding new - raw Prisma queries. - """ - if raw_end_user_id is None: - return None - if not litellm.validate_end_user_id_in_db: - return raw_end_user_id - if prisma_client is None: - return raw_end_user_id - - has_default_budget: Final = bool(litellm.max_end_user_budget_id) or key_end_user_budget_id is not None - cache_key: Final = f"end_user_validation:{raw_end_user_id}" - cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) - if cached == "valid": - return raw_end_user_id - if cached == "invalid": - return raw_end_user_id if has_default_budget else None - - is_valid: Final = await _end_user_id_exists_in_db( - end_user_id=raw_end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, + return await end_user_context_adapter.resolve_and_validate_end_user_id( + raw_end_user_id, + prisma_client, + user_api_key_cache, + parent_otel_span, + proxy_logging_obj, + route, + key_end_user_budget_id, ) - await user_api_key_cache.async_set_cache( - key=cache_key, - value="valid" if is_valid else "invalid", - ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL), - ) - return raw_end_user_id if is_valid or has_default_budget else None - - -async def _end_user_id_exists_in_db( +async def end_user_id_exists_in_db( end_user_id: str, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, @@ -2380,38 +1938,14 @@ async def _end_user_id_exists_in_db( proxy_logging_obj: ProxyLogging | None = None, route: str = "", ) -> bool: - """True when the id matches an EndUser, User, or user_email row.""" - try: - end_user_obj: Final = await get_end_user_object( - end_user_id=end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ) - if end_user_obj is not None: - return True - except Exception as e: - verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", e) + from litellm.proxy.auth import end_user_context_adapter - try: - user_obj: Final = await get_user_object( - user_id=end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_db_only=False, - user_email=end_user_id if "@" in end_user_id else None, - ) - if user_obj is not None: - return True - except Exception as e: - verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", e) + return await end_user_context_adapter.end_user_id_exists_in_db( + end_user_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route + ) - return False + +_end_user_id_exists_in_db: Final = end_user_id_exists_in_db async def _load_tag_registry( @@ -2421,12 +1955,18 @@ async def _load_tag_registry( """The set of tag names that have a row in ``LiteLLM_TagTable``.""" async def fetch_ids() -> tuple[str, ...]: - registry_rows: Final = await _tag_table(TagRepository(prisma_client)).find_many( - take=TAG_REGISTRY_MAX_SIZE + 1, + native: Final = native_auth_repository(prisma_client) + registry_rows: Final = ( + await native.many("tag_registry", [str(TAG_REGISTRY_MAX_SIZE + 1)], TagRegistryRow) + if native is not None + else await _tag_table(TagRepository(prisma_client)).find_many( + take=TAG_REGISTRY_MAX_SIZE + 1, + ) ) return tuple(row.tag_name for row in registry_rows) return await _load_bounded_registry( + native=native_auth_repository(prisma_client) is not None, cache_key=tag_registry_cache_key(), overflow_sentinel=TAG_REGISTRY_OVERFLOW_SENTINEL, max_size=TAG_REGISTRY_MAX_SIZE, @@ -2436,19 +1976,27 @@ async def _load_tag_registry( ) -async def _load_model_access_group_registry( +async def load_model_access_group_registry( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, ) -> frozenset[str] | None: """The set of model access group names that have a row in ``LiteLLM_ModelAccessGroupBudgetTable``.""" async def fetch_ids() -> tuple[str, ...]: - registry_rows: Final = await _model_access_group_budget_table( - ModelAccessGroupBudgetRepository(prisma_client) - ).find_many(take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1) + native: Final = native_auth_repository(prisma_client) + registry_rows: Final = ( + await native.many( + "model_access_group_registry", [str(MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1)], ModelAccessGroupRow + ) + if native is not None + else await _model_access_group_budget_table(ModelAccessGroupBudgetRepository(prisma_client)).find_many( + take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1 + ) + ) return tuple(row.access_group_name for row in registry_rows) return await _load_bounded_registry( + native=native_auth_repository(prisma_client) is not None, cache_key=model_access_group_registry_cache_key(), overflow_sentinel=MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, max_size=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, @@ -2458,6 +2006,9 @@ async def _load_model_access_group_registry( ) +_load_model_access_group_registry: Final = load_model_access_group_registry + + @with_service_target(AUTH_OBJECTS_TARGET) async def _fetch_uncached_model_access_group_budgets( uncached_groups: Sequence[str], @@ -2470,24 +2021,33 @@ async def _fetch_uncached_model_access_group_budgets( ``matched_model_access_groups``, which :func:`collect_matched_model_access_groups` already intersected with the registry, so a name that has no row cannot reach this. """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if not uncached_groups: return () try: - db_rows: Final = await _model_access_group_budget_table( - ModelAccessGroupBudgetRepository(prisma_client) - ).find_many( - where={"access_group_name": {"in": list(uncached_groups)}}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(prisma_client) + db_rows: Final = ( + await native.many("model_access_group", uncached_groups, ModelAccessGroupRow) + if native is not None + else await _model_access_group_budget_table(ModelAccessGroupBudgetRepository(prisma_client)).find_many( + where={"access_group_name": {"in": list(uncached_groups)}}, + include={"litellm_budget_table": True}, + ) ) fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows) for fetched_name, fetched_obj in fetched: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=model_access_group_cache_key(fetched_name), value=fetched_obj, model_type=ModelAccessGroupBudget, ttl=get_management_object_ttl(user_api_key_cache), ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # fail-safe: a budget fetch error must yield "no budget rows", never break auth verbose_proxy_logger.debug("Error batch fetching model access group budgets from database: %s", e) return () @@ -2544,6 +2104,9 @@ async def _fetch_uncached_tags( user_api_key_cache: UserApiKeyCache, ) -> tuple[tuple[str, LiteLLM_TagTable], ...] | Exception: """Rows for the tags a cache probe missed; names absent from the registry never reach the DB.""" + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if not uncached_tags: return () @@ -2558,18 +2121,34 @@ async def _fetch_uncached_tags( return () try: - db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many( - where={"tag_name": {"in": list(tags_to_fetch)}}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(prisma_client) + db_tags: Final = ( + await native.many("tag", tags_to_fetch, LiteLLM_TagTable) + if native is not None + else await _tag_table(TagRepository(prisma_client)).find_many( + where={"tag_name": {"in": list(tags_to_fetch)}}, + include={"litellm_budget_table": True}, + ) + ) + fetched: Final = tuple( + ( + db_tag.tag_name, + LiteLLM_TagTable.model_validate( + db_tag.model_dump() if isinstance(db_tag, BaseModel) else db_tag.dict() + ), + ) + for db_tag in db_tags ) - fetched: Final = tuple((db_tag.tag_name, LiteLLM_TagTable.model_validate(db_tag.dict())) for db_tag in db_tags) for fetched_name, fetched_obj in fetched: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=tag_cache_key(fetched_name), value=fetched_obj, model_type=LiteLLM_TagTable, ttl=get_management_object_ttl(user_api_key_cache), ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # returned as a value: budget checks fail open, the baggage mapper fails closed return e else: @@ -2711,23 +2290,39 @@ async def _fetch_team_membership_from_db( *, use_writer: bool = False, ) -> LiteLLM_TeamMembership | None: - _ = parent_otel_span, proxy_logging_obj - response: Final = await _dictable_table( - TeamMembershipRepository(prisma_client, use_writer=use_writer), "team_membership" - ).find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - include={"litellm_budget_table": True}, + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) + _ = parent_otel_span, proxy_logging_obj + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("member", [user_id, team_id], LiteLLM_TeamMembership, writer=use_writer) + if native is not None + else await _dictable_table( + TeamMembershipRepository(prisma_client, use_writer=use_writer), "team_membership" + ).find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + include={"litellm_budget_table": True}, + ) + ) + membership: Final = ( + None + if response is None + else LiteLLM_TeamMembership.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ) ) - membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) if membership is None: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=_key, value=NO_TEAM_MEMBERSHIP_SENTINEL, ttl=get_management_object_ttl(user_api_key_cache), ) else: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=_key, value=membership, model_type=LiteLLM_TeamMembership, @@ -2828,31 +2423,9 @@ async def get_team_membership( def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool: - from collections import defaultdict + from litellm.proxy.auth import model_helper_adapter - if team_models is None: - return True - if model in team_models: - return True - - access_groups: dict[str, list[str]] = defaultdict(list) - if llm_router: - access_groups = llm_router.registry().access_groups(model_name=model) - - if len(access_groups) > 0: # check if token contains any model access groups - for idx, m in enumerate( - team_models - ): # loop token models, if any of them are an access group add the access group - if m in access_groups: - return True - - # Filter out models that are access_groups - filtered_models: Final = [m for m in team_models if m not in access_groups] - - if model in filtered_models: - return True - - return False + return model_helper_adapter.model_in_access_group(model, team_models, llm_router) def _should_check_db(key: str, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int) -> bool: @@ -2952,18 +2525,27 @@ async def _get_fuzzy_user_object( """ response = None + native: Final = native_auth_repository(prisma_client) if sso_user_id is not None: - response = await _user_table(UserRepository(prisma_client)).find_unique( - where={"sso_user_id": sso_user_id}, - include={"organization_memberships": True}, + response = ( + await native.one("user_sso", [sso_user_id], LiteLLM_UserTable) + if native is not None + else await _user_table(UserRepository(prisma_client)).find_unique( + where={"sso_user_id": sso_user_id}, + include={"organization_memberships": True}, + ) ) if response is None and user_email is not None: # Use case-insensitive query to handle emails with different casing # This matches the pattern used in _check_duplicate_user_email - response = await _user_table(UserRepository(prisma_client)).find_first( - where={"user_email": {"equals": user_email, "mode": "insensitive"}}, - include={"organization_memberships": True}, + response = ( + await native.one("user_email", [user_email], LiteLLM_UserTable) + if native is not None + else await _user_table(UserRepository(prisma_client)).find_first( + where={"user_email": {"equals": user_email, "mode": "insensitive"}}, + include={"organization_memberships": True}, + ) ) if response is not None and sso_user_id is not None: # update sso_user_id @@ -2984,6 +2566,9 @@ async def _backfill_null_user_email( user_row: LiteLLM_UserTable, user_email: str | None, ) -> LiteLLM_UserTable: + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if user_email is None or user_row.user_email is not None or prisma_client is None: return user_row @@ -2992,12 +2577,18 @@ async def _backfill_null_user_email( user_id=user_row.user_id, user_email=user_email, ) - db_row: Final = await user_repo.find_by_id(user_row.user_id) + native: Final = native_auth_repository(prisma_client) + db_row: Final = ( + await native.one("user", [user_row.user_id], LiteLLM_UserTable, writer=True) + if native is not None + else await user_repo.find_by_id(user_row.user_id) + ) if db_row is None: return user_row email_update: Final = {"user_email": db_row.user_email} updated_row: Final = user_row.model_copy(update=email_update) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=user_row.user_id, value=updated_row, model_type=LiteLLM_UserTable, @@ -3025,6 +2616,9 @@ async def get_user_object( - if not, then raise an error """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if user_id is None: return None @@ -3053,8 +2647,13 @@ async def get_user_object( ) if should_check_db: - response = await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).find_unique( - where={"user_id": user_id}, include={"organization_memberships": True} + native: Final = native_auth_repository(prisma_client) + response = ( + await native.one("user", [user_id], LiteLLM_UserTable, writer=bool(check_db_only)) + if native is not None + else await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).find_unique( + where={"user_id": user_id}, include={"organization_memberships": True} + ) ) if response is None: @@ -3134,6 +2733,7 @@ async def get_user_object( # save the user object to cache await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=user_id, value=_response, model_type=LiteLLM_UserTable, @@ -3150,6 +2750,8 @@ async def get_user_object( return _response except UserNotFoundError: raise + except NativeAuthFailure: + raise except Exception as e: _log_budget_lookup_failure("user", e) raise _user_read_failure(user_id=user_id, error=e) @@ -3171,6 +2773,7 @@ async def _cache_management_object( proxy_logging_obj: ProxyLogging | None, *, model_type: type[BaseModel], + fill_generation: int | None = None, ): """ Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis). @@ -3178,6 +2781,7 @@ async def _cache_management_object( ``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned. """ await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=value, model_type=model_type, @@ -3191,6 +2795,8 @@ async def cache_team_object( team_table: LiteLLM_TeamTableCachedObj, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None, + *, + fill_generation: int | None = None, ) -> None: ## CACHE REFRESH TIME! team_table.last_refreshed_at = time.time() @@ -3205,6 +2811,7 @@ async def cache_team_object( # team_id is the table primary key — guaranteed unique, safe to write. await _cache_management_object( + fill_generation=fill_generation, key=key, value=team_table, user_api_key_cache=user_api_key_cache, @@ -3233,6 +2840,8 @@ async def cache_team_object( await user_api_key_cache.async_delete_cache(key=alias_key) else: await pipelined_delete + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the mutation verbose_proxy_logger.warning( "Failed to invalidate cached team alias entry %s; " @@ -3261,6 +2870,8 @@ async def _invalidate_usage_cache_entry( usage_cache.in_memory_cache.delete_cache(key) else: await usage_cache.async_delete_cache(key=key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the write verbose_proxy_logger.warning( "Failed to invalidate internal usage cache entry %s; a stale %s may be served until its TTL expires: %s", @@ -3360,6 +2971,8 @@ async def invalidate_team_member_spend_state( ) try: await spend_counter_store.delete((spend_counter_key,)) + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success verbose_proxy_logger.warning( "Failed to delete stale spend counter %s in Redis after a failed reset write", @@ -3425,6 +3038,8 @@ async def delete_cache_team_object( ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not abort the delete verbose_proxy_logger.warning( "Failed to invalidate cached team entry %s on delete; " @@ -3487,6 +3102,8 @@ async def delete_cache_key_object( ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort: a cache error must not fail a committed write verbose_proxy_logger.warning( "Failed to invalidate cached key entry %s; a stale key object may be served until its TTL expires: %s", @@ -3587,8 +3204,13 @@ class AccessGroupNotFoundError(HTTPException): async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None, *, use_writer: bool = False ) -> "_PrismaTeamRow | None": - response = await _team_table(TeamRepository(prisma_client, use_writer=use_writer)).find_unique( - where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS + native: Final = native_auth_repository(prisma_client) + response = ( + await native.one("team", [team_id], LiteLLM_TeamTable, writer=use_writer) + if native is not None + else await _team_table(TeamRepository(prisma_client, use_writer=use_writer)).find_unique( + where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS + ) ) if response is None and team_id_upsert: @@ -3608,10 +3230,6 @@ async def _get_team_db_check( return response -async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient) -> "_PrismaTeamRow | None": - return await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) - - async def _get_team_object_from_user_api_key_cache( team_id: str, prisma_client: PrismaClient, @@ -3623,6 +3241,9 @@ async def _get_team_object_from_user_api_key_cache( team_id_upsert: bool | None = None, use_writer: bool = False, ) -> LiteLLM_TeamTableCachedObj: + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) db_access_time_key: Final = key should_check_db: Final = _should_check_db( key=db_access_time_key, @@ -3656,6 +3277,8 @@ async def _get_team_object_from_user_api_key_cache( proxy_logging_obj=proxy_logging_obj, check_db_only=use_writer, ) + except NativeAuthFailure: + raise except Exception as e: if use_writer: raise @@ -3668,6 +3291,7 @@ async def _get_team_object_from_user_api_key_cache( # save the team object to cache await cache_team_object( + fill_generation=fill_generation, team_id=team_id, team_table=_response, user_api_key_cache=user_api_key_cache, @@ -3756,6 +3380,8 @@ async def get_team_object( ) except TeamNotFoundError: raise + except NativeAuthFailure: + raise except Exception: raise HTTPException( status_code=404, @@ -3769,9 +3395,12 @@ async def cache_access_object( access_group_table: LiteLLM_AccessGroupTable, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None = None, + *, + fill_generation: int | None = None, ) -> None: key: Final = f"access_group_id:{access_group_id}" await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=access_group_table, model_type=LiteLLM_AccessGroupTable, @@ -3819,6 +3448,9 @@ async def get_access_object( Raises: - HTTPException: If access group doesn't exist in db or cache (status_code=404) """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -3834,17 +3466,25 @@ async def get_access_object( # Not in cache - fetch from DB try: - response: Final = await _dictable_table( - AccessGroupRepository(prisma_client, use_writer=check_db_only), "access_group" - ).find_unique(where={"access_group_id": access_group_id}) + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("access_group", [access_group_id], LiteLLM_AccessGroupTable, writer=check_db_only) + if native is not None + else await _dictable_table( + AccessGroupRepository(prisma_client, use_writer=check_db_only), "access_group" + ).find_unique(where={"access_group_id": access_group_id}) + ) if response is None: raise AccessGroupNotFoundError(access_group_id=access_group_id) - _response: Final = LiteLLM_AccessGroupTable.model_validate(response.dict()) + _response: Final = LiteLLM_AccessGroupTable.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ) # Save to cache await cache_access_object( + fill_generation=fill_generation, access_group_id=access_group_id, access_group_table=_response, user_api_key_cache=user_api_key_cache, @@ -3854,6 +3494,8 @@ async def get_access_object( return _response except HTTPException: raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.exception( "Error getting access group for access_group_id: %s", @@ -3894,6 +3536,9 @@ async def get_team_object_by_alias( Raises: HTTPException: If team doesn't exist or multiple teams have the same alias """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -3911,8 +3556,13 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams: Final = await _team_table(TeamRepository(prisma_client)).find_many( - where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS + native: Final = native_auth_repository(prisma_client) + teams: Final = ( + await native.many("team_alias", [team_alias], LiteLLM_TeamTable) + if native is not None + else await _team_table(TeamRepository(prisma_client)).find_many( + where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS + ) ) if not teams: @@ -3944,6 +3594,8 @@ async def get_team_object_by_alias( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.debug( "Failed to load object_permission for team %s with object_permission_id=%s: %s", @@ -3954,6 +3606,7 @@ async def get_team_object_by_alias( # Cache the result by both alias and team_id await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=team_obj, model_type=LiteLLM_TeamTableCachedObj, @@ -3962,6 +3615,7 @@ async def get_team_object_by_alias( # Also cache by team_id for consistency team_id_cache_key: Final = f"team_id:{team_obj.team_id}" await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=team_id_cache_key, value=team_obj, model_type=LiteLLM_TeamTableCachedObj, @@ -3972,6 +3626,8 @@ async def get_team_object_by_alias( except HTTPException: raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias) raise HTTPException( @@ -4005,6 +3661,9 @@ async def get_org_object_by_alias( Raises: HTTPException: If organization not found or multiple orgs have the same alias """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -4019,8 +3678,13 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many( - where={"organization_alias": org_alias} + native: Final = native_auth_repository(prisma_client) + orgs = ( + await native.many("org_alias", [org_alias], LiteLLM_OrganizationTable) + if native is not None + else await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many( + where={"organization_alias": org_alias} + ) ) if not orgs: @@ -4044,6 +3708,7 @@ async def get_org_object_by_alias( # Cache the result await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=org_obj, model_type=LiteLLM_OrganizationTable, @@ -4051,6 +3716,7 @@ async def get_org_object_by_alias( ) # Also cache by org_id for consistency await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=f"org_id:{org_obj.organization_id}", value=org_obj, model_type=LiteLLM_OrganizationTable, @@ -4061,6 +3727,8 @@ async def get_org_object_by_alias( except HTTPException: raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.exception("Error looking up organization by alias: %s", org_alias) raise HTTPException( @@ -4195,6 +3863,8 @@ class ExperimentalUIJWTToken: return None try: return UserAPIKeyAuth.model_validate(json.loads(decrypted_token)) + except NativeAuthFailure: + raise except Exception as e: raise Exception(f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}") @@ -4238,6 +3908,9 @@ async def _fetch_key_object_from_db_unbounded( *, check_db_only: bool = False, ) -> BaseModel | None: + native: Final = native_auth_repository(prisma_client) + if native is not None: + return await native.key(hashed_token, writer=check_db_only) fetch: Final = partial(prisma_client.get_data, use_writer=True) if check_db_only else prisma_client.get_data async with db_lookup_gate.current(): try: @@ -4247,6 +3920,8 @@ async def _fetch_key_object_from_db_unbounded( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + except NativeAuthFailure: + raise except Exception as e: if PrismaDBExceptionHandler.is_database_transport_error(e): did_reconnect = False @@ -4407,6 +4082,8 @@ async def get_key_object( proxy_logging_obj=proxy_logging_obj, check_db_only=check_db_only, ) + except NativeAuthFailure: + raise except Exception as e: if check_db_only: raise @@ -4459,6 +4136,9 @@ async def get_object_permission( - if valid, return LiteLLM_ObjectPermissionTable object - if not, then raise an error """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -4477,17 +4157,27 @@ async def get_object_permission( # else, check db try: - response: Final = await _dictable_table( - ObjectPermissionRepository(prisma_client, use_writer=check_db_only), "object_permission" - ).find_unique(where={"object_permission_id": object_permission_id}) + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one( + "object_permission", [object_permission_id], LiteLLM_ObjectPermissionTable, writer=check_db_only + ) + if native is not None + else await _dictable_table( + ObjectPermissionRepository(prisma_client, use_writer=check_db_only), "object_permission" + ).find_unique(where={"object_permission_id": object_permission_id}) + ) if response is None: if check_db_only: raise HTTPException(status_code=403, detail="Referenced object permission does not exist") return None - _perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate(response.dict()) + _perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=_perm_obj, model_type=LiteLLM_ObjectPermissionTable, @@ -4495,6 +4185,8 @@ async def get_object_permission( ) return _perm_obj + except NativeAuthFailure: + raise except Exception: if check_db_only: raise @@ -4518,6 +4210,9 @@ async def get_managed_vector_store_rows_by_uuids( Critical-path DB access must go through this helper to avoid raw Prisma calls on the hot request path. """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if not uuids or prisma_client is None: return [] @@ -4538,13 +4233,22 @@ async def get_managed_vector_store_rows_by_uuids( if not cache_misses: return result - rows: Final = await _vector_store_table(ManagedVectorStoresRepository(prisma_client)).find_many( - where={"vector_store_id": {"in": cache_misses}}, - take=len(cache_misses), + native: Final = native_auth_repository(prisma_client) + rows: Final = ( + await native.many("managed_vector_store", cache_misses, LiteLLM_ManagedVectorStoresTable) + if native is not None + else await _vector_store_table(ManagedVectorStoresRepository(prisma_client)).find_many( + where={"vector_store_id": {"in": cache_misses}}, + take=len(cache_misses), + ) ) for row in rows: - row_dict = row.model_dump() if hasattr(row, "model_dump") else (row.dict() if hasattr(row, "dict") else None) + row_dict = ( + row.model_dump() + if isinstance(row, BaseModel) or hasattr(row, "model_dump") + else (row.dict() if hasattr(row, "dict") else None) + ) if not isinstance(row_dict, dict) or not row_dict: row_dict = dict(row) if hasattr(row, "__dict__") else {} if not row_dict: @@ -4552,6 +4256,7 @@ async def get_managed_vector_store_rows_by_uuids( cached_obj = LiteLLM_ManagedVectorStoresTable.model_validate(row_dict) key = f"managed_vector_store_id:{cached_obj.vector_store_id}" await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=cached_obj, model_type=LiteLLM_ManagedVectorStoresTable, @@ -4594,6 +4299,9 @@ async def get_org_object( proxy_logging_obj: Optional proxy logging object include_budget_table: If True, includes litellm_budget_table in the query """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") if not isinstance(org_id, str): @@ -4612,21 +4320,16 @@ async def get_org_object( if deserialized_org is not None: return deserialized_org # else, check db - try: - query_kwargs: Final[dict[str, Mapping[str, object]]] = {"where": {"organization_id": org_id}} - if include_budget_table: - query_kwargs["include"] = {"litellm_budget_table": True} + query_kwargs: Final[dict[str, Mapping[str, object]]] = {"where": {"organization_id": org_id}} + if include_budget_table: + query_kwargs["include"] = {"litellm_budget_table": True} - response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique( - **query_kwargs - ) - except Exception: - # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed - # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them - # apart — a caller that treats absence as "this org places no restriction" then drops a real - # org ceiling during an outage. Propagate the real error; callers that already catch - # Exception are unaffected. - raise + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("org", [org_id], LiteLLM_OrganizationTable) + if native is not None + else await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique(**query_kwargs) + ) if response is None: raise OrganizationNotFoundError( @@ -4636,6 +4339,7 @@ async def get_org_object( _org_obj: Final = LiteLLM_OrganizationTable.model_validate(response.model_dump()) # Cache the result await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=_org_obj, model_type=LiteLLM_OrganizationTable, @@ -4651,7 +4355,10 @@ def _last_known_org_cache_key(org_id: str) -> str: @with_service_target(AUTH_OBJECTS_TARGET) async def _keep_last_known_org( - org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache + org: LiteLLM_OrganizationTable, + org_id: str, + user_api_key_cache: UserApiKeyCache, + fill_generation: int | None = None, ) -> None: cache_key: Final = _last_known_org_cache_key(org_id) held_locally: Final = await user_api_key_cache.async_get_cache( @@ -4660,6 +4367,7 @@ async def _keep_last_known_org( if held_locally is not None: return await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=org, model_type=LiteLLM_OrganizationTable, @@ -4675,6 +4383,9 @@ async def get_org_object_for_request( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, ) -> LiteLLM_OrganizationTable | None: + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) try: org: Final = await get_org_object( org_id=org_id, @@ -4686,6 +4397,8 @@ async def get_org_object_for_request( ) except OrganizationNotFoundError: return None + except NativeAuthFailure: + raise except Exception as e: if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) @@ -4701,7 +4414,7 @@ async def get_org_object_for_request( raise if org is None: return None - await _keep_last_known_org(org, org_id, user_api_key_cache) + await _keep_last_known_org(org, org_id, user_api_key_cache, fill_generation) return org @@ -4749,6 +4462,8 @@ async def get_access_group_tables( raise verbose_proxy_logger.debug("One of the requested access groups does not exist, it grants nothing") return None + except NativeAuthFailure: + raise except Exception: if check_db_only or unreadable == "raise": raise @@ -4895,20 +4610,6 @@ async def get_skill_names_from_access_groups( ) -def _resolve_all_team_model_sentinel_for_auth_check( - models: list[str], - llm_router: RegistryReader | None, - team_id: str | None, -) -> list[str]: - if SpecialModelNames.all_team_models.value not in models or team_id is None or llm_router is None: - return models - proxy_models: Final = llm_router.registry().names() - non_sentinel_models: Final = [model for model in models if model != SpecialModelNames.all_team_models.value] - if not proxy_models: - return non_sentinel_models or models - return list(dict.fromkeys(non_sentinel_models + proxy_models)) - - def check_model_access_helper( model: str, llm_router: RegistryReader | None, @@ -4917,65 +4618,14 @@ def check_model_access_helper( team_id: str | None = None, ) -> bool: ## check if model in allowed model names - from collections import defaultdict + from litellm.proxy.auth import model_helper_adapter - access_groups: dict[str, list[str]] = defaultdict(list) - - if llm_router: - access_groups = llm_router.registry().access_groups(model_name=model, team_id=team_id) - - models = _resolve_all_team_model_sentinel_for_auth_check( - models=models, - llm_router=llm_router, - team_id=team_id, - ) - - if len(access_groups) > 0 and llm_router is not None: # check if token contains any model access groups - for idx, m in enumerate(models): # loop token models, if any of them are an access group add the access group - if m in access_groups: - return True - - # Filter out models that are access_groups - filtered_models: Final = [m for m in models if m not in access_groups] - - if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases): - return True - - if model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models): - return True - - all_model_access: bool = False - - if (len(filtered_models) == 0 and len(models) == 0) or "*" in filtered_models: - all_model_access = True - - if SpecialModelNames.all_proxy_models.value in filtered_models: - all_model_access = True - - if model is not None and model not in filtered_models and all_model_access is False: - return False - return True + return model_helper_adapter.check_model_access_helper(model, llm_router, models, team_model_aliases, team_id) _check_model_access_helper: Final = check_model_access_helper -def _owning_team_names_for_team_deployment( - model: str, - llm_router: RegistryReader | None, - team_id: str | None, - team_model_aliases: Mapping[str, str] | None, -) -> tuple[str, ...]: - if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"): - return () - deployment: Final = llm_router.registry().deployments(model, source="first_in_group") - if deployment is None or deployment.model_info.team_id != team_id: - return () - public_name: Final = deployment.model_info.team_public_model_name - legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model) - return (*((public_name,) if public_name else ()), *legacy_aliases) - - def _can_object_call_model( model: str | list[str], llm_router: RegistryReader | None, @@ -4986,136 +4636,32 @@ def _can_object_call_model( object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"] = "user", fallback_depth: int = 0, ) -> Literal[True]: - """ - Checks if token can call a given model + from litellm.proxy.auth import model_access_adapter - Args: - - model: str - - llm_router: Optional[Router] - - models: List[str] - - team_model_aliases: Optional[Dict[str, str]] - - object_type: Literal["user", "team", "key", "org"]. We use the object type to raise the correct exception type - - Returns: - - True: if token allowed to call model - - Raises: - - Exception: If token not allowed to call model - """ - if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: - raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}") - if isinstance(model, list): - for m in model: - can_object_call_model( - model=m, - llm_router=llm_router, - models=models, - team_model_aliases=team_model_aliases, - team_id=team_id, - key_model_aliases=key_model_aliases, - object_type=object_type, - fallback_depth=fallback_depth + 1, - ) - return True - - from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent - - compaction_parent: Final = native_compaction_parent(model) - global_or_router_alias_target: Final = ( - litellm.model_alias_map[model] - if model in litellm.model_alias_map - else ( - llm_router.registry().resolve(model, kind="alias") - if llm_router is not None and model in llm_router.registry().aliases() - else None - ) - ) - after_team_alias: Final = team_model_aliases.get(model, model) if team_model_aliases else model - after_key_alias: Final = ( - key_model_aliases.get(after_team_alias, after_team_alias) if key_model_aliases else after_team_alias - ) - after_global_alias: Final = litellm.model_alias_map.get(after_key_alias, after_key_alias) - dispatched_model: Final = ( - key_model_aliases.get(after_global_alias, after_global_alias) if key_model_aliases else after_global_alias - ) - key_alias_applied: Final = after_key_alias != after_team_alias or dispatched_model != after_global_alias - potential_models: Final = ( - (dispatched_model,) - if key_alias_applied - else ( - *((model, compaction_parent) if compaction_parent is not None else (model,)), - *((global_or_router_alias_target,) if global_or_router_alias_target else ()), - ) - ) - owning_team_names: Final = ( - () - if key_alias_applied - else _owning_team_names_for_team_deployment( - model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases - ) - ) - - ## check model access for alias + underlying model - allow if either is in allowed models - for m in potential_models: - if check_model_access_helper( - model=m, - llm_router=llm_router, - models=models, - team_model_aliases=team_model_aliases, - team_id=team_id, - ): - return True - if any( - check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id) - for name in owning_team_names - ): - return True - - internal_message: Final = ( - f"{object_type} not allowed to access model. This {object_type} can only access models={models}. " - f"Tried to access {model}" - ) - raise ModelAccessDeniedProxyException( - message=( - customer_model_access_denied_client_message(model=model) - if object_type == "customer" - else model_access_denied_client_message(model=model) - ), - internal_message=internal_message, - type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type=object_type), - param="model", - code=status.HTTP_403_FORBIDDEN, + return model_access_adapter.can_object_call_model( + model=model, + llm_router=llm_router, + models=models, + team_model_aliases=team_model_aliases, + team_id=team_id, + key_model_aliases=key_model_aliases, + object_type=object_type, + fallback_depth=fallback_depth, ) can_object_call_model: Final = _can_object_call_model -def _resolve_team_alias( +def _resolve_team_alias( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint model: str | list[str], team_model_aliases: Mapping[str, str] | None, team_id: str | None, llm_router: Router | None, ) -> str | list[str]: - if not team_model_aliases: - return model - if isinstance(model, str): - return _live_team_alias_target(model, team_model_aliases, team_id, llm_router) - return [_live_team_alias_target(name, team_model_aliases, team_id, llm_router) for name in model] + from litellm.proxy.auth import team_alias_adapter - -def _live_team_alias_target( - model: str, team_model_aliases: Mapping[str, str], team_id: str | None, llm_router: RegistryReader | None -) -> str: - target: Final = team_model_aliases.get(model) - if target is None: - return model - deleted_team_deployment: Final = ( - llm_router is not None - and target.startswith(f"model_name_{team_id}_") - and target not in llm_router.registry().index("model_name") - ) - return model if deleted_team_deployment else target + return team_alias_adapter.resolve_team_alias(model, team_model_aliases, team_id, llm_router) async def can_agent_call_model( @@ -5126,93 +4672,46 @@ async def can_agent_call_model( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: - managed_policy: Final = managed_agent_policy(valid_token) - if model and valid_token is not None and managed_policy is not None: - managed_models: Final = (managed_policy.object_permission or {}).get("models", ()) - if not isinstance(managed_models, (list, tuple)) or not managed_models: - raise HTTPException(403, "This agent has no model grants") - can_object_call_model( - model=_resolve_team_alias( - model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router - ), - llm_router=llm_router, - models=list(managed_models), - team_id=valid_token.team_id, - object_type="agent", - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - await _check_agent_access_group_model_access(model=model, valid_token=valid_token, llm_router=llm_router) - await _check_agent_caller_model_access( - model=model, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.can_agent_call_model( + model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj ) -async def _agent_access_group_ceilings( +async def _agent_access_group_ceilings( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver ) -> tuple[AgentAccessGroupCeiling, ...]: - if not valid_token.agent_id: - return () + from litellm.proxy.auth import managed_model_adapter - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - - managed: Final = managed_agent_policy(valid_token) - if managed is not None: - return await resolve_managed_agent_ceilings(managed) - unmanaged: Final = await resolve_ceiling(valid_token.agent_id) - return (unmanaged,) if unmanaged is not None else () + return await managed_model_adapter.agent_access_group_ceilings(valid_token, resolve_ceiling) -def _managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: - managed: Final = managed_agent_policy(valid_token) - granted: Final = (managed.object_permission or {}).get("models") if managed is not None else None - if not isinstance(granted, (list, tuple)): - return frozenset() - return frozenset(name for name in _OBJECTS_ADAPTER.validate_python(granted) if isinstance(name, str)) +def _managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint + from litellm.proxy.auth import managed_model_adapter + + return managed_model_adapter.managed_agent_model_grants(valid_token) async def agent_model_grants( valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling ) -> frozenset[str]: - ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) - return _managed_agent_model_grants(valid_token).union(*(ceiling.models for ceiling in ceilings)) + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.agent_model_grants(valid_token, resolve_ceiling) -async def _check_agent_access_group_model_access( +async def _check_agent_access_group_model_access( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] valid_token: UserAPIKeyAuth | None, llm_router: Router | None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> Literal[True]: - """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" - if not model or valid_token is None: - return True - ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) - dispatched: Final = _resolve_team_alias( - model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.check_agent_access_group_model_access( + model, valid_token, llm_router, resolve_ceiling ) - for ceiling in ceilings: - if not ceiling.models: - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model), - internal_message=f"agent {valid_token.agent_id} access groups grant no models", - type=ProxyErrorTypes.agent_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) - can_object_call_model( - model=dispatched, - llm_router=llm_router, - models=sorted(ceiling.models), - team_id=valid_token.team_id, - object_type="agent", - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - return True LoadedCallerTeam: TypeAlias = LiteLLM_TeamTable | None @@ -5221,7 +4720,7 @@ CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerT CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerUser]] -async def _check_agent_caller_model_access( +async def _check_agent_caller_model_access( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str] valid_token: UserAPIKeyAuth | None, llm_router: Router | None, @@ -5231,63 +4730,13 @@ async def _check_agent_caller_model_access( load_team: CallerTeamLoader = load_agent_caller_team, load_user: CallerUserLoader = load_agent_caller_user, ) -> None: - """An agent key acting for an invoking user may call only what that user's own key could: the - invoking team's models (and per-member scope) when a team was echoed, else the user's models.""" - if not model or valid_token is None: - return - caller_auth: Final = agent_caller_auth(valid_token) - if caller_auth is None: - return - caller_team: Final = await load_team(valid_token) - caller_key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) - if caller_team is not None: - await can_team_access_model( - model=model, - team_object=caller_team, - llm_router=llm_router, - prisma_client=prisma_client, - key_model_aliases=caller_key_model_aliases, - ) - await check_team_member_model_access( - model=model, - team_object=caller_team, - valid_token=caller_auth, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - key_model_aliases=caller_key_model_aliases, - ) - return - caller_user: Final = await load_user(valid_token) - if caller_user is None: - return - await can_user_call_model( - model=model, - llm_router=llm_router, - user_object=caller_user, - key_model_aliases=caller_key_model_aliases, + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.check_agent_caller_model_access( + model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj, load_team, load_user ) -def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool: - """ - Returns True if `model` being accessed is an alias of a team model - - - `model=gpt-4o` - - `team_model_aliases={"gpt-4o": "gpt-4o-team-1"}` - - returns True - - - `model=gp-4o` - - `team_model_aliases={"o-3": "o3-preview"}` - - returns False - """ - if team_model_aliases: - if model in team_model_aliases: - return True - return False - - def key_model_aliases_for_auth_check(valid_token: UserAPIKeyAuth | None) -> Mapping[str, str] | None: return alias_map(valid_token.aliases) if valid_token is not None and valid_token.aliases else None @@ -5296,7 +4745,7 @@ def team_model_aliases_for_auth_check(valid_token: UserAPIKeyAuth) -> Mapping[st return alias_map(valid_token.team_model_aliases) if valid_token.team_model_aliases else None -def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]: +def resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]: """ Expand key model sentinels before auth checks. @@ -5314,101 +4763,7 @@ def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str] return models -def _model_access_groups_serving_model( - model: str | Sequence[str], - llm_router: Router, - team_id: str | None, -) -> frozenset[str]: - """Every model access group whose deployments serve the requested model(s).""" - requested: Final = (model,) if isinstance(model, str) else tuple(model) - return frozenset( - group - for requested_model in requested - for group in llm_router.registry().access_groups(model_name=requested_model, team_id=team_id) - ) - - -async def _team_member_granted_models( - valid_token: UserAPIKeyAuth, - team_object: LiteLLM_TeamTable | None, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - team_membership: LiteLLM_TeamMembership | None = None, - team_membership_loaded: bool = False, -) -> Sequence[str]: - """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" - if team_object is None or valid_token.user_id is None: - return () - - if not team_membership_loaded: - team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - return () if team_membership is None else _member_allowed_models(team_membership) - - -async def _org_granted_models( - valid_token: UserAPIKeyAuth, - team_object: LiteLLM_TeamTable | None, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> Sequence[str]: - """The org allowlist reached through the key, or through its team when the key names no org.""" - org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None) - if org_id is None: - return () - - try: - org_object: Final = await get_org_object( - org_id=org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: # noqa: BLE001 # fail-safe: attribution degrades to "no org grant", it must never break auth - verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", e) - return () - return org_object.models if org_object is not None else () - - -async def _granted_model_lists( - valid_token: UserAPIKeyAuth, - team_object: LiteLLM_TeamTable | None, - project_object: LiteLLM_ProjectTableCachedObj | None, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - team_membership: LiteLLM_TeamMembership | None = None, - team_membership_loaded: bool = False, -) -> tuple[Sequence[str], ...]: - """One model allowlist per level that participates in authorizing the request.""" - return ( - _resolve_key_models_for_auth_check(valid_token=valid_token), - team_object.models if team_object is not None else (), - await _team_member_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=team_membership_loaded, - ), - project_object.models if project_object is not None else (), - await _org_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - ) +_resolve_key_models_for_auth_check: Final = resolve_key_models_for_auth_check async def collect_matched_model_access_groups( @@ -5423,56 +4778,20 @@ async def collect_matched_model_access_groups( team_membership: LiteLLM_TeamMembership | None = None, team_membership_loaded: bool = False, ) -> tuple[str, ...]: - """ - The budgeted model access groups that authorized this request, sorted and deduplicated. + from litellm.proxy.auth import model_attribution_adapter - A group is charged only when its name appears on an allowlist the caller was granted -- key, - team, team-member scope, project or org -- *and* that group serves the requested model. Asking - for a model that merely belongs to a group attributes nothing, because nothing about the caller - named the group. - - Levels are unioned, never ranked: a team granted ``*`` whose member is scoped to one group is - still a caller gated by that group. An unrestricted allowlist (empty, ``*``) names no group and - so contributes nothing. - - The whole walk is gated on the budget registry, because collecting every match costs a full scan - of each allowlist where the plain access check stops at the first hit. An empty registry means no - group carries a budget, so there is nothing to attribute and no work worth doing. - """ - if model is None or valid_token is None or llm_router is None or prisma_client is None: - return () - - registry: Final = await _load_model_access_group_registry( + return await model_attribution_adapter.collect_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) - if registry is not None and not registry: - return () - - covering_groups: Final = _model_access_groups_serving_model( - model=model, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - budgeted_groups: Final = covering_groups if registry is None else covering_groups & registry - if not budgeted_groups: - return () - - granted: Final = frozenset( - granted_model - for granted_models in await _granted_model_lists( - valid_token=valid_token, - team_object=team_object, - project_object=project_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=team_membership_loaded, - ) - for granted_model in granted_models - ) - return tuple(sorted(budgeted_groups & granted)) async def stamp_matched_model_access_groups( @@ -5487,32 +4806,20 @@ async def stamp_matched_model_access_groups( team_membership: LiteLLM_TeamMembership | None = None, team_membership_loaded: bool = False, ) -> tuple[str, ...]: - """Record the groups that authorized this request on its auth object, for the post-call spend - writer and the reservation counters, and hand them back for the budget check.""" - if valid_token is None: - return () + from litellm.proxy.auth import model_attribution_adapter - try: - matched: Final = await collect_matched_model_access_groups( - model=model, - valid_token=valid_token, - team_object=team_object, - project_object=project_object, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=team_membership_loaded, - ) - except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth - verbose_proxy_logger.debug("model access group attribution failed: %s", e) - return () - if not matched: - return () - matched_groups: Final = list(matched) - valid_token.matched_model_access_groups = matched_groups # rebind-ok: request-scoped carrier for the writer - return matched + return await model_attribution_adapter.stamp_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, + ) async def can_key_call_model( @@ -5522,206 +4829,26 @@ async def can_key_call_model( llm_router: RegistryReader | None, prisma_client: DatabaseClient | None = None, ) -> Literal[True]: - """ - Checks if token can call a given model + from litellm.proxy.auth import model_access_adapter - 1. First checks native key-level model permissions (current implementation) - 2. If not allowed natively, falls back to access_group_ids on the key - - Returns: - - True: if token allowed to call model - - Raises: - - Exception: If token not allowed to call model - """ - key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) - try: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=key_models, - team_model_aliases=valid_token.team_model_aliases, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="key", - ) - except ProxyException: - # Fallback: check key's access_group_ids - key_access_group_ids: Final = valid_token.access_group_ids or [] - if key_access_group_ids: - models_from_groups: Final = await get_models_from_access_groups( - access_group_ids=key_access_group_ids, - prisma_client=prisma_client, - ) - if models_from_groups: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=models_from_groups, - team_model_aliases=valid_token.team_model_aliases, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="key", - ) - raise - - -async def _personal_key_owner( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> LiteLLM_UserTable | None: - if valid_token.user_id is None: - return None - try: - return await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - except Exception as e: # noqa: BLE001 # fail closed: an owner that cannot be loaded grants nothing - verbose_proxy_logger.warning( - "Personal key team model access: owner lookup failed, granting nothing: %s", - e, - ) - return None + return await model_access_adapter.can_key_call_model( + model=model, + llm_model_list=llm_model_list, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + ) async def can_key_call_resolved_model( - model: str, + model: str | Sequence[str], llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None, ) -> None: - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.auth import resolved_model_adapter - skip_key_model_check: Final = valid_token.config or ( - isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models - ) - if not skip_key_model_check: - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - - team_object: LiteLLM_TeamTableCachedObj | None = None - team_object_from_lookup = False - if valid_token.team_id is not None: - try: - team_object = await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - team_object_from_lookup = True - except Exception: - team_object = LiteLLM_TeamTableCachedObj( - team_id=valid_token.team_id, - models=valid_token.team_models, - blocked=valid_token.team_blocked, - team_alias=valid_token.team_alias, - metadata=valid_token.team_metadata, - object_permission_id=valid_token.team_object_permission_id, - object_permission=valid_token.team_object_permission, - ) - - if team_object is not None: - try: - await can_team_access_model( - model=model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=valid_token.team_model_aliases, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - except ProxyException as team_denial: - if team_denial.type != ProxyErrorTypes.team_model_access_denied: - raise - if not await key_access_group_grants_model( - model=model, - valid_token=valid_token, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ): - raise - - if valid_token.user_id is not None and team_object_from_lookup: - await check_team_member_model_access( - model=model, - team_object=team_object, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - if personal_key_team_model_access_applies(valid_token) and valid_token.user_id is not None: - owner: Final = await _personal_key_owner( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - await can_personal_key_call_model_via_teams( - model=model, - user_object=owner, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - if valid_token.project_id is not None: - project_object: Final = await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if project_object is not None and len(project_object.models) > 0: - can_project_access_model( - model=model, - project_object=project_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - if valid_token.end_user_id is not None and prisma_client is not None: - key_metadata: Final = _KEY_METADATA_ADAPTER.validate_python(valid_token.metadata) - end_user_object: Final = await get_end_user_object( - end_user_id=valid_token.end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - token_end_user_max_budget=valid_token.end_user_max_budget, - key_end_user_budget_id=get_key_end_user_budget_id(key_metadata), - ) - if end_user_object is not None and end_user_object.models: - can_customer_access_model( - model=model, - end_user_object=end_user_object, - llm_router=llm_router, - valid_token=valid_token, - ) + return await resolved_model_adapter.can_key_call_resolved_model(model, llm_model_list, valid_token, llm_router) def can_org_access_model( @@ -5730,16 +4857,10 @@ def can_org_access_model( llm_router: Router | None, team_model_aliases: dict[str, str] | None = None, ) -> Literal[True]: - """ - Returns True if the team can access a specific model. + from litellm.proxy.auth import model_access_adapter - """ - return can_object_call_model( - model=model, - llm_router=llm_router, - models=org_object.models if org_object else [], - team_model_aliases=team_model_aliases, - object_type="org", + return model_access_adapter.can_org_access_model( + model=model, org_object=org_object, llm_router=llm_router, team_model_aliases=team_model_aliases ) @@ -5751,41 +4872,16 @@ async def can_team_access_model( key_model_aliases: Mapping[str, str] | None = None, prisma_client: DatabaseClient | None = None, ) -> Literal[True]: - """ - Returns True if the team can access a specific model. + from litellm.proxy.auth import model_access_adapter - 1. First checks native team-level model permissions (current implementation) - 2. If not allowed natively, falls back to access_group_ids on the team - """ - try: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=team_object.models if team_object else [], - team_model_aliases=team_model_aliases, - team_id=team_object.team_id if team_object else None, - key_model_aliases=key_model_aliases, - object_type="team", - ) - except ProxyException: - # Fallback: check team's access_group_ids - team_access_group_ids: Final = (team_object.access_group_ids or []) if team_object else [] - if team_access_group_ids: - models_from_groups: Final = await get_models_from_access_groups( - access_group_ids=team_access_group_ids, - prisma_client=prisma_client, - ) - if models_from_groups: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])), - team_model_aliases=team_model_aliases, - team_id=team_object.team_id if team_object else None, - key_model_aliases=key_model_aliases, - object_type="team", - ) - raise + return await model_access_adapter.can_team_access_model( + model=model, + team_object=team_object, + llm_router=llm_router, + team_model_aliases=team_model_aliases, + key_model_aliases=key_model_aliases, + prisma_client=prisma_client, + ) async def get_authorized_resources_from_key_access_groups( @@ -5795,54 +4891,15 @@ async def get_authorized_resources_from_key_access_groups( check_db_only: bool = False, strict_errors: bool = False, ) -> list[str]: - """ - For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row - and contribute its `resource_field` only if the group authorizes the caller - as an owner — that is, the group's `assigned_team_ids` includes the key's - `team_id`, or the group's `assigned_key_ids` includes the key's token. This - preserves the team-as-owner boundary while still letting a group reach the - key without first being added to the team's `access_group_ids` list. - """ - if valid_token is None: - return [] - key_access_group_ids: Final = list(valid_token.access_group_ids or []) - if not key_access_group_ids: - return [] + from litellm.proxy.auth import model_group_adapter - from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache - - if _prisma_client is None or _user_api_key_cache is None: - return [] - - key_team_id: Final = valid_token.team_id or (team_object.team_id if team_object is not None else None) - key_token: Final = valid_token.token - - authorized_resources: Final[list[str]] = [] - for ag_id in key_access_group_ids: - try: - ag = await get_access_object( - access_group_id=ag_id, - prisma_client=_prisma_client, - user_api_key_cache=_user_api_key_cache, - proxy_logging_obj=_proxy_logging_obj, - check_db_only=check_db_only, - ) - except AccessGroupNotFoundError: - if check_db_only or strict_errors: - raise - continue - except Exception: - if check_db_only or strict_errors: - raise - continue - team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or [])) - key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) - if team_authorized or key_authorized: - authorized_resources.extend(getattr(ag, resource_field, []) or []) - - return list(set(authorized_resources)) + return await model_group_adapter.get_authorized_resources_from_key_access_groups( + valid_token=valid_token, + team_object=team_object, + resource_field=resource_field, + check_db_only=check_db_only, + strict_errors=strict_errors, + ) async def key_access_group_grants_model( @@ -5888,18 +4945,10 @@ def can_project_access_model( llm_router: RegistryReader | None, key_model_aliases: Mapping[str, str] | None = None, ) -> Literal[True]: - """ - Returns True if the project can access a specific model. + from litellm.proxy.auth import model_access_adapter - Raises ProxyException if access is denied. - """ - return can_object_call_model( - model=model, - llm_router=llm_router, - models=project_object.models if project_object else [], - team_id=project_object.team_id if project_object else None, - key_model_aliases=key_model_aliases, - object_type="project", + return model_access_adapter.can_project_access_model( + model=model, project_object=project_object, llm_router=llm_router, key_model_aliases=key_model_aliases ) @@ -5909,27 +4958,11 @@ def can_customer_access_model( llm_router: RegistryReader | None, valid_token: UserAPIKeyAuth | None, ) -> Literal[True]: - team_model_aliases: Final = team_model_aliases_for_auth_check(valid_token) if valid_token is not None else None - team_id: Final = valid_token.team_id if valid_token is not None else None - key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) + from litellm.proxy.auth import model_access_adapter - def check(name: str) -> None: - team_target: Final = ( - _live_team_alias_target(name, team_model_aliases, team_id, llm_router) if team_model_aliases else name - ) - if team_target != name and name in (end_user_object.models or ()): - return - can_object_call_model( - model=team_target, - llm_router=llm_router, - models=end_user_object.models, - key_model_aliases=key_model_aliases, - object_type="customer", - ) - - for name in (model,) if isinstance(model, str) else model: - check(name) - return True + return model_access_adapter.can_customer_access_model( + model=model, end_user_object=end_user_object, llm_router=llm_router, valid_token=valid_token + ) async def can_user_call_model( @@ -5938,235 +4971,26 @@ async def can_user_call_model( user_object: LiteLLM_UserTable | None, key_model_aliases: Mapping[str, str] | None = None, ) -> Literal[True]: - if user_object is None: - return True + from litellm.proxy.auth import model_access_adapter - if SpecialModelNames.no_default_models.value in user_object.models: - internal_message: Final = ( - f"User not allowed to access model. No default model access, only team models allowed. " - f"Tried to access {model}" - ) - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model), - internal_message=internal_message, - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) - - return can_object_call_model( - model=model, - llm_router=llm_router, - models=user_object.models, - key_model_aliases=key_model_aliases, - object_type="user", + return await model_access_adapter.can_user_call_model( + model=model, llm_router=llm_router, user_object=user_object, key_model_aliases=key_model_aliases ) def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) -> bool: - return ( - bool(litellm.personal_key_model_access_from_teams) - and valid_token is not None - and valid_token.team_id is None - and valid_token.user_id is not None - and valid_token.api_key != LITELLM_PROXY_MASTER_KEY_ALIAS - and valid_token.user_role != LitellmUserRoles.PROXY_ADMIN - ) + from litellm.proxy.auth.personal_model_adapter import personal_key_team_model_access_applies as applies + + return applies(valid_token) def personal_key_inherits_team_grants(valid_token: UserAPIKeyAuth | None) -> bool: - """MCP server, MCP tool and agent inheritance is for a virtual key only. A keyless auth with no - team (one grant source of an admitted subject, a JWT user) resolves as the plain caller it is.""" - return ( - personal_key_team_model_access_applies(valid_token) - and valid_token is not None - and valid_token.api_key is not None - ) + from litellm.proxy.auth.personal_model_adapter import personal_key_inherits_team_grants as inherits + + return inherits(valid_token) -@dataclass(frozen=True, slots=True) -class _PersonalKeyTeamLoad: - team: LiteLLM_TeamTableCachedObj | None - unreadable: bool - withheld_team: LiteLLM_TeamTableCachedObj | None - - -async def _load_personal_key_team_with_status( - team_id: str, - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> _PersonalKeyTeamLoad: - try: - team_object: Final = await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - except Exception as e: # noqa: BLE001 # fail closed: a team that cannot be loaded grants nothing - verbose_proxy_logger.warning( - "Personal key team model access: team=%s lookup failed, granting nothing: %s", team_id, e - ) - return _PersonalKeyTeamLoad(team=None, unreadable=True, withheld_team=None) - if team_object.blocked is True: - return _PersonalKeyTeamLoad(team=None, unreadable=False, withheld_team=team_object) - return _PersonalKeyTeamLoad(team=team_object, unreadable=False, withheld_team=None) - - -async def _load_personal_key_team( - team_id: str, - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> LiteLLM_TeamTableCachedObj | None: - loaded: Final = await _load_personal_key_team_with_status( - team_id=team_id, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - return loaded.team - - -async def _load_personal_key_teams( - user_object: LiteLLM_UserTable, - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: - return tuple( - await asyncio.gather( - *( - _load_personal_key_team( - team_id=team_id, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - for team_id in dict.fromkeys(user_object.teams) - ) - ) - ) - - -async def _team_grants_personal_key_model( - model: str, - team_object: LiteLLM_TeamTableCachedObj, - valid_token: UserAPIKeyAuth, - llm_router: RegistryReader | None, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> bool: - key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) - try: - await can_team_access_model( - model=model, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases, - prisma_client=prisma_client, - ) - await _check_team_member_model_access( - model=model, - team_object=team_object, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - key_model_aliases=key_model_aliases, - ) - except ProxyException: - return False - except Exception as e: # noqa: BLE001 # fail closed: a team check that errors grants nothing - verbose_proxy_logger.warning( - "Personal key team model access: team=%s check failed, granting nothing: %s", team_object.team_id, e - ) - return False - return True - - -async def _loaded_teams_allow_model( - model: str, - teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], - valid_token: UserAPIKeyAuth, - llm_router: RegistryReader | None, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> bool: - loaded: Final = tuple(team for team in teams if team is not None) - if not loaded: - return False - grants: Final = await asyncio.gather( - *( - _team_grants_personal_key_model( - model=model, - team_object=team, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - for team in loaded - ) - ) - return any(grants) - - -async def _personal_key_listing_teams( - user_api_key_dict: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache | None, - proxy_logging_obj: ProxyLogging | None, - *, - check_db_only: bool = False, -) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: - if ( - user_api_key_dict.user_id is None - or prisma_client is None - or user_api_key_cache is None - or proxy_logging_obj is None - ): - return () - owner: Final = await _personal_key_owner( - valid_token=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - if owner is None: - return () - return await _load_personal_key_teams( - user_object=owner, - valid_token=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - - -async def _personal_key_team_within_budget( +async def personal_key_team_within_budget( team: LiteLLM_TeamTableCachedObj, valid_token: UserAPIKeyAuth, prisma_client: PrismaClient, @@ -6175,60 +4999,14 @@ async def _personal_key_team_within_budget( *, check_db_only: bool = False, ) -> bool: - source: Final = valid_token.model_copy( - update={"team_id": team.team_id, "org_id": valid_token.org_id or team.organization_id} + from litellm.proxy.auth.personal_budget_adapter import personal_key_team_within_budget + + return await personal_key_team_within_budget( + team, valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only ) - try: - await team_max_budget_check( - team_object=team, - valid_token=source, - proxy_logging_obj=proxy_logging_obj, - ) - await _team_multi_budget_check(team_object=team) - if check_db_only: - team_member_user_id: Final = valid_token.user_id - if team_member_user_id is None: - return False - team_membership: Final = await get_team_membership( - user_id=team_member_user_id, - team_id=team.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=True, - ) - await _check_team_member_budget( - team_object=team, - user_object=None, - valid_token=source, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=True, - ) - else: - await _check_team_member_budget( - team_object=team, - user_object=None, - valid_token=source, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - await organization_max_budget_check( - valid_token=source, - team_object=team, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - except litellm.BudgetExceededError: - return False - except Exception as e: # noqa: BLE001 # fail closed: a team whose budget cannot be checked grants nothing - verbose_proxy_logger.warning("Failed to check budget for personal key team %s, skipping: %s", team.team_id, e) - return False - return True + + +_personal_key_team_within_budget: Final = personal_key_team_within_budget @dataclass(frozen=True, slots=True) @@ -6246,71 +5024,10 @@ async def personal_key_owner_teams_with_status( *, check_db_only: bool = False, ) -> PersonalKeyOwnerTeams: - """The owner's current, loadable, unblocked, within-budget teams: the grant sources a personal key inherits - MCP server, MCP tool and agent access from under ``personal_key_model_access_from_teams``.""" - if valid_token.user_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None: - return PersonalKeyOwnerTeams(teams=(), unreadable=False, withheld=()) - owner: Final = await _personal_key_owner( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - if owner is None: - return PersonalKeyOwnerTeams(teams=(), unreadable=False, withheld=()) - team_ids: Final = tuple(team_id for team_id in dict.fromkeys(owner.teams) if team_id and team_id != UI_TEAM_ID) - team_loads: Final = await asyncio.gather( - *( - _load_personal_key_team_with_status( - team_id=team_id, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - for team_id in team_ids - ) - ) - blocked_teams: Final = tuple( - load.withheld_team - for load in team_loads - if load.withheld_team is not None - and (valid_token.org_id is None or load.withheld_team.organization_id in (None, valid_token.org_id)) - ) - loaded_teams: Final = tuple(load.team for load in team_loads if load.team is not None) - eligible: Final = tuple( - team - for team in loaded_teams - if valid_token.org_id is None or team.organization_id in (None, valid_token.org_id) - ) - if not eligible: - return PersonalKeyOwnerTeams( - teams=(), - unreadable=any(load.unreadable for load in team_loads), - withheld=blocked_teams, - ) - within_budget: Final = await asyncio.gather( - *( - _personal_key_team_within_budget( - team=team, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - for team in eligible - ) - ) - return PersonalKeyOwnerTeams( - teams=tuple(team for team, allowed in zip(eligible, within_budget, strict=True) if allowed), - unreadable=any(load.unreadable for load in team_loads), - withheld=( - *blocked_teams, - *(team for team, allowed in zip(eligible, within_budget, strict=True) if not allowed), - ), + from litellm.proxy.auth import personal_owner_adapter + + return await personal_owner_adapter.personal_key_owner_teams_with_status( + valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only ) @@ -6341,31 +5058,10 @@ async def personal_key_team_visible_models( user_api_key_cache: UserApiKeyCache | None, proxy_logging_obj: ProxyLogging | None, ) -> tuple[str, ...]: - teams: Final = await _personal_key_listing_teams( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - return tuple( - [ - model - for model in (models if teams else ()) - if ( - prisma_client is not None - and user_api_key_cache is not None - and proxy_logging_obj is not None - and await _loaded_teams_allow_model( - model=model, - teams=teams, - valid_token=user_api_key_dict, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - ] + from litellm.proxy.auth import personal_model_listing_adapter + + return await personal_model_listing_adapter.personal_key_team_visible_models( + models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj ) @@ -6378,40 +5074,11 @@ async def can_personal_key_call_model_via_teams( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> Literal[True]: - requested_models: Final = (model,) if isinstance(model, str) else tuple(model) - teams: Final = ( - await _load_personal_key_teams( - user_object=user_object, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if user_object is not None - else () + from litellm.proxy.auth import personal_model_request_adapter + + return await personal_model_request_adapter.can_personal_key_call_model_via_teams( + model, user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj ) - for requested_model in requested_models: - if user_object is None or not await _loaded_teams_allow_model( - model=requested_model, - teams=teams, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ): - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=requested_model), - internal_message=( - f"Personal key model access is derived from the owner's teams. User={valid_token.user_id}, " - f"teams={user_object.teams if user_object is not None else None}, " - f"model={requested_model}" - ), - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) - return True def _search_tool_names_from_object_permission( @@ -6448,11 +5115,7 @@ def _can_object_call_search_tool( ProxyException if access is denied """ # Empty list means all search tools are allowed - if not allowed_search_tools: - return True - - # Check if the search tool is in the allowlist - if search_tool_name in allowed_search_tools: + if access_policy.decide(access_policy.Search(search_tool_name, tuple(allowed_search_tools), False)).denial is None: return True # Access denied @@ -6531,6 +5194,8 @@ def _search_tool_denied(object_type: GrantLayer, message: str) -> ProxyException async def _strict_team_object(load_team_object: TeamObjectLoader) -> LiteLLM_TeamTable | None: try: return await load_team_object() + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # an unresolved team grants nothing under deny-by-default verbose_proxy_logger.debug("Team lookup failed under search_tool_deny_by_default: %s", e) return None @@ -6549,6 +5214,8 @@ async def _strict_user_object( user_id_upsert=False, parent_otel_span=valid_token.parent_otel_span, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # an unresolved user grants nothing under deny-by-default verbose_proxy_logger.debug("User lookup failed under search_tool_deny_by_default: %s", e) return None @@ -6599,7 +5266,14 @@ def check_search_tool_grants(search_tool_name: str, grants: SearchToolGrants) -> for layer in ("key", "team", "user"): grant = grants.grants.get(layer) if layer in grants.strict_layers: - if grant is None or search_tool_name not in (grant.search_tools or ()): + if ( + access_policy.decide( + access_policy.Search( + search_tool_name, tuple(grant.search_tools or ()) if grant is not None else (), True + ) + ).denial + is not None + ): raise _search_tool_denied( layer, f"{layer.capitalize()} not allowed to access search tool: {search_tool_name}. " @@ -6702,22 +5376,51 @@ async def is_valid_fallback_model( _MASKED_KEY_NAME_RE: Final = re.compile(r"sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?") -def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: - """ - Throttle an over-budget key instead of blocking it, when the key opted in - via `throttle_on_budget_exceeded` and a global percentage is configured. +def send_token_budget_alert( + valid_token: UserAPIKeyAuth, + spend: float, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + call_info: Final = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj is not None else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="token_budget", + user_info=call_info, + ) + ) - Records the percentage on the request-scoped `budget_throttle_pct` so the - rate limiter scales the key's TPM/RPM down to it; the persistent limits are - left untouched so the throttle never compounds across requests. Returns True - when the key was throttled (caller skips raising), False when it should still - be hard-blocked. - """ - pct: Final = budget_throttle_percentage() - if pct is None or not should_throttle_budget_exceeded(valid_token): - return False - valid_token.budget_throttle_pct = pct - return True + +def key_budget_exceeded_error( + valid_token: UserAPIKeyAuth, spend: float, max_budget: float +) -> litellm.BudgetExceededError: + # This message is returned to the caller, and key_name has no enforced + # shape (a direct DB write bypasses abbreviate_api_key), so echo it only + # when it still looks masked and fall back to the alias otherwise. + key_label: Final = valid_token.key_alias or "key" + key_descriptor: Final = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) + else key_label + ) + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {max_budget}", + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) async def virtual_key_max_budget_check( @@ -6725,79 +5428,9 @@ async def virtual_key_max_budget_check( proxy_logging_obj: ProxyLogging, user_obj: LiteLLM_UserTable | None = None, ) -> None: - """ - Raises: - BudgetExceededError if the token is over it's max budget. - Triggers a budget alert if the token is over it's max budget. + from litellm.proxy.auth import common_budget_adapter - """ - if valid_token.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - fallback_spend: Final = valid_token.spend or 0.0 - counter_key: Final = key_counter(str(valid_token.token)) - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - spend: Final = await get_current_spend( - counter_key=counter_key, - fallback_spend=fallback_spend, - max_budget=valid_token.max_budget, - ) - - #################################### - # collect information for alerting # - #################################### - - user_email = None - # Check if the token has any user id information - if user_obj is not None: - user_email = user_obj.user_email - - call_info: Final = CallInfo( - token=valid_token.token, - spend=spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - organization_id=valid_token.org_id, - user_email=user_email, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="token_budget", - user_info=call_info, - ) - ) - - #################################### - # collect information for alerting # - #################################### - - # Defense-in-depth (GHSA-2rv4-xv66-fpjg): spend >= NaN is always False, - # so a NaN max_budget would silently disable enforcement. Treat a - # non-finite max_budget as "no configured limit" rather than as a bypass. - if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: - if _apply_budget_exceeded_throttle(valid_token): - return - # This message is returned to the caller, and key_name has no enforced - # shape (a direct DB write bypasses abbreviate_api_key), so echo it only - # when it still looks masked and fall back to the alias otherwise. - key_label: Final = valid_token.key_alias or "key" - key_descriptor: Final = ( - f"{key_label} ({valid_token.key_name})" - if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) - else key_label - ) - raise litellm.BudgetExceededError( - current_cost=spend, - max_budget=valid_token.max_budget, - message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}", - entity_type=Litellm_EntityType.KEY.value, - entity_id=valid_token.token, - ) + return await common_budget_adapter.virtual_key_max_budget_check(valid_token, proxy_logging_obj, user_obj) _virtual_key_max_budget_check: Final = virtual_key_max_budget_check @@ -6806,45 +5439,9 @@ _virtual_key_max_budget_check: Final = virtual_key_max_budget_check async def _virtual_key_multi_budget_check( valid_token: UserAPIKeyAuth, ): - """ - Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded. + from litellm.proxy.auth import budget_windows_adapter - Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}. - Using budget_duration (not list index) keeps counters stable when windows are reordered - or removed during a key update. - - Note: counters are not seeded from DB on Redis cold-start. After a Redis flush, - per-window spend resets to zero within the current window period. This is an acceptable - trade-off: the DB stores reset_at timestamps but not per-window accumulated spend. - """ - if not valid_token.budget_limits: - return - - from litellm.proxy.proxy_server import get_current_spend - - for window in valid_token.budget_limits: - w: dict = window if isinstance(window, dict) else window.model_dump() - counter_key = key_window_counter(str(valid_token.token), str(w["budget_duration"])) - window_spend = await get_current_spend( - counter_key=counter_key, - fallback_spend=0.0, - max_budget=w["max_budget"], - window_entity_type="Key", - window_entity_id=valid_token.token, - window_duration=str(w["budget_duration"]), - window_start=get_budget_window_start(w), - ) - if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: - raise litellm.BudgetExceededError( - current_cost=window_spend, - max_budget=w["max_budget"], - message=( - f"ExceededBudget: Key over {w['budget_duration']} budget. " - f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" - ), - entity_type=Litellm_EntityType.KEY.value, - entity_id=valid_token.token, - ) + return await budget_windows_adapter.virtual_key_multi_budget_check(valid_token) async def virtual_key_soft_budget_check( @@ -6852,164 +5449,58 @@ async def virtual_key_soft_budget_check( proxy_logging_obj: ProxyLogging, user_obj: LiteLLM_UserTable | None = None, ) -> None: - """ - Triggers a budget alert if the token is over it's soft budget. + from litellm.proxy.auth import common_budget_adapter - """ + return await common_budget_adapter.virtual_key_soft_budget_check(valid_token, proxy_logging_obj, user_obj) - if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: - verbose_proxy_logger.debug( - "Crossed Soft Budget for token %s, spend %s, soft_budget %s", - valid_token.token, - valid_token.spend, - valid_token.soft_budget, - ) - call_info: Final = CallInfo( - token=valid_token.token, - spend=valid_token.spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=user_obj.user_email if user_obj else None, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) +def send_soft_budget_alert( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + verbose_proxy_logger.debug( + "Crossed Soft Budget for token %s, spend %s, soft_budget %s", + valid_token.token, + valid_token.spend, + valid_token.soft_budget, + ) + call_info: Final = CallInfo( + token=valid_token.token, + spend=valid_token.spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, ) + ) _virtual_key_soft_budget_check: Final = virtual_key_soft_budget_check -def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: - """Parse emails from a list or comma-separated string.""" - if isinstance(raw, list): - return [e.strip() for e in raw if isinstance(e, str) and e.strip()] - elif isinstance(raw, str): - return [e.strip() for e in raw.split(",") if e.strip()] - return [] - - -def _normalize_alert_emails( - cfg: Mapping[str, str | Sequence[object] | None] | None, -) -> dict[str, list[str]]: - """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. - - Values may legitimately arrive as list, comma-separated string, or None - from YAML/metadata; _parse_email_list tolerates all three. - """ - if not cfg: - return {} - return {k: _parse_email_list(v) for k, v in cfg.items()} - - -def _merge_budget_alert_email_configs( - global_cfg: Mapping[str, str | Sequence[object] | None] | None, - per_key_cfg: Mapping[str, str | Sequence[object] | None] | None, -) -> dict[str, list[str]] | None: - """ - Per-threshold additive merge: each threshold's recipient list is the union - of global + per-key entries (deduped, global-first ordering). Missing - thresholds on one side are inherited from the other. - """ - global_cfg_normalized: Final = _normalize_alert_emails(global_cfg) - per_key_cfg_normalized: Final = _normalize_alert_emails(per_key_cfg) - if not global_cfg_normalized and not per_key_cfg_normalized: - return None - thresholds: Final = set(global_cfg_normalized) | set(per_key_cfg_normalized) - return { - t: list(dict.fromkeys(global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []))) for t in thresholds - } - - async def virtual_key_max_budget_alert_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, user_obj: LiteLLM_UserTable | None = None, ) -> None: - """ - Triggers a budget alert if the token has reached EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - (default 80%) of its max budget. - This is a warning alert before the token actually exceeds the max budget. + from litellm.proxy.auth import budget_alert_adapter - """ - - if valid_token.max_budget is not None and valid_token.spend is not None and valid_token.spend > 0: - owner_email: Final = user_obj.user_email if user_obj else None - alert_email_config: Final[dict[str, list[str]] | None] = _merge_budget_alert_email_configs( - global_cfg=litellm.default_key_max_budget_alert_emails, - per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"), - ) - - if isinstance(alert_email_config, dict) and alert_email_config: - # New path: only create task if spend has crossed the lowest threshold - min_pct: Final = min( - (int(k) for k in alert_email_config if k.isdigit()), - default=None, - ) - if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0): - return - - call_info = CallInfo( - token=valid_token.token, - spend=valid_token.spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=owner_email, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - max_budget_alert_emails=alert_email_config, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, - ) - ) - else: - # Old path: existing single 80% threshold — completely unchanged - alert_threshold: Final = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - - if valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget: - verbose_proxy_logger.debug( - "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", - valid_token.token, - valid_token.spend, - valid_token.max_budget, - alert_threshold, - ) - call_info = CallInfo( - token=valid_token.token, - spend=valid_token.spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=owner_email, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - ) - - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, - ) - ) + return await budget_alert_adapter.virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=user_obj, + ) _virtual_key_max_budget_alert_check: Final = virtual_key_max_budget_alert_check @@ -7019,26 +5510,6 @@ TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY: Final = "team_member_max_budget_alert_e _TEAM_MEMBER_ALERT_CONFIG_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -def _is_valid_alert_threshold_pct(pct: str) -> bool: - return pct.isdigit() and len(pct) <= 3 and 1 <= int(pct) <= 100 - - -def _alert_recipients(raw: object) -> Sequence[str] | None: - if isinstance(raw, (str, Sequence)): - return _parse_email_list(raw) - return None - - -def _valid_alert_threshold_config(raw_config: object) -> Mapping[str, str | Sequence[object] | None] | None: - try: - config: Final = _TEAM_MEMBER_ALERT_CONFIG_ADAPTER.validate_python(raw_config) - except ValidationError: - return None - return MappingProxyType( - {pct: _alert_recipients(emails) for pct, emails in config.items() if _is_valid_alert_threshold_pct(pct)} - ) - - def team_member_max_budget_alert_check( team_id: str, team_alias: str | None, @@ -7050,27 +5521,19 @@ def team_member_max_budget_alert_check( spend: float, max_budget: float, ) -> None: - raw_config: Final = team_metadata.get(TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY) if team_metadata else None - alert_email_config: Final = _merge_budget_alert_email_configs( - global_cfg=None, per_key_cfg=_valid_alert_threshold_config(raw_config) - ) - if not alert_email_config or spend <= 0: - return - min_pct: Final = min(int(pct) for pct in alert_email_config) - if spend < max_budget * (min_pct / 100.0): - return - call_info: Final = CallInfo( - spend=spend, - max_budget=max_budget, - user_id=user_id, + from litellm.proxy.auth import budget_alert_adapter + + return budget_alert_adapter.team_member_max_budget_alert_check( team_id=team_id, team_alias=team_alias, + team_metadata=team_metadata, organization_id=organization_id, + user_id=user_id, user_email=user_email, - event_group=Litellm_EntityType.TEAM_MEMBER, - max_budget_alert_emails=alert_email_config, + proxy_logging_obj=proxy_logging_obj, + spend=spend, + max_budget=max_budget, ) - asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=call_info)) _team_member_max_budget_alert_check: Final = team_member_max_budget_alert_check @@ -7086,79 +5549,18 @@ async def _check_team_member_budget( team_membership: LiteLLM_TeamMembership | None = None, team_membership_loaded: bool = False, ): - """Check if team member is over their max budget within the team.""" - if ( - team_object is not None - and team_object.team_id is not None - and valid_token is not None - and valid_token.user_id is not None - ): - if not team_membership_loaded: - team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - loaded_membership = team_membership + from litellm.proxy.auth import common_budget_lookups - # Per-member override wins; otherwise fall back to the team-level - # default configured via team.metadata["team_member_budget_id"]. - member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None - now: Final = get_utc_datetime() - default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") - team_default_budget: Final = ( - await get_team_member_default_budget( - budget_id=default_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - if (member_budget_row is None or member_budget_row.max_budget is None) - and isinstance(default_budget_id, str) - else None - ) - team_member_budget: Final = resolve_effective_team_member_budget( - member_budget=member_budget_row, - team_default_budget=team_default_budget, - now=now, - ) - - if team_member_budget is not None: - team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 - - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = await get_current_spend( - counter_key=team_member_counter(valid_token.user_id, team_object.team_id), - fallback_spend=team_member_spend, - max_budget=team_member_budget, - ) - - if not math.isfinite(team_member_budget): - return - - team_member_max_budget_alert_check( - team_id=team_object.team_id, - team_alias=team_object.team_alias, - team_metadata=team_object.metadata, - organization_id=team_object.organization_id, - user_id=valid_token.user_id, - user_email=user_object.user_email if user_object is not None else None, - proxy_logging_obj=proxy_logging_obj, - spend=team_member_spend, - max_budget=team_member_budget, - ) - - if team_member_spend >= team_member_budget: - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{team_object.team_id}", - ) + return await common_budget_lookups.check_team_member_budget( + team_object, + user_object, + valid_token, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + ) async def check_team_member_model_access( @@ -7173,54 +5575,20 @@ async def check_team_member_model_access( team_membership_loaded: bool = False, key_model_aliases: Mapping[str, str] | None = None, ) -> None: - """ - Check if a team member's per-member model scope allows access to the requested model. + from litellm.proxy.auth import member_model_adapter - Only enforced when the member's budget table has a non-empty allowed_models list. - If allowed_models is empty or absent, the team-level models list applies (no extra restriction). - """ - if valid_token.user_id is None or team_object.team_id is None: - return - - if not team_membership_loaded: - team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - loaded_membership = team_membership - - if ( - loaded_membership is None - or loaded_membership.litellm_budget_table is None - or not loaded_membership.litellm_budget_table.allowed_models - ): - return # no per-member restriction — inherit team-level check - - member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models - try: - can_object_call_model( - model=model, - llm_router=llm_router, - models=member_allowed_models, - object_type="team", - team_id=team_object.team_id, - key_model_aliases=key_model_aliases, - ) - except ProxyException: - internal_message: Final = ( - f"Team member not allowed to access model. User={valid_token.user_id}, Team={team_object.team_id}, " - f"Model={model}. Allowed member models = {member_allowed_models}" - ) - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model), - internal_message=internal_message, - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) + return await member_model_adapter.check_team_member_model_access( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + key_model_aliases, + ) _check_team_member_model_access: Final = check_team_member_model_access @@ -7231,49 +5599,9 @@ async def team_max_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ) -> None: - """ - Check if the team is over it's max budget. + from litellm.proxy.auth import common_budget_adapter - Raises: - BudgetExceededError if the team is over it's max budget. - Triggers a budget alert if the team is over it's max budget. - """ - if team_object is not None and team_object.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - spend: Final = await get_current_spend( - counter_key=team_counter(team_object.team_id), - fallback_spend=team_object.spend or 0.0, - max_budget=team_object.max_budget, - ) - - if math.isfinite(team_object.max_budget) and spend > team_object.max_budget: - if valid_token: - call_info: Final = CallInfo( - token=valid_token.token, - spend=spend, - max_budget=team_object.max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.TEAM, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="team_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=spend, - max_budget=team_object.max_budget, - message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", - entity_type=Litellm_EntityType.TEAM.value, - entity_id=team_object.team_id, - ) + return await common_budget_adapter.team_max_budget_check(team_object, valid_token, proxy_logging_obj) _team_max_budget_check: Final = team_max_budget_check @@ -7282,41 +5610,9 @@ _team_max_budget_check: Final = team_max_budget_check async def _team_multi_budget_check( team_object: LiteLLM_TeamTable | None, ): - """ - Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded. + from litellm.proxy.auth import budget_windows_adapter - Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}. - Using budget_duration (not list index) keeps counters stable when windows are reordered - or removed during a team update. - """ - if team_object is None or not team_object.budget_limits: - return - - from litellm.proxy.proxy_server import get_current_spend - - for window in team_object.budget_limits: - w: dict = window if isinstance(window, dict) else window.model_dump() - counter_key = team_window_counter(team_object.team_id, str(w["budget_duration"])) - window_spend = await get_current_spend( - counter_key=counter_key, - fallback_spend=0.0, - max_budget=w["max_budget"], - window_entity_type="Team", - window_entity_id=team_object.team_id, - window_duration=str(w["budget_duration"]), - window_start=get_budget_window_start(w), - ) - if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: - raise litellm.BudgetExceededError( - current_cost=window_spend, - max_budget=w["max_budget"], - message=( - f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " - f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" - ), - entity_type=Litellm_EntityType.TEAM.value, - entity_id=team_object.team_id, - ) + return await budget_windows_adapter.team_multi_budget_check(team_object) async def _team_soft_budget_check( @@ -7324,70 +5620,9 @@ async def _team_soft_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ): - """ - Triggers a budget alert if the team is over it's soft budget. - """ - if ( - team_object is not None - and team_object.soft_budget is not None - and team_object.spend is not None - and team_object.spend >= team_object.soft_budget - ): - verbose_proxy_logger.debug( - "Crossed Soft Budget for team %s, spend %s, soft_budget %s", - team_object.team_id, - team_object.spend, - team_object.soft_budget, - ) - if valid_token: - # Extract alert emails from team metadata - alert_emails: list[str] | None = None - if team_object.metadata is not None and isinstance(team_object.metadata, dict): - soft_budget_alert_emails: Final = team_object.metadata.get("soft_budget_alerting_emails") - if soft_budget_alert_emails is not None: - if isinstance(soft_budget_alert_emails, list): - alert_emails = [ - email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip() - ] - elif isinstance(soft_budget_alert_emails, str): - # Handle comma-separated string - alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()] - # Filter out empty strings - if alert_emails: - alert_emails = [email for email in alert_emails if email] - else: - alert_emails = None + from litellm.proxy.auth import common_budget_adapter - # Only send team soft budget alerts if alert_emails are configured - # Team soft budget alerts are sent via metadata.soft_budget_alerting_emails, not global alerting - if alert_emails is None or len(alert_emails) == 0: - verbose_proxy_logger.debug( - "Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails", - team_object.team_id, - ) - return - - call_info: Final = CallInfo( - token=valid_token.token, - spend=team_object.spend, - max_budget=team_object.max_budget, - soft_budget=team_object.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=None, # Team-level alert, no specific user email - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.TEAM, - alert_emails=alert_emails, - ) - - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) - ) + return await common_budget_adapter.team_soft_budget_check(team_object, valid_token, proxy_logging_obj) async def _project_max_budget_check( @@ -7395,57 +5630,9 @@ async def _project_max_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ): - """ - Check if the project is over its max budget. + from litellm.proxy.auth import common_budget_adapter - Raises: - BudgetExceededError if the project is over its max budget. - Triggers a budget alert if the project is over its max budget. - """ - if project_object is None: - return - - max_budget = None - if project_object.litellm_budget_table is not None: - max_budget = project_object.litellm_budget_table.max_budget - - if max_budget is None or not math.isfinite(max_budget): - return - - from litellm.proxy.proxy_server import get_current_spend - - project_spend: Final = await get_current_spend( - counter_key=project_counter(project_object.project_id), - fallback_spend=project_object.spend or 0.0, - max_budget=max_budget, - ) - - if project_spend >= max_budget: - if valid_token: - call_info: Final = CallInfo( - token=valid_token.token, - spend=project_spend, - max_budget=max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.PROJECT, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="project_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=project_spend, - max_budget=max_budget, - message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}", - entity_type=Litellm_EntityType.PROJECT.value, - entity_id=project_object.project_id, - ) + return await common_budget_adapter.project_max_budget_check(project_object, valid_token, proxy_logging_obj) async def _project_soft_budget_check( @@ -7453,43 +5640,9 @@ async def _project_soft_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ): - """ - Triggers a budget alert if the project is over its soft budget. + from litellm.proxy.auth import common_budget_adapter - Mirrors _team_soft_budget_check() pattern. - """ - if project_object is None: - return - - soft_budget = None - if project_object.litellm_budget_table is not None: - soft_budget = project_object.litellm_budget_table.soft_budget - - if soft_budget is not None and project_object.spend is not None and project_object.spend >= soft_budget: - verbose_proxy_logger.debug( - "Crossed Soft Budget for project %s, spend %s, soft_budget %s", - project_object.project_id, - project_object.spend, - soft_budget, - ) - if valid_token: - call_info: Final = CallInfo( - token=valid_token.token, - spend=project_object.spend, - max_budget=None, - soft_budget=soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.PROJECT, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) - ) + return await common_budget_adapter.project_soft_budget_check(project_object, valid_token, proxy_logging_obj) @with_service_target(AUTH_OBJECTS_TARGET) @@ -7506,6 +5659,9 @@ async def get_project_object( Returns LiteLLM_ProjectTableCachedObj or None if not found. """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: return None @@ -7519,9 +5675,14 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row: Final = await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique( - where={"project_id": project_id}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(prisma_client) + project_row: Final = ( + await native.one("project", [project_id], LiteLLM_ProjectTableCachedObj) + if native is not None + else await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique( + where={"project_id": project_id}, + include={"litellm_budget_table": True}, + ) ) if project_row is None: return None @@ -7531,6 +5692,7 @@ async def get_project_object( # Cache with TTL following _cache_management_object pattern project_obj.last_refreshed_at = time.time() await _cache_management_object( + fill_generation=fill_generation, key=cache_key, value=project_obj, user_api_key_cache=user_api_key_cache, @@ -7564,97 +5726,12 @@ async def organization_max_budget_check( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: - """ - Check if the organization is over its max budget. + from litellm.proxy.auth import common_budget_lookups - This function checks the organization budget using: - 1. First, tries to use valid_token.org_id (if key has organization_id set) - 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) - - This ensures organization budget checks work even when keys don't have organization_id - set directly, as long as their team belongs to an organization. - - Raises: - BudgetExceededError if the organization is over its max budget. - Triggers a budget alert if the organization is over its max budget. - """ - if valid_token is None or prisma_client is None: - return - - # Determine organization_id: first try from token, then fallback to team - org_id: str | None = None - if valid_token.org_id is not None: - org_id = valid_token.org_id - elif team_object is not None and team_object.organization_id is not None: - org_id = team_object.organization_id - - # If no organization_id found, skip the check - if org_id is None: - return - - # Get organization object with budget table - use get_org_object so it can be mocked in tests - try: - org_table: Final = await get_org_object( - org_id=org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - include_budget_table=True, - ) - except Exception: - # If organization lookup fails, skip the check - return - - if org_table is None: - return - - carry_organization_budget_state(valid_token=valid_token, org_table=org_table) - - # Get max_budget from organization's budget table - org_max_budget: float | None = None - if org_table.litellm_budget_table is not None: - org_max_budget = org_table.litellm_budget_table.max_budget - - if org_max_budget is None: - return - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - from litellm.proxy.proxy_server import get_current_spend - - org_spend: Final = await get_current_spend( - counter_key=org_counter(org_id), - fallback_spend=org_table.spend or 0.0, - max_budget=org_max_budget, + return await common_budget_lookups.organization_max_budget_check( + valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj ) - # Check if organization spend exceeds max budget - if math.isfinite(org_max_budget) and org_spend >= org_max_budget: - # Trigger budget alert - call_info: Final = CallInfo( - token=valid_token.token, - spend=org_spend, - max_budget=org_max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=org_id, - event_group=Litellm_EntityType.ORGANIZATION, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="organization_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=org_spend, - max_budget=org_max_budget, - message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", - entity_type=Litellm_EntityType.ORGANIZATION.value, - entity_id=org_id, - ) - _organization_max_budget_check: Final = organization_max_budget_check @@ -7691,97 +5768,24 @@ async def tag_max_budget_check_for_tags( proxy_logging_obj: ProxyLogging, valid_token: UserAPIKeyAuth | None, ) -> None: - if prisma_client is None or not tags: - return + from litellm.proxy.auth import common_budget_lookups - tag_objects: Final = await get_tag_objects_batch( - tag_names=tags, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + return await common_budget_lookups.tag_max_budget_check_for_tags( + tags, prisma_client, user_api_key_cache, proxy_logging_obj, valid_token ) - # Check budget for each tag - for tag_name in tags: - tag_object = tag_objects.get(tag_name) - if tag_object is None: - continue - - # Check if tag has budget limits - if tag_object.litellm_budget_table is not None and tag_object.litellm_budget_table.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - tag_spend = await get_current_spend( - counter_key=tag_counter(tag_name), - fallback_spend=tag_object.spend or 0.0, - max_budget=tag_object.litellm_budget_table.max_budget, - fallback_authoritative=True, - ) - if tag_spend <= tag_object.litellm_budget_table.max_budget: - continue - raise litellm.BudgetExceededError( - current_cost=tag_spend, - max_budget=tag_object.litellm_budget_table.max_budget, - message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}", - entity_type=Litellm_EntityType.TAG.value, - entity_id=tag_name, - ) - async def _model_access_group_max_budget_check( matched_model_access_groups: Sequence[str], prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, ) -> None: - """Block the request when a model access group that authorized it is over its max budget. + from litellm.proxy.auth import common_budget_lookups - Only the groups auth already matched are charged and therefore only they are checked, so a - request that no budgeted group authorized costs nothing here. - - Like the tag check this is a plain read with no reservation, so concurrent requests can - overshoot the ceiling slightly. The reservation counters are the precise path; this one covers - the ``disable_budget_reservation`` case. - - The ceiling is exclusive, unlike the tag check it otherwise mirrors: a pool whose recorded - spend has reached ``max_budget`` has nothing left to give, so the next request is refused. - Keys and organizations already draw the line there. A non-positive budget means no budget, - matching what the reservation path treats as unbudgeted. - - Raises: - BudgetExceededError if a matched group is over its max budget. - """ - if prisma_client is None or not matched_model_access_groups: - return - - budgets: Final = await get_model_access_group_budgets_batch( - access_group_names=matched_model_access_groups, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + return await common_budget_lookups.model_access_group_max_budget_check( + matched_model_access_groups, prisma_client, user_api_key_cache ) - from litellm.proxy.proxy_server import get_current_spend - - for group in matched_model_access_groups: - budget = budgets.get(group) - if budget is None or budget.max_budget is None or budget.max_budget <= 0: - continue - - group_spend = await get_current_spend( - counter_key=model_access_group_counter(group), - fallback_spend=budget.spend, - max_budget=budget.max_budget, - fallback_authoritative=True, - ) - if group_spend < budget.max_budget: - continue - raise litellm.BudgetExceededError( - current_cost=group_spend, - max_budget=budget.max_budget, - message=f"Budget has been exceeded! Model access group={group} Current cost: {group_spend}, Max budget: {budget.max_budget}", - entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, - entity_id=group, - ) - def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: """ @@ -7849,6 +5853,8 @@ def _model_custom_llm_provider_matches_wildcard_pattern(model: str, allowed_mode """ try: stripped_model, custom_llm_provider, _, _ = get_llm_provider(model=model) + except NativeAuthFailure: + raise except Exception: return False @@ -7898,13 +5904,18 @@ def _strict_grant_layers( The identities that must each grant an object under a deny-by-default policy: a virtual key and its team, a keyless team member's team, or a keyless user's own grant. The master key and dashboard sessions need none """ - if not deny_by_default or valid_token is None or not _is_strict_grant_identity(valid_token): - return frozenset() - virtual_key: Final = valid_token.via_virtual_key and not valid_token.is_session_token - has_team: Final = team_object is not None or valid_token.team_id is not None - if not virtual_key and not has_team: - return frozenset(("user",)) - return frozenset(layer for layer, required in (("key", virtual_key), ("team", has_team)) if required) + decision: Final = access_policy.decide( + access_policy.GrantLayers( + deny_by_default=deny_by_default, + identity_present=valid_token is not None, + master_key=valid_token is not None and valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS, + dashboard=valid_token is not None and valid_token.team_id == UI_TEAM_ID, + virtual_key=valid_token is not None and valid_token.via_virtual_key, + session_token=valid_token is not None and valid_token.is_session_token, + team_present=team_object is not None or (valid_token is not None and valid_token.team_id is not None), + ) + ) + return frozenset(decision.layers) async def _identity_grants( @@ -8053,7 +6064,21 @@ def _legacy_requested_vector_store_ids( def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, LiteLLM_ObjectPermissionTable | None]]) -> bool: - return any(grant is not None and bool(grant.vector_stores) for _, grant in grants) + return ( + access_policy.decide( + access_policy.VectorRequest( + client_present=True, + strict=False, + ids_read=True, + unreadable=True, + grants=tuple( + tuple(grant.vector_stores) if grant is not None and grant.vector_stores is not None else None + for _, grant in grants + ), + ) + ).denial + is not None + ) def _request_drop_policies( @@ -8069,7 +6094,16 @@ def _require_vector_store_grant( vector_store_ids_to_run: Sequence[str], object_permission: _VectorStorePermissionsRow | None, ) -> None: - if object_permission is None or not object_permission.vector_stores: + decision: Final = access_policy.decide( + access_policy.VectorStores( + requested=tuple(map(_provider_vector_store_id, vector_store_ids_to_run)), + allowed=None + if object_permission is None or object_permission.vector_stores is None + else tuple(map(_provider_vector_store_id, object_permission.vector_stores)), + strict=True, + ) + ) + if decision.denial == "MissingVectorGrant": raise ProxyException( message=f"{object_type.capitalize()} not allowed to access vector store. Tried to access {vector_store_ids_to_run[0]}. vector_store_deny_by_default is enabled and the {object_type} has no vector store grants", type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type), @@ -8152,28 +6186,35 @@ async def vector_store_access_check( ######################################################### # Get the vector store the user is trying to access ######################################################### - if prisma_client is None: + facts: Final = access_policy.VectorRequest(prisma_client is not None, deny_by_default) + initial: Final = access_policy.decide(facts) + if initial.read is None: verbose_proxy_logger.debug("Prisma client not found, skipping vector store access check") return True + assert prisma_client is not None typed_request_body: Final = _typed_request_body(request_body) requested_ids: Final = ( _strict_requested_vector_store_ids(typed_request_body) - if deny_by_default + if initial.read == "StrictVectorIds" else _legacy_requested_vector_store_ids( typed_request_body, lambda: _request_drop_policies(typed_request_body, llm_router, valid_token) ) ) - if isinstance(requested_ids, _UnreadableVectorStoreIds): + ids_plan: Final = access_policy.decide( + replace(facts, ids_read=True, unreadable=isinstance(requested_ids, _UnreadableVectorStoreIds)) + ) + if ids_plan.read == "LegacyVectorGrants": if _restricts_vector_stores( await _identity_grants(valid_token, team_object, None, prisma_client, user_api_key_cache) ): raise _malformed_vector_store_ids() return True + assert not isinstance(requested_ids, _UnreadableVectorStoreIds) rag_vector_store_id: Final = _get_rag_query_vector_store_id(typed_request_body) rag_ids: Final = (rag_vector_store_id,) if rag_vector_store_id is not None else () vector_store_ids_to_run: Final = tuple(dict.fromkeys((*requested_ids, *rag_ids))) - if not vector_store_ids_to_run: + if access_policy.decide(replace(facts, ids_read=True, requested_count=len(vector_store_ids_to_run))).read is None: return True ######################################################### @@ -8213,24 +6254,20 @@ def _can_object_call_vector_stores( """ Raises ProxyException if the object (key, team, org) cannot access the specific vector store. """ - if object_permissions is None: - return True - - if object_permissions.vector_stores is None: - return True - - # If length is 0, then the object has access to all vector stores. - if len(object_permissions.vector_stores) == 0: - return True - - allowed_stores: Final = frozenset(map(_provider_vector_store_id, object_permissions.vector_stores)) - for vector_store_id in vector_store_ids_to_run: - if _provider_vector_store_id(vector_store_id) not in allowed_stores: - raise ProxyException( - message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}", - type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type), - param="vector_store", - code=status.HTTP_401_UNAUTHORIZED, - ) + grants: Final = object_permissions.vector_stores if object_permissions is not None else None + decision: Final = access_policy.decide( + access_policy.VectorStores( + requested=tuple(map(_provider_vector_store_id, vector_store_ids_to_run)), + allowed=tuple(map(_provider_vector_store_id, grants)) if grants is not None else None, + strict=False, + ) + ) + if decision.rejected_index is not None: + raise ProxyException( + message=f"User not allowed to access vector store. Tried to access {vector_store_ids_to_run[decision.rejected_index]}. Only allowed to access {grants}", + type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type), + param="vector_store", + code=status.HTTP_401_UNAUTHORIZED, + ) return True diff --git a/litellm/proxy/auth/auth_checks_organization.py b/litellm/proxy/auth/auth_checks_organization.py index 6820ded0eba..47e2ef6efea 100644 --- a/litellm/proxy/auth/auth_checks_organization.py +++ b/litellm/proxy/auth/auth_checks_organization.py @@ -5,8 +5,6 @@ Auth Checks for Organizations from collections.abc import Awaitable, Callable from typing import Final -from fastapi import status - from litellm.proxy._types import * @@ -15,96 +13,11 @@ def organization_role_based_access_check( user_object: LiteLLM_UserTable | None, route: str, ): - """ - Role based access control checks only run if a user is part of an Organization + from litellm.proxy.auth import org_route_adapter - Organization Checks: - ONLY RUN IF user_object.organization_memberships is not None - - 1. Only Proxy Admins can access /organization/new - 2. IF route is a LiteLLMRoutes.org_admin_only_routes, then check if user is an Org Admin for that organization - - """ - - if user_object is None: - return - - passed_organization_id: Final[str | None] = request_body.get("organization_id", None) - - if route == "/organization/new": - if user_object.user_role != LitellmUserRoles.PROXY_ADMIN.value: - raise ProxyException( - message=f"Only proxy admins can create new organizations. You are {user_object.user_role}", - type=ProxyErrorTypes.auth_error.value, - param="user_role", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if user_object.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return - - # Checks if route is an Org Admin Only Route - if route in LiteLLMRoutes.org_admin_only_routes.value: - ( - _user_organizations, - _user_organization_role_mapping, - ) = get_user_organization_info(user_object) - - if user_object.organization_memberships is None: - raise ProxyException( - message=f"Tried to access route={route} but you are not a member of any organization. Please contact the proxy admin to request access.", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if passed_organization_id is None: - raise ProxyException( - message="Passed organization_id is None, please pass an organization_id in your request", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - user_role: Final[LitellmUserRoles | None] = _user_organization_role_mapping.get(passed_organization_id) - if user_role is None: - raise ProxyException( - message=f"You do not have a role within the selected organization. Passed organization_id: {passed_organization_id}. Please contact the organization admin to request access.", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if user_role != LitellmUserRoles.ORG_ADMIN.value: - raise ProxyException( - message=f"You do not have the required role to perform {route} in Organization {passed_organization_id}. Your role is {user_role} in Organization {passed_organization_id}", - type=ProxyErrorTypes.auth_error.value, - param="user_role", - code=status.HTTP_401_UNAUTHORIZED, - ) - elif route == "/team/new": - # if user is part of multiple teams, then they need to specify the organization_id - ( - _user_organizations, - _user_organization_role_mapping, - ) = get_user_organization_info(user_object) - if user_object.organization_memberships is not None and len(user_object.organization_memberships) > 0: - if passed_organization_id is None: - raise ProxyException( - message=f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {_user_organizations}", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - _user_role_in_passed_org: Final = _user_organization_role_mapping.get(passed_organization_id) - if _user_role_in_passed_org != LitellmUserRoles.ORG_ADMIN.value: - raise ProxyException( - message=f"You do not have the required role to call {route}. Your role is {_user_role_in_passed_org} in Organization {passed_organization_id}", - type=ProxyErrorTypes.auth_error.value, - param="user_role", - code=status.HTTP_401_UNAUTHORIZED, - ) + return org_route_adapter.organization_role_based_access_check( + request_body=request_body, user_object=user_object, route=route + ) def get_user_organization_info( @@ -137,40 +50,9 @@ def user_is_org_admin( request_data: dict, user_object: LiteLLM_UserTable | None = None, ) -> bool: - """ - Helper function to check if user is an org admin for all of the passed organizations. + from litellm.proxy.auth import org_route_adapter - Checks both: - - `organization_id` (singular string) — legacy callers - - `organizations` (list of strings) — used by /user/new - """ - if user_object is None: - return False - - if user_object.organization_memberships is None: - return False - - # Collect candidate org IDs from both fields - candidate_org_ids: Final[list[str]] = [] - singular: Final = request_data.get("organization_id", None) - if singular is not None: - candidate_org_ids.append(singular) - orgs_list: Final = request_data.get("organizations", None) - if isinstance(orgs_list, list): - candidate_org_ids.extend(orgs_list) - - if not candidate_org_ids: - return False - - # Build set of orgs where user is admin - admin_org_ids: Final = { - _membership.organization_id - for _membership in user_object.organization_memberships - if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value and _membership.organization_id is not None - } - - # User must be admin of ALL requested orgs, not just any one - return all(org_id in admin_org_ids for org_id in candidate_org_ids) + return org_route_adapter.user_is_org_admin(request_data=request_data, user_object=user_object) _user_is_org_admin: Final = user_is_org_admin diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py new file mode 100644 index 00000000000..957844269d0 --- /dev/null +++ b/litellm/proxy/auth/auth_denial.py @@ -0,0 +1,152 @@ +from typing import Final, assert_never + +from fastapi import HTTPException, status + +import litellm +from litellm._auth.decision import ( + Denial, + ExpiredKey, + KeyBlocked, + KeyBudgetExceeded, + KeyWithoutToken, + MalformedKey, + MasterKeyOnlyRoute, + MissingKey, + ModelAccessDenied, + ModelBudgetExceeded, + NoDatabase, + NotAVirtualKey, + OwnerDeactivated, + TeamMemberBudgetExceeded, + UserAuthDisabled, +) +from litellm._auth.failure import NativeAuthFailure +from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_KEY_ERROR_MESSAGE +from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import key_budget_exceeded_error +from litellm.proxy.auth.auth_utils import abbreviate_api_key +from litellm.proxy.auth.model_access_adapter import raise_model_denial +from litellm.types.proxy.auth.user_api_key_auth import is_jwt + +_JWT_AUTH_DISABLED_HINT: Final = ( + " This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a" + " virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate" + " with JWTs." +) + + +def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException: + masked: Final = f"{token[:4]}****{token[-4:]}" if len(token) > 8 else "****" + hint: Final = _JWT_AUTH_DISABLED_HINT if not jwt_auth_enabled and is_jwt(token=token) else "" + error: Final = HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={masked}, expected to start with 'sk-'.{hint}", + ) + # Log routing classifies this 401 by where it was raised, never by its message text. + setattr(error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) + return error + + +def denial_exception( + denial: Denial, + token: str | None, + route: str, + jwt_auth_enabled: bool, + key: UserAPIKeyAuth | None, +) -> ProxyException | HTTPException | litellm.BudgetExceededError: + match denial: + case ModelAccessDenied(decision=model_decision): + raise_model_denial(model_decision, model_decision.model or "") + raise NativeAuthFailure("native model denial contained an allow") + case MissingKey(): + return ProxyException( + message="Authentication Error, No api key passed in.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case MalformedKey(): + return ProxyException( + message="Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case UserAuthDisabled(): + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="'allow_user_auth' not set or set to False", + ) + case ExpiredKey(expires=expires, now=now, key_label=key_label): + return ProxyException( + message=f"Authentication Error - Expired Key. Key Expiry time {expires} and current time {now}", + type=ProxyErrorTypes.expired_key, + code=status.HTTP_401_UNAUTHORIZED, + param=abbreviate_api_key(api_key=key_label), + ) + case MasterKeyOnlyRoute(): + return ProxyException( + message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case NoDatabase(): + return ProxyException( + message="No connected db.", + type=ProxyErrorTypes.no_db_connection, + code=400, + param=None, + ) + case NotAVirtualKey(): + return _not_a_virtual_key(token or "", jwt_auth_enabled) + case KeyBlocked(): + return ProxyException( + message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case OwnerDeactivated(user_id=user_id): + return ProxyException( + message=( + f"Authentication Error, User={user_id} has been deactivated via SCIM. " + "Keys owned by this user cannot be used." + ), + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=spend, max_budget=max_budget): + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=( + f"Budget has been exceeded! TeamMember={user_id}:{team_id} " + f"Current cost: {spend}, Max budget: {max_budget}" + ), + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{user_id}:{team_id}", + ) + case KeyBudgetExceeded(spend=spend, max_budget=max_budget): + return key_budget_exceeded_error(key or UserAPIKeyAuth(), spend, max_budget) + case KeyWithoutToken(): + return HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key, no token associated" + ) + case ModelBudgetExceeded( + message=message, + current_cost=current_cost, + max_budget=max_budget, + entity_type=entity_type, + entity_id=entity_id, + ): + return litellm.BudgetExceededError( + current_cost=current_cost, + max_budget=max_budget, + message=message, + entity_type=entity_type, + entity_id=entity_id, + ) + case _: + assert_never(denial) diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index 6257d461cf4..c7f0bea3f17 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -4,6 +4,7 @@ the readers and the fallback, so enforcement never depends on this running.""" from __future__ import annotations +import json import time from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass @@ -12,6 +13,7 @@ from typing import Final, Literal, TypeAlias from pydantic import BaseModel, TypeAdapter, ValidationError +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._internal_context import service_target from litellm._logging import verbose_proxy_logger from litellm.caching.redis_batch import Batch @@ -32,6 +34,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.db.db_span import db_span from litellm.proxy.utils import PrismaClient from litellm.redis.stores.auth_cache import AuthCacheStore +from litellm.repositories.native_auth_repository import native_auth_repository _RowKind: TypeAlias = Literal["user_row", "team_row", "membership_row", "organization_row", "project_row"] @@ -220,6 +223,7 @@ async def _fill_from_redis( """On the request's native pipeline when one is open; a failed read is a miss, like ``async_batch_get_cache``.""" if not entries: return + generation: Final = shared.current_generation() before_read() try: with service_target(AUTH_OBJECTS_TARGET): @@ -229,13 +233,16 @@ async def _fill_from_redis( return for entry, value in ((entry, found.get(entry.cache_key)) for entry in entries): if value is not None: - shared.set_local( - entry.cache_key, value, cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl + shared.set_if_generation( + entry.cache_key, + value, + generation, + cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl, ) def _validate_row( - row_value: object, model_type: type[BaseModel], row: _RowKind, refreshed_at: float + row_value: object, model_type: type[BaseModel], row: _RowKind, refreshed_at: float, *, native: bool = False ) -> BaseModel | None: if row_value is None: return None @@ -246,6 +253,8 @@ def _validate_row( return model_type.model_validate(stamped) return model_type.model_validate(columns) except ValidationError as e: + if native: + raise NativeAuthRuntimeError(f"Invalid native auth {row} projection") from e verbose_proxy_logger.warning("auth prefetch: %s did not validate as %s: %s", row, model_type.__name__, e) return None @@ -253,6 +262,20 @@ def _validate_row( async def _fetch_rows( refs: AuthObjectRefs, kinds: frozenset[_RowKind], prisma_client: PrismaClient ) -> Mapping[str, object]: + native: Final = native_auth_repository(prisma_client) + if native is not None: + refs_json: Final = json.dumps( + [ + refs.user_id if "user_row" in kinds else None, + refs.team_id if kinds & _TEAM_BOUND_ROWS else None, + refs.membership_user_id if "membership_row" in kinds else None, + refs.organization_id if "organization_row" in kinds else None, + refs.project_id if "project_row" in kinds else None, + ] + ) + async with db_span("prefetch_auth_objects", AUTH_OBJECTS_TARGET): + rows: Final = await native.read("prefetch", [refs_json]) + return rows[0] if rows else _NO_ROWS async with db_span("prefetch_auth_objects", AUTH_OBJECTS_TARGET): row: Final[object] = await prisma_client.db.query_first( # pyright: ignore[reportAny] # prisma types query_first as Any _SQL, @@ -265,12 +288,19 @@ async def _fetch_rows( return _RowValues.validate_python(row) if row is not None else _NO_ROWS -async def _write_back(entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: UserApiKeyCache) -> None: +async def _write_back( + entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: UserApiKeyCache, generation: int | None = None +) -> None: payloads: Final = tuple( (entry.cache_key, CacheCodec.serialize(value, model_type=entry.model_type), entry.ttl) for entry, value in entries ) shared: Final = cache.shared + if generation is not None: + assert shared is not None + for cache_key, payload, ttl in payloads: + shared.set_if_generation(cache_key, payload, generation, ttl) + return if shared is not None: with service_target(AUTH_OBJECTS_TARGET): await shared.write_back(payloads, batch=Batch.PRE_CALL) @@ -284,17 +314,22 @@ async def _fill_from_db( ) -> None: if not entries: return + native: Final = native_auth_repository(prisma_client) is not None + generation: Final = cache.current_generation() if native else None model_for: Final[Mapping[_RowKind, type[BaseModel]]] = MappingProxyType( {entry.row: entry.model_type for entry in entries} ) rows: Final = await _fetch_rows(refs, frozenset(model_for), prisma_client) refreshed_at: Final = time.time() objects: Final[Mapping[_RowKind, BaseModel | None]] = MappingProxyType( - {row: _validate_row(rows.get(row), model_type, row, refreshed_at) for row, model_type in model_for.items()} + { + row: _validate_row(rows.get(row), model_type, row, refreshed_at, native=native) + for row, model_type in model_for.items() + } ) writes: Final = tuple((entry, value) for entry in entries if (value := objects[entry.row]) is not None) if writes: - await _write_back(writes, cache) + await _write_back(writes, cache, generation) async def prefetch_auth_objects( @@ -313,6 +348,8 @@ async def prefetch_auth_objects( if prisma_client is None: return await _fill_from_db(refs, _missing_in_memory(missing, user_api_key_cache), user_api_key_cache, prisma_client) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # warm-up only; the getters enforce and fail closed on their own verbose_proxy_logger.warning("auth prefetch skipped, falling back to per-object lookups: %s", e) @@ -338,6 +375,7 @@ async def prefetch_identity_keys( missing: Final = tuple(key for key in dict.fromkeys(cache_keys) if shared.get_local(key) is None) if not missing: return + generation: Final = shared.current_generation() before_read() with service_target(AUTH_OBJECTS_TARGET): found: Final = await shared.read_missing(sorted(missing), batch=Batch.PRE_CALL) @@ -347,4 +385,4 @@ async def prefetch_identity_keys( return for key, value in ((key, found.get(key)) for key in missing): if value is not None: - shared.set_local(key, value, _identity_memory_ttl(value, management_ttl)) + shared.set_if_generation(key, value, generation, _identity_memory_ttl(value, management_ttl)) diff --git a/litellm/proxy/auth/budget_alert_adapter.py b/litellm/proxy/auth/budget_alert_adapter.py new file mode 100644 index 00000000000..2537af0f8b8 --- /dev/null +++ b/litellm/proxy/auth/budget_alert_adapter.py @@ -0,0 +1,178 @@ +import asyncio +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm._auth import budget_alerts +from litellm._auth.budget_alerts import ConfigFacts, Decision, Facts, Failure, ThresholdFacts +from litellm._logging import verbose_proxy_logger +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.proxy._types import CallInfo, Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.common_budget_adapter import metadata +from litellm.proxy.utils import ProxyLogging + +_CONFIG: Final = TypeAdapter(Mapping[str, object]) +_OBJECT: Final = TypeAdapter(object) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def _threshold(key: str, raw: object) -> ThresholdFacts: + text: Final = raw if isinstance(raw, str) else None + items: Final = ( + tuple(item if isinstance(item, str) else None for item in _ITEMS.validate_python(raw)) + if isinstance(raw, list) + else None + ) + initial: Final = ThresholdFacts(key=key, digit=key.isdigit(), text=text, items=items) + try: + integer: Final = int(key) + except ValueError as error: + return replace(initial, integer_failure=Failure("ValueError", str(error))) + parsed: Final = replace(initial, digits=str(integer)) + try: + return replace(parsed, ratio=integer / 100.0) + except OverflowError as error: + return replace(parsed, ratio_failure=Failure("OverflowError", str(error))) + + +def _config(raw: object) -> tuple[tuple[ThresholdFacts, ...], Failure | None]: + if not raw: + return (), None + try: + value: Final = _CONFIG.validate_python(raw) + except ValidationError: + return (), Failure("AttributeError", f"'{type(raw).__name__}' object has no attribute 'items'") + return tuple(_threshold(key, value) for key, value in value.items()), None + + +def read_alert_facts(facts: Facts, global_config: object, local_config: object) -> Facts: + decision: Final = budget_alerts.plan(facts) + if not decision.read_config: + return facts + global_entries, global_failure = _config(global_config) + local_entries, local_failure = _config(local_config) + loaded: Final = replace(facts, config=ConfigFacts(global_entries, local_entries, global_failure, local_failure)) + order: Final = budget_alerts.plan(loaded) + if not order.read_order: + return loaded + assert loaded.config is not None + return replace(loaded, config=replace(loaded.config, order=tuple(set(order.global_keys) | set(order.local_keys)))) + + +def read_key_budget_alert_facts(valid_token: UserAPIKeyAuth) -> Facts: + facts: Final = Facts( + False, valid_token.spend, valid_token.max_budget, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + ) + if not budget_alerts.plan(facts).read_config: + return facts + return read_alert_facts( + facts, + _OBJECT.validate_python(litellm.default_key_max_budget_alert_emails), # pyright: ignore[reportUnknownMemberType] # validate the legacy global config snapshot + metadata(valid_token.metadata).get("max_budget_alert_emails"), # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # legacy metadata is validated at the boundary + ) + + +def read_member_budget_alert_facts( + spend: float, + max_budget: float, + team_metadata: Mapping[str, object] | None, +) -> Facts: + facts: Final = Facts(True, spend, max_budget, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) + return read_alert_facts( + facts, None, team_metadata.get("team_member_max_budget_alert_emails") if team_metadata else None + ) + + +def raise_alert_failure(decision: Decision) -> None: + failure: Final = decision.failure + if failure is None: + return + match failure.kind: + case "AttributeError": + raise AttributeError(failure.message) + case "OverflowError": + raise OverflowError(failure.message) + case _: + raise ValueError(failure.message) + + +def schedule_budget_alert(decision: Decision, call_info: CallInfo, proxy_logging_obj: ProxyLogging) -> None: + raise_alert_failure(decision) + if decision.emit: + projected: Final = ( + call_info + if decision.thresholds is None + else call_info.model_copy(update={"max_budget_alert_emails": alert_emails(decision)}) + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=projected)) + + +def alert_emails( + decision: Decision, +) -> dict[str, list[str]] | None: # mutable-ok: CallInfo requires the legacy dict of lists projection + return None if decision.thresholds is None else {entry.key: list(entry.recipients) for entry in decision.thresholds} + + +async def virtual_key_max_budget_alert_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + decision: Final = budget_alerts.plan(read_key_budget_alert_facts(valid_token)) + raise_alert_failure(decision) + if not decision.emit: + return + if decision.thresholds is None: + assert valid_token.max_budget is not None + verbose_proxy_logger.debug( + "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", + valid_token.token, + valid_token.spend, + valid_token.max_budget, + valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, + ) + call_info: Final = CallInfo( + token=valid_token.token, + spend=valid_token.spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + schedule_budget_alert(decision, call_info, proxy_logging_obj) + + +def team_member_max_budget_alert_check( + team_id: str, + team_alias: str | None, + team_metadata: Mapping[str, object] | None, + organization_id: str | None, + user_id: str, + user_email: str | None, + proxy_logging_obj: ProxyLogging, + spend: float, + max_budget: float, +) -> None: + decision: Final = budget_alerts.plan(read_member_budget_alert_facts(spend, max_budget, team_metadata)) + raise_alert_failure(decision) + if not decision.emit: + return + call_info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + user_id=user_id, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_email=user_email, + event_group=Litellm_EntityType.TEAM_MEMBER, + ) + schedule_budget_alert(decision, call_info, proxy_logging_obj) diff --git a/litellm/proxy/auth/budget_context_adapter.py b/litellm/proxy/auth/budget_context_adapter.py new file mode 100644 index 00000000000..fb9413cbf66 --- /dev/null +++ b/litellm/proxy/auth/budget_context_adapter.py @@ -0,0 +1,385 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import ( + Final, + Literal, + Protocol, + assert_never, + cast, # noqa: TID251 # preserve identity of validated legacy JSON mappings +) + +from opentelemetry.trace import Span + +import litellm +from litellm._auth import budget_context as core +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class TeamModelLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + +class EndUserModelLimiter(Protocol): + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: ... + + +def mapping_value(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return cast(Mapping[str, object], value) # cast-ok: preserve the original JSON map's identity + return None + + +def mapping_count(value: object) -> int | None: + mapping: Final = mapping_value(value) + return len(mapping) if mapping is not None else None + + +def budget_fields(row: LiteLLM_BudgetTable | None) -> core.BudgetFields | None: + if row is None: + return None + return core.BudgetFields( + row.max_budget, + row.tpm_limit, + row.rpm_limit, + row.tpd_limit, + mapping_count(getattr(row, "model_max_budget", None)), + ) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy budget table declares an unparameterized dict + + +async def check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: TeamModelLimiter, + models: Sequence[str], +) -> None: + budgets: Final = valid_token.team_model_max_budget + selected: Final = core.model_checks( + core.ModelCheckFacts("Team", valid_token.team_id, len(budgets) if budgets is not None else 0, len(models)) + ) + for index in selected: + assert valid_token.team_id is not None and budgets is not None + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=budgets, + key_model_max_budget=mapping_value(getattr(valid_token, "model_max_budget", None)), + model=models[index], + ) + + +async def check_end_user_model_budget( + end_user_id: str | None, + end_user_model_max_budget: Mapping[str, object] | None, + model_max_budget_limiter: EndUserModelLimiter, + models: Sequence[str], + team_id: str | None, + *, + scope: Literal["EndUser", "CustomEndUser"] = "EndUser", +) -> None: + selected: Final = core.model_checks( + core.ModelCheckFacts( + scope, + end_user_id, + len(end_user_model_max_budget) if end_user_model_max_budget is not None else 0, + len(models), + ) + ) + for index in selected: + assert end_user_id is not None and end_user_model_max_budget is not None + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=end_user_model_max_budget, + model=models[index], + team_id=team_id, + ) + + +@dataclass(frozen=True, slots=True) +class EndUserIO: + prisma_client: PrismaClient | None + cache: UserApiKeyCache + span: Span | None + + +async def _end_user_row( + facts: core.EndUserFacts, + io: EndUserIO, + own: LiteLLM_EndUserTable | None, + key: LiteLLM_BudgetTable | None = None, + proxy: LiteLLM_BudgetTable | None = None, +) -> LiteLLM_BudgetTable | None: + from litellm.proxy.auth.auth_checks import get_default_end_user_budget + + decision: Final = core.end_user_source(facts) + match decision: + case "Own": + assert own is not None + return own.litellm_budget_table + case "Key": + return key + case "Global": + return proxy + case "None": + return None + case "ReadKey": + row: Final = await get_default_end_user_budget( + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + budget_id=facts.key_budget_id, + ) + return await _end_user_row( + replace(facts, key_read=True, key_budget=budget_fields(row)), io, own, row, proxy + ) + case "ReadGlobalConfig": + return await _end_user_row( + replace(facts, global_config_read=True, global_budget_id=litellm.max_end_user_budget_id), + io, + own, + key, + proxy, + ) + case "ReadGlobal": + row_global: Final = await get_default_end_user_budget( + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + budget_id=facts.global_budget_id, + ) + return await _end_user_row( + replace(facts, global_read=True, global_budget=budget_fields(row_global)), io, own, key, row_global + ) + case _: + assert_never(decision) + + +async def resolve_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + key_end_user_budget_id: str | None, + parent_otel_span: Span | None = None, +) -> LiteLLM_BudgetTable | None: + return await _end_user_row( + core.EndUserFacts("Default", False, None, prisma_client is not None, key_end_user_budget_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + None, + ) + + +async def end_user_model_max_budget( + end_user_id: str | None, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, +) -> Mapping[str, object] | None: + row: Final = await _end_user_row( + core.EndUserFacts( + "Models", end_user_object is not None, end_user_id, prisma_client is not None, key_end_user_budget_id + ), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + end_user_object, + ) + return getattr(row, "model_max_budget", None) if row is not None else None # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy row field + + +async def apply_key_end_user_default_budget_to_token( + valid_token: UserAPIKeyAuth, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + keep_token_limits: bool, +) -> None: + row: Final = await _end_user_row( + core.EndUserFacts("Token", end_user_object is not None, valid_token.end_user_id, True, key_end_user_budget_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + end_user_object, + ) + current: Final = core.BudgetFields( + valid_token.end_user_max_budget, + valid_token.end_user_tpm_limit, + valid_token.end_user_rpm_limit, + valid_token.end_user_tpd_limit, + mapping_count(valid_token.end_user_model_max_budget), + ) + for field in core.token_assignments(core.TokenDefaultsFacts(current, keep_token_limits, row is not None)): + assert row is not None + match field: + case "MaxBudget": + valid_token.end_user_max_budget = ( # rebind-ok: request-local native projection + row.max_budget + ) + case "Tpm": + valid_token.end_user_tpm_limit = ( # rebind-ok: request-local native projection + row.tpm_limit + ) + case "Rpm": + valid_token.end_user_rpm_limit = ( # rebind-ok: request-local native projection + row.rpm_limit + ) + case "Tpd": + valid_token.end_user_tpd_limit = ( # rebind-ok: request-local native projection + row.tpd_limit + ) + case "Models": + valid_token.end_user_model_max_budget = row.model_max_budget # pyright: ignore[reportUnknownMemberType] # preserve the validated row dict and its identity # rebind-ok: native-selected request-local budget assignment + case _: + assert_never(field) + + +@dataclass(frozen=True, slots=True) +class MemberIO: + prisma_client: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging | None + + +@dataclass(frozen=True, slots=True) +class MemberRead: + member: LiteLLM_TeamMembership | None = None + default: LiteLLM_BudgetTable | None = None + error: Exception | None = None + + +async def _read_member(facts: core.MemberFacts, io: MemberIO) -> MemberRead: + from litellm.proxy.auth.auth_checks import get_team_membership + + assert facts.user_id is not None and facts.team_id is not None + try: + row: Final = await get_team_membership( + user_id=facts.user_id, + team_id=facts.team_id, + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # only the database/cache read is caught; native failures rethrow + return MemberRead(error=error) + return MemberRead(member=row) + + +async def _read_member_default(facts: core.MemberFacts, io: MemberIO) -> MemberRead: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget + + assert facts.default_id is not None + try: + row: Final = await get_team_member_default_budget( + budget_id=facts.default_id, + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + raise_on_lookup_error=True, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # only the database/cache read is caught; native failures rethrow + return MemberRead(error=error) + return MemberRead(default=row) + + +async def _member_model_budget( + token: UserAPIKeyAuth, + facts: core.MemberFacts, + io: MemberIO, + reading: MemberRead, +) -> None: + decision: Final = core.member_source(facts) + match decision: + case "Done": + return + case "RaiseLookup": + assert reading.error is not None + raise reading.error + case "SkipLookup": + verbose_proxy_logger.warning( + "Skipping team-member model budget lookup after a lookup error", exc_info=reading.error + ) + return + case "ReadMembership": + member_read: Final = await _read_member(facts, io) + member: Final = member_read.member + member_row: Final = member.litellm_budget_table if member is not None else None + snapshot: Final = ( + core.MembershipSnapshot(member.budget_id, mapping_count(getattr(member_row, "model_max_budget", None))) + if member is not None + else None + ) + await _member_model_budget( + token, + replace( + facts, + membership_read=True, + membership=snapshot, + lookup_failure="Membership" if member_read.error is not None else None, + ), + io, + member_read, + ) + case "ReadDefault": + default_read: Final = await _read_member_default(facts, io) + await _member_model_budget( + token, + replace( + facts, + default_read=True, + default_budget=budget_fields(default_read.default), + lookup_failure="Default" if default_read.error is not None else None, + ), + io, + default_read, + ) + case "Member": + assert reading.member is not None + token.team_member_model_max_budget = mapping_value( # rebind-ok: request-local native projection + getattr(reading.member.litellm_budget_table, "model_max_budget", None) + ) + case "Default": + assert reading.default is not None + token.team_member_model_max_budget = mapping_value( # rebind-ok: request-local native projection + getattr(reading.default, "model_max_budget", None) + ) + case "Clear": + token.team_member_model_max_budget = None # rebind-ok: request-local native projection + case _: + assert_never(decision) + + +async def load_team_member_default_model_budget( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + skip_budget_checks: bool, + proxy_logging_obj: ProxyLogging | None, +) -> None: + metadata: Final = mapping_value(getattr(team_object, "metadata", None)) + default_id: Final = metadata.get("team_member_budget_id") if metadata is not None else None + facts: Final = core.MemberFacts( + user_id=user_api_key_auth_obj.user_id if isinstance(user_api_key_auth_obj.user_id, str) else None, + team_id=user_api_key_auth_obj.team_id if isinstance(user_api_key_auth_obj.team_id, str) else None, + prisma_present=prisma_client is not None, + default_id=default_id if isinstance(default_id, str) else None, + skip_budget_checks=skip_budget_checks, + ) + await _member_model_budget( + user_api_key_auth_obj, facts, MemberIO(prisma_client, user_api_key_cache, proxy_logging_obj), MemberRead() + ) diff --git a/litellm/proxy/auth/budget_lookup_adapter.py b/litellm/proxy/auth/budget_lookup_adapter.py new file mode 100644 index 00000000000..1f12b0eb053 --- /dev/null +++ b/litellm/proxy/auth/budget_lookup_adapter.py @@ -0,0 +1,144 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import datetime +from typing import Final, assert_never + +from litellm._auth import budget_lookup as core +from litellm._auth.failure import NativeAuthFailure +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import ( + LiteLLM_OrganizationTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.common_budget_adapter import metadata +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@dataclass(frozen=True, slots=True) +class LookupIO: + token: UserAPIKeyAuth | None + team: LiteLLM_TeamTable | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + + +@dataclass(frozen=True, slots=True) +class OrganizationReading: + org_id: str + row: LiteLLM_OrganizationTable + + +async def organization( + facts: core.OrganizationFacts, io: LookupIO, org_id: str | None = None +) -> OrganizationReading | None: + from litellm.proxy.auth.auth_checks import get_org_object + + action: Final = core.organization(facts) + match action: + case "Skip": + return None + case "ReadTokenId": + assert io.token is not None + token_id: Final = io.token.org_id + return await organization(replace(facts, token_id=token_id is not None), io, token_id) + case "ReadTeamId": + assert io.team is not None + team_id: Final = io.team.organization_id + return await organization(replace(facts, team_id=team_id is not None), io, team_id) + case "ReadOrganization": + assert org_id is not None and io.prisma is not None + try: + row: Final = await get_org_object( + org_id=org_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + include_budget_table=True, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # pass the raw lookup failure to native Phase 1 policy + return await organization(replace(facts, lookup_failed=True), io, org_id) + result: Final = core.organization(replace(facts, row_present=row is not None)) + assert result in ("Skip", "UseOrganization") + if result == "Skip": + return None + assert row is not None + return OrganizationReading(org_id, row) + case "UseOrganization": + raise AssertionError("organization row was not retained") + case _: + assert_never(action) + + +@dataclass(frozen=True, slots=True) +class MemberReading: + membership: LiteLLM_TeamMembership | None + row: LiteLLM_BudgetTable | None + now: datetime + metadata: Mapping[str, object] + default_id: object + default: LiteLLM_BudgetTable | None = None + + +async def member( + facts: core.MemberFacts, + io: LookupIO, + membership: LiteLLM_TeamMembership | None, + reading: MemberReading | None = None, +) -> MemberReading | None: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget, get_team_membership, get_utc_datetime + + action: Final = core.member(facts) + match action: + case "Skip": + return None + case "ReadTeamId": + assert io.team is not None + return await member(replace(facts, team_id=io.team.team_id is not None), io, membership) # pyright: ignore[reportUnnecessaryComparison] # Phase 1 accepts constructed cache rows with missing IDs + case "ReadUserId": + assert io.token is not None + return await member(replace(facts, user_id=io.token.user_id is not None), io, membership) + case "ReadMembership": + assert io.token is not None and io.token.user_id is not None and io.team is not None + loaded: Final = await get_team_membership( + user_id=io.token.user_id, + team_id=io.team.team_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + ) + return await member(replace(facts, membership_loaded=True), io, loaded) + case "ReadLimits": + assert io.team is not None + row: Final = membership.litellm_budget_table if membership is not None else None + now: Final = get_utc_datetime() + values: Final = metadata(io.team.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # raw legacy metadata boundary + default_id: Final = values.get("team_member_budget_id") + snapshot: Final = MemberReading(membership, row, now, values, default_id) + return await member( + replace( + facts, + limits_read=True, + maximum_present=row is not None and row.max_budget is not None, + default_id_string=isinstance(default_id, str), + ), + io, + membership, + snapshot, + ) + case "ReadDefault": + assert reading is not None and isinstance(reading.default_id, str) + default: Final = await get_team_member_default_budget( + budget_id=reading.default_id, prisma_client=io.prisma, user_api_key_cache=io.cache + ) + return await member(replace(facts, default_read=True), io, membership, replace(reading, default=default)) + case "UseLimits": + assert reading is not None + return reading + case _: + assert_never(action) diff --git a/litellm/proxy/auth/budget_pricing_adapter.py b/litellm/proxy/auth/budget_pricing_adapter.py new file mode 100644 index 00000000000..863294731d0 --- /dev/null +++ b/litellm/proxy/auth/budget_pricing_adapter.py @@ -0,0 +1,180 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from pydantic import TypeAdapter + +import litellm +from litellm._auth import budget_pricing +from litellm._auth.budget_pricing import DeploymentPrice, ModelPriceFacts, RouteBudgetFacts +from litellm._auth.failure import NativeAuthFailure +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.router_api import RegistryReader +from litellm.types.router import DeploymentTypedDict + +_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT: Final = TypeAdapter(object) + + +def _number(raw: object) -> float | None: + return float(raw) if isinstance(raw, (float, int)) and not isinstance(raw, bool) else None + + +def _deployment_price(deployment: DeploymentTypedDict) -> DeploymentPrice: + info: Final = _MAPPING.validate_python(deployment.get("model_info") or {}) # pyright: ignore[reportUnknownMemberType] # deployment TypedDict contains legacy metadata + model_id: Final = info.get("id") + raw: Final = _OBJECT.validate_python(litellm.model_cost.get(model_id)) if model_id is not None else None # pyright: ignore[reportUnknownMemberType] # project one legacy global price entry + entry: Final = as_str_mapping(raw) or {} + params: Final = _MAPPING.validate_python(deployment["litellm_params"]) + provider: Final = params.get("custom_llm_provider") or str(params.get("model")).partition("/")[0] + query_price: Final = decisions_request_price( + {**entry, "litellm_provider": provider}, decisions_pricing_fields(entry) + ) + return DeploymentPrice( + has_input_price="input_cost_per_token" in entry, + has_output_price="output_cost_per_token" in entry, + input_price=_number(entry.get("input_cost_per_token")), + output_price=_number(entry.get("output_cost_per_token")), + query_price=query_price, + ptu_count=0.0 if info.get("ptu_count") is not None else None, + cost_per_ptu_per_hour=0.0 if info.get("cost_per_ptu_per_hour") is not None else None, + ) + + +def read_model_price_facts(model: str, router: RegistryReader, cached: bool | None = None) -> ModelPriceFacts: + from litellm.proxy.auth.auth_checks import ( + _deployments_served_for, # pyright: ignore[reportPrivateUsage] # shared registry projection + ) + + if cached is not None: + return ModelPriceFacts(cached=cached) + try: + group: Final = router.registry().group(model, include_hidden=True) + facts: Final = ModelPriceFacts( + group_found=group is not None, + input_price=group.input_cost_per_token if group is not None else None, + output_price=group.output_cost_per_token if group is not None else None, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve lookup failure as non-exempt without caching + return ModelPriceFacts(lookup_failed=True) + if not budget_pricing.price(facts).read_deployments: + return facts + try: + deployments: Final = tuple( + _deployment_price(deployment) for deployment in _deployments_served_for(model, router) + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve lookup failure as non-exempt without caching + return ModelPriceFacts(lookup_failed=True) + return replace(facts, deployments=deployments) + + +def _read_cached_model_price(model: str, router: RegistryReader, cache: Mapping[str, bool] | None) -> ModelPriceFacts: + cached: Final = cache.get(model) if isinstance(cache, dict) else None + facts: Final = read_model_price_facts(model, router, cached) + decision: Final = budget_pricing.price(facts) + if isinstance(cache, dict) and decision.cache_value is not None: + cache[model] = decision.cache_value # rebind-ok: native-selected registry cache write + return facts + + +def read_model_cost_facts(model: str, router: RegistryReader) -> ModelPriceFacts: + return _read_cached_model_price(model, router, router.registry().zero_cost()) + + +def is_model_cost_zero(model: str | Sequence[str] | None, llm_router: RegistryReader | None) -> bool: + if model is None or llm_router is None: + return budget_pricing.price(ModelPriceFacts()).zero_cost + names: Final = (model,) if isinstance(model, str) else tuple(model) + cache: Final = llm_router.registry().zero_cost() + return all(budget_pricing.price(_read_cached_model_price(name, llm_router, cache)).zero_cost for name in names) + + +def read_route_budget_facts(route: str, body: Mapping[str, object] | None = None) -> RouteBudgetFacts: + from litellm.proxy.auth.auth_checks import BUDGET_ENFORCED_SIDE_EFFECT_ROUTES, MODEL_DISCOVERY_ROUTES + + method: Final = body.get("method") if body is not None else None + return RouteBudgetFacts( + side_effect_route=route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES, + discovery_route=route in MODEL_DISCOVERY_ROUTES, + llm_route=RouteChecks.is_llm_api_route(route=route), + route=route, + method=method if isinstance(method, str) else None, + body_present=body is not None, + ) + + +def _route_decision(route: str, body: Mapping[str, object] | None = None) -> budget_pricing.RouteBudgetDecision: + facts: Final = read_route_budget_facts(route, body) + decision: Final = budget_pricing.route(facts) + if decision != "ReadPassThrough": + return decision + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + registered: Final = InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route) + return budget_pricing.route(replace(facts, registered_pass_through=registered)) + + +def route_skips_budget_checks(route: str) -> bool: + return _route_decision(route) == "Skip" + + +def read_budget_exemption_facts( + route: str, + model: str | Sequence[str] | None, + llm_router: RegistryReader | None, + *, + request_body: Mapping[str, object] | None = None, +) -> budget_pricing.ExemptionFacts: + names: Final = (model,) if isinstance(model, str) else tuple(model or ()) + initial: Final = budget_pricing.ExemptionFacts( + read_route_budget_facts(route, request_body), + model is not None, + llm_router is not None, + len(names), + ) + return _read_exemption(initial, names, llm_router, None) + + +def _read_exemption( + facts: budget_pricing.ExemptionFacts, + models: tuple[str, ...], + router: RegistryReader | None, + cache: Mapping[str, bool] | None, +) -> budget_pricing.ExemptionFacts: + current = facts # rebind-ok: native-selected IO driver state + selected_cache = cache # rebind-ok: cache handle acquired at the first model read + while True: + decision = budget_pricing.exemption(current) + if decision.read_pass_through: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + registered = InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=current.route.route) + current = replace(current, route=replace(current.route, registered_pass_through=registered)) + continue + if decision.read_price_cache: + assert router is not None + selected_cache = router.registry().zero_cost() + current = replace(current, cache_read=True) + continue + if decision.read_model is None: + return current + assert router is not None + price = _read_cached_model_price(models[decision.read_model], router, selected_cache) + current = replace(current, prices=(*current.prices, price)) + + +def request_skips_budget_checks( + route: str, + model: str | Sequence[str] | None, + llm_router: RegistryReader | None, + *, + request_body: Mapping[str, object] | None = None, +) -> bool: + facts: Final = read_budget_exemption_facts(route, model, llm_router, request_body=request_body) + return budget_pricing.exemption(facts).exempt diff --git a/litellm/proxy/auth/budget_throttle.py b/litellm/proxy/auth/budget_throttle.py index 7122c8edd05..8c06454e3b2 100644 --- a/litellm/proxy/auth/budget_throttle.py +++ b/litellm/proxy/auth/budget_throttle.py @@ -10,39 +10,20 @@ so it never compounds across requests. """ import math -from typing import Final -import litellm from litellm.proxy._types import UserAPIKeyAuth def budget_throttle_percentage() -> float | None: - """ - The global throttle percentage, or None when throttling is disabled / - misconfigured (in which case an over-budget key is hard-blocked, the safe - default). - """ - pct: Final = litellm.budget_exceeded_throttle_percentage - if not isinstance(pct, (int, float)) or isinstance(pct, bool): - return None - if not 0 < pct <= 1: - return None - return float(pct) + from litellm.proxy.auth import key_budget_context_adapter + + return key_budget_context_adapter.budget_throttle_percentage() def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: - """ - True when a key that exceeded its own ``max_budget`` should be throttled - rather than blocked: it opted in, a valid global percentage is set, and the - key has a TPM or RPM limit to scale down. A key with neither limit has - nothing to throttle, so it stays hard-blocked (the safe default) rather than - serving unlimited requests past its budget. - """ - if (valid_token.metadata or {}).get("throttle_on_budget_exceeded") is not True: - return False - if valid_token.tpm_limit is None and valid_token.rpm_limit is None: - return False - return budget_throttle_percentage() is not None + from litellm.proxy.auth import key_budget_context_adapter + + return key_budget_context_adapter.should_throttle_budget_exceeded(valid_token) def throttled_limit(limit: int | None, pct: float | None) -> int | None: diff --git a/litellm/proxy/auth/budget_windows_adapter.py b/litellm/proxy/auth/budget_windows_adapter.py new file mode 100644 index 00000000000..647d5bdb75f --- /dev/null +++ b/litellm/proxy/auth/budget_windows_adapter.py @@ -0,0 +1,101 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Final + +from pydantic import BaseModel, TypeAdapter + +from litellm._auth.common_budgets import BudgetFacts, WindowFacts, window_start +from litellm._auth.failure import NativeAuthFailure +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.common_budget_adapter import CounterRead, raise_budget, read_budget + + +@dataclass(frozen=True, slots=True) +class BudgetWindow: + budget_duration: str + max_budget: float + reset_at: object = None + + +_WINDOW: Final = TypeAdapter(BudgetWindow) + + +def _window(raw: object) -> BudgetWindow: + return _WINDOW.validate_python(raw.model_dump() if isinstance(raw, BaseModel) else raw) + + +def _seconds(duration: str) -> int | None: + try: + return duration_in_seconds(duration) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve invalid-duration skip behavior + return None + + +def _reset(value: object) -> datetime | None: + if isinstance(value, datetime): + return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value + if isinstance(value, str): + try: + parsed: Final = datetime.fromisoformat(value.replace("Z", "+00:00")) + except NativeAuthFailure: + raise + except ValueError: + return None + return _reset(parsed) + return None + + +async def check_windows(windows: Sequence[object], *, entity: str, identity: str | None) -> None: + for raw in windows: + await check_window(raw, entity=entity, identity=identity) + + +async def check_window(raw: object, *, entity: str, identity: str | None) -> None: + window: Final = _window(raw) + seconds: Final = _seconds(window.budget_duration) + reset: Final = _reset(window.reset_at) + start_timestamp: Final = window_start( + WindowFacts(seconds, reset.timestamp() if reset is not None else None, datetime.now(timezone.utc).timestamp()) + ) + start: Final = datetime.fromtimestamp(start_timestamp, timezone.utc) if start_timestamp is not None else None + facts, result = await read_budget( + BudgetFacts("KeyWindow" if entity == "Key" else "TeamWindow", window.max_budget), + CounterRead( + f"spend:{entity.lower()}:{identity}:window:{window.budget_duration}", + 0.0, + window.max_budget, + window_entity_type=entity, + window_entity_id=identity, + window_duration=window.budget_duration, + window_start=start, + ), + ) + if result.denied and facts.spend is not None: + label: Final = "Key" if entity == "Key" else f"Team={identity}" + raise_budget( + facts, + Litellm_EntityType.KEY if entity == "Key" else Litellm_EntityType.TEAM, + identity, + f"ExceededBudget: {label} over {window.budget_duration} budget. Spend=${facts.spend:.4f}, Limit=${window.max_budget:.2f}", + ) + + +async def virtual_key_multi_budget_check(valid_token: UserAPIKeyAuth) -> None: + await check_windows( + TypeAdapter(tuple[object, ...]).validate_python(valid_token.budget_limits or ()), # pyright: ignore[reportUnknownMemberType] # legacy input + entity="Key", + identity=valid_token.token, + ) + + +async def team_multi_budget_check(team_object: LiteLLM_TeamTable | None) -> None: + if team_object is not None: + await check_windows( + TypeAdapter(tuple[object, ...]).validate_python(team_object.budget_limits or ()), # pyright: ignore[reportUnknownMemberType] # validate legacy window dictionaries + entity="Team", + identity=team_object.team_id, + ) diff --git a/litellm/proxy/auth/claude_normalization_adapter.py b/litellm/proxy/auth/claude_normalization_adapter.py new file mode 100644 index 00000000000..25275edfa22 --- /dev/null +++ b/litellm/proxy/auth/claude_normalization_adapter.py @@ -0,0 +1,174 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from itertools import accumulate, repeat +from typing import ( + Final, + assert_never, + cast, # noqa: TID251 # checked Mapping view avoids eager reads of unselected values +) + +from starlette.requests import Request + +import litellm +from litellm._auth.claude_normalization import ( + ClaudeMembershipFacts, + ClaudeNormalizationDecision, + ClaudeNormalizationFacts, + claude_membership, + claude_normalization, +) +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.model_policy_read import require_native_fact +from litellm.proxy.common_utils.http_parsing_utils import rewrite_request_model +from litellm.router import Router +from litellm.router_api import RegistryReader, RouterRegistry + +_NORMALIZED: Final = "litellm.claude_model_normalized" + + +def _alias_entry_exists(value: object, name: str) -> bool | None: + if not isinstance(value, Mapping): + return None + return name in value + + +def _settings_aliases(value: object) -> object: + if not isinstance(value, Mapping): + return None + entries: Final = cast(Mapping[object, object], value) # cast-ok: Mapping checked without coercing unrelated values + return entries.get("model_group_alias") + + +@dataclass(frozen=True, slots=True) +class _MembershipRead: + facts: ClaudeMembershipFacts + registry: RouterRegistry | None = None + + +def read_claude_membership_facts( + name: str, router: RegistryReader | None, team_id: str | None, alias_maps: tuple[object, ...] +) -> ClaudeMembershipFacts: + initial: Final = _MembershipRead(ClaudeMembershipFacts(name, router is not None, len(alias_maps))) + + def advance(state: _MembershipRead, _unused: None) -> _MembershipRead: + facts: Final = state.facts + decision: Final = claude_membership(facts) + match decision.kind: + case "ReadGlobal": + return replace(state, facts=replace(facts, global_alias_exists=name in litellm.model_alias_map)) + case "ReadAlias": + alias_exists: Final = _alias_entry_exists(alias_maps[decision.index], name) + return replace(state, facts=replace(facts, alias_entries=(*facts.alias_entries, alias_exists))) + case "ReadRegistry": + return _MembershipRead(replace(facts, registry_loaded=True), require_native_fact(router).registry()) + case "ReadAliases": + alias_in_registry: Final = name in require_native_fact(state.registry).aliases() + return replace(state, facts=replace(facts, registry_alias_exists=alias_in_registry)) + case "ReadDeployment": + return replace( + state, facts=replace(facts, deployment=require_native_fact(state.registry).has_deployment(name)) + ) + case "ReadNames": + names: Final = tuple( + require_native_fact(state.registry).names(name, kind="route_candidates", team_id=team_id) + ) + return replace(state, facts=replace(facts, names=names)) + case "Complete": + raise NativeAuthRuntimeError("completed Claude membership requested IO") + case _: + assert_never(decision.kind) + + return next( + state.facts + for state in accumulate(repeat(None), advance, initial=initial) + if claude_membership(state.facts).kind == "Complete" + ) + + +@dataclass(frozen=True, slots=True) +class _Source: + data: dict[str, object] # mutable-ok: request-local rewrite effect + token: UserAPIKeyAuth + request: Request | None + router: Router | None = None + alias_maps: tuple[object, ...] = () + + +def _membership( + facts: ClaudeNormalizationFacts, source: _Source, decision: ClaudeNormalizationDecision +) -> ClaudeMembershipFacts: + if decision.index != len(facts.membership): + raise NativeAuthRuntimeError("Claude membership read out of order") + return read_claude_membership_facts( + require_native_fact(decision.name), source.router, source.token.team_id, source.alias_maps + ) + + +async def _normalize(facts: ClaudeNormalizationFacts, source: _Source) -> None: + decision: Final = claude_normalization(facts) + match decision.kind: + case "Complete": + return + case "ReadRouter": + from litellm.proxy.proxy_server import llm_router + + return await _normalize( + replace(facts, router_present=llm_router is not None), replace(source, router=llm_router) + ) + case "ReadNormalized": + return await _normalize( + replace(facts, normalized=require_native_fact(source.request).scope.get(_NORMALIZED) is True), source + ) + case "ReadModel": + from litellm.proxy.auth.user_api_key_auth import ( + _get_model_from_request_context, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared raw request extractor has an unparameterized legacy dict argument + ) + + requested: Final = _get_model_from_request_context( + source.data, facts.route, source.request, source.router, source.token.team_id + ) + return await _normalize( + replace( + facts, + model_loaded=True, + requested=requested if isinstance(requested, str) else None, + ), + source, + ) + case "ReadBodyModel": + body: Final = source.data.get("model") + return await _normalize( + replace(facts, body_loaded=True, body_model=body if isinstance(body, str) else None), source + ) + case "ReadSettings": + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=source.token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + aliases: Final = _settings_aliases(settings) + return await _normalize( + replace(facts, settings_loaded=True), + replace(source, alias_maps=(source.token.aliases, source.token.team_model_aliases, aliases)), # pyright: ignore[reportUnknownMemberType] # aliases is an unparameterized legacy mapping, queried only at native-selected names + ) + case "ReadMembership": + membership: Final = _membership(facts, source, decision) + return await _normalize(replace(facts, membership=(*facts.membership, membership)), source) + case "MarkNormalized": + request: Final = require_native_fact(source.request) + request.scope[_NORMALIZED] = True + return await _normalize(replace(facts, marked=True), source) + case "Rewrite": + rewrite_request_model(source.data, source.request, require_native_fact(decision.target)) + case _: + assert_never(decision.kind) + + +async def normalize_claude_model( + request_data: dict[str, object], # mutable-ok: native-selected request-local rewrite + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + await _normalize(ClaudeNormalizationFacts(route, request is not None), _Source(request_data, valid_token, request)) diff --git a/litellm/proxy/auth/codex_model_adapter.py b/litellm/proxy/auth/codex_model_adapter.py new file mode 100644 index 00000000000..25aa3f5e8ea --- /dev/null +++ b/litellm/proxy/auth/codex_model_adapter.py @@ -0,0 +1,168 @@ +from dataclasses import dataclass, replace +from itertools import accumulate, repeat +from typing import Final, assert_never + +from pydantic import BaseModel + +import litellm +from litellm._auth.codex_models import ( + CandidateField, + CodexCandidate, + CodexGrants, + CodexModelFacts, + CodexModelKind, + codex_model, +) +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.model_helpers import ModelHelperFacts +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.model_helper_adapter import read_model_helper_facts +from litellm.proxy.auth.model_policy_read import require_native_fact +from litellm.proxy.common_utils.codex_model_catalog import bundled_codex_models +from litellm.router_api import RegistryReader +from litellm.types.router import Deployment + + +class _ModelGrants(BaseModel, frozen=True): + models: tuple[str, ...] + team_models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _State: + facts: CodexModelFacts + deployments: tuple[Deployment, ...] = () + + +@native_boundary +def _grants(token: UserAPIKeyAuth) -> CodexGrants: + fields: Final = _ModelGrants.model_validate(token, from_attributes=True) + return CodexGrants(fields.models, fields.team_models) + + +def _candidate(field: CandidateField, previous: CodexCandidate, deployment: Deployment) -> CodexCandidate: + match field: + case "Provider": + return replace(previous, provider_loaded=True, provider=deployment.litellm_params.custom_llm_provider) + case "Model": + return replace(previous, model=deployment.litellm_params.model) + case "Name": + return replace(previous, name=deployment.model_name) + case "PublicName": + return replace(previous, public_loaded=True, public_name=deployment.model_info.team_public_model_name) + case _: + assert_never(field) + + +def read_codex_model_facts( + model: object, + router: RegistryReader | None, + team_id: str | None, + *, + kind: CodexModelKind = "Routes", + token: UserAPIKeyAuth | None = None, + route: str | None = None, +) -> CodexModelFacts: + initial: Final = _State( + CodexModelFacts(kind, model if isinstance(model, str) else None, router is not None, team_id, route=route) + ) + + def advance(state: _State, _tick: None) -> _State: + facts: Final = state.facts + decision: Final = codex_model(facts) + match decision.kind: + case "ReadBundled": + return replace(state, facts=replace(facts, bundled=tuple(bundled_codex_models()))) + case "ReadCost": + priced: Final = require_native_fact(decision.model) in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # raw key presence in the dynamically loaded cost registry + return replace(state, facts=replace(facts, priced=priced)) + case "ReadCandidates": + registry: Final = require_native_fact(router).registry() + loaded: Final = tuple( + registry.deployment(model_id) + for model_id in registry.names( + require_native_fact(decision.model), kind="route_candidates", team_id=team_id + ) + ) + deployments: Final = tuple(deployment for deployment in loaded if deployment is not None) + return _State(replace(facts, candidates=tuple(CodexCandidate() for _ in deployments)), deployments) + case "ReadAlias": + return replace( + state, + facts=replace( + facts, + alias_loaded=True, + alias=require_native_fact(router) + .registry() + .resolve(require_native_fact(decision.model), kind="alias"), + ), + ) + case "ReadCandidate": + candidates: Final = require_native_fact(facts.candidates) + updated: Final = _candidate( + require_native_fact(decision.field), candidates[decision.index], state.deployments[decision.index] + ) + return replace( + state, + facts=replace( + facts, + candidates=tuple( + updated if index == decision.index else value for index, value in enumerate(candidates) + ), + ), + ) + case "ReadGrants": + grant_read: Final = _grants(require_native_fact(token)) + return replace( + state, + facts=( + replace(facts, prefixed_grants=grant_read) + if decision.prefixed + else replace(facts, grants=grant_read) + ), + ) + case "ReadHelper": + selected: Final = require_native_fact(facts.prefixed_grants if decision.prefixed else facts.grants) + grants: Final = (selected.key, selected.team)[decision.index] + scope: Final = read_model_helper_facts( + ModelHelperFacts( + "Helper", + require_native_fact(decision.model), + grants, + team_id=team_id, + router_present=bool(router), + ), + router, + ) + updated_facts: Final = ( + replace(facts, prefixed=(*facts.prefixed, scope)) + if decision.prefixed + else replace(facts, bare=(*facts.bare, scope)) + ) + return replace(state, facts=updated_facts) + case "Complete": + raise NativeAuthRuntimeError("completed Codex model decision requested IO") + case _: + assert_never(decision.kind) + + return next( + state.facts + for state in accumulate(repeat(None), advance, initial=initial) + if codex_model(state.facts).kind == "Complete" + ) + + +def routes_to_litellm(model: object, llm_router: RegistryReader | None, team_id: str | None) -> bool: + return codex_model(read_codex_model_facts(model, llm_router, team_id)).allowed + + +def is_codex_native_litellm_turn(route: str, model: object, llm_router: RegistryReader, team_id: str | None) -> bool: + return codex_model(read_codex_model_facts(model, llm_router, team_id, kind="NativeRoutes", route=route)).allowed + + +def key_and_team_allow(model: str, llm_router: RegistryReader, token: UserAPIKeyAuth) -> bool: + return codex_model(read_codex_model_facts(model, llm_router, token.team_id, kind="Grants", token=token)).allowed + + +def native_chatgpt_model(model: str, llm_router: RegistryReader, token: UserAPIKeyAuth) -> str | None: + return codex_model(read_codex_model_facts(model, llm_router, token.team_id, kind="Chatgpt", token=token)).target diff --git a/litellm/proxy/auth/codex_normalization_adapter.py b/litellm/proxy/auth/codex_normalization_adapter.py new file mode 100644 index 00000000000..e207b512c54 --- /dev/null +++ b/litellm/proxy/auth/codex_normalization_adapter.py @@ -0,0 +1,169 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import ( + Final, + cast, # noqa: TID251 # checked Mapping view preserves raw lookup timing without eager validation +) + +from fastapi import Request + +from litellm._auth.codex_normalization import CodexNormalizationDecision, CodexNormalizationFacts, codex_normalization +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._logging import verbose_proxy_logger +from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_model_adapter import read_codex_model_facts +from litellm.proxy.auth.codex_turn_adapter import read_native_turn_facts +from litellm.proxy.auth.model_policy_read import require_native_fact +from litellm.proxy.auth.route_policy_adapter import read_eu_jev_configuration, route_context +from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body, rewrite_request_model +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + data: dict[str, object] # mutable-ok: request-local rewrite is a selected effect + token: UserAPIKeyAuth + request: Request | None + router: RegistryReader | None + + +def _value(value: object, key: str) -> object: + if not isinstance(value, Mapping): + return None + entries: Final = cast(Mapping[object, object], value) # cast-ok: runtime Mapping shape was checked + return entries.get(key) + + +def _rewrite(facts: CodexNormalizationFacts, source: _Source, target: str | None) -> None: + if source.request is None or facts.model is None or target is None: + raise NativeAuthRuntimeError("Codex rewrite without request model") + if facts.kind == "Alias": + verbose_proxy_logger.debug( + "router_settings.model_group_alias resolved %s -> %s before auth", + facts.model.replace("\r", "").replace("\n", ""), + target.replace("\r", "").replace("\n", ""), + ) + source.request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, facts.model) + rewrite_request_model(source.data, source.request, target) + + +async def _alias_scope(facts: CodexNormalizationFacts, source: _Source, decision: CodexNormalizationDecision) -> None: + from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint + + request: Final = require_native_fact(source.request) + match decision.kind: + case "ReadContext": + return await _alias(replace(facts, context=route_context(facts.route)), source) + case "ReadEuJev": + return await _alias(replace(facts, eu_jev=read_eu_jev_configuration()), source) + case "ReadResolved": + return await _alias( + replace(facts, resolved=request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True), source + ) + case "MarkResolved": + request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True + return await _alias(replace(facts, marked=True), source) + case "ReadModel": + model: Final = source.data.get("model") + return await _alias( + replace(facts, model_loaded=True, model=model if isinstance(model, str) else None), source + ) + case "ReadExplicit": + return await _alias(replace(facts, explicit=request_dispatched_to_pass_through_endpoint(request)), source) + case _: + raise NativeAuthRuntimeError("invalid alias scope read stage") + + +async def _alias(facts: CodexNormalizationFacts, source: _Source) -> None: + from litellm.proxy.auth.auth_utils import request_dispatched_to_provider_pass_through + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + decision: Final = codex_normalization(facts) + if decision.kind == "Complete": + return + request: Final = require_native_fact(source.request) + if decision.kind in ("ReadContext", "ReadEuJev", "ReadResolved", "MarkResolved", "ReadModel", "ReadExplicit"): + return await _alias_scope(facts, source, decision) + match decision.kind: + case "ReadProvider": + return await _alias(replace(facts, provider=request_dispatched_to_provider_pass_through(request)), source) + case "ReadRouting": + routing: Final = read_codex_model_facts(facts.model, source.router, source.token.team_id) + return await _alias(replace(facts, routing=routing), source) + case "ReadBody": + return await _alias( + replace(facts, body_present=await read_raw_json_body(request=request) is not None), source + ) + case "ReadSettings": + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=source.token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + entry: Final = _value(_value(settings, "model_group_alias"), require_native_fact(facts.model)) + mapped_entry: Final[object] = entry + target: Final = _value(mapped_entry, "model") if isinstance(entry, Mapping) else entry + return await _alias( + replace(facts, settings_loaded=True, target=target if isinstance(target, str) else None), source + ) + case "ReadProviderAfter": + return await _alias( + replace(facts, provider_after=request_dispatched_to_provider_pass_through(request)), source + ) + case "ReadTurn": + turn: Final = await read_native_turn_facts( + require_native_fact(facts.model), require_native_fact(source.router), request, source.token + ) + return await _alias(replace(facts, turn=turn), source) + case "Rewrite": + _rewrite(facts, source, decision.target) + case _: + raise NativeAuthRuntimeError("invalid alias normalization stage") + + +async def resolve_router_settings_model_group_alias( + request_data: dict[str, object], # mutable-ok: preserve request-local model mutation + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + from litellm.proxy.proxy_server import llm_router + + await _alias( + CodexNormalizationFacts("Alias", route, request is not None, router_present=llm_router is not None), + _Source(request_data, valid_token, request, llm_router), + ) + + +def _prefix(facts: CodexNormalizationFacts, source: _Source) -> None: + decision: Final = codex_normalization(facts) + match decision.kind: + case "Complete": + return + case "ReadRouter": + from litellm.proxy.proxy_server import llm_router + + return _prefix(replace(facts, router_present=llm_router is not None), replace(source, router=llm_router)) + case "ReadPrefix": + prefix: Final = read_codex_model_facts( + facts.model, source.router, source.token.team_id, kind="Chatgpt", token=source.token + ) + return _prefix(replace(facts, prefix=prefix), source) + case "Rewrite": + _rewrite(facts, source, decision.target) + case _: + raise NativeAuthRuntimeError("invalid Codex prefix normalization stage") + + +def resolve_codex_native_chatgpt_model( + request_data: dict[str, object], # mutable-ok: preserve request-local model mutation + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + model: Final = request_data.get("model") + _prefix( + CodexNormalizationFacts( + "Prefix", route, request is not None, model_loaded=True, model=model if isinstance(model, str) else None + ), + _Source(request_data, valid_token, request, None), + ) diff --git a/litellm/proxy/auth/codex_turn_adapter.py b/litellm/proxy/auth/codex_turn_adapter.py new file mode 100644 index 00000000000..b6121ccddc3 --- /dev/null +++ b/litellm/proxy/auth/codex_turn_adapter.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import hmac +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING, Final + +from fastapi import Request + +from litellm._auth.codex_turn import ( + CodexAccessFacts, + CodexCredentials, + CodexKeyFacts, + CodexTurnFacts, + JwtRead, + codex_access, + codex_jwt_token, + codex_key, + codex_turn, +) +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_model_adapter import read_codex_model_facts +from litellm.proxy.auth.model_policy_read import require_native_fact +from litellm.proxy.auth.route_enable_adapter import read_route_enable_facts +from litellm.proxy.auth.route_policy_adapter import read_virtual_route_facts +from litellm.proxy.auth.route_scope_adapter import read_route_scope_facts +from litellm.router_api import RegistryReader + +if TYPE_CHECKING: + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( + _SharedCache, # pyright: ignore[reportPrivateUsage] # existing account-catalog IO protocol + ) + +RESPONSES_PATH: Final = "/v1/responses" +_KEY_FACTS: Final = CodexKeyFacts() + + +def read_codex_key_facts(token: UserAPIKeyAuth, facts: CodexKeyFacts = _KEY_FACTS) -> CodexKeyFacts: + decision: Final = codex_key(facts) + match decision.kind: + case "Complete": + return facts + case "Failure": + raise TypeError(decision.message) + case "ReadProxy": + return read_codex_key_facts(token, replace(facts, proxy=read_route_scope_facts(RESPONSES_PATH, "Proxy"))) + case "ReadEnabled": + return read_codex_key_facts(token, replace(facts, enabled=read_route_enable_facts(RESPONSES_PATH))) + case "ReadVirtual": + return read_codex_key_facts( + token, replace(facts, virtual_route=read_virtual_route_facts(RESPONSES_PATH, token, None)) + ) + case "ReadRole": + return read_codex_key_facts(token, replace(facts, role_loaded=True, role=token.user_role)) + case "ReadAdmin": + return read_codex_key_facts(token, replace(facts, admin=read_route_scope_facts(RESPONSES_PATH, "Admin"))) + case _: + raise NativeAuthRuntimeError("invalid Codex key read stage") + + +async def read_jwt_result(gateway_token: str, model: str, headers: Mapping[str, str]) -> JwtRead: + from litellm.proxy.auth.handle_jwt import JWTAuthManager + from litellm.proxy.proxy_server import ( + general_settings_view, + jwt_handler, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + try: + _ = await JWTAuthManager.auth_builder( # pyright: ignore[reportUnknownMemberType] # unchanged validator accepts legacy dictionaries; its return value is not interpreted here + api_key=gateway_token, + jwt_handler=jwt_handler, + request_data={"model": model}, + general_settings=dict(general_settings_view()), + route=RESPONSES_PATH, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + request_headers=headers, + request_method="POST", + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # the unchanged JWT validator reports refusals by raising + return "Exception" + return "Returned" + + +@dataclass(frozen=True, slots=True) +class CodexAccess: + key: CodexKeyFacts + jwt_token: str | None = field(repr=False) + request_headers: Mapping[str, str] = field(repr=False) + + @property + def key_admits(self) -> bool: + return codex_key(self.key).allowed + + async def admits(self, model: str) -> bool: + return codex_access(await read_codex_access_facts(self, model)).allowed + + +def read_codex_access(token: UserAPIKeyAuth, request: Request) -> CodexAccess: + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( + _gateway_token, # pyright: ignore[reportPrivateUsage] # raw credential extraction, not auth policy + ) + + selection: Final = codex_jwt_token(claims_present=bool(token.jwt_claims), token_present=bool(token.token)) + key: Final = read_codex_key_facts(token) + return CodexAccess(key, _gateway_token(request) if selection.allowed else None, request.headers) + + +async def read_codex_access_facts(source: CodexAccess, model: str) -> CodexAccessFacts: + facts: Final = CodexAccessFacts(source.key, source.jwt_token is not None) + decision: Final = codex_access(facts) + match decision.kind: + case "Complete": + return facts + case "ReadJwt": + if source.jwt_token is None: + raise NativeAuthRuntimeError("Codex JWT read without credential") + return replace(facts, jwt=await read_jwt_result(source.jwt_token, model, source.request_headers)) + case _: + raise NativeAuthRuntimeError("invalid Codex access read stage") + + +@dataclass(frozen=True, slots=True) +class _Source: + router: RegistryReader | None + token: UserAPIKeyAuth + request: Request | None = None + access: CodexAccess | None = None + caller: str | None = field(default=None, repr=False) + headers: Mapping[str, str] | None = field(default=None, repr=False) + cache: _SharedCache | None = None + gateway_token: str | None = field(default=None, repr=False) + + +async def _read(facts: CodexTurnFacts, source: _Source) -> CodexTurnFacts: + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( + _account_slugs, # pyright: ignore[reportPrivateUsage] # unchanged catalog IO + _chatgpt_bearer, # pyright: ignore[reportPrivateUsage] # raw credential parsing + _forwarded_headers, # pyright: ignore[reportPrivateUsage] # transport header projection + _gateway_token, # pyright: ignore[reportPrivateUsage] # raw credential parsing + _models_etag_caller, # pyright: ignore[reportPrivateUsage] # cache identity projection + codex_shared_cache, + ) + + decision: Final = codex_turn(facts) + match decision.kind: + case "Complete": + return facts + case "ReadCredentials": + if source.request is None: + raise NativeAuthRuntimeError("native Codex turn without request") + bearer: Final = _chatgpt_bearer(source.request) + gateway: Final = _gateway_token(source.request) + credentials: Final = CodexCredentials( + bool(bearer), bool(gateway), hmac.compare_digest(bearer.encode(), gateway.encode()) + ) + return await _read(replace(facts, credentials=credentials), replace(source, gateway_token=gateway)) + case "ReadAccess": + if source.request is None or source.gateway_token is None: + raise NativeAuthRuntimeError("native Codex access without request credential") + access: Final = read_codex_access(source.token, source.request) + caller: Final = _models_etag_caller(source.token, source.gateway_token) + headers: Final = _forwarded_headers(source.request) + cache: Final = codex_shared_cache() + return await _read( + replace(facts, access=CodexAccessFacts(access.key, access.jwt_token is not None)), + replace(source, access=access, caller=caller, headers=headers, cache=cache), + ) + case "ReadRouting": + routing: Final = read_codex_model_facts(facts.model, source.router, source.token.team_id) + return await _read(replace(facts, routing=routing), source) + case "ReadJwt": + return await _read( + replace( + facts, + access=await read_codex_access_facts( + require_native_fact(source.access), require_native_fact(facts.model) + ), + ), + source, + ) + case "ReadCatalog": + if source.headers is None or source.caller is None or source.cache is None: + raise NativeAuthRuntimeError("Codex catalog read without IO context") + slugs: Final = await _account_slugs(source.headers, source.caller, source.cache) + return await _read( + replace(facts, catalog_loaded=True, slugs=tuple(slugs) if slugs is not None else None), source + ) + case _: + raise NativeAuthRuntimeError("invalid Codex turn read stage") + + +async def read_native_turn_facts( + model: str, router: RegistryReader, request: Request, token: UserAPIKeyAuth +) -> CodexTurnFacts: + return await _read(CodexTurnFacts(model, native=True), _Source(router, token, request=request)) + + +async def native_turn_serves_through_litellm( + model: str, llm_router: RegistryReader, request: Request, user_api_key_dict: UserAPIKeyAuth +) -> bool: + return codex_turn(await read_native_turn_facts(model, llm_router, request, user_api_key_dict)).allowed + + +async def serves_through_litellm( + model: object, + llm_router: RegistryReader | None, + user_api_key_dict: UserAPIKeyAuth, + access: CodexAccess, + caller: str, + headers: Mapping[str, str], + cache: _SharedCache, +) -> bool: + facts: Final = CodexTurnFacts( + model if isinstance(model, str) else None, access=CodexAccessFacts(access.key, access.jwt_token is not None) + ) + source: Final = _Source(llm_router, user_api_key_dict, access=access, caller=caller, headers=headers, cache=cache) + return codex_turn(await _read(facts, source)).allowed diff --git a/litellm/proxy/auth/common_budget_adapter.py b/litellm/proxy/auth/common_budget_adapter.py new file mode 100644 index 00000000000..cc4d90b26f3 --- /dev/null +++ b/litellm/proxy/auth/common_budget_adapter.py @@ -0,0 +1,432 @@ +import asyncio +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm._auth import budget_lookup, common_budgets +from litellm._auth.common_budgets import AlertFacts, BudgetDecision, BudgetFacts, MemberLimitFacts +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import ( + CallInfo, + LiteLLM_EndUserTable, + Litellm_EntityType, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.utils import ProxyLogging +from litellm.redis.stores.spend_counters import end_user_counter, key_counter, team_counter, user_counter + +_METADATA: Final = TypeAdapter(Mapping[str, object]) +AlertKind: TypeAlias = Literal[ + "token_budget", + "user_budget", + "soft_budget", + "max_budget_alert", + "team_budget", + "organization_budget", + "proxy_budget", + "projected_limit_exceeded", + "project_budget", +] +_ALERT: Final = TypeAdapter[AlertKind](AlertKind) + + +def metadata(raw: object) -> Mapping[str, object]: + return _METADATA.validate_python(raw or {}) + + +@dataclass(frozen=True, slots=True) +class CounterRead: + key: str + fallback: float + max_budget: float | None + authoritative: bool = False + window_entity_type: str | None = None + window_entity_id: str | None = None + window_duration: str | None = None + window_start: datetime | None = None + + +class _CounterOptions(TypedDict, total=False): + fallback_authoritative: ReadOnly[bool] + window_entity_type: ReadOnly[str | None] + window_entity_id: ReadOnly[str | None] + window_duration: ReadOnly[str | None] + window_start: ReadOnly[datetime | None] + + +async def read_counter(counter: CounterRead) -> float: + from litellm.proxy.proxy_server import get_current_spend + + authoritative: Final[_CounterOptions] = {"fallback_authoritative": True} if counter.authoritative else {} + window: Final[_CounterOptions] = ( + { + "window_entity_type": counter.window_entity_type, + "window_entity_id": counter.window_entity_id, + "window_duration": counter.window_duration, + "window_start": counter.window_start, + } + if counter.window_entity_type is not None + else {} + ) + options: Final[_CounterOptions] = {**authoritative, **window} + return await get_current_spend( + counter_key=counter.key, + fallback_spend=counter.fallback, + max_budget=counter.max_budget, + **options, + ) + + +async def read_budget( + facts: BudgetFacts, + counter: CounterRead, + reader: Callable[[CounterRead], Awaitable[float]] = read_counter, +) -> tuple[BudgetFacts, BudgetDecision]: + initial: Final = common_budgets.check(facts) + if not initial.read_spend: + return facts, initial + snapshot: Final = replace(facts, spend=await reader(counter)) + return snapshot, common_budgets.check(snapshot) + + +def emit_alert(logging: ProxyLogging, kind: str, info: CallInfo) -> None: + asyncio.create_task(logging.budget_alerts(type=_ALERT.validate_python(kind), user_info=info)) + + +def key_alert_info( + key: UserAPIKeyAuth, spend: float | None, limit: float | None, entity: Litellm_EntityType +) -> CallInfo: + return CallInfo( + token=key.token, + spend=spend if spend is not None else 0.0, + max_budget=limit, + user_id=key.user_id, + team_id=key.team_id, + team_alias=key.team_alias, + organization_id=key.org_id, + event_group=entity, + ) + + +def raise_budget(facts: BudgetFacts, entity: Litellm_EntityType, identity: str | None, message: str | None) -> None: + assert facts.spend is not None and facts.limit is not None + raise litellm.BudgetExceededError( + current_cost=facts.spend, + max_budget=facts.limit, + message=message, + entity_type=entity.value, + entity_id=identity, + ) + + +async def team_max_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if team_object is None: + return + facts, result = await read_budget( + BudgetFacts("Team", team_object.max_budget, token_present=valid_token is not None), + CounterRead(team_counter(team_object.team_id), team_object.spend or 0.0, team_object.max_budget), + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + key_alert_info(valid_token, facts.spend, facts.limit, Litellm_EntityType.TEAM), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.TEAM, + team_object.team_id, + f"Budget has been exceeded! Team={team_object.team_id} Current cost: {facts.spend}, Max budget: {facts.limit}", + ) + + +async def virtual_key_max_budget_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + from litellm.proxy.auth.auth_checks import key_budget_exceeded_error, send_token_budget_alert + + raw_percentage: Final[object] = litellm.budget_exceeded_throttle_percentage + percentage: Final = ( + float(raw_percentage) + if isinstance(raw_percentage, (int, float)) and not isinstance(raw_percentage, bool) + else None + ) + facts, result = await read_budget( + BudgetFacts( + "Key", + valid_token.max_budget, + throttle_percentage=percentage, + throttle_opt_in=metadata(valid_token.metadata).get("throttle_on_budget_exceeded") is True, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy metadata is validated by metadata() + tpm_limit=valid_token.tpm_limit, + rpm_limit=valid_token.rpm_limit, + ), + CounterRead(key_counter(str(valid_token.token)), valid_token.spend or 0.0, valid_token.max_budget), + ) + if result.alert is not None and facts.spend is not None: + send_token_budget_alert(valid_token, facts.spend, proxy_logging_obj, user_obj) + if result.throttle_percentage is not None: + valid_token.budget_throttle_pct = result.throttle_percentage # rebind-ok: native request-local throttle + if result.denied and facts.spend is not None and facts.limit is not None: + raise key_budget_exceeded_error(valid_token, facts.spend, facts.limit) + + +async def project_max_budget_check( + project_object: LiteLLM_ProjectTableCachedObj | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + from litellm.redis.stores.spend_counters import project_counter + + if project_object is None: + return + limit: Final = ( + project_object.litellm_budget_table.max_budget if project_object.litellm_budget_table is not None else None + ) + facts, result = await read_budget( + BudgetFacts("Project", limit, token_present=valid_token is not None), + CounterRead(project_counter(project_object.project_id), project_object.spend or 0.0, limit), + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + key_alert_info(valid_token, facts.spend, limit, Litellm_EntityType.PROJECT), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.PROJECT, + project_object.project_id, + f"Budget has been exceeded! Project={project_object.project_id} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def user_max_budget_check( + user_object: LiteLLM_UserTable | None, + team_object: LiteLLM_TeamTable | None, + general_settings: Mapping[str, object], + proxy_logging_obj: ProxyLogging, +) -> None: + if user_object is None: + return + facts, result = await read_budget( + BudgetFacts( + "User", + user_object.max_budget, + team_key=team_object is not None and team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # retain the legacy None guard for constructed cache rows + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, + ), + CounterRead(user_counter(user_object.user_id), user_object.spend or 0.0, user_object.max_budget), + ) + if result.alert is not None: + emit_alert( + proxy_logging_obj, + result.alert, + CallInfo( + spend=facts.spend if facts.spend is not None else 0.0, + max_budget=facts.limit, + user_id=user_object.user_id, + user_email=user_object.user_email, + event_group=Litellm_EntityType.USER, + ), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.USER, + user_object.user_id, + f"ExceededBudget: User={user_object.user_id} over budget. Spend={facts.spend}, Budget={facts.limit}", + ) + + +async def check_end_user_budget(end_user_obj: LiteLLM_EndUserTable, route: str) -> None: + from litellm.proxy.auth.route_checks import RouteChecks + + limit: Final = ( + end_user_obj.litellm_budget_table.max_budget if end_user_obj.litellm_budget_table is not None else None + ) + facts, result = await read_budget( + BudgetFacts("EndUser", limit, info_route=RouteChecks.is_info_route(route)), + CounterRead(end_user_counter(end_user_obj.user_id), end_user_obj.spend or 0.0, limit, authoritative=True), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.END_USER, + end_user_obj.user_id, + f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={facts.spend}, Budget={facts.limit}", + ) + + +def global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_checks: bool, route: str) -> None: + from litellm.proxy._types import MODEL_LIST_ROUTES + from litellm.proxy.auth.route_checks import RouteChecks + + facts: Final = BudgetFacts( + "Proxy", + litellm.max_budget, + global_proxy_spend, + skip=skip_budget_checks, + llm_route=RouteChecks.is_llm_api_route(route=route), + model_list_route=route in MODEL_LIST_ROUTES, + ) + if common_budgets.check(facts).denied: + raise_budget(facts, Litellm_EntityType.PROXY, None, None) + + +async def virtual_key_soft_budget_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + from litellm.proxy.auth.auth_checks import send_soft_budget_alert + + if common_budgets.check(BudgetFacts("KeySoft", valid_token.soft_budget, valid_token.spend)).alert is not None: + send_soft_budget_alert(valid_token, proxy_logging_obj, user_obj) + + +def _soft_emails(raw: object) -> tuple[str, ...]: + return budget_lookup.emails(raw) + + +async def team_soft_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if team_object is None: + return + if not budget_lookup.read_soft_emails( + budget_lookup.SoftFacts(team_object.soft_budget, team_object.spend, valid_token is not None) + ): + return + emails: Final = _soft_emails(budget_lookup.soft_email_value(team_object.metadata)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # raw legacy metadata projection + result: Final = common_budgets.check( + BudgetFacts( + "TeamSoft", + team_object.soft_budget, + team_object.spend, + token_present=valid_token is not None, + alert_recipients=emails, + ) + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + CallInfo( + token=valid_token.token, + spend=team_object.spend if team_object.spend is not None else 0.0, + max_budget=team_object.max_budget, + soft_budget=team_object.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.TEAM, + alert_emails=list(emails), + ), + ) + + +async def project_soft_budget_check( + project_object: LiteLLM_ProjectTableCachedObj | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if project_object is None: + return + limit: Final = ( + project_object.litellm_budget_table.soft_budget if project_object.litellm_budget_table is not None else None + ) + result: Final = common_budgets.check( + BudgetFacts("ProjectSoft", limit, project_object.spend, token_present=valid_token is not None) + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + CallInfo( + token=valid_token.token, + spend=project_object.spend, + max_budget=None, + soft_budget=limit, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.PROJECT, + ), + ) + + +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + return common_budgets.member_limit( + read_member_limit_facts( + member_budget, team_default_budget.max_budget if team_default_budget is not None else None, now + ) + ) + + +def read_member_limit_facts( + member_budget: LiteLLM_BudgetTable | None, default_limit: float | None, now: datetime +) -> MemberLimitFacts: + expiry: Final = member_budget.temp_budget_expiry if member_budget is not None else None + aware_expiry: Final = ( + expiry.replace(tzinfo=timezone.utc) if expiry is not None and expiry.tzinfo is None else expiry + ) + expiry_number: Final = ( + aware_expiry.replace(tzinfo=timezone.utc).timestamp() + if aware_expiry is not None and aware_expiry.tzinfo is now.tzinfo + else aware_expiry.timestamp() + if aware_expiry is not None + else None + ) + now_number: Final = ( + now.replace(tzinfo=timezone.utc).timestamp() + if aware_expiry is not None and aware_expiry.tzinfo is now.tzinfo + else now.timestamp() + ) + return MemberLimitFacts( + member_limit=member_budget.max_budget if member_budget is not None else None, + default_limit=default_limit, + temporary_increase=member_budget.temp_budget_increase if member_budget is not None else None, + expires_at=expiry_number, + now=now_number, + now_naive=now.utcoffset() is None, + ) + + +def threshold_alert_decision( + spend: float | None, + max_budget: float | None, + config: Mapping[str, Sequence[str]] | None, + *, + team_member: bool = False, +) -> bool: + minimum: Final = min((int(key) for key in (config or {}) if team_member or key.isdigit()), default=None) + return common_budgets.alert( + AlertFacts(spend, max_budget, bool(config), minimum, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, team_member) + ) diff --git a/litellm/proxy/auth/common_budget_lookups.py b/litellm/proxy/auth/common_budget_lookups.py new file mode 100644 index 00000000000..e7b24ab8d70 --- /dev/null +++ b/litellm/proxy/auth/common_budget_lookups.py @@ -0,0 +1,240 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Final + +from litellm._auth import budget_lookup, common_budgets +from litellm._auth.common_budgets import BudgetFacts +from litellm.proxy._types import ( + Litellm_EntityType, + LiteLLM_TagTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth import budget_lookup_adapter +from litellm.proxy.auth.common_budget_adapter import ( + CounterRead, + emit_alert, + key_alert_info, + raise_budget, + read_budget, + resolve_effective_team_member_budget, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.redis.stores.spend_counters import ( + model_access_group_counter, + org_counter, + tag_counter, + team_member_counter, +) +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget + + +async def organization_max_budget_check( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + reading: Final = await budget_lookup_adapter.organization( + budget_lookup.OrganizationFacts(valid_token is not None, prisma_client is not None, team_object is not None), + budget_lookup_adapter.LookupIO(valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + if reading is None: + return + assert valid_token is not None + org_id: Final = reading.org_id + org: Final = reading.row + carry_organization_budget_state(valid_token=valid_token, org_table=org) + limit: Final = org.litellm_budget_table.max_budget if org.litellm_budget_table is not None else None + facts, result = await read_budget( + BudgetFacts("Organization", limit), CounterRead(org_counter(org_id), org.spend or 0.0, limit) + ) + if result.alert is not None: + info: Final = key_alert_info(valid_token, facts.spend, limit, Litellm_EntityType.ORGANIZATION) + info.organization_id = org_id + emit_alert(proxy_logging_obj, result.alert, info) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.ORGANIZATION, + org_id, + f"Budget has been exceeded! Organization={org_id} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def tag_max_budget_check_for_tags( + tags: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + valid_token: UserAPIKeyAuth | None, +) -> None: + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + if prisma_client is None or not tags: + return + rows: Final = await get_tag_objects_batch( + tag_names=tags, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for name in tags: + await check_tag_row(name, rows.get(name)) + + +async def check_tag_row(name: str, row: LiteLLM_TagTable | None) -> None: + limit: Final = ( + row.litellm_budget_table.max_budget if row is not None and row.litellm_budget_table is not None else None + ) + facts, result = await read_budget( + BudgetFacts("Tag", limit), + CounterRead(tag_counter(name), (row.spend or 0.0) if row is not None else 0.0, limit, authoritative=True), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.TAG, + name, + f"Budget has been exceeded! Tag={name} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def model_access_group_max_budget_check( + matched_model_access_groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + if prisma_client is None or not matched_model_access_groups: + return + rows: Final = await get_model_access_group_budgets_batch( + access_group_names=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + for name in matched_model_access_groups: + await check_access_group_row(name, rows.get(name)) + + +async def check_access_group_row(name: str, row: ModelAccessGroupBudget | None) -> None: + limit: Final = row.max_budget if row is not None else None + facts, result = await read_budget( + BudgetFacts("AccessGroup", limit), + CounterRead( + model_access_group_counter(name), + row.spend if row is not None else 0.0, + limit, + authoritative=True, + ), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.MODEL_ACCESS_GROUP, + name, + f"Budget has been exceeded! Model access group={name} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def check_team_member_budget( + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> None: + from litellm.proxy.auth.budget_alert_adapter import team_member_max_budget_alert_check + + reading: Final = await budget_lookup_adapter.member( + budget_lookup.MemberFacts(team_object is not None, valid_token is not None, team_membership_loaded), + budget_lookup_adapter.LookupIO(valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj), + team_membership, + ) + if reading is None: + return + assert team_object is not None and valid_token is not None and valid_token.user_id is not None + membership: Final = reading.membership + team_metadata: Final = reading.metadata + limit: Final = resolve_effective_team_member_budget(reading.row, reading.default, reading.now) + facts, result = await read_budget( + BudgetFacts("TeamMember", limit), + CounterRead( + team_member_counter(valid_token.user_id, team_object.team_id), + (membership.spend or 0.0) if membership is not None else 0.0, + limit, + ), + ) + if result.alert is not None and facts.spend is not None and limit is not None: + team_member_max_budget_alert_check( + team_id=team_object.team_id, + team_alias=team_object.team_alias, + team_metadata=dict(team_metadata), + organization_id=team_object.organization_id, + user_id=valid_token.user_id, + user_email=user_object.user_email if user_object is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=facts.spend, + max_budget=limit, + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.TEAM_MEMBER, + f"{valid_token.user_id}:{team_object.team_id}", + f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +@dataclass(frozen=True, slots=True) +class AccessGroupBudgetReading: + name: str + facts: BudgetFacts + + +async def read_access_group_budget_row(name: str, row: ModelAccessGroupBudget | None) -> AccessGroupBudgetReading: + limit: Final = row.max_budget if row is not None else None + facts, _decision = await read_budget( + BudgetFacts("AccessGroup", limit), + CounterRead(model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True), + ) + return AccessGroupBudgetReading(name, facts) + + +async def read_access_group_budget_facts( + groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> tuple[AccessGroupBudgetReading, ...]: + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + if prisma_client is None or not groups: + return () + rows: Final = await get_model_access_group_budgets_batch( + access_group_names=groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return await _read_access_group_rows(tuple(groups), rows) + + +async def _read_access_group_rows( + groups: tuple[str, ...], + rows: Mapping[str, ModelAccessGroupBudget], +) -> tuple[AccessGroupBudgetReading, ...]: + readings: tuple[AccessGroupBudgetReading, ...] = () # rebind-ok: native-selected IO driver accumulator + for name in groups: + reading = await read_access_group_budget_row(name, rows.get(name)) + readings = (*readings, reading) + if common_budgets.check(reading.facts).denied: + return readings + return readings diff --git a/litellm/proxy/auth/end_user_context_adapter.py b/litellm/proxy/auth/end_user_context_adapter.py new file mode 100644 index 00000000000..aae5d73f770 --- /dev/null +++ b/litellm/proxy/auth/end_user_context_adapter.py @@ -0,0 +1,406 @@ +from collections.abc import Mapping, MutableMapping +from dataclasses import replace +from types import MappingProxyType +from typing import Final, assert_never + +from opentelemetry.trace import Span + +import litellm +from litellm._auth import end_user_context as core +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_EndUserTable, UserAPIKeyAuth +from litellm.proxy.auth.budget_context_adapter import mapping_value, resolve_default_end_user_budget +from litellm.proxy.auth.end_user_context_io import EndUserIO, read_default_budget +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +_PARAM_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType( + { + "EndUserId": "end_user_id", + "Tpm": "end_user_tpm_limit", + "Rpm": "end_user_rpm_limit", + "Tpd": "end_user_tpd_limit", + "MaxBudget": "end_user_max_budget", + "Models": "end_user_model_max_budget", + "Region": "allowed_model_region", + } +) +_BUDGET_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType( + { + "Tpm": "tpm_limit", + "Rpm": "rpm_limit", + "Tpd": "tpd_limit", + "MaxBudget": "max_budget", + "Models": "model_max_budget", + } +) + + +def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: + raw: Final = key_metadata.get("end_user_budget_id") if key_metadata is not None else None + return core.normalize_id(core.IdFacts(raw if isinstance(raw, str) else None)) + + +def select_default_id(facts: core.DefaultFacts) -> str | None: + decision: Final = core.default_id(facts) + if decision.read_global: + return select_default_id(replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id)) + return decision.budget_id + + +async def get_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + budget_id: str | None = None, +) -> LiteLLM_BudgetTable | None: + selected: Final = select_default_id(core.DefaultFacts(prisma_client is not None, budget_id, None)) + if selected is None: + return None + return await read_default_budget(selected, EndUserIO(prisma_client, user_api_key_cache, parent_otel_span)) + + +async def _apply_default( + facts: core.ApplyFacts, row: LiteLLM_EndUserTable, io: EndUserIO, default: LiteLLM_BudgetTable | None = None +) -> LiteLLM_EndUserTable: + decision: Final = core.apply_default(facts) + match decision: + case "Keep": + return row + case "ReadGlobal": + return await _apply_default( + replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), row, io + ) + case "ReadDefault": + budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span) + return await _apply_default( + replace(facts, default_read=True, default_present=budget is not None), row, io, budget + ) + case "Copy": + assert default is not None + verbose_proxy_logger.debug("Applied default budget %s to end user %s", default.budget_id, row.user_id) + return row.model_copy(update=MappingProxyType({"litellm_budget_table": default})) + case _: + assert_never(decision) + + +async def apply_default_budget_to_end_user( + end_user_obj: LiteLLM_EndUserTable, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable: + return await _apply_default( + core.ApplyFacts(end_user_obj.budget_id, end_user_obj.litellm_budget_table is not None, key_end_user_budget_id), + end_user_obj, + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + ) + + +async def _registry(facts: core.RegistryFacts, io: EndUserIO) -> bool: + from litellm.proxy.auth import auth_checks + + decision: Final = core.registry(facts) + match decision: + case "Restricted": + return False + case "Unrestricted": + return True + case "Read": + assert io.prisma is not None + ids: Final = await auth_checks._load_end_user_restricted_registry( # pyright: ignore[reportPrivateUsage] # raw cached registry IO + prisma_client=io.prisma, user_api_key_cache=io.cache + ) + return await _registry( + replace(facts, registry_read=True, registry=tuple(ids) if ids is not None else None), io + ) + case _: + assert_never(decision) + + +async def end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, + key_end_user_budget_id: str | None = None, +) -> bool: + return await _registry( + core.RegistryFacts( + end_user_id, + litellm.max_end_user_budget_id, + key_end_user_budget_id, + litellm.validate_end_user_id_in_db, + token_end_user_max_budget is not None, + ), + EndUserIO(prisma_client, user_api_key_cache), + ) + + +def _budget_fields(row: LiteLLM_BudgetTable | None) -> core.Fields: + return core.Fields( + tpm=getattr(row, "tpm_limit", None) is not None, + rpm=getattr(row, "rpm_limit", None) is not None, + tpd=getattr(row, "tpd_limit", None) is not None, + max_budget=getattr(row, "max_budget", None) is not None, + models=getattr(row, "model_max_budget", None) is not None, + ) + + +def _budget_params(row: LiteLLM_BudgetTable | None) -> Mapping[str, object]: + selected: Final = core.assignments(core.AssignmentFacts("Params", _budget_fields(row))) + return MappingProxyType({_PARAM_NAMES[field]: getattr(row, _BUDGET_NAMES[field]) for field in selected}) + + +def apply_budget_limits_to_end_user_params( + end_user_params: MutableMapping[str, object], # mutable-ok: legacy request-local output parameter + budget_info: LiteLLM_BudgetTable, + end_user_id: str | None, +) -> None: + for name, value in _budget_params(budget_info).items(): + end_user_params[name] = value # rebind-ok: native-selected legacy request parameter projection + verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) + + +def _token_fields(end_user_params: Mapping[str, object]) -> tuple[core.Field, ...]: + fields: Final = core.Fields( + tpm=end_user_params.get("end_user_tpm_limit") is not None, + rpm=end_user_params.get("end_user_rpm_limit") is not None, + tpd=end_user_params.get("end_user_tpd_limit") is not None, + models=end_user_params.get("end_user_model_max_budget") is not None, + region=end_user_params.get("allowed_model_region") is not None, + ) + return core.assignments(core.AssignmentFacts("Token", fields)) + + +def _project_token( + valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object], fields: tuple[core.Field, ...] +) -> UserAPIKeyAuth: + for field in fields: + setattr(valid_token, _PARAM_NAMES[field], end_user_params.get(_PARAM_NAMES[field])) + return valid_token + + +def update_valid_token_with_end_user_params( + valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object] +) -> UserAPIKeyAuth: + return _project_token(valid_token, end_user_params, _token_fields(end_user_params)) + + +async def _validation(facts: core.ValidationFacts, io: EndUserIO) -> core.ValidationDecision: + from litellm.proxy.auth import auth_checks + + decision: Final = core.validation(facts) + match decision.read: + case None: + return decision + case "Cache": + cached: Final[object] = await auth_checks._raw_cache(io.cache).async_get_cache( # pyright: ignore[reportPrivateUsage] # shared raw cache IO + key=f"end_user_validation:{facts.end_user_id}" + ) + return await _validation( + replace(facts, cache_read=True, cached=cached if isinstance(cached, str) else None), io + ) + case "EndUser": + try: + row: Final = await auth_checks.get_end_user_object( + end_user_id=facts.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + route=io.route, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve end-user validation lookup failure handling + verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", error) + return await _validation(replace(facts, end_user_read=True), io) + return await _validation(replace(facts, end_user_read=True, end_user_present=row is not None), io) + case "User": + assert facts.end_user_id is not None and io.prisma is not None + try: + user: Final = await auth_checks.get_user_object( + user_id=facts.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + user_id_upsert=False, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + check_db_only=False, + user_email=decision.user_email, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve secondary identity lookup failure handling + verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", error) + return await _validation(replace(facts, user_read=True), io) + return await _validation(replace(facts, user_read=True, user_present=user is not None), io) + case "WriteCache": + await io.cache.async_set_cache( + key=f"end_user_validation:{facts.end_user_id}", value=decision.cache_value, ttl=decision.cache_ttl + ) + return await _validation(replace(facts, cache_written=True), io) + case _: + assert_never(decision.read) + + +async def resolve_and_validate_end_user_id( + raw_end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", + key_end_user_budget_id: str | None = None, +) -> str | None: + decision: Final = await _validation( + core.ValidationFacts( + "Resolve", + raw_end_user_id, + litellm.validate_end_user_id_in_db, + prisma_client is not None, + litellm.max_end_user_budget_id, + key_end_user_budget_id, + ), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) + return decision.end_user_id + + +async def end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", +) -> bool: + decision: Final = await _validation( + core.ValidationFacts("Exists", end_user_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) + return decision.valid + + +def _complete_lookup( + decision: core.LookupDecision, + token: UserAPIKeyAuth, + row: LiteLLM_EndUserTable | None, + default: LiteLLM_BudgetTable | None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + if decision.set_max_budget: + assert default is not None + try: + token.end_user_max_budget = default.max_budget # rebind-ok: request-local native projection + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve request-local setter failure handling + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return token, row + + +async def _lookup( + facts: core.LookupFacts, + token: UserAPIKeyAuth, + io: EndUserIO, + row: LiteLLM_EndUserTable | None = None, + default: LiteLLM_BudgetTable | None = None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + from litellm.proxy.auth.auth_checks import get_end_user_object + + decision: Final = core.lookup(facts) + match decision.action: + case "Done": + return _complete_lookup(decision, token, row, default) + case "ReadGlobal": + return await _lookup( + replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), token, io, row + ) + case "ReadEndUser": + try: + end_user: Final = await get_end_user_object( + end_user_id=token.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + route=io.route, + token_end_user_max_budget=token.end_user_max_budget, + key_end_user_budget_id=facts.key_id, + ) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve the legacy lookup IO failure path + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace( + facts, + end_user_read=True, + end_user_present=end_user is not None, + own_budget_present=end_user is not None and end_user.litellm_budget_table is not None, + ), + token, + io, + end_user, + ) + case "ReadDefault": + try: + budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve the legacy default lookup IO failure path + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace(facts, default_read=True, default_present=budget is not None), token, io, row, budget + ) + case "ApplyOwn" | "ApplyDefault": + budget_row: Final = ( + row.litellm_budget_table if decision.action == "ApplyOwn" and row is not None else default + ) + params: Final = MappingProxyType( + { + "end_user_id": token.end_user_id, + "allowed_model_region": getattr(row, "allowed_model_region", None), + **_budget_params(budget_row), + } + ) + fields: Final = _token_fields(params) + try: + if decision.project_budget: + verbose_proxy_logger.debug("Applied budget limits to end user %s", token.end_user_id or "") + _project_token(token, params, fields) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve request-local setter failure handling + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace(facts, applied=True, token_max_present=token.end_user_max_budget is not None), + token, + io, + row, + default, + ) + case _: + assert_never(decision.action) + + +async def lookup_end_user_and_apply_budget( + valid_token: UserAPIKeyAuth, + route: str, + parent_otel_span: Span | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + key_id: Final = get_key_end_user_budget_id(mapping_value(getattr(valid_token, "metadata", None))) + return await _lookup( + core.LookupFacts(key_id), + valid_token, + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) diff --git a/litellm/proxy/auth/end_user_context_io.py b/litellm/proxy/auth/end_user_context_io.py new file mode 100644 index 00000000000..e515f9f6c5d --- /dev/null +++ b/litellm/proxy/auth/end_user_context_io.py @@ -0,0 +1,64 @@ +from dataclasses import dataclass +from typing import Final + +from opentelemetry.trace import Span + +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_BudgetTable +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.native_auth_repository import native_auth_repository + + +@dataclass(frozen=True, slots=True) +class EndUserIO: + prisma: PrismaClient | None + cache: UserApiKeyCache + span: Span | None = None + logging: ProxyLogging | None = None + route: str = "" + + +async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTable | None: + from litellm.proxy.auth import auth_checks + + assert io.prisma is not None + key: Final = f"default_end_user_budget:{budget_id}" + cached: Final = await io.cache.async_get_cache(key=key, model_type=LiteLLM_BudgetTable) + if cached is not None: + return cached + try: + native: Final = native_auth_repository(io.prisma) + fill_generation: Final = io.cache.current_generation() if native is not None else None + record: Final = ( + await native.one("budget", [budget_id], LiteLLM_BudgetTable) + if native is not None + else await auth_checks._dictable_table(BudgetRepository(io.prisma), "budget").find_unique( # pyright: ignore[reportPrivateUsage] # shared raw repository projection + where={"budget_id": budget_id} + ) + ) + if record is None: + verbose_proxy_logger.warning( + "Default end user budget not found in database: %s", budget_id.replace("\r", "").replace("\n", "") + ) + return None + row: Final = LiteLLM_BudgetTable.model_validate( + record.model_dump() if isinstance(record, LiteLLM_BudgetTable) else record.dict() + ) + await io.cache.async_set_cache( + fill_generation=fill_generation, + key=key, + value=row, + model_type=LiteLLM_BudgetTable, + ttl=get_management_object_ttl(io.cache), + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve the default budget DB/cache IO failure contract + verbose_proxy_logger.error("Error fetching default end user budget: %s", error) + return None + + else: + return row diff --git a/litellm/proxy/auth/fallback_access_adapter.py b/litellm/proxy/auth/fallback_access_adapter.py new file mode 100644 index 00000000000..26e7530c924 --- /dev/null +++ b/litellm/proxy/auth/fallback_access_adapter.py @@ -0,0 +1,154 @@ +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope, fallback_authorization +from litellm._auth.model_access import ModelAccessFacts, key_model_grants +from litellm.proxy._types import LiteLLM_ProjectTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_group_adapter import get_authorized_resources_from_key_access_groups +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Sources: + model: str + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: DatabaseClient | None + team: LiteLLM_TeamTable | None + project: LiteLLM_ProjectTable | None + + +async def load_fallback_authorization_facts( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: DatabaseClient | None = None, + team_membership: LiteLLM_TeamMembership | None = None, + project_object: LiteLLM_ProjectTable | None = None, + refuse_agent_keys: bool = False, + team_object: LiteLLM_TeamTable | None = None, +) -> FallbackAuthorizationFacts: + facts: Final = FallbackAuthorizationFacts( + refuse_agent_keys=refuse_agent_keys, + agent_id=valid_token.agent_id, + team_present=team_object is not None, + embedded_team_grants=tuple(valid_token.team_models or ()), + member_grants=tuple(team_membership.litellm_budget_table.allowed_models or ()) + if team_membership is not None and team_membership.litellm_budget_table is not None + else (), + project_grants=tuple(project_object.models) if project_object is not None else (), + ) + return await _read(facts, _Sources(model, valid_token, llm_router, prisma_client, team_object, project_object)) + + +async def _read(facts: FallbackAuthorizationFacts, source: _Sources) -> FallbackAuthorizationFacts: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + decision: Final = fallback_authorization(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadKey": + grants: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + key: Final = GroupedModelScope( + project_model_facts( + source.model, + source.router, + grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ), + tuple(source.token.access_group_ids or ()), + ) + return await _read(replace(facts, key=key), source) + case "ReadTeam": + if source.team is None: + raise NativeAuthRuntimeError("missing fallback team source") + team: Final = GroupedModelScope( + project_model_facts( + source.model, + source.router, + source.team.models, + source.token.team_model_aliases, + source.team.team_id, + key_model_aliases_for_auth_check(source.token), + "team", + ), + tuple(source.team.access_group_ids or ()), + ) + return await _read(replace(facts, team=team), source) + case "ReadEmbeddedTeam": + embedded: Final = project_model_facts( + source.model, + source.router, + facts.embedded_team_grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "team", + ) + return await _read(replace(facts, embedded_team=embedded), source) + case "ReadMember": + member: Final = project_model_facts( + source.model, + source.router, + facts.member_grants, + team_id=source.token.team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="team", + ) + return await _read(replace(facts, member=member), source) + case "ReadProject": + project: Final = _project_scope(facts, source) + return await _read(replace(facts, project=project), source) + case "ReadKeyGroups": + if facts.key is None: + raise NativeAuthRuntimeError("missing fallback key scope") + key_scope: Final = await read_grouped_scope(facts.key, source.router, source.prisma, merge=False) + return await _read(replace(facts, key=key_scope), source) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing fallback team scope") + team_scope: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=team_scope), source) + case "ReadAssignedGroups": + assigned: Final = await get_authorized_resources_from_key_access_groups( + source.token, source.team, "access_model_names" + ) + return await _read(replace(facts, assigned_models=tuple(assigned)), source) + case "ReadAssignedScope": + if facts.assigned_models is None: + raise NativeAuthRuntimeError("missing fallback assigned models") + scope: Final = project_model_facts( + source.model, + source.router, + facts.assigned_models, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ) + return await _read(replace(facts, assigned_scope=scope), source) + + +def _project_scope(facts: FallbackAuthorizationFacts, source: _Sources) -> ModelAccessFacts: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + if source.project is None: + raise NativeAuthRuntimeError("missing fallback project source") + return project_model_facts( + source.model, + source.router, + facts.project_grants, + team_id=source.project.team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="project", + ) diff --git a/litellm/proxy/auth/fallback_candidate_adapter.py b/litellm/proxy/auth/fallback_candidate_adapter.py new file mode 100644 index 00000000000..32b0870798b --- /dev/null +++ b/litellm/proxy/auth/fallback_candidate_adapter.py @@ -0,0 +1,320 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Final, Protocol + +from pydantic import BaseModel + +import litellm +from litellm._auth.decision import FallbackTarget +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_candidate import ( + CandidateBudgetRead, + CandidateBudgetScope, + FallbackCandidateDecision, + FallbackCandidateFacts, + fallback_candidate, +) +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.model_budget_data import ModelBudgetEntity +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.budget_pricing_adapter import read_model_cost_facts +from litellm.proxy.auth.common_budget_lookups import read_access_group_budget_facts +from litellm.proxy.auth.fallback_access_adapter import load_fallback_authorization_facts +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_attribution_adapter import read_model_group_attribution_facts +from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading +from litellm.proxy.auth.route_policy_adapter import route_context +from litellm.proxy.auth.unpriced_model_adapter import read_unpriced_model_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.hooks.model_max_budget_limiter import team_member_budget_entity_id +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router import Router + + +class ModelBudgetReader(Protocol): + async def read_model_budget_raw( + self, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> ModelBudgetReading: ... + + +class _BudgetFields(BaseModel, frozen=True): + model_max_budget: object = None + user_model_max_budget: object = None + end_user_model_max_budget: object = None + team_model_max_budget: object = None + team_member_model_max_budget: object = None + + +@dataclass(frozen=True, slots=True) +class _BudgetSource: + entity: ModelBudgetEntity + identity: str | None + budgets: Mapping[str, object] | None + + +@dataclass(frozen=True, slots=True) +class _Dependencies: + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + + +@dataclass(frozen=True, slots=True) +class _Sources: + target: FallbackTarget + token: UserAPIKeyAuth + limiter: ModelBudgetReader + router: Router | None + route: str + project: LiteLLM_ProjectTable | None + team: LiteLLM_TeamTable | None + customer: LiteLLM_EndUserTable | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + dependencies: _Dependencies | None = None + membership: LiteLLM_TeamMembership | None = None + budgets: tuple[_BudgetSource, ...] = () + + +def _mapping(raw: object) -> Mapping[str, object] | None: + from litellm.proxy.auth.budget_context_adapter import mapping_value + + return mapping_value(raw) + + +def _budget_sources(token: UserAPIKeyAuth) -> tuple[_BudgetSource, ...]: + fields: Final = _BudgetFields.model_validate(token, from_attributes=True) + member_id: Final = ( + team_member_budget_entity_id(token.user_id, token.team_id) + if token.user_id is not None and token.team_id is not None + else None + ) + return ( + _BudgetSource("Key", token.token, _mapping(fields.model_max_budget)), + _BudgetSource("User", token.user_id, _mapping(fields.user_model_max_budget)), + _BudgetSource("EndUser", token.end_user_id, _mapping(fields.end_user_model_max_budget)), + _BudgetSource("Team", token.team_id, _mapping(fields.team_model_max_budget)), + _BudgetSource("TeamMember", member_id, _mapping(fields.team_member_model_max_budget)), + ) + + +async def read_fallback_candidate_facts( + *, + target: FallbackTarget, + valid_token: UserAPIKeyAuth, + limiter: ModelBudgetReader, + llm_router: Router | None, + route: str, + project_object: LiteLLM_ProjectTable | None = None, + team_object: LiteLLM_TeamTable | None = None, + end_user_object: LiteLLM_EndUserTable | None = None, + prisma_client: PrismaClient | None = None, + user_api_key_cache: UserApiKeyCache | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> FallbackCandidateFacts: + facts: Final = FallbackCandidateFacts( + router_present=llm_router is not None, + user_present=valid_token.user_id is not None, + team_present=valid_token.team_id is not None, + block_unpriced=litellm.block_requests_for_models_without_pricing, + customer_present=end_user_object is not None, + ) + return await _read( + facts, + _Sources( + target, + valid_token, + limiter, + llm_router, + route, + project_object, + team_object, + end_user_object, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + ), + ) + + +def _warn(source: _Sources, *, exception: bool) -> None: + verbose_proxy_logger.warning( + "Skipping fallback candidate %s after a budget or authorization check error", + source.target.model.replace("\r", "").replace("\n", ""), + exc_info=exception, + ) + + +async def _read(facts: FallbackCandidateFacts, source: _Sources) -> FallbackCandidateFacts: + decision: Final = fallback_candidate(facts) + if decision.kind == "Complete": + if decision.log_failure: + _warn(source, exception=False) + return facts + if decision.kind == "ReadMembership": + return await _read_membership(facts, source) + try: + updated, updated_source = await _effect(decision, facts, source) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # legacy candidate IO errors become failure facts, never authorization verdicts + _warn(source, exception=True) + return replace(facts, io_failed=True) + return await _read(updated, updated_source) + + +async def _read_membership(facts: FallbackCandidateFacts, source: _Sources) -> FallbackCandidateFacts: + from litellm.proxy import proxy_server + from litellm.proxy.auth.auth_checks import get_team_membership + + if source.token.user_id is None or source.token.team_id is None: + raise NativeAuthRuntimeError("candidate membership read requested without identity") + dependencies: Final = _Dependencies( + source.prisma if source.prisma is not None else proxy_server.prisma_client, + source.cache if source.cache is not None else proxy_server.user_api_key_cache, + source.logging if source.logging is not None else proxy_server.proxy_logging_obj, + ) + try: + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=source.token.team_id, + prisma_client=dependencies.prisma, + user_api_key_cache=dependencies.cache, + proxy_logging_obj=dependencies.logging, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve silent membership IO failure as a native candidate fact + return replace(facts, membership="Failed") + return await _read( + replace(facts, membership="Loaded"), replace(source, dependencies=dependencies, membership=membership) + ) + + +async def _effect( + decision: FallbackCandidateDecision, facts: FallbackCandidateFacts, source: _Sources +) -> tuple[FallbackCandidateFacts, _Sources]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check, team_model_aliases_for_auth_check + + router: Final = source.router + dependencies: Final = source.dependencies + if router is None or dependencies is None: + raise NativeAuthRuntimeError("candidate read requested without router or dependencies") + match decision.kind: + case "ReadRoute": + return replace(facts, route=route_context(source.route)), source + case "ReadUnpriced": + return replace(facts, unpriced=read_unpriced_model_facts(source.target.served_model, router)), source + case "ReadCustomerGrants": + if source.customer is None: + raise NativeAuthRuntimeError("candidate customer grants requested without source") + return replace(facts, customer_grants=object_model_grants(source.customer)), source + case "ReadCustomer": + if facts.customer_grants is None: + raise NativeAuthRuntimeError("candidate customer scope requested before grants") + customer: Final = project_model_facts( + source.target.model, + router, + facts.customer_grants, + team_model_aliases_for_auth_check(source.token), + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "customer", + ) + return replace(facts, customer=customer), source + case "ReadPrice": + return replace(facts, price=read_model_cost_facts(source.target.served_model, router)), source + case "ReadBudgetModels": + models: Final = ( + source.target.requested, + source.target.model, + source.target.served_model, + router.registry().resolve(source.target.served_model), + ) + return replace(facts, budget_models=models), source + case "ReadBudgetScopes": + budgets: Final = _budget_sources(source.token) + scopes: Final = tuple( + CandidateBudgetScope(budget.entity, budget.identity is not None, bool(budget.budgets)) + for budget in budgets + ) + return replace(facts, budget_scopes=scopes), replace(source, budgets=budgets) + case "ReadBudget": + if facts.budget_models is None: + raise NativeAuthRuntimeError("candidate budget read requested without model names") + budget: Final = next(budget for budget in source.budgets if budget.entity == decision.entity) + key_budget: Final = next(budget.budgets for budget in source.budgets if budget.entity == "Key") + reading: Final = await source.limiter.read_model_budget_raw( + decision.entity, + budget.identity, + source.token.team_id, + facts.budget_models[decision.model_index], + budget.budgets, + key_budget if decision.entity == "Team" else None, + ) + answer: Final = CandidateBudgetRead(decision.model_index, decision.entity, reading.snapshot) + return replace(facts, budgets=(*facts.budgets, answer)), source + case "ReadAuthorization": + authorization: Final = await load_fallback_authorization_facts( + source.target.model, + source.token, + router, + source.prisma, + source.membership, + source.project, + True, + source.team, + ) + return replace(facts, authorization=authorization), source + case "ResetAttribution": + source.token.matched_model_access_groups = [] # rebind-ok: preserve native-selected request-local reset before attribution IO + return replace(facts, attribution_reset=True), source + case "ReadAttribution": + return replace(facts, attribution=await _attribution(source, dependencies, router)), source + case "StampAttribution": + source.token.matched_model_access_groups = list(decision.groups) # rebind-ok: native attribution effect + return replace(facts, attribution_stamped=True), source + case "ReadGroupBudgets": + readings: Final = await read_access_group_budget_facts( + decision.groups, dependencies.prisma, dependencies.cache + ) + return replace( + facts, group_budgets_read=True, group_budgets=tuple(reading.facts for reading in readings) + ), source + case "Complete" | "ReadMembership": + raise NativeAuthRuntimeError("unexpected candidate data read") + + +async def _attribution(source: _Sources, dependencies: _Dependencies, router: Router) -> ModelGroupAttributionFacts: + try: + return await read_model_group_attribution_facts( + source.target.model, + source.token, + source.team, + source.project, + router, + dependencies.prisma, + dependencies.cache, + dependencies.logging, + source.membership, + True, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve attribution telemetry IO failure separately from candidate IO failure + verbose_proxy_logger.debug("model access group attribution failed: %s", error) + return ModelGroupAttributionFacts(io_failed=True) diff --git a/litellm/proxy/auth/key_budget_context_adapter.py b/litellm/proxy/auth/key_budget_context_adapter.py new file mode 100644 index 00000000000..3f7d245704b --- /dev/null +++ b/litellm/proxy/auth/key_budget_context_adapter.py @@ -0,0 +1,95 @@ +from collections.abc import Mapping +from dataclasses import replace +from datetime import datetime, timezone +from typing import Final, assert_never + +import litellm +from litellm._auth import budget_context as core +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.budget_context_adapter import mapping_value + + +def numeric_facts(value: object) -> core.NumericFacts: + return core.NumericFacts(float(value) if isinstance(value, (int, float)) else None, isinstance(value, bool)) + + +def budget_throttle_percentage() -> float | None: + return core.percentage(numeric_facts(litellm.budget_exceeded_throttle_percentage)) + + +def read_key_throttle_facts(valid_token: UserAPIKeyAuth) -> core.ThrottleFacts: + metadata: Final = mapping_value(getattr(valid_token, "metadata", None)) + facts: Final = core.ThrottleFacts( + (metadata or {}).get("throttle_on_budget_exceeded") is True, + valid_token.tpm_limit, + valid_token.rpm_limit, + ) + decision: Final = core.throttle(facts) + return ( + replace(facts, percentage=numeric_facts(litellm.budget_exceeded_throttle_percentage)) + if decision.read_percentage + else facts + ) + + +def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: + return core.throttle(read_key_throttle_facts(valid_token)).percentage is not None + + +def inspect_key_budget_throttle(valid_token: UserAPIKeyAuth) -> float | None: + decision: Final = core.throttle(read_key_throttle_facts(valid_token)) + return budget_throttle_percentage() if decision.percentage is not None else None + + +def _microseconds(value: datetime) -> int: + delta: Final = value - datetime(1970, 1, 1, tzinfo=timezone.utc) + return (delta.days * 86400 + delta.seconds) * 1_000_000 + delta.microseconds + + +def _temporary( + token: UserAPIKeyAuth, + facts: core.TemporaryFacts, + metadata: Mapping[str, object] | None, +) -> UserAPIKeyAuth: + decision: Final = core.temporary(facts) + if decision.invalid_increase_type is not None: + raise TypeError( + f"unsupported operand type(s) for +: '{type(token.max_budget).__name__}' and '{decision.invalid_increase_type}'" + ) + if decision.copy: + return token.model_copy(update={"max_budget": decision.max_budget}) + match decision.read: + case None: + return token + case "Metadata": + raw: Final = mapping_value(getattr(token, "metadata", None)) + return _temporary( + token, + replace( + facts, + metadata_read=True, + has_increase=raw is not None and "temp_budget_increase" in raw, + has_expiry=raw is not None and "temp_budget_expiry" in raw, + ), + raw, + ) + case "Expiry": + assert metadata is not None + expiry: Final = datetime.fromisoformat(metadata["temp_budget_expiry"]) # pyright: ignore[reportArgumentType] # preserve datetime's public TypeError for malformed metadata + aware: Final = expiry.replace(tzinfo=timezone.utc) if expiry.tzinfo is None else expiry + return _temporary(token, replace(facts, expires_at=_microseconds(aware)), metadata) + case "Now": + return _temporary(token, replace(facts, now=_microseconds(datetime.now(timezone.utc))), metadata) + case "Increase": + assert metadata is not None + value: Final = metadata["temp_budget_increase"] + raw_increase: Final = core.IncreaseFacts( + float(value) if isinstance(value, (int, float)) else None, bool(value), type(value).__name__ + ) + return _temporary(token, replace(facts, increase=raw_increase), metadata) + case _: + assert_never(decision.read) + + +def update_key_budget_with_temp_budget_increase(valid_token: UserAPIKeyAuth) -> UserAPIKeyAuth: + return _temporary(valid_token, core.TemporaryFacts(valid_token.max_budget), None) diff --git a/litellm/proxy/auth/key_model_access_adapter.py b/litellm/proxy/auth/key_model_access_adapter.py new file mode 100644 index 00000000000..be311767faf --- /dev/null +++ b/litellm/proxy/auth/key_model_access_adapter.py @@ -0,0 +1,161 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from fastapi import Request +from pydantic import BaseModel, Field + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts, key_and_fallback_model_access +from litellm._auth.model_access import key_model_grants +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_normalization_adapter import ( + resolve_codex_native_chatgpt_model, + resolve_router_settings_model_group_alias, +) +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.router import Router + + +class _Config(BaseModel, frozen=True): + config: Mapping[str, object] = Field(default_factory=dict[str, object]) + + +@dataclass(frozen=True, slots=True) +class _Sources: + token: UserAPIKeyAuth + request_data: dict[str, object] # mutable-ok: alias resolution mutates the caller's request + route: str + request: Request | None + router: Router | None + + +async def read_key_and_fallback_model_access_facts( + *, + valid_token: UserAPIKeyAuth, + request_data: dict[str, object], # mutable-ok: preserve request-local alias mutations + route: str, + request: Request | None, + llm_model_list: Sequence[object] | None, + llm_router: Router | None, +) -> KeyAndFallbackModelAccessFacts: + from litellm.proxy.auth.claude_normalization_adapter import normalize_claude_model + + await normalize_claude_model(request_data, valid_token, request, route) + await resolve_router_settings_model_group_alias(request_data, valid_token, request, route) + resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) + facts: Final = KeyAndFallbackModelAccessFacts( + config_nonempty=bool(_Config.model_validate(valid_token, from_attributes=True).config), + key_models=object_model_grants(valid_token), + ) + return await _read(facts, _Sources(valid_token, request_data, route, request, llm_router), ()) + + +def _scope(model: str | Sequence[str], source: _Sources) -> GroupedModelScope: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + grants: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + return GroupedModelScope( + project_model_facts( + model, + source.router, + grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ), + tuple(source.token.access_group_ids or ()), + ) + + +async def _groups(scope: GroupedModelScope, source: _Sources) -> GroupedModelScope: + return await read_grouped_scope(scope, source.router, merge=False) + + +async def _route(name: str, source: _Sources) -> Exception | None: + from litellm.proxy.route_llm_request import ( + route_request, # pyright: ignore[reportUnknownVariableType] # preserve the existing routing probe and exception identity + ) + + try: + await route_request( + data={"model": name, "messages": [{"role": "user", "content": "Who was Alexander?"}]}, + llm_router=source.router, + user_model=None, + route_type="acompletion", + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # native policy sequences a legacy routing IO failure before rethrowing + return error + return None + + +async def _read( + facts: KeyAndFallbackModelAccessFacts, source: _Sources, route_errors: tuple[Exception | None, ...] +) -> KeyAndFallbackModelAccessFacts: + from litellm.proxy.auth.auth_utils import fallback_target_model_name, iter_request_fallback_targets + from litellm.proxy.auth.user_api_key_auth import ( + _get_model_from_request_context, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # read the existing mutable request model + ) + + decision: Final = key_and_fallback_model_access(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadPrimary": + model: Final = _get_model_from_request_context( + request_data=source.request_data, + route=source.route, + request=source.request, + llm_router=source.router, + team_id=source.token.team_id, + ) + return await _read( + replace(facts, primary_loaded=True, primary=_scope(model, source) if model is not None else None), + source, + route_errors, + ) + case "ReadPrimaryGroups": + if facts.primary is None: + raise NativeAuthRuntimeError("missing primary model scope") + return await _read(replace(facts, primary=await _groups(facts.primary, source)), source, route_errors) + case "ReadFallbackNames": + names: Final = tuple( + dict.fromkeys( + name + for target in iter_request_fallback_targets(source.request_data) + if (name := fallback_target_model_name(target)) is not None + ) + ) + return await _read(replace(facts, fallback_names=names), source, route_errors) + case "ReadFallbackScope": + if facts.fallback_names is None: + raise NativeAuthRuntimeError("missing fallback names") + scope: Final = _scope(facts.fallback_names[decision.index], source) + return await _read(replace(facts, fallback_scopes=(*facts.fallback_scopes, scope)), source, route_errors) + case "ReadFallbackGroups": + grouped: Final = await _groups(facts.fallback_scopes[decision.index], source) + scopes: Final = tuple( + grouped if index == decision.index else scope for index, scope in enumerate(facts.fallback_scopes) + ) + return await _read(replace(facts, fallback_scopes=scopes), source, route_errors) + case "RouteFallback": + if facts.fallback_names is None: + raise NativeAuthRuntimeError("missing fallback names") + error: Final = await _route(facts.fallback_names[decision.index], source) + return await _read( + replace(facts, routing=(*facts.routing, "Routed" if error is None else "Failed")), + source, + (*route_errors, error), + ) + case "RoutingFailure": + failure: Final = route_errors[decision.index] + if failure is None: + raise NativeAuthRuntimeError("missing fallback routing failure") + raise failure diff --git a/litellm/proxy/auth/managed_model_adapter.py b/litellm/proxy/auth/managed_model_adapter.py new file mode 100644 index 00000000000..069e739615a --- /dev/null +++ b/litellm/proxy/auth/managed_model_adapter.py @@ -0,0 +1,366 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.managed_models import ( + AgentCeilingSourceFacts, + ManagedModelFacts, + agent_ceiling_source, + managed_model_grants, + managed_models, +) +from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, + resolve_agent_access_group_ceiling, + resolve_managed_agent_ceilings, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import ( + agent_caller_auth, + load_agent_caller_team, + load_agent_caller_user, +) +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy +from litellm.proxy.auth.member_model_adapter import read_member_model_facts +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.proxy.auth.team_alias_adapter import resolve_team_alias +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader +from litellm.types.agents import AgentResponse + +CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] +CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]] +CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_UserTable | None]] +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] | None + token: UserAPIKeyAuth | None + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + ceiling: CeilingResolver + team_loader: CallerTeamLoader + user_loader: CallerUserLoader + managed_grants: tuple[object, ...] = () + dispatched: str | Sequence[str] = () + caller: UserAPIKeyAuth | None = None + team: LiteLLM_TeamTable | None = None + user: LiteLLM_UserTable | None = None + key_aliases: Mapping[str, str] | None = None + + +async def agent_access_group_ceilings( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> tuple[AgentAccessGroupCeiling, ...]: + return await _ceilings(AgentCeilingSourceFacts(bool(valid_token.agent_id)), valid_token, resolve_ceiling) + + +async def _ceilings( + facts: AgentCeilingSourceFacts, + token: UserAPIKeyAuth, + resolver: CeilingResolver, + managed: AgentResponse | None = None, + rows: tuple[AgentAccessGroupCeiling, ...] = (), +) -> tuple[AgentAccessGroupCeiling, ...]: + decision: Final = agent_ceiling_source(facts) + match decision.kind: + case "Complete": + return tuple(rows[index] for index in decision.indices) + case "ReadManaged": + policy: Final = managed_agent_policy(token) + return await _ceilings( + replace(facts, managed_loaded=True, managed_present=policy is not None), token, resolver, policy + ) + case "ReadManagedCeilings": + if managed is None: + raise NativeAuthRuntimeError("missing managed agent ceiling source") + loaded: Final = await resolve_managed_agent_ceilings(managed) + return await _ceilings(replace(facts, loaded=True, count=len(loaded)), token, resolver, managed, loaded) + case "ReadUnmanagedCeiling": + if token.agent_id is None: + raise NativeAuthRuntimeError("missing agent ceiling identity") + row: Final = await resolver(token.agent_id) + projected: Final = (row,) if row is not None else () + return await _ceilings( + replace(facts, loaded=True, count=len(projected)), token, resolver, managed, projected + ) + + +async def read_managed_model_facts( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + purpose: Literal["All", "Ceilings", "Caller"] = "All", + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> ManagedModelFacts: + facts, _source = await _read( + ManagedModelFacts(purpose, bool(model), valid_token is not None), + _Source( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + resolve_ceiling, + load_team, + load_user, + ), + ) + return facts + + +async def can_agent_call_model( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + purpose: Literal["All", "Ceilings", "Caller"] = "All", + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + facts, source = await _read( + ManagedModelFacts(purpose, bool(model), valid_token is not None), + _Source( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + resolve_ceiling, + load_team, + load_user, + ), + ) + decision: Final = managed_models(facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None or model is None: + raise NativeAuthRuntimeError("incomplete managed model authorization") + raise_model_denial( + decision.denial, + model, + user_id=source.caller.user_id + if decision.member_models is not None and source.caller is not None + else valid_token.agent_id + if valid_token is not None + else None, + team_id=source.caller.team_id if source.caller is not None else None, + member_models=decision.member_models, + ) + + +@native_boundary +def _managed_grants(raw: tuple[object, ...]) -> tuple[str, ...]: + return _NAMES.validate_python(raw) + + +def _declared_grants(token: UserAPIKeyAuth) -> tuple[str | None, ...] | None: + managed: Final = managed_agent_policy(token) + raw: Final = (managed.object_permission or {}).get("models") if managed is not None else None + if not isinstance(raw, (list, tuple)): + return None + return tuple(value if isinstance(value, str) else None for value in _ITEMS.validate_python(raw)) + + +def managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: + return managed_model_grants(_declared_grants(valid_token), ()) + + +async def agent_model_grants( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> frozenset[str]: + ceilings: Final = await agent_access_group_ceilings(valid_token, resolve_ceiling) + return managed_model_grants(_declared_grants(valid_token), tuple(tuple(ceiling.models) for ceiling in ceilings)) + + +async def check_agent_access_group_model_access( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> Literal[True]: + await can_agent_call_model( + model, valid_token, llm_router, None, None, None, purpose="Ceilings", resolve_ceiling=resolve_ceiling + ) + return True + + +async def check_agent_caller_model_access( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + await can_agent_call_model( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + purpose="Caller", + load_team=load_team, + load_user=load_user, + ) + + +def _model(source: _Source) -> str | Sequence[str]: + if source.model is None: + raise NativeAuthRuntimeError("managed model read without a request model") + return source.model + + +def _token(source: _Source) -> UserAPIKeyAuth: + if source.token is None: + raise NativeAuthRuntimeError("managed model read without a token") + return source.token + + +def _dispatch(source: _Source) -> str | Sequence[str]: + from litellm.proxy.auth.auth_checks import team_model_aliases_for_auth_check + + token: Final = _token(source) + return resolve_team_alias(_model(source), team_model_aliases_for_auth_check(token), token.team_id, source.router) + + +async def _read(facts: ManagedModelFacts, source: _Source) -> tuple[ManagedModelFacts, _Source]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + decision: Final = managed_models(facts) + match decision.kind: + case "Allow" | "Deny": + return facts, source + case "ReadManaged": + policy: Final = managed_agent_policy(source.token) + raw: Final = (policy.object_permission or {}).get("models", ()) if policy is not None else () + items: Final = _ITEMS.validate_python(raw) if isinstance(raw, (list, tuple)) else () + return await _read( + replace( + facts, + managed_loaded=True, + managed_present=policy is not None, + managed_models_valid=isinstance(raw, (list, tuple)), + managed_models_count=len(items), + ), + replace(source, managed_grants=items), + ) + case "ReadManagedScope": + dispatched: Final = _dispatch(source) + managed: Final = project_model_facts( + dispatched, + source.router, + _managed_grants(source.managed_grants), + team_id=_token(source).team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="agent", + ) + return await _read(replace(facts, managed=managed), source) + case "ReadCeilings": + ceilings: Final = await agent_access_group_ceilings(_token(source), source.ceiling) + return await _read( + replace(facts, ceilings=tuple(tuple(sorted(ceiling.models)) for ceiling in ceilings)), source + ) + case "ReadAliases": + return await _read(replace(facts, aliases_loaded=True), replace(source, dispatched=_dispatch(source))) + case "ReadCeilingScope": + if facts.ceilings is None: + raise NativeAuthRuntimeError("managed model scope without ceilings") + ceiling_scope: Final = project_model_facts( + source.dispatched, + source.router, + facts.ceilings[decision.index], + team_id=_token(source).team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="agent", + ) + return await _read(replace(facts, scopes=(*facts.scopes, ceiling_scope)), source) + case "ReadCaller": + caller: Final = agent_caller_auth(_token(source)) + return await _read( + replace(facts, caller_loaded=True, caller_present=caller is not None), replace(source, caller=caller) + ) + case "ReadTeam": + team: Final = await source.team_loader(_token(source)) + aliases: Final = key_model_aliases_for_auth_check(source.token) + return await _read( + replace(facts, team_loaded=True, team_present=team is not None), + replace(source, team=team, key_aliases=aliases), + ) + case "ReadTeamScope": + if source.team is None: + raise NativeAuthRuntimeError("missing caller team for model scope") + scope: Final = project_model_facts( + _model(source), + source.router, + object_model_grants(source.team), + team_id=source.team.team_id, + key_model_aliases=source.key_aliases, + object_type="team", + ) + return await _read( + replace(facts, team=GroupedModelScope(scope, tuple(source.team.access_group_ids or ()))), source + ) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing caller team group scope") + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=grouped), source) + case "ReadMember": + if source.team is None or source.caller is None or source.cache is None or source.logging is None: + raise NativeAuthRuntimeError("missing caller member model context") + member: Final = await read_member_model_facts( + _model(source), + source.team, + source.caller, + source.router, + source.prisma, + source.cache, + source.logging, + key_model_aliases=source.key_aliases, + ) + return await _read(replace(facts, member=member), source) + case "ReadUser": + user: Final = await source.user_loader(_token(source)) + return await _read( + replace( + facts, + user_loaded=True, + user_present=user is not None, + user_grants=object_model_grants(user) if user is not None else (), + ), + replace(source, user=user), + ) + case "ReadUserScope": + user_scope: Final = project_model_facts( + _model(source), source.router, facts.user_grants, key_model_aliases=source.key_aliases + ) + return await _read(replace(facts, user=user_scope), source) diff --git a/litellm/proxy/auth/member_model_adapter.py b/litellm/proxy/auth/member_model_adapter.py new file mode 100644 index 00000000000..08c2af532d6 --- /dev/null +++ b/litellm/proxy/auth/member_model_adapter.py @@ -0,0 +1,132 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.member_models import MemberModelFacts, member_models +from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import project_model_facts, raise_model_denial +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] + team: LiteLLM_TeamTable + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + membership: LiteLLM_TeamMembership | None + aliases: Mapping[str, str] | None + + +async def read_member_model_facts( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, + key_model_aliases: Mapping[str, str] | None = None, +) -> MemberModelFacts: + facts: Final = MemberModelFacts( + user_present=valid_token.user_id is not None, + team_present=team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # preserve unvalidated cache row identity guard + membership_loaded=team_membership_loaded, + ) + return await _read( + facts, + _Source( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + key_model_aliases, + ), + ) + + +async def check_team_member_model_access( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, + key_model_aliases: Mapping[str, str] | None = None, +) -> None: + facts: Final = await read_member_model_facts( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + key_model_aliases, + ) + decision: Final = member_models(facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete member model access facts") + raise_model_denial( + decision.denial, + model, + user_id=valid_token.user_id, + team_id=team_object.team_id, + member_models=decision.member_models, + ) + + +async def _read(facts: MemberModelFacts, source: _Source) -> MemberModelFacts: + from litellm.proxy.auth.auth_checks import get_team_membership + + decision: Final = member_models(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadMembership": + if source.token.user_id is None: + raise NativeAuthRuntimeError("member model lookup without user identity") + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=source.team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + return await _read(replace(facts, membership_loaded=True), replace(source, membership=membership)) + case "ReadGrants": + row: Final = source.membership.litellm_budget_table if source.membership is not None else None + return await _read( + replace(facts, grants=tuple(row.allowed_models or ()) if row is not None else ()), source + ) + case "ReadScope": + if facts.grants is None: + raise NativeAuthRuntimeError("member model scope read without grants") + scope: Final = project_model_facts( + source.model, + source.router, + facts.grants, + team_id=source.team.team_id, + key_model_aliases=source.aliases, + object_type="team", + ) + return await _read(replace(facts, scope=scope), source) diff --git a/litellm/proxy/auth/model_access_adapter.py b/litellm/proxy/auth/model_access_adapter.py new file mode 100644 index 00000000000..f5322635870 --- /dev/null +++ b/litellm/proxy/auth/model_access_adapter.py @@ -0,0 +1,310 @@ +from collections.abc import Mapping, Sequence +from typing import Final, Literal + +from fastapi import HTTPException +from pydantic import BaseModel + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ( + ModelAccessDecision, + ModelAccessFacts, + ModelAliases, + ModelObject, + ModelRegistry, + grouped_model_access, + key_model_grants, + model_access, + user_model_precheck, +) +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, + ModelAccessDeniedProxyException, + ProxyErrorTypes, + UserAPIKeyAuth, +) +from litellm.proxy.auth.model_access_denied import ( + customer_model_access_denied_client_message, + model_access_denied_client_message, +) +from litellm.proxy.auth.model_projection_adapter import project_scope, read_grouped_scope +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader + + +class ModelLists(BaseModel, frozen=True): + models: tuple[str, ...] = () + + +@native_boundary +def object_model_grants(source: object) -> tuple[str, ...]: + return ModelLists.model_validate(source, from_attributes=True).models + + +@native_boundary +def _initial_model_facts( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: ModelObject = "user", + fallback_depth: int = 0, +) -> ModelAccessFacts: + requested: Final = (model,) if isinstance(model, str) else tuple(model) + aliases: Final = ModelAliases( + team=dict(team_model_aliases or {}), + key=dict(key_model_aliases or {}), + global_=dict(litellm.model_alias_map), + team_order=tuple(team_model_aliases or ()), + ) + registry: Final = ModelRegistry( + present=llm_router is not None, + helper_present=bool(llm_router), + evaluations=(), + ) + initial: Final = ModelAccessFacts( + requested=requested, + grants=tuple(models), + team_id=team_id, + aliases=aliases, + registry=registry, + object_type=object_type, + fallback_depth=fallback_depth, + max_depth=DEFAULT_MAX_RECURSE_DEPTH, + is_list=not isinstance(model, str), + ) + return initial + + +def project_model_facts( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: ModelObject = "user", + fallback_depth: int = 0, +) -> ModelAccessFacts: + initial: Final = _initial_model_facts( + model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type, fallback_depth + ) + return project_scope(initial, llm_router) + + +def raise_model_denial( + decision: ModelAccessDecision, + requested: str | Sequence[str], + *, + user_id: str | None = None, + team_id: str | None = None, + teams: Sequence[str] | None = None, + member_models: Sequence[str] | None = None, +) -> Literal[True]: + model: Final = requested if isinstance(requested, str) else list(requested) + match decision.kind: + case "Allow": + return True + case "Depth": + raise Exception( # noqa: TRY002 # preserve the legacy depth failure type and envelope + f"Unable to parse model, max fallback depth exceeded - received model: {decision.model if decision.model is not None else model}" + ) + case "NoAgentModels": + raise HTTPException(403, "This agent has no model grants") + case "ReadGroups" | "ReadScope" | "ReadGroupScope": + raise NativeAuthRuntimeError("model grant lookup was not answered") + case "Deny" | "NoDefaultModels" | "EmptyAgentCeiling" | "PersonalTeam": + error_model: Final = decision.model if decision.model is not None else model + object_type: Final = decision.object_type + message: Final = ( + customer_model_access_denied_client_message(error_model) + if object_type == "customer" + else model_access_denied_client_message(error_model) + ) + internal: Final = ( + f"Team member not allowed to access model. User={user_id}, Team={team_id}, Model={model}. Allowed member models = {list(member_models)}" + if member_models is not None + else ( + f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}" + if decision.kind == "NoDefaultModels" + else ( + f"agent {user_id} access groups grant no models" + if decision.kind == "EmptyAgentCeiling" + else ( + f"Personal key model access is derived from the owner's teams. User={user_id}, teams={list(teams) if teams is not None else None}, model={error_model}" + if decision.kind == "PersonalTeam" + else ( + f"{object_type} not allowed to access model. This {object_type} can only access models={list(decision.grants)}. Tried to access {error_model}" + ) + ) + ) + ) + ) + error_type: Final = ( + ProxyErrorTypes.team_model_access_denied + if member_models is not None + else ( + ProxyErrorTypes.key_model_access_denied + if decision.kind in ("NoDefaultModels", "PersonalTeam") + else ( + ProxyErrorTypes.agent_model_access_denied + if decision.kind == "EmptyAgentCeiling" + else ProxyErrorTypes.get_model_access_error_type_for_object(object_type) + ) + ) + ) + raise ModelAccessDeniedProxyException( + message=message, internal_message=internal, type=error_type, param="model", code=403 + ) + + +def can_object_call_model( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: ModelObject = "user", + fallback_depth: int = 0, +) -> Literal[True]: + facts: Final = project_model_facts( + model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type, fallback_depth + ) + return raise_model_denial(model_access(facts), model) + + +async def _grouped_model_access( + facts: ModelAccessFacts, + group_ids: Sequence[str], + prisma_client: DatabaseClient | None, + *, + merge: bool, + router: RegistryReader | None = None, +) -> ModelAccessDecision: + scope: Final = await read_grouped_scope( + GroupedModelScope(facts, tuple(group_ids)), router, prisma_client, merge=merge + ) + return grouped_model_access(scope.facts, scope.group_ids, scope.group_models, merge=merge) + + +async def can_key_call_model( + model: str | Sequence[str], + llm_model_list: Sequence[object] | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: DatabaseClient | None = None, +) -> Literal[True]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + grants: Final = key_model_grants( + object_model_grants(valid_token), valid_token.team_id, tuple(valid_token.team_models or ()) + ) + facts: Final = project_model_facts( + model, + llm_router, + grants, + valid_token.team_model_aliases, + valid_token.team_id, + key_model_aliases_for_auth_check(valid_token), + "key", + ) + decision: Final = await _grouped_model_access( + facts, valid_token.access_group_ids or (), prisma_client, merge=False, router=llm_router + ) + return raise_model_denial(decision, model) + + +async def can_team_access_model( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> Literal[True]: + facts: Final = project_model_facts( + model, + llm_router, + team_object.models if team_object is not None else (), + team_model_aliases, + team_object.team_id if team_object is not None else None, + key_model_aliases, + "team", + ) + decision: Final = await _grouped_model_access( + facts, + (team_object.access_group_ids or ()) if team_object is not None else (), + prisma_client, + merge=True, + router=llm_router, + ) + return raise_model_denial(decision, model) + + +def can_org_access_model( + model: str, + org_object: LiteLLM_OrganizationTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, +) -> Literal[True]: + return can_object_call_model( + model, llm_router, org_object.models if org_object is not None else (), team_model_aliases, object_type="org" + ) + + +def can_project_access_model( + model: str | Sequence[str], + project_object: LiteLLM_ProjectTable, + llm_router: RegistryReader | None, + key_model_aliases: Mapping[str, str] | None = None, +) -> Literal[True]: + return can_object_call_model( + model, + llm_router, + project_object.models, + team_id=project_object.team_id, + key_model_aliases=key_model_aliases, + object_type="project", + ) + + +def can_customer_access_model( + model: str | Sequence[str], + end_user_object: LiteLLM_EndUserTable, + llm_router: RegistryReader | None, + valid_token: UserAPIKeyAuth | None, +) -> Literal[True]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check, team_model_aliases_for_auth_check + + facts: Final = project_model_facts( + model, + llm_router, + end_user_object.models or (), + team_model_aliases_for_auth_check(valid_token) if valid_token is not None else None, + valid_token.team_id if valid_token is not None else None, + key_model_aliases_for_auth_check(valid_token), + "customer", + ) + return raise_model_denial(model_access(facts, kind="Customer"), model) + + +async def can_user_call_model( + model: str | Sequence[str], + llm_router: RegistryReader | None, + user_object: LiteLLM_UserTable | None, + key_model_aliases: Mapping[str, str] | None = None, +) -> Literal[True]: + grants: Final = object_model_grants(user_object) if user_object is not None else () + precheck: Final = user_model_precheck(present=user_object is not None, grants=grants) + if precheck.kind != "ReadScope": + return raise_model_denial(precheck, model) + return can_object_call_model(model, llm_router, grants, key_model_aliases=key_model_aliases) diff --git a/litellm/proxy/auth/model_attribution_adapter.py b/litellm/proxy/auth/model_attribution_adapter.py new file mode 100644 index 00000000000..f7ea8e588e5 --- /dev/null +++ b/litellm/proxy/auth/model_attribution_adapter.py @@ -0,0 +1,213 @@ +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.model_access import key_model_grants +from litellm._auth.model_attribution import ModelGroupAttributionFacts, model_group_attribution +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_ProjectTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Sources: + model: tuple[str, ...] + token: UserAPIKeyAuth + team: LiteLLM_TeamTable | None + project: LiteLLM_ProjectTable | None + router: RegistryReader + prisma: PrismaClient + cache: UserApiKeyCache + logging: ProxyLogging + membership: LiteLLM_TeamMembership | None + + +async def read_model_group_attribution_facts( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTable | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> ModelGroupAttributionFacts: + facts: Final = ModelGroupAttributionFacts( + model_present=model is not None, + token_present=valid_token is not None, + router_present=llm_router is not None, + database_present=prisma_client is not None, + team_present=team_object is not None, + user_present=valid_token is not None and valid_token.user_id is not None, + membership_loaded=team_membership_loaded, + token_org=valid_token.org_id if valid_token is not None else None, + team_org=team_object.organization_id if team_object is not None else None, + ) + if model_group_attribution(facts).kind == "Complete": + return facts + if model is None or valid_token is None or llm_router is None or prisma_client is None: + raise NativeAuthRuntimeError("attribution read requested without its source") + return await _read( + facts, + _Sources( + (model,) if isinstance(model, str) else tuple(model), + valid_token, + team_object, + project_object, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + ), + ) + + +def _covering_groups(source: _Sources) -> Iterator[str]: + for model in source.model: + yield from source.router.registry().access_groups(model_name=model, team_id=source.token.team_id) + + +async def _read(facts: ModelGroupAttributionFacts, source: _Sources) -> ModelGroupAttributionFacts: + from litellm.proxy.auth.auth_checks import ( + _load_model_access_group_registry, # pyright: ignore[reportPrivateUsage] # reuse the existing registry IO without its attribution policy + ) + + decision: Final = model_group_attribution(facts) + match decision.kind: + case "Complete": + return facts + case "ReadRegistry": + registry: Final = await _load_model_access_group_registry(source.prisma, source.cache) + return await _read( + replace(facts, registry_loaded=True, registry=tuple(registry) if registry is not None else None), source + ) + case "ReadCovering": + return await _read(replace(facts, covering=tuple(_covering_groups(source))), source) + case "ReadKey": + key: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + return await _read(replace(facts, grants=(*facts.grants, key)), source) + case "ReadTeam": + team: Final = object_model_grants(source.team) if source.team is not None else () + return await _read(replace(facts, grants=(*facts.grants, team)), source) + case "ReadMember": + membership: Final = await _read_membership(source) if decision.lookup else source.membership + member: Final = ( + tuple(membership.litellm_budget_table.allowed_models or ()) + if decision.include and membership is not None and membership.litellm_budget_table is not None + else () + ) + return await _read(replace(facts, grants=(*facts.grants, member)), source) + case "ReadProject": + project: Final = object_model_grants(source.project) if source.project is not None else () + return await _read(replace(facts, grants=(*facts.grants, project)), source) + case "ReadOrg": + org, failed = await _org_grants(decision.id, source) + return await _read(replace(facts, grants=(*facts.grants, org), org_failed=failed), source) + + +async def _read_membership(source: _Sources) -> LiteLLM_TeamMembership | None: + from litellm.proxy.auth.auth_checks import get_team_membership + + if source.team is None or source.token.user_id is None: + raise NativeAuthRuntimeError("member attribution lookup requested without identity") + return await get_team_membership( + user_id=source.token.user_id, + team_id=source.team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + + +async def _org_grants(org_id: str | None, source: _Sources) -> tuple[tuple[str, ...], bool]: + from litellm.proxy.auth.auth_checks import get_org_object + + if org_id is None: + return (), False + try: + org: Final = await get_org_object( + org_id=org_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve attribution's org IO failure behavior, never swallow native failure + verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", error) + return (), True + return (tuple(org.models) if org is not None else ()), False + + +async def collect_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTable | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> tuple[str, ...]: + facts: Final = await read_model_group_attribution_facts( + model, + valid_token, + team_object, + project_object, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + ) + return model_group_attribution(facts).groups + + +async def stamp_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTable | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> tuple[str, ...]: + try: + facts: Final = await read_model_group_attribution_facts( + model, + valid_token, + team_object, + project_object, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # native attribution decides the legacy telemetry IO failure result + verbose_proxy_logger.debug("model access group attribution failed: %s", error) + return model_group_attribution(ModelGroupAttributionFacts(io_failed=True)).groups + decision: Final = model_group_attribution(facts) + if decision.stamp: + if valid_token is None: + raise NativeAuthRuntimeError("attribution stamping requested without token") + valid_token.matched_model_access_groups = list(decision.groups) # rebind-ok: native attribution effect + return decision.groups diff --git a/litellm/proxy/auth/model_budget_adapter.py b/litellm/proxy/auth/model_budget_adapter.py new file mode 100644 index 00000000000..9c1833f2919 --- /dev/null +++ b/litellm/proxy/auth/model_budget_adapter.py @@ -0,0 +1,306 @@ +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from typing import Final + +from pydantic import TypeAdapter + +import litellm +from litellm._auth import model_budget_aliases, model_budget_data +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.model_budget_data import ( + ModelBudgetEntity, + ModelBudgetEntry, + ModelBudgetQueryFacts, + ModelBudgetResolution, + ModelBudgetResolveFacts, + ModelBudgetSnapshot, +) +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry +from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.hooks.model_budget_shared_windows import SharedWindows +from litellm.proxy.hooks.model_max_budget_limiter import ModelAccessGroups, model_budget_spend_cache_key +from litellm.router_strategy.budget_limiter import budget_window_seconds +from litellm.types.utils import BudgetConfig + +_FLOAT: Final = TypeAdapter(float) +_OBJECT: Final = TypeAdapter(object) +_ENTITIES: Final[Mapping[ModelBudgetEntity, Litellm_EntityType]] = { + "Key": Litellm_EntityType.KEY, + "User": Litellm_EntityType.USER, + "Team": Litellm_EntityType.TEAM, + "TeamMember": Litellm_EntityType.TEAM_MEMBER, + "EndUser": Litellm_EntityType.END_USER, +} + + +def _entry(name: str, raw: object) -> ModelBudgetEntry: + try: + config: Final = BudgetConfig.model_validate(raw) + seconds: Final = duration_in_seconds(config.budget_duration) if config.budget_duration is not None else None + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve malformed entry fallthrough + return ModelBudgetEntry(name, None, None, None) + return ModelBudgetEntry(name, config.max_budget, config.budget_duration, seconds) + + +def _spend(raw: object) -> float: + try: + return _FLOAT.validate_python(raw or 0.0) + except NativeAuthFailure: + raise + except ValueError: + return 0.0 + + +def _price_entry(model: str) -> object: + return _OBJECT.validate_python(litellm.model_cost.get(model)) # pyright: ignore[reportUnknownMemberType] # raw catalog IO + + +@dataclass(frozen=True, slots=True) +class ModelBudgetReading: + snapshot: ModelBudgetSnapshot + budget_model: str | None + named_model: bool + + +@dataclass(frozen=True, slots=True) +class ModelBudgetIO: + cache: DualCache + model_access_groups: ModelAccessGroups + shared_windows: SharedWindows + process_started_at: float + clock: Callable[[], float] = time.time + monotonic: Callable[[], float] = time.monotonic + price_entry: Callable[[str], object] = _price_entry + + async def counter(self, key: str) -> float: + redis: Final = self.cache.redis_cache + raw: Final = _OBJECT.validate_python( + await self.cache.async_get_cache(key=key) if redis is None else await redis.async_get_cache(key=key) # pyright: ignore[reportUnknownMemberType] # cache APIs remain legacy IO + ) + return _spend(raw) + + +@dataclass(frozen=True, slots=True) +class AliasReading: + row: model_budget_aliases.Row + source: AliasedCostEntry | None + + +def read_alias_models(model: str, io: ModelBudgetIO, readings: tuple[AliasReading, ...] = ()) -> tuple[str, ...]: + decision: Final = model_budget_aliases.decide( + model_budget_aliases.Facts(model, tuple(reading.row for reading in readings)) + ) + match decision.read: + case None: + return decision.models + case "Entry": + raw: Final = io.price_entry(decision.key) + source: Final = raw if isinstance(raw, AliasedCostEntry) else None + entry: Final = AliasReading( + model_budget_aliases.Row(decision.key, source.names if source is not None else ()), source + ) + return read_alias_models(model, io, (*readings, entry)) + case "Prefix": + prefix_known: Final = decision.key.partition("/")[0] in litellm.provider_list + updated: Final = tuple( + replace(reading, row=replace(reading.row, prefix_known=prefix_known)) + if reading.row.key == decision.key + else reading + for reading in readings + ) + return read_alias_models(model, io, updated) + case "Provider": + selected: Final = next(reading for reading in readings if reading.row.key == decision.key) + assert selected.source is not None + provider: Final = selected.source.get("litellm_provider") + provided: Final = replace( + selected, + row=replace(selected.row, provider_read=True, provider=provider if isinstance(provider, str) else None), + ) + return read_alias_models( + model, io, tuple(provided if reading.row.key == decision.key else reading for reading in readings) + ) + + +def _resolve_facts( + model: str, + budgets: Mapping[str, object], + io: ModelBudgetIO, + key_budgets: Mapping[str, object] | None = None, + *, + team: bool = False, + excluded_entries: tuple[str, ...] = (), +) -> ModelBudgetResolveFacts: + base: Final = get_bedrock_base_model(model) + raw_price: Final = as_str_mapping(io.price_entry(base)) or {} + return ModelBudgetResolveFacts( + model=model, + bedrock_base=base, + bedrock_provider=str(raw_price.get("litellm_provider") or ""), + entries=tuple(_entry(name, config) for name, config in budgets.items()), + key_entries=tuple(_entry(name, config) for name, config in (key_budgets or {}).items()), + team=team, + excluded_entries=excluded_entries, + ) + + +def resolve_model_budget( + model: str, + budgets: Mapping[str, object], + io: ModelBudgetIO, + team_id: str | None, + key_budgets: Mapping[str, object] | None = None, + *, + team: bool = False, + excluded_entries: tuple[str, ...] = (), +) -> ModelBudgetResolution: + facts: Final = _resolve_facts(model, budgets, io, key_budgets, team=team, excluded_entries=excluded_entries) + initial: Final = model_budget_data.resolve(facts) + if not initial.read_groups: + return initial + groups: Final = tuple(io.model_access_groups(model, team_id)) + return model_budget_data.resolve(replace(facts, access_groups=groups)) + + +async def read_model_budget( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, +) -> ModelBudgetReading: + models: Final = read_alias_models(model, io) + exclusions: Final = ( + model_budget_data.team_exclusions( + tuple(_resolve_facts(candidate, {}, io, key_budgets, team=True) for candidate in models) + ) + if entity == "Team" + else () + ) + scope_models: Final = read_alias_models(model, io) if entity == "Team" else models + return await _read_scope(io, entity, entity_id, team_id, scope_models, budgets, exclusions) + + +async def _read_scope( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + models: tuple[str, ...], + budgets: Mapping[str, object] | None, + excluded_entries: tuple[str, ...], + readings: tuple[ModelBudgetReading, ...] = (), +) -> ModelBudgetReading: + decision: Final = model_budget_aliases.scope( + model_budget_aliases.ScopeFacts(models, tuple(reading.snapshot for reading in readings), bool(budgets)) + ) + if decision.read_index is not None: + reading: Final = await _read_single( + io, entity, entity_id, team_id, models[decision.read_index], budgets, excluded_entries + ) + return await _read_scope( + io, entity, entity_id, team_id, models, budgets, excluded_entries, (*readings, reading) + ) + return ( + readings[decision.result_index] + if decision.result_index is not None + else ModelBudgetReading(ModelBudgetSnapshot(), None, False) + ) + + +async def _read_single( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + excluded_entries: tuple[str, ...], +) -> ModelBudgetReading: + resolution: Final = resolve_model_budget(model, budgets or {}, io, team_id, excluded_entries=excluded_entries) + entry: Final = resolution.entry + if entry is None: + return ModelBudgetReading(ModelBudgetSnapshot(), None, False) + spend_key: Final = model_budget_spend_cache_key(_ENTITIES[entity], entity_id, entry.name, entry.duration) + legacy_key: Final = model_budget_spend_cache_key(_ENTITIES[entity], entity_id, model, entry.duration) + facts: Final = ModelBudgetQueryFacts( + entity=entity, + snapshot=ModelBudgetSnapshot(max_budget=entry.max_budget), + duration_seconds=entry.duration_seconds, + process_uptime=io.monotonic() - io.process_started_at, + same_counter=spend_key == legacy_key, + redis_present=io.cache.redis_cache is not None, + ) + snapshot: Final = await _read_snapshot(io, facts, spend_key, legacy_key, entry.duration) + return ModelBudgetReading(snapshot, entry.name, resolution.named_model) + + +async def _read_snapshot( + io: ModelBudgetIO, + facts: ModelBudgetQueryFacts, + key: str, + legacy_key: str, + duration: str | None, +) -> ModelBudgetSnapshot: + decision: Final = model_budget_data.query(facts) + match decision.read: + case None: + return facts.snapshot + case "Local": + snapshot: Final = replace(facts.snapshot, local_spend=await io.counter(key)) + return await _read_snapshot(io, replace(facts, snapshot=snapshot), key, legacy_key, duration) + case "Legacy": + legacy_snapshot: Final = replace(facts.snapshot, legacy_spend=await io.counter(legacy_key)) + return await _read_snapshot(io, replace(facts, snapshot=legacy_snapshot), key, legacy_key, duration) + case "Shared": + now: Final = io.clock() + shared: Final = ( + await io.shared_windows.read(spend_key=key, budget_duration=duration) if duration is not None else None + ) + shared_snapshot: Final = replace( + facts.snapshot, + shared_spend=shared.spend if shared is not None else None, + shared_started_at=shared.started_at if shared is not None else None, + shared_window_seconds=budget_window_seconds(shared.started_at, duration) + if shared is not None and duration is not None + else None, + now=now, + ) + return await _read_snapshot( + io, replace(facts, snapshot=shared_snapshot, shared_read=True), key, legacy_key, duration + ) + + +async def enforce_model_budget( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + exceeded_message: str, + key_budgets: Mapping[str, object] | None = None, +) -> bool: + reading: Final = await read_model_budget(io, entity, entity_id, team_id, model, budgets, key_budgets) + decision: Final = model_budget_data.check(reading.snapshot) + if decision.denied: + assert reading.snapshot.max_budget is not None + raise litellm.BudgetExceededError( + message=exceeded_message + if reading.named_model + else f"{exceeded_message}, budget_model={reading.budget_model}", + current_cost=decision.current_spend, + max_budget=reading.snapshot.max_budget, + entity_type=_ENTITIES[entity].value, + entity_id=entity_id, + ) + return True diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 7f2aa33246d..2646b28e8a1 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -14,14 +14,16 @@ from types import MappingProxyType from typing import Final, TypeVar import litellm +from litellm._auth.failure import NativeAuthFailure from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping -from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialModelNames, UserAPIKeyAuth from litellm.proxy.common_utils.model_listing_utils import CallerAliases, alias_map, dispatched_model from litellm.proxy.common_utils.token_count_settings import cost_map_entry +from litellm.repositories.native_auth_repository import native_auth_repository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.responses.main import strip_responses_routing_prefix from litellm.router import Router @@ -132,12 +134,21 @@ async def get_mcp_server_ids( # Make a direct SQL query to get just the mcp_servers try: - result: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( - where={"object_permission_id": user_api_key_dict.object_permission_id}, + native: Final = native_auth_repository(prisma_client) + result: Final = ( + await native.one( + "object_permission", [user_api_key_dict.object_permission_id], LiteLLM_ObjectPermissionTable + ) + if native is not None + else await ObjectPermissionRepository(prisma_client).table.find_unique( + where={"object_permission_id": user_api_key_dict.object_permission_id}, + ) ) if result and result.mcp_servers: return result.mcp_servers return [] + except NativeAuthFailure: + raise except Exception: return [] diff --git a/litellm/proxy/auth/model_group_adapter.py b/litellm/proxy/auth/model_group_adapter.py new file mode 100644 index 00000000000..c6e0ed4f060 --- /dev/null +++ b/litellm/proxy/auth/model_group_adapter.py @@ -0,0 +1,94 @@ +from dataclasses import dataclass, replace +from typing import Final, Literal, TypeAlias + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.model_groups import GroupLookup, GroupResources, ResourceGroupFacts, resource_group_decision +from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth + +ResourceField: TypeAlias = Literal[ + "access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names" +] + + +@dataclass(frozen=True, slots=True) +class GroupRead: + lookup: GroupLookup + error: Exception | None = None + + +async def get_authorized_resources_from_key_access_groups( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + resource_field: ResourceField, + check_db_only: bool = False, + strict_errors: bool = False, +) -> list[str]: # mutable-ok: retain the legacy public return contract + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + facts: Final = ResourceGroupFacts( + subject_present=valid_token is not None, + database_present=prisma_client is not None, + cache_present=_present(user_api_key_cache), + group_ids=tuple(valid_token.access_group_ids or ()) if valid_token is not None else (), + team_id=(valid_token.team_id if valid_token is not None else None) + or (team_object.team_id if team_object is not None else None), + key_id=valid_token.token if valid_token is not None else None, + check_db_only=check_db_only, + strict_errors=strict_errors, + ) + return list(set(await _drive(facts, resource_field, ()))) + + +async def _read(group_id: str, resource_field: ResourceField, check_db_only: bool) -> GroupRead: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + group: Final = await get_access_object( + access_group_id=group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + resources: Final = { + "access_model_names": group.access_model_names, + "access_mcp_server_ids": group.access_mcp_server_ids, + "access_agent_ids": group.access_agent_ids, + "access_skill_names": group.access_skill_names, + }[resource_field] + return GroupRead( + GroupLookup( + "Found", + GroupResources( + tuple(group.assigned_team_ids or ()), tuple(group.assigned_key_ids or ()), tuple(resources or ()) + ), + ) + ) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError as error: + return GroupRead(GroupLookup("Missing"), error) + except Exception as error: # noqa: BLE001 # native strict-error policy decides whether to rethrow this lookup failure + return GroupRead(GroupLookup("Failed"), error) + + +async def _drive( + facts: ResourceGroupFacts, resource_field: ResourceField, reads: tuple[GroupRead, ...] +) -> tuple[str, ...]: + decision: Final = resource_group_decision(facts) + match decision.kind: + case "Complete": + return decision.resources + case "Raise": + error: Final = reads[decision.index].error + if error is None: + raise NativeAuthRuntimeError("missing access group lookup failure") + raise error + case "Read": + result: Final = await _read(facts.group_ids[decision.index], resource_field, facts.check_db_only) + return await _drive(replace(facts, loaded=(*facts.loaded, result.lookup)), resource_field, (*reads, result)) + + +def _present(value: object) -> bool: + return value is not None diff --git a/litellm/proxy/auth/model_helper_adapter.py b/litellm/proxy/auth/model_helper_adapter.py new file mode 100644 index 00000000000..b6b02d27281 --- /dev/null +++ b/litellm/proxy/auth/model_helper_adapter.py @@ -0,0 +1,55 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.model_helpers import ModelHelperFacts, model_helper +from litellm.proxy.auth.model_projection_adapter import provider_facts +from litellm.router_api import RegistryReader + + +def read_model_helper_facts(facts: ModelHelperFacts, router: RegistryReader | None) -> ModelHelperFacts: + if facts.provider_reads is None and not facts.provider_loaded: + return read_model_helper_facts(replace(facts, provider_reads=()), router) + decision: Final = model_helper(facts) + match decision.kind: + case "Complete": + return facts + case "ReadGroups": + if router is None: + raise NativeAuthRuntimeError("model helper registry read without router") + groups: Final = tuple(router.registry().access_groups(model_name=facts.model, team_id=facts.team_id)) + return read_model_helper_facts(replace(facts, groups=groups), router) + case "ReadNames": + if router is None: + raise NativeAuthRuntimeError("model helper names read without router") + return read_model_helper_facts(replace(facts, names=tuple(router.registry().names())), router) + case "ReadProvider": + return read_model_helper_facts( + replace( + facts, + provider_loaded=True, + provider_reads=(*(facts.provider_reads or ()), provider_facts(facts.model)), + ), + router, + ) + + +def check_model_access_helper( + model: str, + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, +) -> bool: + facts: Final = ModelHelperFacts( + "Helper", model, tuple(models), team_id, tuple(team_model_aliases or ()), bool(llm_router) + ) + return model_helper(read_model_helper_facts(facts, llm_router)).allowed + + +def model_in_access_group(model: str, team_models: Sequence[str] | None, llm_router: RegistryReader | None) -> bool: + facts: Final = ModelHelperFacts( + "Group", model, tuple(team_models) if team_models is not None else None, router_present=bool(llm_router) + ) + return model_helper(read_model_helper_facts(facts, llm_router)).allowed diff --git a/litellm/proxy/auth/model_policy_read.py b/litellm/proxy/auth/model_policy_read.py new file mode 100644 index 00000000000..77e0a6c13db --- /dev/null +++ b/litellm/proxy/auth/model_policy_read.py @@ -0,0 +1,11 @@ +from typing import TypeVar + +from litellm._auth.failure import NativeAuthRuntimeError + +_T = TypeVar("_T") # rebind-ok: TypeVar must remain a bare static type-variable binding + + +def require_native_fact(value: _T | None) -> _T: + if value is None: + raise NativeAuthRuntimeError("missing fact for native model authorization read") + return value diff --git a/litellm/proxy/auth/model_projection_adapter.py b/litellm/proxy/auth/model_projection_adapter.py new file mode 100644 index 00000000000..b95bf682867 --- /dev/null +++ b/litellm/proxy/auth/model_projection_adapter.py @@ -0,0 +1,121 @@ +from collections.abc import Iterable +from dataclasses import dataclass, replace +from itertools import accumulate, repeat +from typing import Final, assert_never + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ( + ModelAccessFacts, + ModelEvaluation, + ModelProvider, + ModelReading, + ModelReadKind, + grouped_model_access, +) +from litellm._auth.model_projection import ModelProjectionDecision, model_projection +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader +from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent + + +def provider_facts(model: str) -> ModelProvider | None: + try: + stripped, provider, _, _ = litellm.get_llm_provider(model=model) + return ModelProvider(stripped, provider) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # provider inference failure is an absent provider, as in the legacy helper + return None + + +def project_scope( + facts: ModelAccessFacts, router: RegistryReader | None, reads: tuple[ModelReading, ...] = () +) -> ModelAccessFacts: + initial: Final = _Projection(reads, model_projection(facts, reads)) + + def advance(state: _Projection, _tick: None) -> _Projection: + decision: Final = state.decision + if decision.read is None or decision.model is None: + raise NativeAuthRuntimeError("incomplete model registry read") + reading: Final = read_registry(decision.read, decision.model, decision.team_id, router) + updated: Final = (*state.reads, reading) + return _Projection(updated, model_projection(facts, updated)) + + complete: Final = next( + state for state in accumulate(repeat(None), advance, initial=initial) if state.decision.kind == "Complete" + ) + return replace( + facts, + registry=replace( + facts.registry, + evaluations=( + *(facts.registry.evaluations or ()), + ModelEvaluation(facts.grants, complete.reads, facts.aliases), + ), + ), + ) + + +@dataclass(frozen=True, slots=True) +class _Projection: + reads: tuple[ModelReading, ...] + decision: ModelProjectionDecision + + +@native_boundary +def _strings(values: Iterable[str]) -> tuple[str, ...]: + return tuple(values) + + +def read_registry(kind: ModelReadKind, model: str, team_id: str | None, router: RegistryReader | None) -> ModelReading: + if kind == "Compaction": + return ModelReading(kind, model, value=native_compaction_parent(model)) + if kind == "Provider": + return ModelReading(kind, model, provider=provider_facts(model)) + if router is None: + raise NativeAuthRuntimeError("model registry read without router") + match kind: + case "Aliases": + return ModelReading(kind, model, values=_strings(router.registry().aliases())) + case "AliasTarget": + return ModelReading(kind, model, value=router.registry().resolve(model, kind="alias")) + case "Names": + return ModelReading(kind, model, values=_strings(router.registry().names())) + case "Groups": + return ModelReading( + kind, model, values=_strings(router.registry().access_groups(model_name=model, team_id=team_id)) + ) + case "Deployment": + deployment: Final = router.registry().deployments(model, source="first_in_group") + return ModelReading( + kind, + model, + value=deployment.model_info.team_public_model_name if deployment is not None else None, + team=deployment.model_info.team_id if deployment is not None else None, + ) + case _: + assert_never(kind) + + +async def read_grouped_scope( + scope: GroupedModelScope, router: RegistryReader | None, prisma: DatabaseClient | None = None, *, merge: bool +) -> GroupedModelScope: + from litellm.proxy.auth.auth_checks import get_models_from_access_groups + + decision: Final = grouped_model_access(scope.facts, scope.group_ids, scope.group_models, merge=merge) + if decision.kind == "ReadGroups": + models: Final = await get_models_from_access_groups(scope.group_ids, prisma_client=prisma) + return await read_grouped_scope(replace(scope, group_models=tuple(models)), router, prisma, merge=merge) + if decision.kind == "ReadGroupScope": + scoped: Final = project_scope( + replace( + scope.facts, + grants=decision.grants, + aliases=replace(scope.facts.aliases, global_=dict(litellm.model_alias_map)), + ), + router, + ) + return replace(scope, facts=replace(scoped, grants=scope.facts.grants)) + return scope diff --git a/litellm/proxy/auth/org_route_adapter.py b/litellm/proxy/auth/org_route_adapter.py new file mode 100644 index 00000000000..ff008f66552 --- /dev/null +++ b/litellm/proxy/auth/org_route_adapter.py @@ -0,0 +1,130 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.org_routes import ( + AdminFacts, + Identifier, + OrganizationMembership, + OrganizationRouteFacts, + admin, + organization_route, +) +from litellm.proxy._types import LiteLLM_UserTable, LiteLLMRoutes, ProxyErrorTypes, ProxyException + +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) + + +def memberships(user: LiteLLM_UserTable | None) -> tuple[OrganizationMembership, ...] | None: + return ( + tuple(OrganizationMembership(item.organization_id, item.user_role) for item in user.organization_memberships) + if user is not None and user.organization_memberships is not None + else None + ) + + +def user_is_org_admin(request_data: Mapping[str, object], user_object: LiteLLM_UserTable | None = None) -> bool: + facts = AdminFacts(memberships(user_object)) # rebind-ok: native-directed read stage + singular: object = None # rebind-ok: raw request field loaded only after the native directive + plural: tuple[object, ...] = () # rebind-ok: raw request field loaded only after the native directive + while True: + decision = admin(facts) + match decision.action: + case "ReadRequest": + singular = request_data.get("organization_id") + raw_plural = request_data.get("organizations") + plural = _OBJECTS.validate_python(raw_plural) if isinstance(raw_plural, list) else () + facts = replace( + facts, request_read=True, singular_present=singular is not None, plural_count=len(plural) + ) + case "ReadSingular" | "ReadPlural": + raw = singular if decision.action == "ReadSingular" else plural[decision.plural_index] + facts = replace(facts, candidate_index=decision.candidate_index, candidate=_identifier(raw)) + case "Allow": + return True + case "Deny": + return False + case "HashFailure": + raise TypeError(decision.hash_failure) + + +def _identifier(value: object) -> Identifier: + try: + hash(value) + except NativeAuthFailure: + raise + except TypeError as error: + return Identifier(None, str(error)) + return Identifier(value if isinstance(value, str) else None) + + +def organization_role_based_access_check( + request_body: Mapping[str, object], user_object: LiteLLM_UserTable | None, route: str +) -> None: + facts = OrganizationRouteFacts( # rebind-ok: native-directed read driver + route, + user_object is not None, + None, + None, + None, + tuple(LiteLLMRoutes.org_admin_only_routes.value), + request_pending=True, + ) + passed: object = None # rebind-ok: raw field is read only at the native-selected stage + while True: + decision = organization_route(facts) + if decision.read_request: + assert user_object is not None + passed = request_body.get("organization_id") + facts = replace( + facts, + user_role=user_object.user_role, + organization_id=passed if isinstance(passed, str) else None, + organization_id_present=passed is not None, + request_pending=False, + memberships_pending=True, + identifier_pending=True, + ) + elif decision.read_memberships: + facts = replace(facts, memberships=memberships(user_object), memberships_pending=False) + elif decision.read_identifier: + facts = replace(facts, identifier_pending=False, hash_failure=_identifier(passed).hash_failure) + else: + break + if decision.hash_failure is not None: + raise TypeError(decision.hash_failure) + organization_id: Final = passed + role: Final = facts.user_role + messages: Final = { + "CreateOrganization": (f"Only proxy admins can create new organizations. You are {role}", "user_role"), + "NoMemberships": ( + f"Tried to access route={route} but you are not a member of any organization. Please contact the proxy admin to request access.", + "organization_id", + ), + "MissingOrganization": ( + "Passed organization_id is None, please pass an organization_id in your request", + "organization_id", + ), + "MissingRole": ( + f"You do not have a role within the selected organization. Passed organization_id: {organization_id}. Please contact the organization admin to request access.", + "organization_id", + ), + "WrongRole": ( + f"You do not have the required role to perform {route} in Organization {organization_id}. Your role is {decision.organization_role} in Organization {organization_id}", + "user_role", + ), + "TeamMissingOrganization": ( + f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {list(decision.organization_ids)}", + "organization_id", + ), + "TeamWrongRole": ( + f"You do not have the required role to call {route}. Your role is {decision.organization_role} in Organization {organization_id}", + "user_role", + ), + } + if decision.denial is not None: + message, param = messages[decision.denial] + raise ProxyException(message=message, type=ProxyErrorTypes.auth_error.value, param=param, code=401) diff --git a/litellm/proxy/auth/owner_team_adapter.py b/litellm/proxy/auth/owner_team_adapter.py new file mode 100644 index 00000000000..e0ea089269b --- /dev/null +++ b/litellm/proxy/auth/owner_team_adapter.py @@ -0,0 +1,104 @@ +import asyncio +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.owner_teams import OwnerTeamBudget, OwnerTeamFacts, OwnerTeamLoad, owner_teams +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.personal_budget_adapter import read_personal_team_budget_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@dataclass(frozen=True, slots=True) +class OwnerTeamReading: + facts: OwnerTeamFacts + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + token: UserAPIKeyAuth + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + check_db_only: bool + + +async def read_owner_team_facts( + user_object: LiteLLM_UserTable, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + purpose: Literal["Models", "Grants"], + check_db_only: bool = False, +) -> OwnerTeamReading: + facts: Final = OwnerTeamFacts( + purpose, + tuple(user_object.teams), + UI_TEAM_ID, + user_present=valid_token.user_id is not None, + database_present=prisma_client is not None, + org=valid_token.org_id, + ) + return await _read( + OwnerTeamReading(facts), + _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only), + ) + + +async def _load(team_id: str, source: _Source) -> tuple[OwnerTeamLoad, LiteLLM_TeamTableCachedObj | None]: + from litellm.proxy.auth.auth_checks import get_team_object + + try: + team: Final = await get_team_object( + team_id=team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + check_db_only=source.check_db_only, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # raw team lookup failures grant nothing; native failure is never converted + verbose_proxy_logger.warning( + "Personal key team model access: team=%s lookup failed, granting nothing: %s", team_id, error + ) + return OwnerTeamLoad(False, failed=True), None + return OwnerTeamLoad(True, blocked=team.blocked is True, org=team.organization_id), team + + +async def _budget(index: int, reading: OwnerTeamReading, source: _Source) -> OwnerTeamBudget: + team: Final = reading.teams[index] + if team is None or source.prisma is None: + raise NativeAuthRuntimeError("owner team budget read requested without its source") + return OwnerTeamBudget( + index, + await read_personal_team_budget_facts( + team, source.token, source.prisma, source.cache, source.logging, check_db_only=source.check_db_only + ), + ) + + +async def _read(reading: OwnerTeamReading, source: _Source) -> OwnerTeamReading: + decision: Final = owner_teams(reading.facts) + match decision.kind: + case "Complete": + return reading + case "ReadTeams": + rows: Final = await asyncio.gather(*(_load(team_id, source) for team_id in decision.ids)) + updated: Final = OwnerTeamReading( + replace(reading.facts, loaded=True, loads=tuple(row[0] for row in rows)), tuple(row[1] for row in rows) + ) + return await _read(updated, source) + case "ReadBudgets": + budgets: Final = tuple( + await asyncio.gather(*(_budget(index, reading, source) for index in decision.indices)) + ) + return await _read( + replace(reading, facts=replace(reading.facts, budgets=(*reading.facts.budgets, *budgets))), source + ) diff --git a/litellm/proxy/auth/personal_budget_adapter.py b/litellm/proxy/auth/personal_budget_adapter.py new file mode 100644 index 00000000000..a251c162fbc --- /dev/null +++ b/litellm/proxy/auth/personal_budget_adapter.py @@ -0,0 +1,280 @@ +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from typing import Final + +from pydantic import BaseModel + +import litellm +from litellm._auth.common_budgets import BudgetFacts, MemberLimitFacts, WindowFacts, window_start +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.personal_budget import PersonalTeamBudgetDecision, PersonalTeamBudgetFacts, personal_team_budget +from litellm._logging import verbose_proxy_logger +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import Litellm_EntityType, LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj, UserAPIKeyAuth +from litellm.proxy.auth.budget_alert_adapter import team_member_max_budget_alert_check +from litellm.proxy.auth.budget_windows_adapter import ( + _reset, # pyright: ignore[reportPrivateUsage] # shared raw window timestamp projection + _seconds, # pyright: ignore[reportPrivateUsage] # shared raw duration projection + _window, # pyright: ignore[reportPrivateUsage] # shared typed window decoding +) +from litellm.proxy.auth.common_budget_adapter import CounterRead, emit_alert, key_alert_info, metadata, read_budget +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class _Windows(BaseModel, frozen=True): + budget_limits: tuple[object, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class _Sources: + team: LiteLLM_TeamTableCachedObj + token: UserAPIKeyAuth + prisma: PrismaClient + cache: UserApiKeyCache + logging: ProxyLogging + windows: tuple[object, ...] = () + membership: LiteLLM_TeamMembership | None = None + + +async def read_personal_team_budget_facts( + team: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, +) -> PersonalTeamBudgetFacts: + facts: Final = PersonalTeamBudgetFacts( + user_present=valid_token.user_id is not None, + team_id_present=team.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # unvalidated cache rows + check_db_only=check_db_only, + token_org=valid_token.org_id, + team_org=team.organization_id, + ) + decision: Final = personal_team_budget(facts) + if decision.kind != "Source": + raise NativeAuthRuntimeError("personal budget source was not selected") + source: Final = valid_token.model_copy(update={"team_id": team.team_id, "org_id": decision.org_id}) + return await _read( + replace(facts, source_ready=True), _Sources(team, source, prisma_client, user_api_key_cache, proxy_logging_obj) + ) + + +async def personal_key_team_within_budget( + team: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, +) -> bool: + facts: Final = await read_personal_team_budget_facts( + team, valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only + ) + return personal_team_budget(facts).allowed + + +async def _read(facts: PersonalTeamBudgetFacts, source: _Sources) -> PersonalTeamBudgetFacts: + decision: Final = personal_team_budget(facts) + if decision.kind == "Complete": + if decision.log_failure: + verbose_proxy_logger.warning( + "Failed to check budget for personal key team %s, skipping: incompatible datetimes", source.team.team_id + ) + return facts + try: + updated, updated_source = await _effect(decision, facts, source) + except NativeAuthFailure: + raise + except litellm.BudgetExceededError: + return replace(facts, io_failed=True) + except Exception as error: # noqa: BLE001 # preserve fail-closed personal budget IO behavior without catching native failures + verbose_proxy_logger.warning( + "Failed to check budget for personal key team %s, skipping: %s", source.team.team_id, error + ) + return replace(facts, io_failed=True) + return await _read(updated, updated_source) + + +def _member_limit(row: LiteLLM_BudgetTable | None, now: datetime) -> MemberLimitFacts: + expiry: Final = row.temp_budget_expiry if row is not None else None + aware: Final = expiry.replace(tzinfo=timezone.utc) if expiry is not None and expiry.tzinfo is None else expiry + expiry_number: Final = ( + aware.replace(tzinfo=timezone.utc).timestamp() + if aware is not None and aware.tzinfo is now.tzinfo + else aware.timestamp() + if aware is not None + else None + ) + now_number: Final = ( + now.replace(tzinfo=timezone.utc).timestamp() + if aware is not None and aware.tzinfo is now.tzinfo + else now.timestamp() + ) + return MemberLimitFacts( + row.max_budget if row is not None else None, + None, + row.temp_budget_increase if row is not None else None, + expiry_number, + now_number, + now.utcoffset() is None, + ) + + +async def _effect( + decision: PersonalTeamBudgetDecision, facts: PersonalTeamBudgetFacts, source: _Sources +) -> tuple[PersonalTeamBudgetFacts, _Sources]: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget, get_team_membership, get_utc_datetime + + team: Final = source.team + match decision.kind: + case "ReadTeam": + budget, _decision = await read_budget( + BudgetFacts("Team", team.max_budget, token_present=True), + CounterRead(f"spend:team:{team.team_id}", team.spend or 0.0, team.max_budget), + ) + return replace(facts, team=budget), source + case "AlertTeam": + if facts.team is None: + raise NativeAuthRuntimeError("missing personal team budget alert facts") + emit_alert( + source.logging, + "team_budget", + key_alert_info(source.token, facts.team.spend, facts.team.limit, Litellm_EntityType.TEAM), + ) + return replace(facts, team_alerted=True), source + case "ReadWindows": + windows: Final = _Windows.model_validate(team, from_attributes=True).budget_limits or () + return replace(facts, window_count=len(windows)), replace(source, windows=windows) + case "ReadWindow": + window: Final = _window(source.windows[decision.index]) + seconds: Final = _seconds(window.budget_duration) + reset: Final = _reset(window.reset_at) + start: Final = window_start( + WindowFacts( + seconds, reset.timestamp() if reset is not None else None, datetime.now(timezone.utc).timestamp() + ) + ) + window_facts, _window_decision = await read_budget( + BudgetFacts("TeamWindow", window.max_budget), + CounterRead( + f"spend:team:{team.team_id}:window:{window.budget_duration}", + 0.0, + window.max_budget, + window_entity_type="Team", + window_entity_id=team.team_id, + window_duration=window.budget_duration, + window_start=datetime.fromtimestamp(start, timezone.utc) if start is not None else None, + ), + ) + return replace(facts, windows=(*facts.windows, window_facts)), source + case "ReadMembership": + if source.token.user_id is None: + raise NativeAuthRuntimeError("missing personal team membership identity") + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + check_db_only=facts.check_db_only, + ) + return replace(facts, membership_loaded=True), replace(source, membership=membership) + case "ReadMemberLimit": + row: Final = source.membership.litellm_budget_table if source.membership is not None else None + now: Final = get_utc_datetime() + default_id: Final = metadata(team.metadata).get("team_member_budget_id") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy metadata projected as an opaque mapping + return replace( + facts, + member_limit=_member_limit(row, now), + default_id=default_id if isinstance(default_id, str) else None, + ), source + case "ReadDefault": + if facts.member_limit is None: + raise NativeAuthRuntimeError("missing personal member limit facts") + default: Final = await get_team_member_default_budget( + budget_id=decision.id, prisma_client=source.prisma, user_api_key_cache=source.cache + ) + return replace( + facts, + default_loaded=True, + member_limit=replace( + facts.member_limit, default_limit=default.max_budget if default is not None else None + ), + ), source + case "ReadMember": + member, _member_decision = await read_budget( + BudgetFacts("TeamMember", decision.limit), + CounterRead( + f"spend:team_member:{source.token.user_id}:{team.team_id}", + (source.membership.spend or 0.0) if source.membership is not None else 0.0, + decision.limit, + ), + ) + return replace(facts, member=member), source + case "AlertMember": + alert_member: Final = facts.member + if ( + alert_member is None + or alert_member.spend is None + or alert_member.limit is None + or source.token.user_id is None + ): + raise NativeAuthRuntimeError("missing personal member alert facts") + team_member_max_budget_alert_check( + team_id=team.team_id, + team_alias=team.team_alias, + team_metadata=team.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # helper validates legacy alert metadata + organization_id=team.organization_id, + user_id=source.token.user_id, + user_email=None, + proxy_logging_obj=source.logging, + spend=alert_member.spend, + max_budget=alert_member.limit, + ) + return replace(facts, member_alerted=True), source + case "ReadOrg": + org_budget, org_failed = await _read_org(decision.id, source) + return replace(facts, org_loaded=True, org=org_budget, org_failed=org_failed), source + case "AlertOrg": + if facts.org is None: + raise NativeAuthRuntimeError("missing personal org budget alert facts") + emit_alert( + source.logging, + "organization_budget", + key_alert_info( + source.token, facts.org.spend, facts.org.limit, Litellm_EntityType.ORGANIZATION + ).model_copy(update={"organization_id": decision.id}), + ) + return replace(facts, org_alerted=True), source + case "Source" | "Complete": + raise NativeAuthRuntimeError("unexpected personal budget read") + + +async def _read_org(org_id: str, source: _Sources) -> tuple[BudgetFacts | None, bool]: + from litellm.proxy.auth.auth_checks import get_org_object + + try: + org: Final = await get_org_object( + org_id=org_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + include_budget_table=True, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # legacy org lookup failure skips only the org budget, not the other personal-team ceilings + return None, True + if org is None: + return None, False + carry_organization_budget_state(valid_token=source.token, org_table=org) + limit: Final = org.litellm_budget_table.max_budget if org.litellm_budget_table is not None else None + budget, _decision = await read_budget( + BudgetFacts("Organization", limit), CounterRead(f"spend:org:{org_id}", org.spend or 0.0, limit) + ) + return budget, False diff --git a/litellm/proxy/auth/personal_model_adapter.py b/litellm/proxy/auth/personal_model_adapter.py new file mode 100644 index 00000000000..9fa87d43502 --- /dev/null +++ b/litellm/proxy/auth/personal_model_adapter.py @@ -0,0 +1,192 @@ +import asyncio +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Final + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import personal_model_access_applies +from litellm._auth.personal_models import ( + LoadedTeamModelFacts, + PersonalModelDecision, + PersonalModelScope, + loaded_team_models, + personal_model_scope, +) +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +def _eligibility(token: UserAPIKeyAuth | None, *, inheritance: bool) -> bool: + return personal_model_access_applies( + enabled=bool(litellm.personal_key_model_access_from_teams), + token_present=token is not None, + team_id=token.team_id if token is not None else None, + user_id=token.user_id if token is not None else None, + master=token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS if token is not None else False, + role=token.user_role if token is not None else None, + api_key_present=token.api_key is not None if token is not None else False, + inheritance=inheritance, + ) + + +def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) -> bool: + return _eligibility(valid_token, inheritance=False) + + +def personal_key_inherits_team_grants(valid_token: UserAPIKeyAuth | None) -> bool: + return _eligibility(valid_token, inheritance=True) + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str + team: LiteLLM_TeamTableCachedObj + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + aliases: Mapping[str, str] | None + + +async def read_personal_model_scope( + model: str, + team_object: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> PersonalModelScope: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + aliases: Final = key_model_aliases_for_auth_check(valid_token) + facts: Final = PersonalModelScope( + user_present=valid_token.user_id is not None, + team_id_present=team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # unvalidated cache rows + ) + return await _read_scope( + facts, + _Source( + model, team_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj, aliases + ), + ) + + +async def _read_scope(facts: PersonalModelScope, source: _Source) -> PersonalModelScope: + decision: Final = personal_model_scope(facts) + if decision.kind == "Complete": + return facts + try: + updated: Final = await _scope_effect(decision, facts, source) + except NativeAuthFailure: + raise + except ProxyException: + return replace(facts, io_failed=True) + except Exception as error: # noqa: BLE001 # preserve personal team lookup failures as raw failure facts, not policy callbacks + verbose_proxy_logger.warning( + "Personal key team model access: team=%s check failed, granting nothing: %s", source.team.team_id, error + ) + return replace(facts, io_failed=True) + return await _read_scope(updated, source) + + +async def _scope_effect( + decision: PersonalModelDecision, facts: PersonalModelScope, source: _Source +) -> PersonalModelScope: + from litellm.proxy.auth.auth_checks import get_team_membership + + match decision.kind: + case "ReadTeam": + team: Final = project_model_facts( + source.model, + source.router, + object_model_grants(source.team), + team_id=source.team.team_id, + key_model_aliases=source.aliases, + object_type="team", + ) + return replace(facts, team=GroupedModelScope(team, tuple(source.team.access_group_ids or ()))) + case "ReadGroups": + if facts.team is None: + raise NativeAuthRuntimeError("personal model group read without team facts") + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return replace(facts, team=grouped) + case "ReadMembership": + if source.token.user_id is None: + raise NativeAuthRuntimeError("personal model membership read without identity") + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=source.team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + row: Final = membership.litellm_budget_table if membership is not None else None + grants: Final = tuple(row.allowed_models or ()) if row is not None else () + return replace(facts, member_grants=grants) + case "ReadMember": + if facts.member_grants is None: + raise NativeAuthRuntimeError("personal model member scope read without grants") + member: Final = project_model_facts( + source.model, + source.router, + facts.member_grants, + team_id=source.team.team_id, + key_model_aliases=source.aliases, + object_type="team", + ) + return replace(facts, member=member) + case "ReadScopes" | "Complete": + raise NativeAuthRuntimeError("unexpected personal model scope read") + + +async def read_loaded_team_model_facts( + model: str, + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LoadedTeamModelFacts: + facts: Final = LoadedTeamModelFacts(tuple(team is not None for team in teams)) + decision: Final = loaded_team_models(facts) + if decision.kind == "Complete": + return facts + if decision.kind != "ReadScopes": + raise NativeAuthRuntimeError("unexpected loaded team model read") + + async def read(index: int) -> PersonalModelScope: + team: Final = teams[index] + if team is None: + raise NativeAuthRuntimeError("missing selected personal model team") + return await read_personal_model_scope( + model, team, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + + scopes: Final = tuple(await asyncio.gather(*(read(index) for index in decision.indices))) + return replace(facts, scopes=scopes) + + +async def loaded_teams_allow_model( + model: str, + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> bool: + facts: Final = await read_loaded_team_model_facts( + model, teams, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + return loaded_team_models(facts).allowed diff --git a/litellm/proxy/auth/personal_model_listing_adapter.py b/litellm/proxy/auth/personal_model_listing_adapter.py new file mode 100644 index 00000000000..5e79f062535 --- /dev/null +++ b/litellm/proxy/auth/personal_model_listing_adapter.py @@ -0,0 +1,111 @@ +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.personal_model_listing import PersonalModelListingFacts, personal_model_listing +from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, UserAPIKeyAuth +from litellm.proxy.auth.personal_model_adapter import read_loaded_team_model_facts +from litellm.proxy.auth.personal_owner_adapter import read_personal_owner_facts, selected_team +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + names: Sequence[str] + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +async def personal_key_listing_teams( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: + reading: Final = await read_personal_owner_facts( + user_api_key_dict, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + check_db_only=check_db_only, + purpose="Models", + ) + decision: Final = personal_owner(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete listing team facts") + return tuple(team if index in decision.selected else None for index, team in enumerate(reading.teams)) + + +async def read_personal_model_listing_facts( + models: Sequence[str], + user_api_key_dict: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> PersonalModelListingFacts: + owner: Final = PersonalOwnerFacts( + user_api_key_dict.user_id is not None, + prisma_client is not None, + user_api_key_cache is not None, + proxy_logging_obj is not None, + ) + return await _read( + PersonalModelListingFacts(owner), + _Source(models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + + +async def personal_key_team_visible_models( + models: Sequence[str], + user_api_key_dict: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[str, ...]: + facts: Final = await read_personal_model_listing_facts( + models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = personal_model_listing(facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete personal model listing facts") + return tuple((facts.names or ())[index] for index in decision.selected) + + +async def _read(facts: PersonalModelListingFacts, source: _Source) -> PersonalModelListingFacts: + decision: Final = personal_model_listing(facts) + match decision.kind: + case "Complete": + return facts + case "ReadOwner": + owner: Final = await read_personal_owner_facts( + source.token, source.prisma, source.cache, source.logging, purpose="Models" + ) + return await _read(replace(facts, owner=owner.facts), replace(source, teams=owner.teams)) + case "ReadNames": + return await _read(replace(facts, names=tuple(source.names)), source) + case "ReadModel": + if source.prisma is None or source.cache is None or source.logging is None or facts.names is None: + raise NativeAuthRuntimeError("personal listing scope read without dependencies") + teams: Final = tuple(selected_team(source.teams, index) for index in decision.teams) + model: Final = await read_loaded_team_model_facts( + facts.names[decision.index], + teams, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + ) + return await _read(replace(facts, models=(*facts.models, model)), source) diff --git a/litellm/proxy/auth/personal_model_request_adapter.py b/litellm/proxy/auth/personal_model_request_adapter.py new file mode 100644 index 00000000000..f96da4834b1 --- /dev/null +++ b/litellm/proxy/auth/personal_model_request_adapter.py @@ -0,0 +1,93 @@ +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.model_access import ModelAccessDecision +from litellm._auth.personal_model_request import PersonalModelRequestFacts, personal_model_request +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import raise_model_denial +from litellm.proxy.auth.owner_team_adapter import read_owner_team_facts +from litellm.proxy.auth.personal_model_adapter import read_loaded_team_model_facts +from litellm.proxy.auth.personal_owner_adapter import selected_team +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + owner: LiteLLM_UserTable | None + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +async def read_personal_model_request_facts( + model: str | Sequence[str], + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> PersonalModelRequestFacts: + requested: Final = (model,) if isinstance(model, str) else tuple(model) + return await _read( + PersonalModelRequestFacts(requested, user_object is not None), + _Source(user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + + +async def can_personal_key_call_model_via_teams( + model: str | Sequence[str], + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Literal[True]: + facts: Final = await read_personal_model_request_facts( + model, user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = personal_model_request(facts) + if decision.kind == "Allow": + return True + if decision.kind != "Deny": + raise NativeAuthRuntimeError("incomplete personal model request facts") + return raise_model_denial( + ModelAccessDecision("PersonalTeam", model=decision.model), + model, + user_id=valid_token.user_id, + teams=user_object.teams if user_object is not None else None, + ) + + +async def _read(facts: PersonalModelRequestFacts, source: _Source) -> PersonalModelRequestFacts: + decision: Final = personal_model_request(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadTeams": + if source.owner is None: + raise NativeAuthRuntimeError("personal model team read without owner") + reading: Final = await read_owner_team_facts( + source.owner, source.token, source.prisma, source.cache, source.logging, purpose="Models" + ) + return await _read(replace(facts, teams=reading.facts), replace(source, teams=reading.teams)) + case "ReadModel": + teams: Final = tuple(selected_team(source.teams, index) for index in decision.teams) + model: Final = await read_loaded_team_model_facts( + facts.requested[decision.index], + teams, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + ) + return await _read(replace(facts, models=(*facts.models, model)), source) diff --git a/litellm/proxy/auth/personal_owner_adapter.py b/litellm/proxy/auth/personal_owner_adapter.py new file mode 100644 index 00000000000..01c3ea5ddd1 --- /dev/null +++ b/litellm/proxy/auth/personal_owner_adapter.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Final, Literal + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.owner_team_adapter import read_owner_team_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +if TYPE_CHECKING: + from litellm.proxy.auth.auth_checks import PersonalKeyOwnerTeams + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerReading: + facts: PersonalOwnerFacts + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + token: UserAPIKeyAuth + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + check_db_only: bool + owner: LiteLLM_UserTable | None = None + purpose: Literal["Models", "Grants"] = "Grants" + + +async def read_personal_owner_facts( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, + purpose: Literal["Models", "Grants"] = "Grants", +) -> PersonalOwnerReading: + facts: Final = PersonalOwnerFacts( + valid_token.user_id is not None, + prisma_client is not None, + user_api_key_cache is not None, + proxy_logging_obj is not None, + ) + return await _read( + PersonalOwnerReading(facts), + _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only, purpose=purpose), + ) + + +async def personal_key_owner_teams_with_status( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> PersonalKeyOwnerTeams: + from litellm.proxy.auth.auth_checks import PersonalKeyOwnerTeams + + reading: Final = await read_personal_owner_facts( + valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only + ) + decision: Final = personal_owner(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete personal owner team facts") + return PersonalKeyOwnerTeams( + teams=tuple(selected_team(reading.teams, index) for index in decision.selected), + withheld=tuple(selected_team(reading.teams, index) for index in decision.withheld), + unreadable=decision.unreadable, + ) + + +def selected_team(teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], index: int) -> LiteLLM_TeamTableCachedObj: + team: Final = teams[index] + if team is None: + raise NativeAuthRuntimeError("missing selected owner team row") + return team + + +async def _read(reading: PersonalOwnerReading, source: _Source) -> PersonalOwnerReading: + from litellm.proxy.auth.auth_checks import get_user_object + + decision: Final = personal_owner(reading.facts) + if decision.kind == "Complete": + return reading + if source.token.user_id is None or source.prisma is None or source.cache is None or source.logging is None: + raise NativeAuthRuntimeError("owner read requested without required dependencies") + match decision.kind: + case "ReadOwner": + try: + owner: Final = await get_user_object( + user_id=source.token.user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + user_id_upsert=False, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + check_db_only=source.check_db_only, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve legacy owner lookup failure as raw data + verbose_proxy_logger.warning( + "Personal key team model access: owner lookup failed, granting nothing: %s", error + ) + return PersonalOwnerReading(replace(reading.facts, owner_loaded=True, owner_failed=True)) + return await _read( + PersonalOwnerReading(replace(reading.facts, owner_loaded=True, owner_present=owner is not None)), + replace(source, owner=owner), + ) + case "ReadTeams": + if source.owner is None: + raise NativeAuthRuntimeError("personal owner team read without owner row") + teams: Final = await read_owner_team_facts( + source.owner, + source.token, + source.prisma, + source.cache, + source.logging, + purpose=source.purpose, + check_db_only=source.check_db_only, + ) + return PersonalOwnerReading(replace(reading.facts, teams=teams.facts), teams.teams) diff --git a/litellm/proxy/auth/resolved_model_adapter.py b/litellm/proxy/auth/resolved_model_adapter.py new file mode 100644 index 00000000000..8c3cb246f0b --- /dev/null +++ b/litellm/proxy/auth/resolved_model_adapter.py @@ -0,0 +1,286 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from pydantic import BaseModel, Field, TypeAdapter + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import key_model_grants +from litellm._auth.resolved_models import ResolvedModelFacts, ResolvedTeamRead, resolved_models +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.member_model_adapter import read_member_model_facts +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.proxy.auth.personal_model_request_adapter import read_personal_model_request_facts +from litellm.proxy.auth.team_override_adapter import read_team_override_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +class _Config(BaseModel, frozen=True): + config: Mapping[str, object] = Field(default_factory=dict[str, object]) + + +_METADATA: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + +@dataclass(frozen=True, slots=True) +class ResolvedModelReading: + facts: ResolvedModelFacts + team: LiteLLM_TeamTableCachedObj | None = None + owner: LiteLLM_UserTable | None = None + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + team: LiteLLM_TeamTableCachedObj | None = None + owner: LiteLLM_UserTable | None = None + project: LiteLLM_ProjectTable | None = None + customer: LiteLLM_EndUserTable | None = None + + +async def read_resolved_model_facts( + model: str | Sequence[str], + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> ResolvedModelReading: + facts: Final = ResolvedModelFacts( + model=str(model), + config_nonempty=bool(_Config.model_validate(valid_token, from_attributes=True).config), + key_models=object_model_grants(valid_token), + team_id=valid_token.team_id, + user_id=valid_token.user_id, + project_present=valid_token.project_id is not None, + customer_present=valid_token.end_user_id is not None, + database_present=prisma_client is not None, + personal_enabled=bool(litellm.personal_key_model_access_from_teams), + master=valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS, + role=valid_token.user_role, + ) + return await _read( + facts, _Source(model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj) + ) + + +async def can_key_call_resolved_model( + model: str | Sequence[str], + llm_model_list: Sequence[object] | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, +) -> None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + reading: Final = await read_resolved_model_facts( + model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = resolved_models(reading.facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete resolved model access facts") + raise_model_denial( + decision.denial, + model, + user_id=valid_token.user_id, + team_id=reading.team.team_id if reading.team is not None else None, + teams=reading.owner.teams if reading.owner is not None else None, + member_models=decision.member_models, + ) + + +async def _read(facts: ResolvedModelFacts, source: _Source) -> ResolvedModelReading: + from litellm.proxy.auth.auth_checks import ( + get_end_user_object, + get_key_end_user_budget_id, + get_project_object, + get_team_object, + get_user_object, + key_model_aliases_for_auth_check, + team_model_aliases_for_auth_check, + ) + + decision: Final = resolved_models(facts) + match decision.kind: + case "Allow" | "Deny": + return ResolvedModelReading(facts, source.team, source.owner) + case "ReadKey": + grants: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + key: Final = project_model_facts( + source.model, + source.router, + grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ) + return await _read( + replace(facts, key=GroupedModelScope(key, tuple(source.token.access_group_ids or ()))), source + ) + case "ReadKeyGroups": + if facts.key is None: + raise NativeAuthRuntimeError("resolved key group read without scope") + scope: Final = await read_grouped_scope(facts.key, source.router, merge=False) + return await _read(replace(facts, key=scope), source) + case "ReadTeamObject": + if source.token.team_id is None: + raise NativeAuthRuntimeError("resolved team read without identity") + try: + team: Final = await get_team_object( + team_id=source.token.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # legacy resolved-model team lookup falls back to the embedded token row + return await _read(replace(facts, team_read=ResolvedTeamRead(failed=True)), source) + return await _read( + replace(facts, team_read=ResolvedTeamRead(found=team is not None)), # pyright: ignore[reportUnnecessaryComparison] # preserve the legacy missing-row result + replace(source, team=team), + ) + case "ReadEmbeddedTeam": + embedded: Final = LiteLLM_TeamTableCachedObj( + team_id=source.token.team_id, # pyright: ignore[reportArgumentType] # native selects this read only for a present team ID + models=source.token.team_models, + blocked=source.token.team_blocked, + team_alias=source.token.team_alias, + metadata=source.token.team_metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy typed row validates embedded metadata + object_permission_id=source.token.team_object_permission_id, + object_permission=source.token.team_object_permission, + ) + return await _read(replace(facts, embedded_ready=True), replace(source, team=embedded)) + case "ReadTeamScope": + team_scope: Final = await read_team_override_facts( + source.model, + source.token, + source.team, + source.router, + source.token.team_model_aliases, + key_model_aliases_for_auth_check(source.token), + ) + return await _read(replace(facts, team=team_scope), source) + case "ReadMember": + if source.team is None: + raise NativeAuthRuntimeError("resolved member model read without team") + member: Final = await read_member_model_facts( + source.model, + source.team, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + ) + return await _read(replace(facts, member=member), source) + case "ReadOwner": + if source.token.user_id is None: + raise NativeAuthRuntimeError("resolved personal owner read without identity") + try: + owner: Final = await get_user_object( + user_id=source.token.user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + user_id_upsert=False, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve raw owner lookup failure and native fail-closed model selection + verbose_proxy_logger.warning( + "Personal key team model access: owner lookup failed, granting nothing: %s", error + ) + return await _read(replace(facts, owner_loaded=True, owner_failed=True), source) + return await _read( + replace(facts, owner_loaded=True, owner_present=owner is not None), replace(source, owner=owner) + ) + case "ReadPersonal": + personal: Final = await read_personal_model_request_facts( + source.model, source.owner, source.token, source.router, source.prisma, source.cache, source.logging + ) + return await _read(replace(facts, personal=personal), source) + case "ReadProjectObject": + if source.token.project_id is None: + raise NativeAuthRuntimeError("resolved project read without identity") + project: Final = await get_project_object( + project_id=source.token.project_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + return await _read( + replace( + facts, + project_loaded=True, + project_grants=object_model_grants(project) if project is not None else (), + ), + replace(source, project=project), + ) + case "ReadProjectScope": + if source.project is None: + raise NativeAuthRuntimeError("resolved project scope read without row") + project_scope: Final = project_model_facts( + source.model, + source.router, + facts.project_grants, + team_id=source.project.team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="project", + ) + return await _read(replace(facts, project=project_scope), source) + case "ReadCustomerObject": + if source.token.end_user_id is None or source.prisma is None: + raise NativeAuthRuntimeError("resolved customer read without identity or database") + customer: Final = await get_end_user_object( + end_user_id=source.token.end_user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + token_end_user_max_budget=source.token.end_user_max_budget, + key_end_user_budget_id=get_key_end_user_budget_id(_METADATA.validate_python(source.token.metadata)), # pyright: ignore[reportUnknownMemberType] # validate the legacy metadata field at the IO boundary + ) + return await _read( + replace( + facts, + customer_loaded=True, + customer_grants=tuple(customer.models or ()) if customer is not None else (), + ), + replace(source, customer=customer), + ) + case "ReadCustomerScope": + customer_scope: Final = project_model_facts( + source.model, + source.router, + facts.customer_grants, + team_model_aliases_for_auth_check(source.token), + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "customer", + ) + return await _read(replace(facts, customer=customer_scope), source) diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py index 0c8d05e0b94..2924aa1372b 100644 --- a/litellm/proxy/auth/resolvers/store.py +++ b/litellm/proxy/auth/resolvers/store.py @@ -1,10 +1,12 @@ from __future__ import annotations +import time from collections.abc import Sequence from typing import TYPE_CHECKING, Final from pydantic import BaseModel +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth @@ -36,7 +38,12 @@ from litellm.proxy.auth.resolvers.models import ( UserIdentity, ) from litellm.proxy.auth.roles import TeamRole, map_role, team_role -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET +from litellm.proxy.common_utils.user_api_key_cache import ( + AUTH_OBJECTS_TARGET, + UserApiKeyCache, + get_management_object_ttl, +) +from litellm.repositories.native_auth_repository import NativeAuthDatabaseError if TYPE_CHECKING: from litellm.caching.caching import DualCache @@ -108,10 +115,25 @@ class IdentityStore: async def _resolve_key(self, hashed_token: str) -> UserAPIKeyAuth: if self._prisma is None: raise NoDatabaseConnectionError() + first: Final = await self._resolve_key_once(hashed_token) + if first is not None: + return first + retried: Final = await self._resolve_key_once(hashed_token) + if retried is not None: + return retried + raise NativeAuthRuntimeError("Auth identity changed during lookup") + async def _resolve_key_once(self, hashed_token: str) -> UserAPIKeyAuth | None: + assert self._prisma is not None + cache: Final = self._cache if isinstance(self._cache, UserApiKeyCache) else None + generation: Final = cache.begin_fill(hashed_token) if cache is not None else None cached: Final = await self._cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth) if cached is not None: - return copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + return ( + copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + if cache is None or cache.current_generation() == generation + else None + ) if self._check_cache_only: raise KeyNotInCacheError(hashed_token) @@ -137,12 +159,21 @@ class IdentityStore: proxy_logging_obj=self._proxy_logging_obj, ) except Exception as e: + if isinstance(e, NativeAuthFailure) and not isinstance(e, NativeAuthDatabaseError): + raise verbose_proxy_logger.debug( "Failed to load object_permission for key with object_permission_id=%s: %s", key.object_permission_id, e, ) + if cache is not None and generation is not None: + refreshed: Final = key.model_copy(update={"last_refreshed_at": time.time()}) + stored: Final = cache.set_if_generation( + hashed_token, refreshed, generation, ttl=get_management_object_ttl(cache) + ) + return refreshed if stored or cache.current_generation() == generation else None + await cache_key_object( hashed_token=hashed_token, user_api_key_obj=key, diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 0aa2528d46b..4d822f22300 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -1,15 +1,12 @@ -import itertools -import re from collections.abc import Collection, Iterable, Mapping from typing import Final, cast from fastapi import HTTPException, Request, status -from litellm._logging import verbose_proxy_logger +from litellm._auth import routes from litellm.proxy._types import ( MODEL_LIST_ROUTES, UI_TEAM_ID, - CommonProxyErrors, KeyManagementRoutes, LiteLLM_UserTable, LiteLLMRoutes, @@ -18,10 +15,13 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.auth import route_policy_adapter from .auth_checks_organization import ( # noqa: F401 # legacy module exports _user_is_org_admin, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - user_is_org_admin, +) +from .auth_checks_organization import ( + user_is_org_admin as user_is_org_admin, # noqa: PLC0414 # preserve the public compatibility re-export ) # Management write routes denied to PROXY_ADMIN_VIEW_ONLY. Adding a new write @@ -112,14 +112,9 @@ _TOKEN_COUNTING_ACTION: Final = "countTokens" class RouteChecks: @staticmethod def check_route_enabled(route: str) -> None: - try: - from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + from litellm.proxy.auth.route_enable_adapter import check_route_enabled - EnterpriseRouteChecks.should_call_route(route=route) - except HTTPException as e: - raise e - except Exception: - pass + check_route_enabled(route) @staticmethod def should_call_route( @@ -141,120 +136,7 @@ class RouteChecks: valid_token: UserAPIKeyAuth, request: Request | None = None, ) -> bool: - """ - Raises Exception if Virtual Key is not allowed to call the route - """ - - # Only check if valid_token.allowed_routes is set and is a list with at least one item - if valid_token.allowed_routes is None: - return True - if not isinstance(valid_token.allowed_routes, list): - return True - if len(valid_token.allowed_routes) == 0: - return True - - denied_auth_enforced_pass_through_route = False - - # explicit check for allowed routes (exact match or prefix match) - for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): - return True - - ## check if 'allowed_route' is a field name in LiteLLMRoutes - if any(allowed_route in LiteLLMRoutes._member_names_ for allowed_route in valid_token.allowed_routes): - for allowed_route in valid_token.allowed_routes: - if allowed_route in LiteLLMRoutes._member_names_: - if RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes[allowed_route].value, - ): - if ( - allowed_route in _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS - and RouteChecks.is_auth_enforced_pass_through_route( - route=route, - method=RouteChecks.get_request_method(request=request), - ) - ): - if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): - return True - denied_auth_enforced_pass_through_route = True - else: - return True - - ################################################ - # For llm_api_routes, also check registered pass-through endpoints - ################################################ - if allowed_route in _LLM_API_ROUTE_GROUPS: - if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET": - return True - - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - InitPassThroughEndpointHelpers, - ) - - if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): - if RouteChecks.is_auth_enforced_pass_through_route( - route=route, - method=RouteChecks.get_request_method(request=request), - ): - if RouteChecks.check_passthrough_route_access( - route=route, user_api_key_dict=valid_token - ): - return True - denied_auth_enforced_pass_through_route = True - else: - return True - - # Method-aware carve-out: allow GET on the two - # read-only MCP-server discovery endpoints - # (`/v1/mcp/server` and `/v1/mcp/server/{server_id}`) - # so virtual keys with allowed_routes=["llm_api_routes"] - # can list/inspect MCP servers. The GET handlers in - # mcp_management_endpoints.py sanitize the response - # for restricted virtual keys (stripping url, - # headers, env, credentials). POST/PUT/DELETE on - # these paths are admin-only management writes and - # are intentionally not covered. - if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route( - route=route, request=request - ): - return True - - # Agent registry CRUD moved from llm_api_routes into - # management_routes so DISABLE_LLM_API_ENDPOINTS stops - # blocking it. Keys configured with - # allowed_routes=["llm_api_routes"] before that split - # could reach these paths, so keep them reachable here; - # the handlers in agent_endpoints/endpoints.py still - # enforce proxy-admin on writes and scope reads by role. - if allowed_route == "llm_api_routes" and RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes.agent_management_routes.value, - ): - return True - - # check if wildcard pattern is allowed - for allowed_route in valid_token.allowed_routes: - if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): - return True - - if denied_auth_enforced_pass_through_route: - raise RouteChecks._auth_pass_through_denied_exception(route=route) - - if valid_token.metadata.get("password_reset_required") is True: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=( - "This account's password must be changed before the session can be used: " - "it was either found in a known data breach or set by an admin. " - "Change it via POST /user/password/change (UI: /ui/change-password), then log in again." - ), - ) - - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {valid_token.allowed_routes}. Tried to call route: {route}", - ) + return route_policy_adapter.check_virtual_route(route, valid_token, request) @staticmethod def _mask_user_id(user_id: str) -> str: @@ -318,209 +200,29 @@ class RouteChecks: valid_token: UserAPIKeyAuth, request_data: dict, ): - """ - Checks if Non Proxy Admin User is allowed to access the route - """ - - # Check user has defined custom admin routes - RouteChecks.custom_admin_only_route_check( - route=route, + return route_policy_adapter.check_non_admin_route( + user_obj, _user_role, route, request, valid_token, request_data ) - if RouteChecks.is_auth_enforced_pass_through_route( - route=route, - method=RouteChecks.get_request_method(request=request), - ): - RouteChecks._require_auth_pass_through_access( - route=route, - valid_token=valid_token, - jwt_team_allowed_routes=RouteChecks._jwt_team_allowed_routes(valid_token=valid_token), - ) - elif RouteChecks.is_llm_api_route(route=route): - pass - elif RouteChecks.is_info_route(route=route): - # check if user allowed to call an info route - if route == "/key/info": - # handled by function itself - pass - elif route == "/user/info": - # check if user can access this route - query_params: Final = request.query_params - user_id: Final = query_params.get("user_id") - verbose_proxy_logger.debug("user_id: %s & valid_token.user_id: %s", user_id, valid_token.user_id) - if ( - user_id - and user_id != valid_token.user_id - and _user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value - ): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"key not allowed to access this user's info. user_id={user_id}, key's user_id={valid_token.user_id}", - ) - elif route == "/v2/user/info": - # handled by the endpoint itself (full RBAC in handler) - pass - elif route == "/model/info": - # /model/info just shows models user has access to - pass - elif route == "/team/info": - pass # handled by function itself - elif ( - route in LiteLLMRoutes.global_spend_tracking_routes.value - and getattr(valid_token, "permissions", None) is not None - and "get_spend_routes" in getattr(valid_token, "permissions", []) - ): - pass - elif route == "/lens" or route.startswith("/lens/"): - pass # Lens authorizes the delegated role, including its POST-based read endpoints. - elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: - RouteChecks._check_proxy_admin_viewer_access( - route=route, - _user_role=_user_role, - request_data=request_data, - request=request, - ) - elif ( - _user_role == LitellmUserRoles.INTERNAL_USER.value - and ( - RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value) - or RouteChecks.is_self_scoped_credential_route(route=route, request=request) - ) - or user_is_org_admin(request_data=request_data, user_object=user_obj) - and RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value) - or _user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - and RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes.internal_user_view_only_routes.value, - ) - or RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.self_managed_routes.value) - ): - pass - elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"): - pass # authN/authZ handled by api itself - elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token) or ( - valid_token.is_team_service_account - and RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.team_service_account_key_routes.value - ) - ): - pass - elif valid_token.allowed_routes is not None: - # check if route is in allowed_routes (exact match or prefix match) - route_allowed = False - for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): - route_allowed = True - break - - if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): - route_allowed = True - break - - if not route_allowed: - RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) - else: - RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) - @staticmethod def custom_admin_only_route_check(route: str): - from litellm.proxy.proxy_server import general_settings, premium_user + from litellm.proxy.auth.route_scope_adapter import custom_admin_only_route_check - if "admin_only_routes" in general_settings: - if premium_user is not True: - verbose_proxy_logger.error( - "Trying to use 'admin_only_routes' this is an Enterprise only feature. %s", - CommonProxyErrors.not_premium_user.value, - ) - return - if route in general_settings["admin_only_routes"]: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route. Route={route} is an admin only route", - ) + return custom_admin_only_route_check(route) @staticmethod def is_allowed_on_proxy(route: str) -> bool: - from litellm.proxy.proxy_server import general_settings, premium_user + from litellm.proxy.auth.route_scope_adapter import is_allowed_on_proxy - if "allowed_routes" not in general_settings: - return True - if premium_user is not True: - verbose_proxy_logger.error( - "Trying to set allowed_routes. This is an Enterprise feature. %s", - CommonProxyErrors.not_premium_user.value, - ) - return route in general_settings["allowed_routes"] + return is_allowed_on_proxy(route) @staticmethod def is_unbilled_route(route: str) -> bool: - resource, _, action = route.rsplit("/", 1)[-1].partition(":") - return route in _UNBILLED_ROUTES or resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION + return route_policy_adapter.classify(route, "Unbilled") @staticmethod def is_llm_api_route(route: str) -> bool: - """ - Helper to checks if provided route is an OpenAI route - - - Returns: - - True: if route is an OpenAI route - - False: if route is not an OpenAI route - """ - # Ensure route is a string before performing checks - if not isinstance(route, str): - return False - - if route in LiteLLMRoutes.openai_routes.value: - return True - - if route in LiteLLMRoutes.anthropic_routes.value: - return True - - if route in LiteLLMRoutes.google_routes.value: - return True - - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): - return True - - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_inference_routes.value): - return True - - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.litellm_native_routes.value): - return True - - # fuzzy match routes like "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" - # Check for routes with placeholders or wildcard patterns - for openai_route in LiteLLMRoutes.openai_routes.value: - # Replace placeholders with regex pattern - # placeholders are written as "/threads/{thread_id}" - if "{" in openai_route: - if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): - return True - # Check for wildcard patterns like "/containers/*" - if RouteChecks.is_wildcard_pattern(pattern=openai_route): - if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): - return True - - # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" - for google_route in LiteLLMRoutes.google_routes.value: - if "{" in google_route: - if RouteChecks._route_matches_pattern(route=route, pattern=google_route): - return True - - # Check for Anthropic routes with placeholders - for anthropic_route in LiteLLMRoutes.anthropic_routes.value: - if "{" in anthropic_route: - if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): - return True - - if RouteChecks._is_azure_openai_route(route=route): - return True - - for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: - if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): - return True - return RouteChecks.is_builtin_eu_jev_decision_route(route=route) + return isinstance(route, str) and route_policy_adapter.classify(route, "Llm") @staticmethod def is_builtin_eu_jev_decision_route(route: str) -> bool: @@ -536,26 +238,7 @@ class RouteChecks: @staticmethod def _is_get_mcp_server_discovery_route(route: str, request: Request | None) -> bool: - """ - Returns True if `request` is a GET against one of the two read-only - MCP-server discovery paths: - - - GET `/v1/mcp/server` (list) - - GET `/v1/mcp/server/{server_id}` (single server, single segment) - - Multi-segment paths (`/v1/mcp/server/{id}/approve`, etc.) and any - non-GET method return False, so admin-only management writes on the - same path prefix are not reachable through this carve-out. - """ - if request is None or request.method.upper() != "GET": - return False - if route == "/v1/mcp/server": - return True - prefix: Final = "/v1/mcp/server/" - if not route.startswith(prefix): - return False - remainder: Final = route[len(prefix) :] - return bool(remainder) and "/" not in remainder + return route_policy_adapter.classify(route, "Discovery", RouteChecks.get_request_method(request)) @staticmethod def is_management_route(route: str) -> bool: @@ -580,67 +263,11 @@ class RouteChecks: @staticmethod def _is_azure_openai_route(route: str) -> bool: - """ - Check if route is a route from AzureOpenAI SDK client - - eg. - route='/openai/deployments/vertex_ai/gemini-1.5-flash/chat/completions' - """ - # Ensure route is a string before attempting regex matching - if not isinstance(route, str): - return False - # Add support for deployment and engine model paths - deployment_pattern: Final = r"^/openai/deployments/[^/]+/[^/]+/chat/completions$" - engine_pattern: Final = r"^/engines/[^/]+/chat/completions$" - - if re.match(deployment_pattern, route) or re.match(engine_pattern, route): - return True - return False + return isinstance(route, str) and route_policy_adapter.classify(route, "Azure") @staticmethod def _route_matches_pattern(route: str, pattern: str) -> bool: - """ - Check if route matches the pattern placed in proxy/_types.py - - Example: - - pattern: "/threads/{thread_id}" - - route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" - - returns: True - - - - pattern: "/key/{token_id}/regenerate" - - route: "/key/regenerate/82akk800000000jjsk" - - returns: False, pattern is "/key/{token_id}/regenerate" - """ - # Ensure route is a string before attempting regex matching - if not isinstance(route, str): - return False - - def _placeholder_to_regex(match: re.Match) -> str: - placeholder: Final = match.group(0).strip("{}") - if not placeholder.endswith(":path"): - return r"[^/]+" - # A ":path" placeholder takes whatever the router's own path - # converter takes, slashes and colons alike, so an id spelled with - # either (or both) still matches the template it was mounted under. - # - # Unless the template puts a ":" literal of its own after the - # placeholder: the Google routes end in ":generateContent" and - # friends, and there the value has to stop before that suffix - # rather than swallow it and match a different verb. - # - # "[\s\S]" rather than ".", because "." stops at a newline and the - # path converter does not: a %0A anywhere in the value would leave - # the route unmatched here while still reaching the handler, which - # turns this gate into a bypass for the lists built on it. - return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" - - pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) - # Anchor the pattern to match the entire string - pattern = f"^{pattern}$" - if re.match(pattern, route): - return True - return False + return isinstance(route, str) and routes.match_route(route, (pattern,), "Template") @staticmethod def is_wildcard_pattern(pattern: str) -> bool: @@ -653,100 +280,15 @@ class RouteChecks: @staticmethod def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: - """ - Check if route matches the wildcard pattern - - eg. - - pattern: "/scim/v2/*" - route: "/scim/v2/Users" - - returns: True - - pattern: "/scim/v2/*" - route: "/chat/completions" - - returns: False - - - pattern: "/scim/v2/*" - route: "/scim/v2/Users/123" - - returns: True - - """ - if pattern.endswith("*"): - # Get the prefix (everything before the wildcard) - prefix: Final = pattern[:-1] - return route.startswith(prefix) - else: - # If there's no wildcard, the pattern and route should match exactly - return route == pattern + return routes.match_route(route, (pattern,), "Wildcard") @staticmethod def _route_matches_allowed_route(route: str, allowed_route: str) -> bool: - """ - Check if route matches the allowed_route pattern. - Supports both exact match and prefix match. - - Examples: - - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6" -> True (exact match) - - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6/v1/chat/completions" -> True (prefix match) - - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-600" -> False (not a valid prefix) - - Args: - route: The actual route being accessed - allowed_route: The allowed route pattern - - Returns: - bool: True if route matches (exact or prefix), False otherwise - """ - # Exact match - if route == allowed_route: - return True - # Prefix match - ensure we add "/" to prevent false matches like /fake-openai-proxy-600 - if route.startswith(allowed_route + "/"): - return True - return False + return routes.match_route(route, (allowed_route,), "Prefix") @staticmethod def check_route_access(route: str, allowed_routes: Collection[str]) -> bool: - """ - Check if a route has access by checking both exact matches and patterns - - Args: - route (str): The route to check - allowed_routes (Sequence): Allowed routes/patterns - - Returns: - bool: True if route is allowed, False otherwise - """ - ######################################################### - # exact match route is in allowed_routes - ######################################################### - if route in allowed_routes: - return True - - ######################################################### - # wildcard match route is in allowed_routes - # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* - ######################################################### - if any( - RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) - for allowed_route in allowed_routes - if RouteChecks.is_wildcard_pattern(pattern=allowed_route) - ): - return True - - ######################################################### - # pattern match route is in allowed_routes - # pattern: "/threads/{thread_id}" - # route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" - # returns: True - ######################################################### - if any( # Check pattern match - RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes - ): - return True - - return False + return routes.match_route(route, tuple(allowed_routes)) @staticmethod def get_request_method(request: Request | None) -> str | None: @@ -768,13 +310,7 @@ class RouteChecks: @staticmethod def is_self_scoped_credential_route(route: str, request: Request | None) -> bool: - method: Final = RouteChecks.get_request_method(request) - if method is None: - return False - return any( - method == expected_method and RouteChecks.check_route_access(route=route, allowed_routes=(pattern,)) - for expected_method, pattern in _SELF_SCOPED_CREDENTIAL_ROUTES - ) + return route_policy_adapter.classify(route, "SelfCredential", RouteChecks.get_request_method(request)) @staticmethod def is_auth_enforced_pass_through_route(route: str, method: str | None = None) -> bool: @@ -806,48 +342,13 @@ class RouteChecks: @staticmethod def _route_matches_denied_route(route: str, denied_route: str) -> bool: - """A `/` entry denies every route, since every route sits under the root.""" - normalized_denied_route: Final = denied_route.rstrip("/") or "/" - return ( - normalized_denied_route == "/" - or RouteChecks._route_matches_allowed_route(route=route, allowed_route=normalized_denied_route) - or RouteChecks.route_matches_wildcard_pattern(route=route, pattern=denied_route) - ) + return routes.match_route(route, (denied_route,), "Denied") @staticmethod def matching_denied_passthrough_route( route: str, metadata_sources: Iterable[Mapping[str, object] | None] ) -> str | None: - """ - First ``denied_passthrough_routes`` entry across ``metadata_sources`` that matches ``route``. - Unlike the allowlist (key list, else team list), every source's deny list applies. - """ - denied_routes: Final = tuple( - itertools.chain.from_iterable( - cast( # cast-ok: management endpoints validate this metadata key as a list of route strings on write - "list[str]", (metadata or {}).get("denied_passthrough_routes") or [] - ) - for metadata in metadata_sources - ) - ) - if not denied_routes: - return None - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - InitPassThroughEndpointHelpers, - ) - - forwarded_routes: Final = InitPassThroughEndpointHelpers.forwarded_routes(route) - return next( - ( - denied_route - for denied_route in denied_routes - if any( - RouteChecks._route_matches_denied_route(route=candidate, denied_route=denied_route) - for candidate in forwarded_routes - ) - ), - None, - ) + return route_policy_adapter.denied_passthrough(route, metadata_sources) @staticmethod def passthrough_route_denied_exception(route: str, denied_route: str) -> HTTPException: @@ -870,16 +371,7 @@ class RouteChecks: @staticmethod def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool: - """ - Explicit paths and trailing-wildcard prefixes grant auth=true pass-through. Blanket grants never do: - a named route group like ``openai_routes`` is only ever compared as a path, and an entry that names - no path segment (``*``, ``/*``) is skipped. - """ - return any( - RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) - for allowed_route in team_allowed_routes - if allowed_route.rstrip("*").strip("/") - ) + return routes.match_route(route, tuple(team_allowed_routes), "Jwt") @staticmethod def _jwt_team_allowed_routes(valid_token: UserAPIKeyAuth) -> Collection[str]: @@ -897,46 +389,18 @@ class RouteChecks: valid_token: UserAPIKeyAuth, jwt_team_allowed_routes: Collection[str] = (), ) -> None: - """ - Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the - key or team, or an explicit JWT ``team_allowed_routes`` entry. A key or team - ``denied_passthrough_routes`` match blocks the route even when one of those grants it. - """ - RouteChecks._raise_if_passthrough_route_denied(route=route, valid_token=valid_token) - if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): - return - if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes): - return - raise RouteChecks._auth_pass_through_denied_exception(route=route) + decision: Final = routes.passthrough_route( + routes.RouteContext(route), + route_policy_adapter.passthrough_facts(route, valid_token, jwt_team_allowed_routes, include_denied=True), + required=True, + ) + route_policy_adapter.raise_route_denial(decision, route, token=valid_token) @staticmethod def check_passthrough_route_access(route: str, user_api_key_dict: UserAPIKeyAuth) -> bool: - """ - Check if route is a passthrough route. - Supports both exact match and prefix match. - """ - metadata: Final = user_api_key_dict.metadata - team_metadata: Final = user_api_key_dict.team_metadata or {} - if metadata is None and team_metadata is None: - return False - if "allowed_passthrough_routes" not in metadata and "allowed_passthrough_routes" not in team_metadata: - return False - if ( - metadata.get("allowed_passthrough_routes") is None - and team_metadata.get("allowed_passthrough_routes") is None - ): - return False - - allowed_passthrough_routes: Final = ( - metadata.get("allowed_passthrough_routes") or team_metadata.get("allowed_passthrough_routes") or [] - ) - - # Check if route matches any allowed passthrough route (exact or prefix match) - for allowed_route in allowed_passthrough_routes: - if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): - return True - - return False + return routes.passthrough_route( + routes.RouteContext(route), route_policy_adapter.passthrough_facts(route, user_api_key_dict) + ).allowed @staticmethod def is_assistants_api_request(request: Request) -> bool: @@ -1017,117 +481,4 @@ class RouteChecks: request_data: dict, request: Request | None = None, ) -> None: - """ - Check access for PROXY_ADMIN_VIEW_ONLY role. - - Admin Viewer follows a read-parity-with-Proxy-Admin rule: anything Proxy - Admin can read/list/get, Admin Viewer can read/list/get. The only - exclusions are cost-incurring inference routes (Playground, /chat/ - completions, etc.) and any state-mutating request. - - Implementation: - 1. LLM/inference routes → 403 (cost-incurring). - 2. Safe HTTP method (GET/HEAD/OPTIONS) → allow by default. This is - the read-parity guarantee — every new GET endpoint added anywhere - in the codebase is automatically readable by Admin Viewer - without needing to remember to add it to an allowlist. - 3. Unsafe HTTP method (POST/PUT/PATCH/DELETE): - - Allow `/user/update` only when restricted to user_email. - - Allow `/user/password/change` (endpoint only writes the caller's own row). - - Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. - - Otherwise allow only if the route is in admin_viewer_routes / - global_spend_tracking_routes (legacy explicit-allow set). - - Else 403. - """ - if RouteChecks.is_llm_api_route(route=route): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", - ) - - # Check if this is a write operation on management routes - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value): - # For management routes, only allow read operations or specific allowed updates - if route == "/user/update": - # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY - if request_data is not None and isinstance(request_data, dict): - _params_updated: Final = request_data.keys() - for param in _params_updated: - if param != "user_email": - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email can be updated", - ) - elif RouteChecks.check_route_access(route=route, allowed_routes=_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES) or ( - route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) - ): - # Block write operations for PROXY_ADMIN_VIEW_ONLY - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", - ) - # Allow read operations on management routes (like /user/info, /team/info, /model/info) - method: Final = request.method.upper() if request is not None else "GET" - is_safe_method: Final = method in RouteChecks._SAFE_HTTP_METHODS - - # ── Safe HTTP method: default-allow ────────────────────────────── - if is_safe_method: - return - - # ── Unsafe HTTP method: explicit checks ────────────────────────── - # Allow `/user/update` for self-service email change. - if route == "/user/update": - if request_data is not None and isinstance(request_data, dict): - for param in request_data: - if param != "user_email": - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=( - f"user not allowed to access this route, role= {_user_role}. " - f"Trying to access: {route} and updating invalid param: {param}. " - "only user_email can be updated" - ), - ) - return - - # Self-service password change; the endpoint only writes the caller's own row. - if route == "/user/password/change": - return - - # Self-service logout; the endpoint only revokes the caller's own session key. - if route == "/session/logout": - return - - if RouteChecks.is_self_scoped_credential_route(route=route, request=request): - return - - # Hard-block known write routes regardless of HTTP method (defensive - # — these are POSTs in practice, but pinning them here protects - # against future GET-shaped writes). - if RouteChecks.check_route_access( - route=route, allowed_routes=RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES - ) or (route.startswith("/key/") and route.endswith("/regenerate")): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", - ) - - # Legacy explicit-allow sets (kept for routes that are POST but - # semantically read-only, e.g. /spend/calculate). Both admin_viewer_routes - # and global_spend_tracking_routes are reads/listings. - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value): - return - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value): - return - - # NOTE: We intentionally do NOT fall back to allowing all - # `management_routes`. That set is a mix of reads (info/list — handled - # via the safe-method branch above) and writes (`/team/block`, - # `/team/permissions_update`, `/jwt/key/mapping/{new,update,delete}`, - # `/key/bulk_update`, `/key/{id}/reset_spend`). A blanket allow would - # let Admin Viewer POST these write endpoints — violating the - # "no writes, ever" rule. Default-deny instead. - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", - ) + return route_policy_adapter.check_viewer_route(route, _user_role, request_data, request) diff --git a/litellm/proxy/auth/route_enable_adapter.py b/litellm/proxy/auth/route_enable_adapter.py new file mode 100644 index 00000000000..91358dbb73b --- /dev/null +++ b/litellm/proxy/auth/route_enable_adapter.py @@ -0,0 +1,130 @@ +import os +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Final, assert_never + +from fastapi import HTTPException +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteFeature, RouteSetting, route_enable +from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.auth.route_policy_adapter import read_eu_jev_configuration, route_context + +_KEYS: Final = {"Admin": "DISABLE_ADMIN_ENDPOINTS", "Llm": "DISABLE_LLM_API_ENDPOINTS"} + + +_EXEMPTIONS: Final = TypeAdapter(tuple[str, ...]) + + +def read_enterprise() -> object: + from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + + return EnterpriseRouteChecks + + +def read_setting(feature: RouteFeature) -> RouteSetting: + from litellm.proxy.proxy_server import premium_user + from litellm.secret_managers.main import ( + get_secret_bool as _get_secret_bool, # noqa: F401 # pyright: ignore[reportUnusedImport] # preserve the Phase 1 import before environment and license reads + ) + + return RouteSetting(_KEYS[feature] in os.environ, bool(premium_user)) + + +def read_secret(feature: RouteFeature) -> bool: + from litellm.secret_managers.main import get_secret_bool + + return get_secret_bool(_KEYS[feature]) is True + + +def read_exemptions(enterprise: object) -> tuple[str, ...]: + return _EXEMPTIONS.validate_python( + getattr(enterprise, "LLM_API_EXEMPT_ROUTES") # noqa: B009 # optional class attributes are validated at this boundary + ) + + +@dataclass(frozen=True, slots=True) +class RouteEnableReaders: + enterprise: Callable[[], object] = read_enterprise + setting: Callable[[RouteFeature], RouteSetting] = read_setting + secret: Callable[[RouteFeature], bool] = read_secret + exemptions: Callable[[object], tuple[str, ...]] = read_exemptions + eu_jev: Callable[[], bool] = read_eu_jev_configuration + + +_READERS: Final = RouteEnableReaders() + + +def read_route_enable_facts(route: str) -> RouteEnableFacts: + facts, _error = _read(RouteEnableFacts(route_context(route)), _READERS) + return facts + + +def _answer( + facts: RouteEnableFacts, decision: RouteEnableDecision, readers: RouteEnableReaders, enterprise: object +) -> tuple[RouteEnableFacts, object]: + match decision.kind: + case "ReadImport": + return replace(facts, imported=True), readers.enterprise() + case "ReadExemptions": + return replace(facts, exemptions=readers.exemptions(enterprise)), enterprise + case "ReadEuJev": + return replace(facts, eu_jev_configured=readers.eu_jev()), enterprise + case "ReadSettings" | "ReadSecret": + if decision.feature is None: + raise NativeAuthRuntimeError("route setting read without feature") + previous: Final = facts.admin if decision.feature == "Admin" else facts.llm + if decision.kind == "ReadSecret" and previous is None: + raise NativeAuthRuntimeError("route secret read without settings") + setting: Final = ( + readers.setting(decision.feature) + if decision.kind == "ReadSettings" or previous is None + else replace(previous, secret=readers.secret(decision.feature)) + ) + return ( + replace(facts, admin=setting) if decision.feature == "Admin" else replace(facts, llm=setting) + ), enterprise + case "Allow" | "Deny" | "PropagateHttp": + raise NativeAuthRuntimeError("completed route enable decision requested IO") + case _: + assert_never(decision.kind) + + +def _read( + facts: RouteEnableFacts, + readers: RouteEnableReaders, + enterprise: object = None, + http_error: HTTPException | None = None, +) -> tuple[RouteEnableFacts, HTTPException | None]: + decision: Final = route_enable(facts) + if decision.kind in ("Allow", "Deny", "PropagateHttp"): + return facts, http_error + try: + updated, imported = _answer(facts, decision, readers, enterprise) + except NativeAuthFailure: + raise + except HTTPException as error: + return _read(replace(facts, failed="Http"), readers, enterprise, error) + except Exception: # noqa: BLE001 # optional enterprise import and configuration IO failures were ignored in Phase 1 + return _read(replace(facts, failed="Other"), readers, enterprise) + return _read(updated, readers, imported) + + +def check_route_enabled(route: str, *, readers: RouteEnableReaders = _READERS) -> None: + facts, error = _read(RouteEnableFacts(route_context(route)), readers) + decision: Final = route_enable(facts) + if decision.kind == "Allow": + return + if decision.kind == "PropagateHttp" and error is not None: + raise error + if decision.kind != "Deny" or decision.feature is None or decision.reason is None: + raise NativeAuthRuntimeError("incomplete route enable decision") + if decision.reason == "License": + label: Final = "ADMIN" if decision.feature == "Admin" else "LLM API" + raise HTTPException( + 500, + f"\U0001f6a8\U0001f6a8\U0001f6a8 DISABLING {label} ENDPOINTS is an Enterprise feature\n\U0001f6a8 {CommonProxyErrors.not_premium_user.value}", + ) + disabled_label: Final = "Management" if decision.feature == "Admin" else "LLM API" + raise HTTPException(403, f"{disabled_label} routes are disabled for this instance.") diff --git a/litellm/proxy/auth/route_policy_adapter.py b/litellm/proxy/auth/route_policy_adapter.py new file mode 100644 index 00000000000..de7309c7ad4 --- /dev/null +++ b/litellm/proxy/auth/route_policy_adapter.py @@ -0,0 +1,256 @@ +from collections.abc import Collection, Iterable, Mapping +from dataclasses import replace +from typing import Final, Literal + +from fastapi import HTTPException, Request +from pydantic import BaseModel, Field, TypeAdapter + +from litellm._auth import routes +from litellm.proxy._types import LiteLLM_UserTable, LiteLLMRoutes, UserAPIKeyAuth + +_STRINGS: Final = TypeAdapter(tuple[str, ...]) + + +class _TokenFields(BaseModel, frozen=True): + allowed_routes: object = None + metadata: Mapping[str, object] | None = Field(default_factory=dict[str, object]) + permissions: Mapping[str, object] | None = None + + +class _ProxyFields(BaseModel, frozen=True): + general_settings: Mapping[str, object] + premium_user: bool + + +def _allowed_routes(value: object) -> tuple[str, ...] | None: + return _STRINGS.validate_python(value) if isinstance(value, list) else None + + +def route_catalog() -> routes.RouteCatalog: + from litellm.proxy.agent_endpoints.auth.managed_authorization import ( + _MANAGED_LIVE_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_MCP_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_MODEL_PATHS, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_MODEL_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_REALTIME_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + ) + from litellm.proxy.auth.route_checks import ( + _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _SELF_SCOPED_CREDENTIAL_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _UNBILLED_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + RouteChecks, + ) + + return routes.RouteCatalog( + groups={name: tuple(member.value) for name, member in LiteLLMRoutes.__members__.items()}, + viewer_blocked=tuple(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES), + viewer_write_blocked=tuple(RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES), # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + self_credentials=_SELF_SCOPED_CREDENTIAL_ROUTES, + managed_models=tuple(_MANAGED_MODEL_ROUTES), + managed_paths=_MANAGED_MODEL_PATHS, + managed_realtime=tuple(_MANAGED_REALTIME_ROUTES), + managed_live=tuple(_MANAGED_LIVE_ROUTES), + managed_mcp=_MANAGED_MCP_ROUTES, + unbilled=tuple(_UNBILLED_ROUTES), + ) + + +def route_context(route: str, method: str | None = None, *, registration: bool = False) -> routes.RouteContext: + if not registration: + return routes.RouteContext(route=route, method=method, catalog=route_catalog()) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + registered: Final = InitPassThroughEndpointHelpers.get_registered_pass_through_route(route=route, method=method) + return routes.RouteContext( + route=route, + method=method, + catalog=route_catalog(), + registered=InitPassThroughEndpointHelpers.is_registered_pass_through_route(route), + registered_auth=registered.get("auth") is True if registered is not None else None, + configured_eu_jev=InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash( + "/eu_jev/v1/systemone" + ), + ) + + +def passthrough_facts( + route: str, token: UserAPIKeyAuth, jwt_team_routes: Collection[str] = (), *, include_denied: bool = False +) -> routes.PassThroughFacts: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + metadata: Final = _TokenFields.model_validate(token, from_attributes=True).metadata or {} + return routes.PassThroughFacts( + key_allowed=_STRINGS.validate_python(metadata.get("allowed_passthrough_routes") or ()), + team_allowed=_STRINGS.validate_python((token.team_metadata or {}).get("allowed_passthrough_routes") or ()), + denied=( + *_STRINGS.validate_python(metadata.get("denied_passthrough_routes") or ()), + *_STRINGS.validate_python((token.team_metadata or {}).get("denied_passthrough_routes") or ()), + ) + if include_denied + else (), + forwarded=InitPassThroughEndpointHelpers.forwarded_routes(route) if include_denied else (), + jwt_team_routes=tuple(jwt_team_routes), + ) + + +def raise_route_denial( + decision: routes.RouteDecision, + route: str, + *, + token: UserAPIKeyAuth | None = None, + role: str | None = None, + user: LiteLLM_UserTable | None = None, + query_user_id: str | None = None, +) -> None: + from litellm.proxy.auth.route_checks import RouteChecks + + match decision.denial: + case None: + return + case "AdminOnly": + RouteChecks._raise_admin_only_route_exception(user_obj=user, route=route) # pyright: ignore[reportPrivateUsage] # preserve the public exception envelope + case "PassThrough": + raise RouteChecks._auth_pass_through_denied_exception(route) # pyright: ignore[reportPrivateUsage] # preserve the public exception envelope + case "DeniedPassThrough": + raise RouteChecks.passthrough_route_denied_exception(route, decision.detail or "") + case "VirtualRoute": + raise HTTPException( + 403, + f"Virtual key is not allowed to call this route. Only allowed to call routes: {_TokenFields.model_validate(token, from_attributes=True).allowed_routes if token is not None else None}. Tried to call route: {route}", + ) + case "PasswordReset": + raise HTTPException( + 403, + "This account's password must be changed before the session can be used: " + "it was either found in a known data breach or set by an admin. " + "Change it via POST /user/password/change (UI: /ui/change-password), then log in again.", + ) + case "CustomAdmin": + raise HTTPException(403, f"user not allowed to access this route. Route={route} is an admin only route") + case "UserInfo": + raise HTTPException( + 403, + f"key not allowed to access this user's info. user_id={query_user_id}, key's user_id={token.user_id if token is not None else None}", + ) + case "ViewerInference": + raise HTTPException(403, f"user not allowed to access this OpenAI routes, role= {role}") + case "ViewerWrite": + raise HTTPException(403, f"user not allowed to access this route, role= {role}. Trying to access: {route}") + case "ViewerParameter": + raise HTTPException( + 403, + f"user not allowed to access this route, role= {role}. Trying to access: {route} and updating invalid param: {decision.detail}. only user_email can be updated", + ) + + +def read_virtual_route_facts(route: str, token: UserAPIKeyAuth, request: Request | None) -> routes.VirtualRouteFacts: + from litellm.proxy.auth.route_checks import RouteChecks + + fields: Final = _TokenFields.model_validate(token, from_attributes=True) + return routes.VirtualRouteFacts( + context=route_context(route, RouteChecks.get_request_method(request), registration=True), + allowed_routes=_allowed_routes(fields.allowed_routes), + passthrough=passthrough_facts(route, token), + password_reset_required=(fields.metadata or {}).get("password_reset_required") is True, + ) + + +def check_virtual_route(route: str, token: UserAPIKeyAuth, request: Request | None) -> bool: + raise_route_denial(routes.virtual_route(read_virtual_route_facts(route, token, request)), route, token=token) + return True + + +def check_viewer_route(route: str, role: str, request_data: Mapping[str, object], request: Request | None) -> None: + from litellm.proxy.auth.route_checks import RouteChecks + + context: Final = route_context(route, RouteChecks.get_request_method(request), registration=True) + facts: Final = routes.ViewerFacts( + replace(context, method="") if request is not None and context.method is None else context, tuple(request_data) + ) + raise_route_denial(routes.viewer_route(facts), route, role=role) + + +def check_non_admin_route( + user: LiteLLM_UserTable | None, + role: str | None, + route: str, + request: Request, + token: UserAPIKeyAuth, + request_data: Mapping[str, object], +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth.route_checks import RouteChecks + + proxy: Final = _ProxyFields.model_validate(vars(proxy_server)) + fields: Final = _TokenFields.model_validate(token, from_attributes=True) + singular: Final = request_data.get("organization_id") + plural: Final = request_data.get("organizations") + organization_ids: Final = ( + *((singular,) if isinstance(singular, str) else ()), + *(_STRINGS.validate_python(plural) if isinstance(plural, list) else ()), + ) + context: Final = route_context(route, RouteChecks.get_request_method(request), registration=True) + facts: Final = routes.NonAdminFacts( + viewer=routes.ViewerFacts( + replace(context, method="") if context.method is None else context, tuple(request_data) + ), + role=role, + custom_admin_routes=_STRINGS.validate_python(proxy.general_settings.get("admin_only_routes") or ()), + premium=proxy.premium_user is True, + query_user_id=request.query_params.get("user_id") if route == "/user/info" else None, + token_user_id=token.user_id, + permissions=tuple(fields.permissions or ()), + organization_ids=organization_ids, + admin_organization_ids=tuple( + membership.organization_id + for membership in (user.organization_memberships or ()) + if membership.user_role == "org_admin" + ) + if user is not None + else (), + team_service_account=bool(token.is_team_service_account), + allowed_routes=_allowed_routes(fields.allowed_routes), + passthrough=passthrough_facts(route, token, RouteChecks._jwt_team_allowed_routes(token), include_denied=True), # pyright: ignore[reportPrivateUsage] # project legacy JWT team route metadata + ) + raise_route_denial( + routes.non_admin_route(facts), route, token=token, role=role, user=user, query_user_id=facts.query_user_id + ) + + +def classify( + route: str, + kind: Literal["Llm", "Unbilled", "Discovery", "SelfCredential", "Azure", "Managed"], + method: str | None = None, +) -> bool: + from litellm._auth.route_enable import llm_route + + context: Final = route_context(route, method) + if kind != "Llm": + return routes.classify_route(context, kind) + decision: Final = llm_route(context) + if decision == "ReadEuJev": + return llm_route(context, read_eu_jev_configuration()) == "Allow" + return decision == "Allow" + + +def read_eu_jev_configuration() -> bool: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + return InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash("/eu_jev/v1/systemone") + + +def denied_passthrough(route: str, metadata_sources: Iterable[Mapping[str, object] | None]) -> str | None: + from itertools import chain + + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + facts: Final = routes.PassThroughFacts( + denied=tuple( + chain.from_iterable( + _STRINGS.validate_python((source or {}).get("denied_passthrough_routes") or ()) + for source in metadata_sources + ) + ), + forwarded=InitPassThroughEndpointHelpers.forwarded_routes(route), + ) + return routes.denied_passthrough_route(route, facts).detail diff --git a/litellm/proxy/auth/route_role_adapter.py b/litellm/proxy/auth/route_role_adapter.py new file mode 100644 index 00000000000..238a5a53d1c --- /dev/null +++ b/litellm/proxy/auth/route_role_adapter.py @@ -0,0 +1,163 @@ +from collections.abc import Iterable +from dataclasses import replace +from itertools import chain +from typing import Final, NoReturn, assert_never + +from fastapi import Request +from pydantic import TypeAdapter +from starlette.routing import compile_path + +from litellm._auth import route_roles as core +from litellm._auth.failure import NativeAuthFailure +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_UserTable, + LiteLLMRoutes, + LitellmUserRoles, + UserAPIKeyAuth, +) + +_STRINGS: Final = TypeAdapter(tuple[str, ...]) +_END: Final = object() + + +def _role(user: LiteLLM_UserTable | None) -> core.RoleFacts: + return core.RoleFacts(user is not None, user.user_role if user is not None else None) + + +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + return LitellmUserRoles(core.role(core.RoleFacts(True, user_role)).effective) + + +def get_user_role(user_obj: LiteLLM_UserTable | None) -> LitellmUserRoles | None: + role: Final = core.role(_role(user_obj)).user_role + return LitellmUserRoles(role) if role is not None else None + + +def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: + return core.role(_role(user_obj)).admin + + +def is_api_route_allowed( + route: str, + request: Request, + request_data: dict[str, object], # mutable-ok: pass the original request dictionary to the shared route wrapper + valid_token: UserAPIKeyAuth | None, + user_obj: LiteLLM_UserTable | None = None, +) -> bool: + from litellm.proxy.auth.route_checks import RouteChecks + + role: Final = get_user_role(user_obj) + initial: Final = core.ApiFacts(valid_token is not None, False, core.RoleFacts(False, None)) + decision: Final = core.api(initial) + if decision == "InvalidToken": + raise Exception("Invalid proxy server token passed. valid_token=None.") # noqa: TRY002 # preserve the public legacy exception class + result: Final = core.api(replace(initial, admin_read=True, user=_role(user_obj))) + if result == "NonAdmin": + assert valid_token is not None + RouteChecks.non_proxy_admin_allowed_routes_check( # pyright: ignore[reportUnknownMemberType] # the shared wrapper still declares an unparameterized request dictionary + user_obj=user_obj, + _user_role=role, + route=route, + request=request, + request_data=request_data, + valid_token=valid_token, + ) + return True + + +def _entry(raw: object) -> core.Entry: + try: + hash(raw) + except NativeAuthFailure: + raise + except TypeError as error: + return core.Entry(None, type(raw).__name__, hash_failure=str(error)) + text: Final = raw if isinstance(raw, str) else None + group: Final = LiteLLMRoutes.__members__.get(text) if text is not None else None + return core.Entry( + text, + type(raw).__name__, + templates=_STRINGS.validate_python(group.value) if group is not None else None, # pyright: ignore[reportUnknownMemberType] # enum groups are legacy route collections + ) + + +def _raise_failure(failure: core.Failure | None) -> NoReturn: + assert failure is not None + match failure.kind: + case "TypeError": + raise TypeError(failure.message) + case "AttributeError": + raise AttributeError(failure.message) + case _: + assert_never(failure.kind) + + +def allowed_routes_list_check(user_route: str, allowed_routes: Iterable[object]) -> bool: + entries: Final = iter(allowed_routes) + facts = core.ListedFacts(user_route) # rebind-ok: native-selected data-read driver state + while True: + decision = core.listed(facts) + match decision.action: + case "ReadEntry": + raw = next(entries, _END) + facts = core.ListedFacts(user_route, None if raw is _END else _entry(raw), raw is _END) + case "ReadTemplate": + assert facts.entry is not None and facts.entry.templates is not None + compiled, _, _ = compile_path(facts.entry.templates[decision.template_index]) + facts = replace(facts, template_index=decision.template_index, compiled=compiled.pattern) + case "Allow": + return True + case "Deny": + return False + case "Failure": + _raise_failure(decision.failure) + case _: + raise AssertionError(f"Unexpected listed-route directive: {decision.action}") + + +def allowed_routes_check( + user_role: LitellmUserRoles, + user_route: str, + litellm_proxy_roles: LiteLLM_JWTAuth, +) -> bool: + facts: Final = core.JwtFacts(user_role) + decision: Final = core.jwt(facts) + match decision: + case "Deny": + return False + case "ReadAdmin": + return allowed_routes_list_check(user_route, litellm_proxy_roles.admin_allowed_routes) + case "ReadTeam": + configured: Final = litellm_proxy_roles.team_allowed_routes + selected: Final = core.jwt(replace(facts, config_read=True, config_missing=configured is None)) # pyright: ignore[reportUnnecessaryComparison] # constructed legacy JWT configs retain the None default branch + return allowed_routes_list_check( + user_route, ("openai_routes", "info_routes") if selected == "DefaultTeam" else configured + ) + case _: + raise AssertionError(f"Unexpected JWT route directive: {decision}") + + +def allowed_route_check_inside_route(user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None) -> bool: + return core.inside(core.InsideFacts(user_api_key_dict.user_role, user_api_key_dict.user_id, requested_user_id)) + + +def _expanded(raw: object) -> tuple[object, ...]: + facts: Final = _entry(raw) + decision: Final = core.expand(facts) + match decision.action: + case "Expand": + assert facts.templates is not None + return facts.templates + case "Keep": + return (raw,) + case "Failure": + _raise_failure(decision.failure) + case _: + raise AssertionError(f"Unexpected route expansion directive: {decision.action}") + + +def get_actual_routes( + allowed_routes: Iterable[object], +) -> list[object]: # mutable-ok: preserve the public list projection + return list(chain.from_iterable(_expanded(raw) for raw in allowed_routes)) diff --git a/litellm/proxy/auth/route_scope_adapter.py b/litellm/proxy/auth/route_scope_adapter.py new file mode 100644 index 00000000000..a4c5de07178 --- /dev/null +++ b/litellm/proxy/auth/route_scope_adapter.py @@ -0,0 +1,73 @@ +from collections.abc import Iterable, Mapping +from dataclasses import replace +from typing import Final, assert_never + +from fastapi import HTTPException +from pydantic import BaseModel, TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.route_scope import RouteScopeFacts, RouteScopeKind, RouteValues, route_scope +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors + +_OBJECTS: Final = TypeAdapter(Iterable[object]) + + +class _Settings(BaseModel, frozen=True): + general_settings: Mapping[str, object] + premium_user: object + + +def route_values(value: object) -> RouteValues: + if isinstance(value, str): + return RouteValues(text=value) + if isinstance(value, (bytes, bytearray)): + return RouteValues(failure="a bytes-like object is required, not 'str'") + if isinstance(value, Iterable): + return RouteValues(names=tuple(item for item in _OBJECTS.validate_python(value) if isinstance(item, str))) + return RouteValues(failure=f"argument of type '{type(value).__name__}' is not iterable") + + +def read_route_scope_facts(route: str, kind: RouteScopeKind) -> RouteScopeFacts: + from litellm.proxy import proxy_server + + settings: Final = _Settings.model_validate(vars(proxy_server)) + setting: Final = "allowed_routes" if kind == "Proxy" else "admin_only_routes" + initial: Final = RouteScopeFacts(kind, route, setting in settings.general_settings, settings.premium_user is True) + + def advance(facts: RouteScopeFacts) -> RouteScopeFacts: + decision: Final = route_scope(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "Warn": + verbose_proxy_logger.error( + "Trying to set allowed_routes. This is an Enterprise feature. %s" + if kind == "Proxy" + else "Trying to use 'admin_only_routes' this is an Enterprise only feature. %s", + CommonProxyErrors.not_premium_user.value, + ) + return advance(replace(facts, warned=True)) + case "ReadValues": + return advance(replace(facts, values=route_values(settings.general_settings[setting]))) + case "Failure": + raise TypeError(decision.message) + case _: + assert_never(decision.kind) + + return advance(initial) + + +def is_allowed_on_proxy(route: str) -> bool: + return route_scope(read_route_scope_facts(route, "Proxy")).kind == "Allow" + + +def custom_admin_only_route_check(route: str) -> None: + decision: Final = route_scope(read_route_scope_facts(route, "Admin")) + match decision.kind: + case "Allow": + return + case "Deny": + raise HTTPException(403, f"user not allowed to access this route. Route={route} is an admin only route") + case _: + raise NativeAuthRuntimeError("incomplete custom admin route decision") diff --git a/litellm/proxy/auth/spend_counter_adapter.py b/litellm/proxy/auth/spend_counter_adapter.py new file mode 100644 index 00000000000..9985d5c5ab0 --- /dev/null +++ b/litellm/proxy/auth/spend_counter_adapter.py @@ -0,0 +1,34 @@ +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, replace +from typing import Final, NoReturn + +from litellm._auth import spend_counter_policy as core +from litellm._auth.failure import NativeAuthRuntimeError + + +@dataclass(frozen=True, slots=True) +class SpendIO: + read_strict: Callable[[], bool] + read_floor: Callable[[], Awaitable[float | None]] + repair: Callable[[float], Awaitable[None]] + unavailable: Callable[[], NoReturn] + + +async def resolve_spend(facts: core.Spend, io: SpendIO) -> float: + decision: Final = core.decide(facts) + match decision: + case core.SpendValue(kind="Value", spend=spend): + return spend + case core.SpendValue(kind="Repair", spend=spend): + await io.repair(spend) + return spend + case core.Action(kind="ReadStrictBefore"): + return await resolve_spend(replace(facts, strict_before=io.read_strict()), io) + case core.Action(kind="ReadStrictAfter"): + return await resolve_spend(replace(facts, strict_after=io.read_strict()), io) + case core.Action(kind="ReadFloor"): + return await resolve_spend(replace(facts, authoritative_read=True, authoritative=await io.read_floor()), io) + case core.Action(kind="Unverifiable"): + io.unavailable() + case _: + raise NativeAuthRuntimeError("Unexpected spend-counter policy decision") diff --git a/litellm/proxy/auth/team_alias_adapter.py b/litellm/proxy/auth/team_alias_adapter.py new file mode 100644 index 00000000000..1c83caf1dba --- /dev/null +++ b/litellm/proxy/auth/team_alias_adapter.py @@ -0,0 +1,33 @@ +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.team_aliases import TeamAliasFacts, team_aliases +from litellm.router_api import RegistryReader + + +def resolve_team_alias( + model: str | Sequence[str], + team_model_aliases: Mapping[str, str] | None, + team_id: str | None, + llm_router: RegistryReader | None, +) -> str | list[str]: # mutable-ok: the existing public wrapper preserves the input scalar/list shape + initial: Final = TeamAliasFacts( + (model,) if isinstance(model, str) else tuple(model), team_model_aliases or {}, team_id, llm_router is not None + ) + decision: Final = team_aliases(initial) + if decision.kind == "Complete": + return decision.models[0] if isinstance(model, str) else list(decision.models) + if llm_router is None: + raise NativeAuthRuntimeError("team aliases requested names without a router") + names: Final = tuple(_names(llm_router.registry().index("model_name")) for _ in range(decision.count)) + completed: Final = team_aliases(replace(initial, names=names)) + if completed.kind != "Complete": + raise NativeAuthRuntimeError("incomplete team alias resolution") + return completed.models[0] if isinstance(model, str) else list(completed.models) + + +@native_boundary +def _names(names: Iterable[str]) -> tuple[str, ...]: + return tuple(names) diff --git a/litellm/proxy/auth/team_override_adapter.py b/litellm/proxy/auth/team_override_adapter.py new file mode 100644 index 00000000000..ee0ed0c1d3b --- /dev/null +++ b/litellm/proxy/auth/team_override_adapter.py @@ -0,0 +1,102 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.team_override import TeamOverrideFacts, team_override +from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.model_group_adapter import get_authorized_resources_from_key_access_groups +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] + token: UserAPIKeyAuth | None + team: LiteLLM_TeamTable | None + router: RegistryReader | None + team_aliases: Mapping[str, str] | None + key_aliases: Mapping[str, str] | None + prisma: DatabaseClient | None + + +async def read_team_override_facts( + model: str | Sequence[str], + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> TeamOverrideFacts: + return await _read( + TeamOverrideFacts(), + _Source(model, valid_token, team_object, llm_router, team_model_aliases, key_model_aliases, prisma_client), + ) + + +async def can_team_call_model_with_override( + model: str | Sequence[str], + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> Literal[True]: + facts: Final = await read_team_override_facts( + model, valid_token, team_object, llm_router, team_model_aliases, key_model_aliases, prisma_client + ) + decision: Final = team_override(facts) + if decision.kind == "Allow": + return True + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete team override authorization") + return raise_model_denial(decision.denial, model) + + +async def _read(facts: TeamOverrideFacts, source: _Source) -> TeamOverrideFacts: + + decision: Final = team_override(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadTeam": + scope: Final = project_model_facts( + source.model, + source.router, + object_model_grants(source.team) if source.team is not None else (), + source.team_aliases, + source.team.team_id if source.team is not None else None, + source.key_aliases, + "team", + ) + groups: Final = tuple(source.team.access_group_ids or ()) if source.team is not None else () + return await _read(replace(facts, team=GroupedModelScope(scope, groups)), source) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing team override scope") + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=grouped), source) + case "ReadAssignedGroups": + assigned: Final = await get_authorized_resources_from_key_access_groups( + source.token, source.team, "access_model_names" + ) + return await _read(replace(facts, assigned_models=tuple(assigned)), source) + case "ReadAssignedScope": + if facts.assigned_models is None: + raise NativeAuthRuntimeError("missing assigned model override grants") + assigned_scope: Final = project_model_facts( + source.model, + source.router, + facts.assigned_models, + source.token.team_model_aliases if source.token is not None else None, + source.token.team_id if source.token is not None else None, + source.key_aliases, + "key", + ) + return await _read(replace(facts, assigned_scope=assigned_scope), source) diff --git a/litellm/proxy/auth/unpriced_model_adapter.py b/litellm/proxy/auth/unpriced_model_adapter.py new file mode 100644 index 00000000000..7ed766df168 --- /dev/null +++ b/litellm/proxy/auth/unpriced_model_adapter.py @@ -0,0 +1,128 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from pydantic import BaseModel, Field, TypeAdapter + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.unpriced_models import DeploymentPricing, PriceEntry, PriceMetric, UnpricedModelFacts, unpriced_model +from litellm.router_api import RegistryReader + +_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OPTIONAL_STRING: Final = TypeAdapter[str | None](str | None) + + +class _ModelInfo(BaseModel, frozen=True): + id: str | None = None + + +class _Deployment(BaseModel, frozen=True): + litellm_params: Mapping[str, object] = Field(default_factory=dict[str, object]) + model_info: _ModelInfo = Field(default_factory=_ModelInfo) + + +class _CostMap(BaseModel, frozen=True): + model_cost: Mapping[str, object] + + +def _number(value: object) -> float | None: + return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None + + +def _metric(name: str, value: object) -> PriceMetric: + return PriceMetric( + name, + _number(value), + tuple(_number(item) for item in _MAPPING.validate_python(value).values()) if isinstance(value, dict) else (), + ) + + +def price_entry(entry: Mapping[str, object]) -> PriceEntry: + return PriceEntry( + entry.get("tiered_pricing") is not None, tuple(_metric(key, value) for key, value in entry.items()) + ) + + +def read_unpriced_model_facts(model: str | None, llm_router: RegistryReader | None) -> UnpricedModelFacts: + return _read(UnpricedModelFacts(model, llm_router is not None), llm_router, ()) + + +def _read(facts: UnpricedModelFacts, router: RegistryReader | None, rows: Sequence[_Deployment]) -> UnpricedModelFacts: + decision: Final = unpriced_model(facts) + if decision.kind == "Complete": + return facts + if router is None or facts.model is None: + raise NativeAuthRuntimeError("unpriced model read without a router or model") + match decision.kind: + case "ReadGroup": + found: Final = router.registry().group(facts.model) is not None + return _read(replace(facts, group_found=found), router, rows) + case "ReadDeployments" | "ReadCostDeployments": + deployments: Final = tuple( + _Deployment.model_validate(row) for row in router.registry().deployments(facts.model) or () + ) + projected: Final = tuple( + DeploymentPricing(tuple(row.litellm_params), row.model_info.id is not None) for row in deployments + ) + snapshot: Final = ( + replace(facts, deployments=projected) + if decision.kind == "ReadDeployments" + else replace(facts, cost_deployments=projected) + ) + return _read(snapshot, router, deployments) + case "ReadInfo" | "ReadCost": + if facts.deployments is None: + raise NativeAuthRuntimeError("missing unpriced deployment facts") + row: Final = rows[decision.index] + model_id: Final = row.model_info.id + if model_id is None: + raise NativeAuthRuntimeError("missing unpriced deployment identifier") + current: Final = ( + facts.deployments[decision.index] + if decision.kind == "ReadInfo" + else (facts.cost_deployments or ())[decision.index] + ) + if decision.kind == "ReadInfo": + info: Final = router.registry().deployment_model_info( + model_id, _OPTIONAL_STRING.validate_python(row.litellm_params.get("model")) or "" + ) + updated: Final = replace( + current, + info_loaded=True, + info=price_entry(_MAPPING.validate_python(info)) if info is not None else None, + ) + return _read( + replace( + facts, + deployments=tuple( + updated if index == decision.index else item for index, item in enumerate(facts.deployments) + ), + ), + router, + rows, + ) + keys: Final = tuple( + _MAPPING.validate_python(_CostMap.model_validate(vars(litellm)).model_cost.get(model_id, {})) + ) + cost: Final = replace(current, cost_keys=keys) + return _read( + replace( + facts, + cost_deployments=tuple( + cost if index == decision.index else item + for index, item in enumerate(facts.cost_deployments or ()) + ), + ), + router, + rows, + ) + + +def model_has_no_cost_mapping(model: str | None, llm_router: RegistryReader | None) -> bool: + return unpriced_model(read_unpriced_model_facts(model, llm_router)).unpriced + + +def unpriced_models_in_request(model: str | Sequence[str] | None, llm_router: RegistryReader | None) -> tuple[str, ...]: + names: Final = (model,) if isinstance(model, str) else tuple(model or ()) + return tuple(name for name in names if model_has_no_cost_mapping(name, llm_router)) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a52ad14e293..f1b624cb35b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -9,13 +9,12 @@ Returns a UserAPIKeyAuth object if the API key is valid import asyncio import fnmatch -import re -import secrets from collections.abc import Mapping, Sequence +from dataclasses import replace from datetime import datetime, timezone from enum import Enum from types import MappingProxyType -from typing import Final, NamedTuple, Protocol, Union, cast +from typing import Final, Literal, NamedTuple, Protocol, Union, assert_never, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -24,17 +23,55 @@ from pydantic import TypeAdapter from starlette.exceptions import WebSocketException import litellm +from litellm._auth import access as access_policy +from litellm._auth import admission, auth_context, common_flow +from litellm._auth.budget_pricing import ExemptionFacts +from litellm._auth.credential import ( + PresentedKeys, + WebSocketDenial, + WebSocketKeys, +) +from litellm._auth.decision import ( + AdminKey, + Allow, + AuthCore, + CachedAdminKey, + Deny, + FallbackTarget, + Grant, + KeyFacts, + MasterKey, + ModelBudgetRead, + ModelBudgetScope, + NoMasterKey, + NonAdminKey, + Outcome, + OwnerFacts, + PresentedKey, + RequestFacts, + Settings, + Step, + TeamMemberBudget, + TeamMemberModelBudgetScope, + UserAuthRoute, + VirtualKey, + WithinModelBudgets, +) +from litellm._auth.driver import drive, drive_model_budgets, drive_team_member_model_budgets +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts, key_and_fallback_model_access +from litellm._auth.model_budget_data import ModelBudgetEntity, ModelBudgetSnapshot +from litellm._auth.runtime import get_auth_core, scheme_credential, select_api_key, select_websocket_api_key +from litellm._auth.rust_core import expired from litellm._internal_context import service_target, set_request_parent_otel_span from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.constants import ( - CLIENT_REQUESTED_MODEL_SCOPE_KEY, GLOBAL_PROXY_SPEND_CACHE_KEY, - INVALID_VIRTUAL_KEY_ERROR_MARKER, - INVALID_VIRTUAL_KEY_ERROR_MESSAGE, LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS, - MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY, ) from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_event, phase_span, seed_request_identity @@ -54,49 +91,39 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo _get_user_role, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _is_user_proxy_admin, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - _model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # reuse common-check group budget enforcement _team_member_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - _unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # reuse common-check pricing enforcement _virtual_key_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _virtual_key_max_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _virtual_key_soft_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export cache_key_object, - can_customer_access_model, - can_key_call_model, - can_object_call_model, - can_project_access_model, - can_team_access_model, check_end_user_budget, common_checks, delete_cache_key_object, + get_default_end_user_budget, get_end_user_object, get_jwt_key_mapping_object, get_key_end_user_budget_id, get_object_permission, get_org_object_for_request, get_project_object, - get_team_member_default_budget, get_team_membership, get_team_object, get_user_object, get_user_role, is_model_cost_zero, is_user_proxy_admin, - is_valid_fallback_model, jwt_key_mapping_cache_key, - key_access_group_grants_model, - key_model_aliases_for_auth_check, - model_max_budget_of, request_skips_budget_checks, resolve_and_validate_end_user_id, - resolve_default_end_user_budget, route_skips_budget_checks, - stamp_matched_model_access_groups, + send_soft_budget_alert, + send_token_budget_alert, team_member_max_budget_alert_check, virtual_key_max_budget_alert_check, virtual_key_max_budget_check, - virtual_key_soft_budget_check, + virtual_key_soft_budget_check, # pyright: ignore[reportUnusedImport] # backwards-compatible package export ) +from litellm.proxy.auth.auth_denial import denial_exception from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_method import AuthMethod from litellm.proxy.auth.auth_object_prefetch import ( @@ -107,21 +134,22 @@ from litellm.proxy.auth.auth_object_prefetch import ( ) from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, - fallback_target_model_name, get_end_user_id_from_request_body, get_model_from_request, get_request_route, get_request_route_template, is_invalid_virtual_key_error, - iter_request_fallback_targets, log_model_access_denial, normalize_request_route, pre_db_read_auth_checks, - request_dispatched_to_pass_through_endpoint, - request_dispatched_to_provider_pass_through, route_in_additonal_public_routes, ) +from litellm.proxy.auth.budget_throttle import budget_throttle_percentage, should_throttle_budget_exceeded +from litellm.proxy.auth.fallback_access_adapter import load_fallback_authorization_facts from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy.auth.key_model_access_adapter import read_key_and_fallback_model_access_facts +from litellm.proxy.auth.model_access_adapter import raise_model_denial +from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request @@ -154,7 +182,6 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_query_params, safe_set_request_parsed_body, ) -from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401 # legacy module exports _realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export realtime_request_body, @@ -274,50 +301,21 @@ def _get_model_from_request_context( ) -_CLAUDE_MODEL_ROUTES: Final = frozenset( - f"/{prefix}{endpoint}" for prefix in ("", "v1/") for endpoint in ("messages", "chat/completions", "responses") -) -_CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized" _GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_PASS_THROUGH_HEADERS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( + Mapping[str, object] | None +) async def _normalize_claude_model( - request_data: dict, valid_token: UserAPIKeyAuth, request: Request | None, route: str + request_data: dict[str, object], # mutable-ok: normalization rewrites the request body for downstream readers + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, ) -> None: - from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + from litellm.proxy.auth.claude_normalization_adapter import normalize_claude_model - if route not in _CLAUDE_MODEL_ROUTES or llm_router is None: - return - if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True: - return - requested: Final = _get_model_from_request_context(request_data, route, request, llm_router, valid_token.team_id) - if not isinstance(requested, str) or requested != request_data.get("model"): - return - if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"): - return - settings: Final = await proxy_config.get_hierarchical_router_settings( - user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) - aliases: Final = settings.get("model_group_alias") if isinstance(settings, Mapping) else None - source: Final = claude_code_requested_group( - requested, llm_router, valid_token.team_id, (valid_token.aliases, valid_token.team_model_aliases, aliases) - ) - if request is not None: - request.scope[_CLAUDE_MODEL_NORMALIZED] = True - if source is None: - return - rewrite_request_model(request_data, request, source) - - -_CODEX_NATIVE_RESPONSES_ROUTE: Final = "/backend-api/codex/responses" - - -def _is_codex_native_litellm_turn(route: str, model: object, llm_router: litellm.Router, team_id: str | None) -> bool: - if route != _CODEX_NATIVE_RESPONSES_ROUTE: - return False - from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import routes_to_litellm - - return routes_to_litellm(model, llm_router, team_id) + await normalize_claude_model(request_data, valid_token, request, route) async def _resolve_router_settings_model_group_alias( @@ -326,44 +324,9 @@ async def _resolve_router_settings_model_group_alias( request: Request | None, route: str, ) -> None: - """Rewrite the requested model through the key's or team's ``router_settings.model_group_alias`` - before the allowlist checks, so they authorize the model group the request is routed to. - """ - from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + from litellm.proxy.auth.codex_normalization_adapter import resolve_router_settings_model_group_alias - if request is None or llm_router is None or not RouteChecks.is_llm_api_route(route=route): - return - if request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True: - return - request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True - requested: Final = request_data.get("model") - if request_dispatched_to_pass_through_endpoint(request) or ( - request_dispatched_to_provider_pass_through(request) - and not _is_codex_native_litellm_turn(route, requested, llm_router, valid_token.team_id) - ): - return - if not isinstance(requested, str) or await read_raw_json_body(request=request) is None: - return - settings: Final = await proxy_config.get_hierarchical_router_settings( - user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) - if not isinstance(settings, Mapping): - return - target: Final = resolve_model_group_alias(settings.get("model_group_alias"), requested) - if target is None or target == requested: - return - if request_dispatched_to_provider_pass_through(request): - from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import native_turn_serves_through_litellm - - if not await native_turn_serves_through_litellm(requested, llm_router, request, valid_token): - return - verbose_proxy_logger.debug( - "router_settings.model_group_alias resolved %s -> %s before auth", - requested.replace("\r", "").replace("\n", ""), - target.replace("\r", "").replace("\n", ""), - ) - request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, requested) - rewrite_request_model(request_data, request, target) + await resolve_router_settings_model_group_alias(request_data, valid_token, request, route) def _resolve_codex_native_chatgpt_model( @@ -372,19 +335,9 @@ def _resolve_codex_native_chatgpt_model( request: Request | None, route: str, ) -> None: - requested: Final = request_data.get("model") - if route != _CODEX_NATIVE_RESPONSES_ROUTE or request is None or not isinstance(requested, str): - return - from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import native_chatgpt_model - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.auth.codex_normalization_adapter import resolve_codex_native_chatgpt_model - if llm_router is None: - return - target: Final = native_chatgpt_model(requested, llm_router, valid_token) - if target is None: - return - request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, requested) - rewrite_request_model(request_data, request, target) + resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) def _get_model_names_for_budget_checks( @@ -409,13 +362,21 @@ def _budget_fallback_target_model(target: object) -> str | None: return model if isinstance(model, str) else None -class _KeyModelBudgetLimiter(Protocol): +class _RawModelBudgetLimiter(Protocol): + async def read_model_budget_raw( + self, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> ModelBudgetReading: ... + + +class _KeyAndUserModelBudgetLimiter(_RawModelBudgetLimiter, Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... - async def get_fallback_model_within_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> str | None: ... - - -class _UserModelBudgetLimiter(Protocol): async def is_user_within_model_budget( self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None ) -> bool: ... @@ -441,7 +402,7 @@ class _EndUserModelBudgetLimiter(Protocol): ) -> bool: ... -class _EveryModelBudgetLimiter(Protocol): +class _EveryModelBudgetLimiter(_RawModelBudgetLimiter, Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... async def is_user_within_model_budget( @@ -493,21 +454,6 @@ def _token_team_models(valid_token: _TokenTeamModels) -> list[str]: return valid_token.team_models -class _TokenModelBudgetFields(Protocol): - @property - def model_max_budget(self) -> Mapping[str, object] | None: ... - - @property - def end_user_model_max_budget(self) -> Mapping[str, object] | None: ... - - @property - def team_model_aliases(self) -> dict[str, str] | None: ... # mutable-ok: mirrors _can_object_call_model's contract - - -def _token_model_budget_fields(valid_token: _TokenModelBudgetFields) -> _TokenModelBudgetFields: - return valid_token - - def _router_fallbacks(router: litellm.Router) -> Sequence[object] | None: fallbacks: Final[Sequence[object] | None] = router.settings().fallbacks return fallbacks @@ -536,51 +482,22 @@ async def _read_user_model_max_budget( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # mirrors the main path's tolerance verbose_logger.debug("Unable to read user for the per-model budget check: %s", e) return None return user_obj.model_max_budget if user_obj is not None else None -async def _check_user_model_budget( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _UserModelBudgetLimiter, - models: list[str], -) -> None: - """Enforce the internal user's own `model_max_budget` across the request's models. - - Separate from the key check: a user's per-model budget caps every key they - own, so a caller cannot escape it by minting another key. - """ - user_model_max_budget: Final = valid_token.user_model_max_budget - if valid_token.user_id is None or not isinstance(user_model_max_budget, Mapping) or not user_model_max_budget: - return - for model_name in models: - await model_max_budget_limiter.is_user_within_model_budget( - user_id=valid_token.user_id, - user_model_max_budget=user_model_max_budget, - model=model_name, - team_id=valid_token.team_id, - ) - - async def _check_team_model_budget( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _TeamModelBudgetLimiter, models: list[str], ) -> None: - """Enforce the team's `model_max_budget` for every requested model the key does not override.""" - team_model_max_budget: Final = valid_token.team_model_max_budget - if valid_token.team_id is None or not team_model_max_budget: - return - key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget - for model_name in models: - await model_max_budget_limiter.is_team_within_model_budget( - team_id=valid_token.team_id, - team_model_max_budget=team_model_max_budget, - key_model_max_budget=key_model_max_budget, - model=model_name, - ) + from litellm.proxy.auth import budget_context_adapter + + return await budget_context_adapter.check_team_model_budget(valid_token, model_max_budget_limiter, models) async def _end_user_model_max_budget( @@ -591,21 +508,11 @@ async def _end_user_model_max_budget( user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, ) -> Mapping[str, object] | None: - budget_table: Final = ( - end_user_object.litellm_budget_table - if end_user_object is not None - else ( - await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, - ) - if end_user_id and prisma_client is not None - else None - ) + from litellm.proxy.auth import budget_context_adapter + + return await budget_context_adapter.end_user_model_max_budget( + end_user_id, end_user_object, key_end_user_budget_id, prisma_client, user_api_key_cache, parent_otel_span ) - return model_max_budget_of(budget_table) if budget_table is not None else None async def _check_end_user_model_budget( @@ -615,169 +522,333 @@ async def _check_end_user_model_budget( models: Sequence[str], team_id: str | None, ) -> None: - if not end_user_id or not end_user_model_max_budget: - return - for model_name in models: - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=end_user_id, - end_user_model_max_budget=end_user_model_max_budget, - model=model_name, - team_id=team_id, - ) + from litellm.proxy.auth import budget_context_adapter + + return await budget_context_adapter.check_end_user_model_budget( + end_user_id, end_user_model_max_budget, model_max_budget_limiter, models, team_id + ) -async def _is_within_every_model_budget( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - model: str, -) -> bool: - budget_fields: Final = _token_model_budget_fields(valid_token) - key_model_max_budget: Final = budget_fields.model_max_budget - user_model_max_budget: Final = valid_token.user_model_max_budget - end_user_model_max_budget: Final = budget_fields.end_user_model_max_budget - team_model_max_budget: Final = valid_token.team_model_max_budget - team_member_model_max_budget: Final = valid_token.team_member_model_max_budget - try: - if valid_token.token is not None and isinstance(key_model_max_budget, Mapping) and key_model_max_budget: - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=model, - ) - if valid_token.user_id is not None and isinstance(user_model_max_budget, Mapping) and user_model_max_budget: - await model_max_budget_limiter.is_user_within_model_budget( - user_id=valid_token.user_id, - user_model_max_budget=user_model_max_budget, - model=model, - team_id=valid_token.team_id, - ) - if ( - valid_token.end_user_id is not None - and isinstance(end_user_model_max_budget, Mapping) - and end_user_model_max_budget - ): - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_model_max_budget, - model=model, - team_id=valid_token.team_id, - ) - if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget: - await model_max_budget_limiter.is_team_within_model_budget( - team_id=valid_token.team_id, - team_model_max_budget=team_model_max_budget, - key_model_max_budget=key_model_max_budget, - model=model, - ) - if ( - valid_token.user_id is not None - and valid_token.team_id is not None - and isinstance(team_member_model_max_budget, Mapping) - and team_member_model_max_budget - ): - await model_max_budget_limiter.is_team_member_within_model_budget( - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_member_model_max_budget=team_member_model_max_budget, - model=model, - ) - except litellm.BudgetExceededError: - return False - return True +def _model_budget_read( + reading: ModelBudgetReading, message: str, entity_type: str, entity_id: str | None +) -> ModelBudgetRead: + return ModelBudgetRead( + reading.snapshot, + message if reading.named_model else f"{message}, budget_model={reading.budget_model}", + entity_type, + entity_id, + ) -async def _is_fallback_model_authorized( - model: str, - valid_token: UserAPIKeyAuth, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router | None, - prisma_client: PrismaClient | None = None, - team_membership: LiteLLM_TeamMembership | None = None, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - refuse_agent_keys: bool = False, - team_object: LiteLLM_TeamTable | None = None, -) -> bool: - if refuse_agent_keys and valid_token.agent_id is not None: - return False - try: - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - ) - if team_object is not None: - try: - await can_team_access_model( - model=model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - prisma_client=prisma_client, +def _model_budget_scope_enabled(scope: Literal["Key", "User", "Member"], token: UserAPIKeyAuth, budget: object) -> bool: + facts = common_flow.ModelBudgets(scope) # rebind-ok: native-selected raw field read stages + while True: + match common_flow.read(facts): + case "ReadModelBudgetPresence": + facts = replace(facts, budget_present=bool(budget)) + case "ReadModelBudgetShape": + facts = replace(facts, mapping=isinstance(budget, Mapping)) + case "ReadModelBudgetToken": + facts = replace(facts, entity_present=token.token is not None) + case "ReadModelBudgetUser": + facts = replace(facts, entity_present=token.user_id is not None) + case "ReadModelBudgetTeam": + facts = replace(facts, team_present=token.team_id is not None) + case "CheckModelBudgets": + return True + case "Continue": + return False + case _: + raise NativeAuthRuntimeError("Unexpected model-budget scope read") + + +def _model_budget_fallback_enabled(router_present: bool, data: Mapping[str, object]) -> bool: + initial: Final = common_flow.FallbackAllowed(router_present) + facts: Final = ( + replace(initial, disabled=fallbacks_disabled(data)) + if common_flow.read(initial) == "ReadFallbackDisabled" + else initial + ) + return common_flow.read(facts) == "AllowFallback" + + +class _ModelBudgetEffects: + def __init__( + self, + *, + request: Request, + request_data: dict[str, object], # mutable-ok: budget fallbacks rewrite the request model in place + valid_token: UserAPIKeyAuth, + route: str, + limiter: _KeyAndUserModelBudgetLimiter, + llm_router: litellm.Router | None, + llm_model_list: Sequence[object] | None, + ) -> None: + self._request: Final = request + self._request_data: Final = request_data + self._valid_token: Final = valid_token + self._route: Final = route + self._limiter: Final = limiter + self._llm_router: Final = llm_router + self._llm_model_list: Final = llm_model_list + + @property + def valid_token(self) -> UserAPIKeyAuth: + return self._valid_token + + @property + def route(self) -> str: + return self._route + + async def read_model_budget_scope(self) -> ModelBudgetScope: + requested: Final = self._request_data.get("model") + user_model_max_budget: Final = self._valid_token.user_model_max_budget + budget_fallbacks: Final[Mapping[str, list[str]]] = self._valid_token.budget_fallbacks or {} + return ModelBudgetScope( + models=tuple( + _get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=self._request_data, + route=self._route, + request=self._request, + llm_router=self._llm_router, + team_id=self._valid_token.team_id, + ) ) - except ProxyException as team_denial: - if team_denial.type != ProxyErrorTypes.team_model_access_denied: - raise - if not await key_access_group_grants_model( - model=model, - valid_token=valid_token, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ): - raise - else: - team_models: Final = _token_team_models(valid_token) - if team_models: - can_object_call_model( - model=model, - llm_router=llm_router, - models=team_models, - team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="team", - ) - member_allowed_models: Final = ( - team_membership.litellm_budget_table.allowed_models - if team_membership is not None and team_membership.litellm_budget_table is not None - else None + ), + rewritable_model=requested if isinstance(requested, str) else None, + key_has_model_budgets=_model_budget_scope_enabled( + "Key", self._valid_token, self._valid_token.model_max_budget + ), + key_budget_fallbacks=MappingProxyType({model: tuple(chain) for model, chain in budget_fallbacks.items()}), + user_has_model_budgets=_model_budget_scope_enabled("User", self._valid_token, user_model_max_budget), ) - if member_allowed_models: - can_object_call_model( - model=model, - llm_router=llm_router, - models=member_allowed_models, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="team", - ) - if project_object is not None and len(project_object.models) > 0: - can_project_access_model( - model=model, - project_object=project_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - except ProxyException: - return False - return True + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: + token: Final = self._valid_token + reading: Final = await self._limiter.read_model_budget_raw( + "Key", token.token, token.team_id, model, token.model_max_budget + ) + return _model_budget_read( + reading, + f"LiteLLM Virtual Key: {token.token}, key_alias: {token.key_alias}, exceeded budget for model={model}", + "key", + token.token, + ) -async def _is_within_all_fallback_model_budgets( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - models: Sequence[str], -) -> bool: - for model in models: - if not await _is_within_every_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: + return await load_fallback_authorization_facts( model=model, - ): - return False - return True + valid_token=self._valid_token, + llm_router=self._llm_router, + ) + + async def rewrite_request_model(self, model: str) -> None: + _reroute_request(self._request_data, self._request, model) + + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: + user_id: Final = self._valid_token.user_id + user_model_max_budget: Final = self._valid_token.user_model_max_budget + if user_id is None or not isinstance(user_model_max_budget, Mapping): + return ModelBudgetRead(ModelBudgetSnapshot()) + reading: Final = await self._limiter.read_model_budget_raw( + "User", user_id, self._valid_token.team_id, model, user_model_max_budget + ) + return _model_budget_read( + reading, f"LiteLLM User: {user_id}, exceeded budget for model={model}", "user", user_id + ) + + +async def _enforce_model_budgets(step: Step, effects: _ModelBudgetEffects) -> None: + match await drive_model_budgets(step, effects): + case Deny(denial=denial): + raise denial_exception( + denial, + token=effects.valid_token.token, + route=effects.route, + jwt_auth_enabled=False, + key=effects.valid_token, + ) + case Allow(): + return + + +def _reroute_request( + request_data: dict[str, object], # mutable-ok: the fallback model replaces the requested one in place + request: Request, + model: str, +) -> None: + rewrite_request_model(request_data=request_data, request=request, model=model) + if "path_params" in request.scope: + request.scope["path_params"]["model"] = model # rebind-ok: route handlers use the fallback model + + +def _budget_fallback_candidates(targets: Sequence[object]) -> tuple[str, ...]: + return tuple(model for target in targets if (model := _budget_fallback_target_model(target)) is not None) + + +class _TeamMemberModelBudgetEffects: + def __init__( + self, + *, + valid_token: UserAPIKeyAuth, + limiter: _EveryModelBudgetLimiter, + models: Sequence[str], + request_data: dict[str, object], # mutable-ok: fallback rewrites the request body + request: Request, + llm_model_list: Sequence[object] | None, + llm_router: litellm.Router | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + team_object: LiteLLM_TeamTable | None, + end_user_object: LiteLLM_EndUserTable | None, + route: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + ) -> None: + self._valid_token: Final = valid_token + self._limiter: Final = limiter + self._models: Final = tuple(models) + self._request_data: Final = request_data + self._request: Final = request + self._llm_model_list: Final = llm_model_list + self._llm_router: Final = llm_router + self._project_object: Final = project_object + self._team_object: Final = team_object + self._end_user_object: Final = end_user_object + self._route: Final = route + self._prisma_client: Final = prisma_client + self._user_api_key_cache: Final = user_api_key_cache + self._proxy_logging_obj: Final = proxy_logging_obj + + def _resolved_dependencies(self) -> tuple[PrismaClient | None, UserApiKeyCache, ProxyLogging]: + from litellm.proxy import proxy_server + + return ( + self._prisma_client if self._prisma_client is not None else proxy_server.prisma_client, + self._user_api_key_cache if self._user_api_key_cache is not None else proxy_server.user_api_key_cache, + self._proxy_logging_obj if self._proxy_logging_obj is not None else proxy_server.proxy_logging_obj, + ) + + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: + budget: Final = self._valid_token.team_member_model_max_budget + requested: Final = self._request_data.get("model") + return TeamMemberModelBudgetScope( + models=self._models, + rewritable_model=requested if isinstance(requested, str) else None, + has_model_budgets=_model_budget_scope_enabled("Member", self._valid_token, budget), + can_fall_back=_model_budget_fallback_enabled(self._llm_router is not None, self._request_data), + ) + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetRead: + user_id: Final = self._valid_token.user_id + team_id: Final = self._valid_token.team_id + budget: Final = self._valid_token.team_member_model_max_budget + if user_id is None or team_id is None or not isinstance(budget, Mapping): + return ModelBudgetRead(ModelBudgetSnapshot()) + from litellm.proxy.hooks.model_max_budget_limiter import team_member_budget_entity_id + + identity: Final = team_member_budget_entity_id(user_id, team_id) + reading: Final = await self._limiter.read_model_budget_raw("TeamMember", identity, team_id, model, budget) + return _model_budget_read( + reading, + f"LiteLLM Team Member: user={user_id}, team={team_id}, exceeded budget for model={model}", + "team_member", + identity, + ) + + async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: + llm_router: Final = self._llm_router + initial: Final = common_flow.MemberFallback(llm_router is not None) + content: Final = ( + replace( + initial, + multipart=self._request.headers.get("content-type", "").lower().startswith("multipart/form-data"), + ) + if common_flow.read(initial) == "ReadContentType" + else initial + ) + resource: Final = ( + replace(content, resource_bound=references_provider_scoped_resource(self._request_data)) + if common_flow.read(content) == "ReadResourceBinding" + else content + ) + if common_flow.read(resource) != "ReadFallbackSettings": + return () + assert llm_router is not None + from litellm.proxy.proxy_server import proxy_config + + prisma_client, _, proxy_logging_obj = self._resolved_dependencies() + try: + router_settings: Final[dict[str, object] | None] = ( + cast( # cast-ok: proxy settings return an unparameterized dict + dict[str, object] | None, + await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=self._valid_token, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ), + ) + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback. + return () + + configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None + router_fallbacks: Final[Sequence[object] | None] = ( + cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below + if isinstance(configured_fallbacks, list) + else _router_fallbacks(llm_router) + ) + if router_fallbacks is None: + return () + fallback_models, _ = request_fallback_chain(list(router_fallbacks), self._request_data, model) + model_group_aliases: Final[object] = ( + router_settings.get("model_group_alias") if router_settings is not None else None + ) + key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime + object, self._valid_token.aliases + ) + resolved: Final[tuple[tuple[str, str], ...]] = tuple( + ( + candidate, + target + if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None + else candidate, + ) + for candidate in _budget_fallback_candidates(fallback_models or ()) + ) + return tuple( + FallbackTarget( + requested=candidate, + model=fallback_model, + served_model=resolve_served_model( + fallback_model, litellm.model_alias_map, key_aliases, model_group_aliases + ), + ) + for candidate, fallback_model in resolved + ) + + async def vet_team_member_fallback(self, target: FallbackTarget) -> FallbackCandidateFacts: + from litellm.proxy.auth.fallback_candidate_adapter import read_fallback_candidate_facts + + return await read_fallback_candidate_facts( + target=target, + valid_token=self._valid_token, + limiter=self._limiter, + llm_router=self._llm_router, + route=self._route, + project_object=self._project_object, + team_object=self._team_object, + end_user_object=self._end_user_object, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def reroute_team_member_request(self, model: str) -> None: + _reroute_request(self._request_data, self._request, model) async def _check_team_member_model_budget_with_fallback( @@ -796,352 +867,29 @@ async def _check_team_member_model_budget_with_fallback( user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> None: - team_member_model_max_budget: Final = valid_token.team_member_model_max_budget - if ( - valid_token.user_id is None - or valid_token.team_id is None - or not isinstance(team_member_model_max_budget, Mapping) - or not team_member_model_max_budget - ): - return - - user_id: Final = valid_token.user_id - team_id: Final = valid_token.team_id - for model_name in models: - await _check_team_member_model_budget_for_model( - user_id=user_id, - team_id=team_id, - team_member_model_max_budget=team_member_model_max_budget, - model_name=model_name, - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - request_data=request_data, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - project_object=project_object, - team_object=team_object, - end_user_object=end_user_object, - route=route, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - -async def _check_team_member_model_budget_for_model( - user_id: str, - team_id: str, - team_member_model_max_budget: Mapping[str, object], - model_name: str, - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router | None, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - team_object: LiteLLM_TeamTable | None = None, - end_user_object: LiteLLM_EndUserTable | None = None, - route: str = "", - prisma_client: PrismaClient | None = None, - user_api_key_cache: UserApiKeyCache | None = None, - proxy_logging_obj: ProxyLogging | None = None, -) -> None: - try: - await model_max_budget_limiter.is_team_member_within_model_budget( - user_id=user_id, - team_id=team_id, - team_member_model_max_budget=team_member_model_max_budget, - model=model_name, - ) - except litellm.BudgetExceededError: - if request_data.get("model") != model_name or llm_router is None or fallbacks_disabled(request_data): - raise - fallback_succeeded: Final = await _try_team_member_model_budget_fallback( - user_id=user_id, - team_id=team_id, - model_name=model_name, - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - request_data=request_data, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - project_object=project_object, - team_object=team_object, - end_user_object=end_user_object, - route=route, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if not fallback_succeeded: - raise - - -async def _try_team_member_model_budget_fallback( - user_id: str, - team_id: str, - model_name: str, - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - team_object: LiteLLM_TeamTable | None = None, - end_user_object: LiteLLM_EndUserTable | None = None, - route: str = "", - prisma_client: PrismaClient | None = None, - user_api_key_cache: UserApiKeyCache | None = None, - proxy_logging_obj: ProxyLogging | None = None, -) -> bool: - if request.headers.get("content-type", "").lower().startswith( - "multipart/form-data" - ) or references_provider_scoped_resource(request_data): - return False - - from litellm.proxy import proxy_server - from litellm.proxy.proxy_server import proxy_config - - resolved_prisma_client: Final = prisma_client if prisma_client is not None else proxy_server.prisma_client - resolved_user_api_key_cache: Final = ( - user_api_key_cache if user_api_key_cache is not None else proxy_server.user_api_key_cache + effects: Final = _TeamMemberModelBudgetEffects( + valid_token=valid_token, + limiter=model_max_budget_limiter, + models=models, + request_data=request_data, + request=request, + llm_model_list=llm_model_list, + llm_router=llm_router, + project_object=project_object, + team_object=team_object, + end_user_object=end_user_object, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) - resolved_proxy_logging_obj: Final = ( - proxy_logging_obj if proxy_logging_obj is not None else proxy_server.proxy_logging_obj - ) - - try: - router_settings: Final[dict[str, object] | None] = ( - cast( # cast-ok: proxy settings return an unparameterized dict - dict[str, object] | None, - await proxy_config.get_hierarchical_router_settings( - user_api_key_dict=valid_token, - prisma_client=resolved_prisma_client, - proxy_logging_obj=resolved_proxy_logging_obj, - ), + match await drive_team_member_model_budgets(_AUTH_CORE.check_team_member_model_budgets(), effects): + case Deny(denial=denial): + raise denial_exception( + denial, token=valid_token.token, route=route, jwt_auth_enabled=False, key=valid_token ) - ) - except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback. - return False - - configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None - router_fallbacks: Final[Sequence[object] | None] = ( - cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below - if isinstance(configured_fallbacks, list) - else _router_fallbacks(llm_router) - ) - if router_fallbacks is None: - return False - - fallback_models, _ = request_fallback_chain(list(router_fallbacks), request_data, model_name) - fallback_targets: Final[Sequence[object]] = fallback_models or () - fallback_candidates: Final[tuple[str, ...]] = tuple( - model for target in fallback_targets if (model := _budget_fallback_target_model(target)) is not None - ) - if not fallback_candidates: - return False - - model_group_aliases: Final[object] = ( - router_settings.get("model_group_alias") if router_settings is not None else None - ) - key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime - object, valid_token.aliases - ) - resolved_fallback_candidates: Final[tuple[tuple[str, str], ...]] = tuple( - ( - candidate, - target if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None else candidate, - ) - for candidate in fallback_candidates - ) - fallback_candidates_with_served_models: Final = ( - ( - fallback_candidate, - fallback_model, - resolve_served_model( - fallback_model, - litellm.model_alias_map, - key_aliases, - model_group_aliases, - ), - ) - for fallback_candidate, fallback_model in resolved_fallback_candidates - ) - - try: - team_membership: Final = await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - proxy_logging_obj=resolved_proxy_logging_obj, - ) - except Exception: # noqa: BLE001 # membership lookup failures must not authorize a fallback - return False - - for fallback_candidate, fallback_model, served_model in fallback_candidates_with_served_models: - if fallback_model == model_name: - continue - try: - if ( - litellm.block_requests_for_models_without_pricing - and RouteChecks.is_llm_api_route(route=route) - and _unpriced_models_in_request(model=served_model, llm_router=llm_router) - ): - continue - if end_user_object is not None and end_user_object.models: - can_customer_access_model( - model=fallback_model, - end_user_object=end_user_object, - llm_router=llm_router, - valid_token=valid_token, - ) - if not is_model_cost_zero( - model=served_model, llm_router=llm_router - ) and not await _is_within_all_fallback_model_budgets( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=tuple( - dict.fromkeys( - ( - fallback_candidate, - fallback_model, - served_model, - llm_router.registry().resolve(served_model), - ) - ) - ), - ): - continue - if not await _is_fallback_model_authorized( - model=fallback_model, - valid_token=valid_token, - llm_model_list=llm_model_list, - llm_router=llm_router, - prisma_client=prisma_client, - team_object=team_object, - team_membership=team_membership, - project_object=project_object, - refuse_agent_keys=True, - ): - continue - valid_token.matched_model_access_groups = [] # rebind-ok: clear prior groups before evaluating this candidate - matched_groups = await stamp_matched_model_access_groups( - model=fallback_model, - valid_token=valid_token, - team_object=team_object, - project_object=project_object, - llm_router=llm_router, - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - proxy_logging_obj=resolved_proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=True, - ) - valid_token.matched_model_access_groups = list(matched_groups) # rebind-ok: spend uses the candidate groups - if matched_groups: - await _model_access_group_max_budget_check( - matched_model_access_groups=list(matched_groups), - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - ) - except Exception: # noqa: BLE001 # A candidate check failure must not escape the fail-closed walk. - verbose_proxy_logger.warning( - "Skipping fallback candidate %s after a budget or authorization check error", - fallback_model.replace("\r", "").replace("\n", ""), - exc_info=True, - ) - continue - rewrite_request_model(request_data=request_data, request=request, model=fallback_model) - if "path_params" in request.scope: - request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model - return True - return False - - -async def _check_key_model_budget_with_fallback( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _KeyModelBudgetLimiter, - model_name: str, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None = None, - llm_router: litellm.Router | None = None, -) -> None: - """ - Enforce the key's per-model budget for `model_name`. If exceeded and the - key has a `budget_fallbacks` chain configured for `model_name`, reroute - the request to the first fallback model still within its own budget - instead of rejecting the request. - - The selected fallback is validated against the key's model-access - allowlist and the team's model restrictions so that budget_fallbacks - cannot bypass model authorization. The rewrite is persisted to the - parsed-body cache, Starlette's JSON cache (``request._json``), and - path parameters so that downstream handlers see the final model - regardless of whether they consume ``_read_request_body()``, - ``request.json()``, or the path ``model`` parameter. - - Fallback is only attempted when ``model_name`` matches the top-level - ``request_data["model"]``; models extracted from nested fields - (``session.model``, ``completion.model``, etc.) are not rewritable - and raise immediately. - - Raises: - BudgetExceededError: if `model_name` is over budget and no configured - fallback is within budget either (or the fallback is not authorized). - """ - try: - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=model_name, - ) - except litellm.BudgetExceededError as e: - if request_data.get("model") != model_name: - raise e - fallback_model: Final = await model_max_budget_limiter.get_fallback_model_within_budget( - user_api_key_dict=valid_token, - model=model_name, - ) - if fallback_model is None: - raise e - if not await _is_fallback_model_authorized( - model=fallback_model, - valid_token=valid_token, - llm_model_list=llm_model_list, - llm_router=llm_router, - ): - raise - rewrite_request_model(request_data=request_data, request=request, model=fallback_model) - if "path_params" in request.scope: - request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model - - -def _get_bearer_token_or_received_api_key(api_key: str) -> str: - if api_key.startswith("Bearer "): # ensure Bearer token passed in - api_key = api_key.replace("Bearer ", "") # extract the token - elif api_key.startswith("Basic "): - api_key = api_key.replace("Basic ", "") # handle langfuse input - elif api_key.startswith("bearer "): - api_key = api_key.replace("bearer ", "") - elif api_key.startswith("AWS4-HMAC-SHA256"): - # Handle AWS Signature V4 format from LangChain - # Format: AWS4-HMAC-SHA256 Credential=Bearer $LITELLM_MASTER_KEY/date/region/service/aws4_request, SignedHeaders=..., Signature=... - # Extract the Bearer token from the Credential field - match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r"Credential=([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - - return api_key + case Allow(): + return def _routing_selector_matches_claim( @@ -1236,29 +984,7 @@ def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> def get_bearer_token( api_key: str, ) -> str: - if api_key.startswith("Bearer "): # ensure Bearer token passed in - api_key = api_key.replace("Bearer ", "") # extract the token - elif api_key.startswith("Basic "): - api_key = api_key.replace("Basic ", "") # handle langfuse input - elif api_key.startswith("bearer "): - api_key = api_key.replace("bearer ", "") - elif api_key.startswith("AWS4-HMAC-SHA256"): - # Handle AWS Signature V4 format from LangChain - # Format: AWS4-HMAC-SHA256 Credential=Bearer $LITELLM_MASTER_KEY/date/region/service/aws4_request, SignedHeaders=..., Signature=... - # Extract the Bearer token from the Credential field - match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r"Credential=([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - api_key = "" - else: - api_key = "" - return api_key + return scheme_credential(api_key) or "" _get_bearer_token: Final = get_bearer_token @@ -1269,30 +995,17 @@ def _apply_budget_limits_to_end_user_params( budget_info: LiteLLM_BudgetTable, end_user_id: str | None, ) -> None: - """ - Helper function to apply budget limits to end user parameters. + from litellm.proxy.auth import end_user_context_adapter - Args: - end_user_params: Dictionary to update with budget parameters - budget_info: Budget table object containing limits - end_user_id: ID of the end user for logging - """ - if budget_info.tpm_limit is not None: - end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit + return end_user_context_adapter.apply_budget_limits_to_end_user_params(end_user_params, budget_info, end_user_id) - if budget_info.rpm_limit is not None: - end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit - if budget_info.tpd_limit is not None: - end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit - - if budget_info.max_budget is not None: - end_user_params["end_user_max_budget"] = budget_info.max_budget - - if budget_info.model_max_budget is not None: - end_user_params["end_user_model_max_budget"] = budget_info.model_max_budget - - verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) +_WEBSOCKET_DENIAL_DETAIL: Final = MappingProxyType( + { + WebSocketDenial.NO_API_KEY: "No API key provided", + WebSocketDenial.INVALID_AUTHORIZATION: "Invalid Authorization header format", + } +) async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: @@ -1355,27 +1068,17 @@ async def user_api_key_auth_websocket_for_model( request.body = return_body - authorization: Final = websocket.headers.get("authorization") - # If no Authorization header, try the api-key header - if not authorization: - api_key = websocket.headers.get("api-key") - if not api_key: - # Try extracting from WebSocket subprotocol (browser clients) - for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): - protocol = protocol.strip() - if protocol.startswith("openai-insecure-api-key."): - api_key = protocol[len("openai-insecure-api-key.") :] - break - if not api_key: - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="No API key provided") - else: - # Extract the API key from the Bearer token - if not authorization.startswith("Bearer "): - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="Invalid Authorization header format") - - api_key = authorization[len("Bearer ") :].strip() + selected: Final = select_websocket_api_key( + WebSocketKeys( + authorization=websocket.headers.get("authorization"), + api_key=websocket.headers.get("api-key"), + subprotocols=websocket.headers.get("sec-websocket-protocol"), + ) + ) + if isinstance(selected, WebSocketDenial): + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + raise HTTPException(status_code=403, detail=_WEBSOCKET_DENIAL_DETAIL[selected]) + api_key: Final = selected # Call user_api_key_auth with the extracted API key # Note: You'll need to modify this to work with WebSocket context if needed @@ -1385,6 +1088,8 @@ async def user_api_key_auth_websocket_for_model( raise HTTPException( 403, "Required baggage tag mappings need HTTP inference, WebSocket inference is disabled" ) + except NativeAuthFailure: + raise except Exception as e: if not close_on_auth_error: raise @@ -1408,22 +1113,9 @@ async def user_api_key_auth_accepted_websocket_for_model(websocket: WebSocket, m def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: - valid_token.end_user_id = end_user_params.get("end_user_id") - # Only overwrite token fields when the DB-derived value is not None. - # This prevents DB lookups (where the budget table has no value set) - # from silently clearing values that a custom auth function may have - # already set on the token. - if end_user_params.get("end_user_tpm_limit") is not None: - valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] - if end_user_params.get("end_user_rpm_limit") is not None: - valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] - if end_user_params.get("end_user_tpd_limit") is not None: - valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] - if end_user_params.get("allowed_model_region") is not None: - valid_token.allowed_model_region = end_user_params["allowed_model_region"] - if end_user_params.get("end_user_model_max_budget") is not None: - valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"] - return valid_token + from litellm.proxy.auth import end_user_context_adapter + + return end_user_context_adapter.update_valid_token_with_end_user_params(valid_token, end_user_params) # Reusable coordinator for global spend to prevent cache stampede @@ -1444,6 +1136,11 @@ async def fetch_global_spend_with_event_coordination( """ async def _load_global_spend() -> float | None: + from litellm.repositories.native_auth_repository import native_auth_repository + + native: Final = native_auth_repository(prisma_client) + if native is not None: + return await bounded_db_lookup(native.spend("user", LITELLM_PROXY_BUDGET_NAME), name="proxy_budget") proxy_budget_row: Final = await bounded_db_lookup( UserRepository(prisma_client).table.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), name="proxy_budget", @@ -1503,6 +1200,43 @@ def get_rbac_role(jwt_handler: JWTHandler, scopes: list[str]) -> str: return LitellmUserRoles.TEAM +def _query_api_key(route: str, request: Request) -> str | None: + from litellm.proxy.auth.route_checks import RouteChecks + from litellm.proxy.common_utils.http_parsing_utils import ( + safe_get_request_query_params, + ) + + if not RouteChecks.is_generate_content_route(route=route): + return None + return safe_get_request_query_params(request).get("key") + + +def _pass_through_header_api_key(endpoint: Mapping[str, object], route: str, request: Request) -> str | None: + if endpoint.get("path", "") != route: + return None + headers: Final = _PASS_THROUGH_HEADERS_ADAPTER.validate_python(endpoint.get("headers")) + header_name: Final = headers.get("litellm_user_api_key", "") if headers is not None else None + return request.headers.get(header_name) if isinstance(header_name, str) else None + + +def _pass_through_api_key( + pass_through_endpoints: Sequence[Mapping[str, object]] | None, route: str, request: Request +) -> str | None: + header_values: Final = tuple( + value + for value in ( + _pass_through_header_api_key(endpoint, route, request) for endpoint in pass_through_endpoints or () + ) + if value is not None + ) + return header_values[-1] if header_values else None + + +def _header_or_none(value: object) -> str | None: + """Direct callers of ``user_api_key_auth`` leave unpassed headers as their ``fastapi.Security`` defaults.""" + return value if isinstance(value, str) else None + + def get_api_key( custom_litellm_key_header: str | None, api_key: str, @@ -1510,57 +1244,29 @@ def get_api_key( anthropic_api_key_header: str | None, google_ai_studio_api_key_header: str | None, azure_apim_header: str | None, - pass_through_endpoints: list[dict] | None, + pass_through_endpoints: Sequence[Mapping[str, object]] | None, route: str, request: Request, + *, + custom_header: str | None = None, ) -> tuple[str, str | None]: - """ - Returns: - Tuple[Optional[str], Optional[str]]: Tuple of the api_key and the passed_in_key - """ - from litellm.proxy.auth.route_checks import RouteChecks - from litellm.proxy.common_utils.http_parsing_utils import ( - safe_get_request_query_params, + selected: Final = select_api_key( + PresentedKeys( + custom_header=custom_header, + litellm_header=_header_or_none(custom_litellm_key_header), + authorization=_header_or_none(api_key), + api_key=_header_or_none(azure_api_key_header), + x_api_key=_header_or_none(anthropic_api_key_header), + google=_header_or_none(google_ai_studio_api_key_header), + apim=_header_or_none(azure_apim_header), + query=_query_api_key(route, request), + pass_through=_pass_through_api_key(pass_through_endpoints, route, request), + ) ) - - api_key = api_key - passed_in_key: str | None = None - if isinstance(custom_litellm_key_header, str): - passed_in_key = custom_litellm_key_header - api_key = _get_bearer_token_or_received_api_key(custom_litellm_key_header) - elif isinstance(api_key, str) and len(api_key) > 0: - passed_in_key = api_key - api_key = get_bearer_token(api_key=api_key) - elif isinstance(azure_api_key_header, str): - passed_in_key = azure_api_key_header - api_key = azure_api_key_header - elif isinstance(anthropic_api_key_header, str): - passed_in_key = anthropic_api_key_header - api_key = anthropic_api_key_header - elif isinstance(google_ai_studio_api_key_header, str): - passed_in_key = google_ai_studio_api_key_header - api_key = google_ai_studio_api_key_header - elif isinstance(azure_apim_header, str): - passed_in_key = azure_apim_header - api_key = azure_apim_header - elif ( - RouteChecks.is_generate_content_route(route=route) - and request is not None - and safe_get_request_query_params(request).get("key") - ): - google_auth_key: Final[str] = safe_get_request_query_params(request).get("key") or "" - passed_in_key = google_auth_key - api_key = google_auth_key - elif pass_through_endpoints is not None: - for endpoint in pass_through_endpoints: - if endpoint.get("path", "") == route: - headers: dict | None = endpoint.get("headers", None) - if headers is not None: - header_key: str = headers.get("litellm_user_api_key", "") - if request.headers.get(header_key) is not None: - api_key = request.headers.get(header_key) or "" - passed_in_key = api_key - return api_key, passed_in_key + if selected is None: + return api_key, None + verbose_proxy_logger.debug("Selected API key from %s", selected.source.value) + return selected.token, selected.presented async def check_api_key_for_custom_headers_or_pass_through_endpoints( @@ -1618,12 +1324,6 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( # the lookup and return None (caller proceeds to auth_builder). _JWT_PROXY_ADMIN_SENTINEL: Final = "__JWT_PROXY_ADMIN__" -_JWT_AUTH_DISABLED_HINT = ( - " This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a" - " virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate" - " with JWTs." -) - class _JWTKeyMappingFailure(Enum): NO_ELIGIBLE_KEY = "JWT Key Mapping: No eligible existing virtual key for this identity. Access denied." @@ -1759,6 +1459,8 @@ async def _auto_register_jwt_mapping( "updated_by": behavior.value, } ) + except NativeAuthFailure: + raise except Exception as e: error_str: Final = str(e).lower() if "unique" in error_str or "p2002" in error_str: @@ -1777,6 +1479,8 @@ async def _auto_register_jwt_mapping( try: async with db_span("delete_orphaned_jwt_key", "LiteLLM_VerificationToken"): await VerificationTokenRepository(prisma_client).table.delete(where={"token": token_hash}) + except NativeAuthFailure: + raise except Exception as delete_err: # Don't fail the request if cleanup fails — the orphan is # unmapped and inert. Log so an operator can prune it later. @@ -2080,6 +1784,8 @@ def _ensure_litellm_received_at_on_request_state(request: Request) -> datetime: received_at: Final = datetime.now(timezone.utc) try: request.state.litellm_received_at = received_at + except NativeAuthFailure: + raise except Exception: pass return received_at @@ -2162,6 +1868,8 @@ async def _record_unparsable_body_failure( error_type=ProxyErrorTypes.bad_request_error, route=route, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # any logging failure must leave the caller's 400 untouched verbose_proxy_logger.exception("Failed to log the request rejected for an unparsable body: %s", e) @@ -2236,10 +1944,543 @@ def _custom_auth_token_for_this_request(response: object) -> UserAPIKeyAuth: return UserAPIKeyAuth.model_validate(response).model_copy() +_AUTH_CORE: Final[AuthCore] = get_auth_core() + + +class _ResolvedEndUser(NamedTuple): + end_user_id: str | None + end_user_object: LiteLLM_EndUserTable | None + budget: LiteLLM_BudgetTable | None + + @property + def params(self) -> dict[str, object]: + if not self.end_user_id: + return {} + region: Final = ( + {} if self.end_user_object is None else {"allowed_model_region": self.end_user_object.allowed_model_region} + ) + params: Final[dict[str, object]] = {"end_user_id": self.end_user_id, **region} + if self.budget is not None: + _apply_budget_limits_to_end_user_params( + end_user_params=params, budget_info=self.budget, end_user_id=self.end_user_id + ) + return params + + +async def _resolve_end_user( + request: Request, + request_data: Mapping[str, object], + token_end_user_id: str | None, + api_key: str, + key_is_resolved: bool, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + route: str, +) -> _ResolvedEndUser: + raw_end_user_id: Final = token_end_user_id or get_end_user_id_from_request_body( + request_data, safe_get_request_headers(request) + ) + end_user_id: Final = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if prisma_client is not None: + await prefetch_identity_keys( + _identity_cache_keys(api_key, end_user_id=end_user_id, key_is_resolved=key_is_resolved), + user_api_key_cache=user_api_key_cache, + before_read=arm_routing_cooldowns, + ) + if not end_user_id: + return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=None) + try: + with tracer.trace("litellm.proxy.auth.get_end_user_object"): + end_user_object: Final = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if end_user_object is not None: + return _ResolvedEndUser( + end_user_id=end_user_id, end_user_object=end_user_object, budget=end_user_object.litellm_budget_table + ) + if litellm.max_end_user_budget_id is None: + return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=None) + default_budget: Final = await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) + return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=default_budget) + except litellm.BudgetExceededError: + raise + except NativeAuthFailure: + raise + except Exception as e: + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) + return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=None) + + +def _proxy_model_list() -> list[object] | None: + from litellm.proxy.proxy_server import llm_model_list + + return llm_model_list + + +_MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value) + + +def _core_settings(master_key: str | None, allow_user_auth: bool, prisma_client: PrismaClient | None) -> Settings: + return Settings( + master_key=master_key, + allow_user_auth=allow_user_auth, + has_db=prisma_client is not None, + master_key_only_routes=_MASTER_KEY_ONLY_ROUTES, + proxy_max_budget=litellm.max_budget, + ) + + +def _aware_expiry(expires: datetime | str) -> datetime: + expiry: Final = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) + return ( + expiry if expiry.tzinfo is not None and expiry.utcoffset() is not None else expiry.replace(tzinfo=timezone.utc) + ) + + +def _presented_key(key: UserAPIKeyAuth) -> PresentedKey: + if key.user_role != LitellmUserRoles.PROXY_ADMIN: + return NonAdminKey() + return AdminKey(expires=None if key.expires is None else _aware_expiry(key.expires)) + + +def _with_end_user(key: UserAPIKeyAuth, end_user: _ResolvedEndUser, parent_otel_span: Span | None) -> UserAPIKeyAuth: + with_end_user: Final = update_valid_token_with_end_user_params(valid_token=key, end_user_params=end_user.params) + with_end_user.parent_otel_span = parent_otel_span + if end_user.end_user_object is not None: + with_end_user.end_user_object_permission = end_user.end_user_object.object_permission + return with_end_user + + +def _reset_end_user_fields(key: UserAPIKeyAuth, end_user: _ResolvedEndUser) -> UserAPIKeyAuth: + key.end_user_id = end_user.end_user_id or None + key.end_user_tpm_limit = end_user.budget.tpm_limit if end_user.budget is not None else None + key.end_user_rpm_limit = end_user.budget.rpm_limit if end_user.budget is not None else None + key.end_user_tpd_limit = end_user.budget.tpd_limit if end_user.budget is not None else None + key.allowed_model_region = ( + end_user.end_user_object.allowed_model_region if end_user.end_user_object is not None else None + ) + return key + + +async def _refresh_team_fields_from_cache( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if valid_token.team_id is None or valid_token.team_id == UI_TEAM_ID: + return + ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token + try: + team_obj: Final[LiteLLM_TeamTableCachedObj] = await get_team_object( + team_id=valid_token.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=True, + ) + + if ( + team_obj.last_refreshed_at is not None + and valid_token.last_refreshed_at is not None + and team_obj.last_refreshed_at > valid_token.last_refreshed_at + ): + team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) + + for k, v in team_obj_dict.items(): + field_name = f"team_{k}" + if field_name in valid_token.__fields__: + setattr(valid_token, field_name, v) + except NativeAuthFailure: + raise + except Exception as e: + verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. + + +class _BuilderEffects: + def __init__( + self, + *, + request: Request, + request_data: dict[str, object], # mutable-ok: model access checks rewrite the request model in place + token: str, + preresolved_key: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + route: str, + ) -> None: + self._request: Final = request + self._request_data: Final = request_data + self._token: Final = token + self._preresolved_key: Final = preresolved_key + self._prisma_client: Final = prisma_client + self._user_api_key_cache: Final = user_api_key_cache + self._parent_otel_span: Final = parent_otel_span + self._proxy_logging_obj: Final = proxy_logging_obj + self._route: Final = route + self.end_user: _ResolvedEndUser = _ResolvedEndUser(end_user_id=None, end_user_object=None, budget=None) + self.resolved_key: UserAPIKeyAuth | None = preresolved_key + self.owner: LiteLLM_UserTable | None = None + + @property + def token(self) -> str: + return self._token + + @property + def parent_otel_span(self) -> Span | None: + return self._parent_otel_span + + @property + def user_api_key_cache(self) -> UserApiKeyCache: + return self._user_api_key_cache + + @property + def proxy_logging_obj(self) -> ProxyLogging: + return self._proxy_logging_obj + + @property + def key(self) -> UserAPIKeyAuth: + if self.resolved_key is None: + raise RuntimeError("the auth core granted a key it never resolved") + return self.resolved_key + + def _identity_store(self, *, check_cache_only: bool) -> IdentityStore: + return IdentityStore( + self._prisma_client, + self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + check_cache_only=check_cache_only, + ) + + async def resolve_end_user(self, prefetch_key: bool) -> None: + self.end_user = await _resolve_end_user( + request=self._request, + request_data=self._request_data, + token_end_user_id=self._preresolved_key.end_user_id if self._preresolved_key is not None else None, + api_key=self._token, + key_is_resolved=not prefetch_key or self._preresolved_key is not None, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + route=self._route, + ) + + async def _cached_key(self, hashed_token: str) -> UserAPIKeyAuth | None: + try: + with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): + cached: Final = IdentityStore.key_from_principal( + await self._identity_store(check_cache_only=True).resolve(hashed_token=hashed_token) + ) + except NativeAuthFailure: + raise + except Exception: + verbose_logger.debug("api key not found in cache.") + return None + # Key-cache entries are written only after the proxy validated a virtual key or the master key, but + # via_virtual_key is exclude=True so serialization drops it. Restore it at this trusted boundary. + cached.via_virtual_key = True + return cached + + def _ui_login_key(self) -> UserAPIKeyAuth | None: + # Decryption fails closed for anything that is not a genuine UI login blob. + if self._token.startswith("sk-") or get_secret_bool("EXPERIMENTAL_UI_LOGIN") is False: + return None + return ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(self._token) + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: + found: Final = self._preresolved_key or await self._cached_key(hashed_token) or self._ui_login_key() + if found is None: + return None + self.resolved_key = ( + await _refresh_session_token_grants( + valid_token=found, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + if found.is_session_token and self._prisma_client is not None + else found + ) + return _presented_key(self.resolved_key) + + async def evict_presented_key(self, hashed_token: str) -> None: + await delete_cache_key_object( + hashed_token=hashed_token, + user_api_key_cache=self._user_api_key_cache, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def load_key(self, hashed_token: str) -> None: + with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): + self.resolved_key = IdentityStore.key_from_principal( + await self._identity_store(check_cache_only=False).resolve(hashed_token=hashed_token) + ) + + async def prepare_key(self, loaded_from_db: bool) -> None: + if loaded_from_db: + _reset_end_user_fields(self.key, self.end_user) + return + await _refresh_team_fields_from_cache( + valid_token=self.key, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def inspect_key(self) -> KeyFacts: + key: Final = _update_key_budget_with_temp_budget_increase(self.key) + self.resolved_key = key + return KeyFacts( + has_token=key.token is not None, + blocked=key.blocked is True, + expires=None if key.expires is None else _aware_expiry(key.expires), + spend=key.spend, + soft_budget=key.soft_budget, + max_budget=key.max_budget, + team_member_spend=key.team_member_spend, + user_id=key.user_id, + team_id=key.team_id, + budget_throttle_pct=budget_throttle_percentage() if should_throttle_budget_exceeded(key) else None, + ) + + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: + from litellm.proxy.proxy_server import llm_router + + return await read_key_and_fallback_model_access_facts( + valid_token=self.key, + request_data=self._request_data, + route=self._route, + request=self._request, + llm_model_list=_proxy_model_list(), + llm_router=llm_router, + ) + + async def prefetch_auth_objects(self) -> None: + await _prefetch_referenced_auth_objects( + self.key, + end_user_id=self.end_user.end_user_id, + user_api_key_cache=self._user_api_key_cache, + prisma_client=self._prisma_client, + ) + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: + try: + with tracer.trace("litellm.proxy.auth.get_user_object"): + owner: Final = await get_user_object( + user_id=user_id, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + user_id_upsert=False, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + except NativeAuthFailure: + raise + except Exception as e: + verbose_logger.debug("Unable to get user %s from db/cache, continuing without it: %s", user_id, e) + return None + self.owner = owner + if owner is None: + return None + key: Final = self.key + # The joint verification-token view carries the key's columns only, so the user's own per-model budget + # reaches enforcement and the post-call increment through the row fetched here. + key.user_model_max_budget = owner.model_max_budget + return OwnerFacts(deactivated=isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False) + + async def check_budget_exemption(self) -> ExemptionFacts: + from litellm.proxy.auth.budget_pricing_adapter import read_budget_exemption_facts + from litellm.proxy.proxy_server import llm_router + + model: Final = _get_model_from_request_context( + request_data=self._request_data, + route=self._route, + request=self._request, + llm_router=llm_router, + team_id=self.key.team_id, + ) + return read_budget_exemption_facts( + route=self._route, model=model, llm_router=llm_router, request_body=self._request_data + ) + + async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: + from litellm._auth import reservation + from litellm.proxy.auth.common_budget_adapter import resolve_effective_team_member_budget + from litellm.proxy.proxy_server import get_current_spend + + if self._prisma_client is None: + return None + cache_key: Final = team_membership_auth_cache_key(team_id=team_id, user_id=user_id) + cached: Final = await self._user_api_key_cache.async_get_cache(key=cache_key, model_type=LiteLLM_TeamMembership) + membership: Final = cached if cached is not None else await self._load_team_membership(user_id, team_id) + if membership is None or membership.litellm_budget_table is None: + return None + max_budget: Final = resolve_effective_team_member_budget( + membership.litellm_budget_table, None, datetime.now(timezone.utc) + ) + if not reservation.counter("Member", max_budget): + return None + assert max_budget is not None + spend: Final = await get_current_spend( + counter_key=team_member_counter(user_id, team_id), + fallback_spend=self.key.team_member_spend or 0.0, + max_budget=max_budget, + ) + return TeamMemberBudget(spend=spend, max_budget=max_budget) + + async def _load_team_membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None: + from litellm.repositories.native_auth_repository import native_auth_repository + + if self._prisma_client is None: + return None + native: Final = native_auth_repository(self._prisma_client) + fill_generation: Final = self._user_api_key_cache.current_generation() if native is not None else None + row: Final = ( + await native.one("member", [user_id, team_id], LiteLLM_TeamMembership) + if native is not None + else await TeamMembershipRepository(self._prisma_client).table.find_first( + where={"user_id": user_id, "team_id": team_id}, + include={"litellm_budget_table": True}, + ) + ) + if row is None: + return None + membership: Final = LiteLLM_TeamMembership.model_validate(row.model_dump()) + await self._user_api_key_cache.async_set_cache( + fill_generation=fill_generation, + key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + value=membership, + model_type=LiteLLM_TeamMembership, + ttl=5, + ) + return membership + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: + # common_checks sends this alert on requests that get past here, so only the rejected request sends it here. + team_member_max_budget_alert_check( + team_id=team_id, + team_alias=self.key.team_alias, + team_metadata=self.key.team_metadata, + organization_id=self.key.org_id, + user_id=user_id, + user_email=self.owner.user_email if self.owner is not None else None, + proxy_logging_obj=self._proxy_logging_obj, + spend=budget.spend, + max_budget=budget.max_budget, + ) + + async def send_key_budget_alerts(self) -> None: + await virtual_key_max_budget_alert_check( + valid_token=self.key, proxy_logging_obj=self._proxy_logging_obj, user_obj=self.owner + ) + + async def read_key_spend(self, max_budget: float) -> float: + from litellm.proxy.proxy_server import get_current_spend + + # Redis-first cross-pod counter, falling back to the key object's recorded spend. + return await get_current_spend( + counter_key=f"spend:key:{self.key.token}", + fallback_spend=self.key.spend or 0.0, + max_budget=max_budget, + ) + + async def alert_key_budget(self, spend: float) -> None: + send_token_budget_alert(self.key, spend, self._proxy_logging_obj, self.owner) + + async def alert_key_soft_budget(self) -> None: + send_soft_budget_alert(self.key, self._proxy_logging_obj, self.owner) + + def _model_budgets(self) -> _ModelBudgetEffects: + from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter + + return _ModelBudgetEffects( + request=self._request, + request_data=self._request_data, + valid_token=self.key, + route=self._route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=_proxy_model_list(), + ) + + async def read_model_budget_scope(self) -> ModelBudgetScope: + return await self._model_budgets().read_model_budget_scope() + + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: + return await self._model_budgets().check_key_model_budget(model) + + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: + return await self._model_budgets().authorize_budget_fallback(model) + + async def rewrite_request_model(self, model: str) -> None: + await self._model_budgets().rewrite_request_model(model) + + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: + return await self._model_budgets().check_user_model_budget(model) + + async def load_key_context(self) -> None: + await _load_team_and_project( + valid_token=self.key, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def read_proxy_spend(self) -> float | None: + if self._prisma_client is None: + return None + with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): + return await fetch_global_spend_with_event_coordination( + cache_key=GLOBAL_PROXY_SPEND_CACHE_KEY, + user_api_key_cache=self._user_api_key_cache, + prisma_client=self._prisma_client, + ) + + async def alert_proxy_budget(self, spend: float) -> None: + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + call_info: Final = CallInfo( + token=self.key.token, + spend=spend, + max_budget=litellm.max_budget, + user_id=litellm_proxy_admin_name, + team_id=self.key.team_id, + event_group=Litellm_EntityType.PROXY, + ) + asyncio.create_task(self._proxy_logging_obj.budget_alerts(type="proxy_budget", user_info=call_info)) + + async def user_api_key_auth_builder( request: Request, api_key: str, - azure_api_key_header: str, + azure_api_key_header: str | None, anthropic_api_key_header: str | None, google_ai_studio_api_key_header: str | None, azure_apim_header: str | None, @@ -2267,6 +2508,8 @@ async def user_api_key_auth_builder( start_time: Final = getattr(request.state, "litellm_received_at", None) or datetime.now(timezone.utc) try: request.state.litellm_received_at = start_time + except NativeAuthFailure: + raise except Exception: pass route: Final[str] = get_request_route(request=request) @@ -2281,8 +2524,13 @@ async def user_api_key_auth_builder( route=route, ) pass_through_endpoints: Final[list[dict] | None] = general_settings.get("pass_through_endpoints", None) - ## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header - api_key, passed_in_key = get_api_key( + custom_litellm_key_header_name: Final = general_settings.get("litellm_key_header_name") + api_key, _ = get_api_key( + custom_header=( + request.headers.get(custom_litellm_key_header_name) + if custom_litellm_key_header_name is not None + else None + ), custom_litellm_key_header=custom_litellm_key_header, api_key=api_key, azure_api_key_header=azure_api_key_header, @@ -2293,12 +2541,6 @@ async def user_api_key_auth_builder( route=route, request=request, ) - custom_litellm_key_header_name: Final = general_settings.get("litellm_key_header_name") - if custom_litellm_key_header_name is not None and custom_litellm_key_header_name in request.headers: - api_key = get_api_key_from_custom_header( - request=request, - custom_litellm_key_header_name=custom_litellm_key_header_name, - ) if open_telemetry_logger is not None: # Reuse the span created by user_api_key_auth (before body parse) @@ -2458,6 +2700,8 @@ async def user_api_key_auth_builder( proxy_logging_obj=proxy_logging_obj, user_email=mapped_user_email, ) + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.debug("JWT mapped-key user_email backfill skipped: %s", e) else: @@ -2613,17 +2857,16 @@ async def user_api_key_auth_builder( # zero-cost model would be refused here and allowed there, # while the log above claims all budget checks were skipped. if not skip_budget_checks: - await _check_user_model_budget( - valid_token=cast(UserAPIKeyAuth, valid_token), - model_max_budget_limiter=model_max_budget_limiter, - models=_get_model_names_for_budget_checks( - model=_get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) + await _enforce_model_budgets( + _AUTH_CORE.check_user_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=cast(UserAPIKeyAuth, valid_token), + route=route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=_proxy_model_list(), ), ) @@ -2642,344 +2885,49 @@ async def user_api_key_auth_builder( api_key = response elif isinstance(response, UserAPIKeyAuth): return response - if master_key is None: - if isinstance(api_key, str): - return UserAPIKeyAuth( - api_key=api_key, - user_role=LitellmUserRoles.INTERNAL_USER, - parent_otel_span=parent_otel_span, - ) - else: - return UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - parent_otel_span=parent_otel_span, - ) - elif api_key is None: # only require api key if master key is set - raise ProxyException( - message="Authentication Error, No api key passed in.", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, + if master_key is not None and not isinstance(master_key, str): + raise HTTPException( + status_code=500, + detail={f"Master key must be a valid string. Current type={type(master_key)}"}, ) - elif api_key == "": - # missing 'Bearer ' prefix - raise ProxyException( - message="Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix.", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - if route == "/user/auth": - if general_settings.get("allow_user_auth", False) is True: - return UserAPIKeyAuth() - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="'allow_user_auth' not set or set to False", - ) - - ## Check END-USER OBJECT - _end_user_object = None - end_user_params: Final = {} - - raw_end_user_id: Final = ( - valid_token.end_user_id if valid_token is not None else None - ) or get_end_user_id_from_request_body(request_data, safe_get_request_headers(request)) - end_user_id = await resolve_and_validate_end_user_id( - raw_end_user_id=raw_end_user_id, + effects: Final = _BuilderEffects( + request=request, + request_data=request_data, + token=api_key, + preresolved_key=valid_token, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, ) - if prisma_client is not None: - await prefetch_identity_keys( - _identity_cache_keys(api_key, end_user_id=end_user_id, key_is_resolved=valid_token is not None), - user_api_key_cache=user_api_key_cache, - before_read=arm_routing_cooldowns, - ) - if end_user_id: - try: - end_user_params["end_user_id"] = end_user_id - - with tracer.trace("litellm.proxy.auth.get_end_user_object"): - _end_user_object = await get_end_user_object( - end_user_id=end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ) - if _end_user_object is not None: - end_user_params["allowed_model_region"] = _end_user_object.allowed_model_region - if _end_user_object.litellm_budget_table is not None: - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=_end_user_object.litellm_budget_table, - end_user_id=end_user_id, - ) - elif litellm.max_end_user_budget_id is not None: - # End user doesn't exist yet, but apply default budget limits if configured - from litellm.proxy.auth.auth_checks import ( - get_default_end_user_budget, - ) - - default_budget: Final = await get_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - ) - if default_budget is not None: - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=default_budget, - end_user_id=end_user_id, - ) - except Exception as e: - if isinstance(e, litellm.BudgetExceededError): - raise e - verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) - - ### CHECK IF ADMIN ### - # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead - ### CHECK IF ADMIN ### - # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead - if valid_token is None: - ## Check CACHE - try: - with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): - valid_token = IdentityStore.key_from_principal( - await IdentityStore( - prisma_client, - user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, - ).resolve(hashed_token=hash_token(api_key)) - ) - # Key-cache entries are written only after the proxy validated a - # virtual key or the master key, but via_virtual_key is exclude=True - # so serialization drops it; restore it at this trusted boundary. - # The UI-login JWT fallback below constructs its token from a - # decrypted blob, not this cache, and stays unmarked. - if isinstance(valid_token, UserAPIKeyAuth): - valid_token.via_virtual_key = True - except Exception: - verbose_logger.debug("api key not found in cache.") - valid_token = None - - ## Check UI/CLI Hash Key - # Attempt decryption for non-sk- tokens unless the operator has - # explicitly set EXPERIMENTAL_UI_LOGIN=false to disable it. - # Unset (None) keeps the new default of always attempting decryption; - # decryption fails closed for anything that is not a genuine blob. - if ( - valid_token is None - and not api_key.startswith("sk-") - and get_secret_bool("EXPERIMENTAL_UI_LOGIN") is not False - ): - valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) - - if valid_token is not None and valid_token.is_session_token and prisma_client is not None: - valid_token = await _refresh_session_token_grants( # rebind-ok: later checks read this name - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - if ( - valid_token is not None - and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.user_role == LitellmUserRoles.PROXY_ADMIN - ): - if valid_token.expires is not None: - current_time = datetime.now(timezone.utc) - if isinstance(valid_token.expires, datetime): - expiry_time = valid_token.expires - else: - expiry_time = datetime.fromisoformat(valid_token.expires) - if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: - expiry_time = expiry_time.replace(tzinfo=timezone.utc) - if expiry_time < current_time: - await delete_cache_key_object( - hashed_token=hash_token(api_key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - raise ProxyException( - message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", - type=ProxyErrorTypes.expired_key, - code=status.HTTP_401_UNAUTHORIZED, - param=abbreviate_api_key(api_key=api_key), - ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user_params - ) - valid_token.parent_otel_span = parent_otel_span - if _end_user_object is not None: - valid_token.end_user_object_permission = _end_user_object.object_permission - - return valid_token - - if ( - valid_token is not None - and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.team_id is not None - and valid_token.team_id != UI_TEAM_ID - ): - ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token - try: - team_obj: Final[LiteLLM_TeamTableCachedObj] = await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, - ) - - if ( - team_obj.last_refreshed_at is not None - and valid_token.last_refreshed_at is not None - and team_obj.last_refreshed_at > valid_token.last_refreshed_at - ): - team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) - - for k, v in team_obj_dict.items(): - field_name = f"team_{k}" - if field_name in valid_token.__fields__: - setattr(valid_token, field_name, v) - except Exception as e: - verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. - try: - is_master_key_valid = secrets.compare_digest(api_key, master_key) - except Exception: - is_master_key_valid = False - - ## VALIDATE MASTER KEY ## - if not isinstance(master_key, str): - raise HTTPException( - status_code=500, - detail={f"Master key must be a valid string. Current type={type(master_key)}"}, + outcome: Final = await drive( + _AUTH_CORE.start( + RequestFacts( + route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route) + ), + _core_settings( + master_key=master_key, + allow_user_auth=general_settings.get("allow_user_auth", False) is True, + prisma_client=prisma_client, + ), + datetime.now(timezone.utc), + ), + effects, ) - - if is_master_key_valid: - # Substitute a stable alias for the raw master key so neither the - # master key nor its hash propagates into spend logs, Prometheus - # /metrics labels, audit trails, rate-limit buckets, or any other - # downstream consumer of UserAPIKeyAuth.api_key. - _user_api_key_obj = await _return_user_api_key_auth_obj( - user_obj=None, - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, - parent_otel_span=parent_otel_span, - valid_token_dict={ - **end_user_params, - "user_id": litellm_proxy_admin_name, - }, - route=route, - start_time=start_time, - ) - asyncio.create_task( - cache_key_object( - hashed_token=hash_token(master_key), - user_api_key_obj=_user_api_key_obj, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - - _user_api_key_obj = update_valid_token_with_end_user_params( - valid_token=_user_api_key_obj, end_user_params=end_user_params - ) - _user_api_key_obj.via_virtual_key = True - - return _user_api_key_obj - - ## IF it's not a master key - ## Route should not be in master_key_only_routes - if route in LiteLLMRoutes.master_key_only_routes.value: - raise ProxyException( - message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - ## Check DB - - if ( - prisma_client is None - ): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error - raise ProxyException( - message="No connected db.", - type=ProxyErrorTypes.no_db_connection, - code=400, - param=None, - ) - - if valid_token is None: - if isinstance(api_key, str): # if generated token, make sure it starts with sk-. - _masked_key: Final = f"{api_key[:4]}****{api_key[-4:]}" if len(api_key) > 8 else "****" - if not api_key.startswith("sk-"): - _hint = _JWT_AUTH_DISABLED_HINT if not enable_jwt_auth and JWTHandler.is_jwt(token=api_key) else "" - _malformed_key_error = HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=( - f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={_masked_key}, " - f"expected to start with 'sk-'.{_hint}" - ), - ) # prevent token hashes from being used - # Stamp provenance here so log routing classifies this 401 by - # where it was raised, never by its message text. - setattr(_malformed_key_error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) - raise _malformed_key_error - else: - verbose_logger.warning( - "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( - type(api_key) if api_key is not None else "None" - ) - ) - if api_key.startswith("sk-"): - api_key = hash_token(token=api_key) - - with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): - valid_token = IdentityStore.key_from_principal( - await IdentityStore( - prisma_client, - user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ).resolve(hashed_token=api_key) - ) - # update end-user params on valid token - # These can change per request - it's important to update them here - valid_token.end_user_id = end_user_params.get("end_user_id") - valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") - valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") - valid_token.end_user_tpd_limit = end_user_params.get("end_user_tpd_limit") - valid_token.allowed_model_region = end_user_params.get("allowed_model_region") - - return await validate_resolved_virtual_key( - request=request, - request_data=request_data, - valid_token=valid_token, - api_key=api_key, + finally: + valid_token = effects.resolved_key # rebind-ok: the exception handler below reports the resolved identity + return await _authenticated( + outcome, + effects, + admin_user_id=litellm_proxy_admin_name, route=route, start_time=start_time, - parent_otel_span=parent_otel_span, - end_user_id=end_user_id, - end_user_params=cast( # cast-ok: builder assembles this dict from validated end-user fields - dict[str, object], end_user_params - ), - _end_user_object=_end_user_object, + jwt_auth_enabled=enable_jwt_auth, ) + except NativeAuthFailure: + raise except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( e=e, @@ -2995,417 +2943,194 @@ async def user_api_key_auth_builder( _user_api_key_auth_builder: Final = user_api_key_auth_builder -async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of existing shared authorization checks - request: Request, - request_data: dict[str, object], - valid_token: UserAPIKeyAuth | None, - api_key: str, - route: str, - start_time: datetime, +async def _load_team_and_project( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, - end_user_id: str | None, - end_user_params: dict[str, object], - _end_user_object: LiteLLM_EndUserTable | None, -) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - litellm_proxy_admin_name, - llm_model_list, - llm_router, - model_max_budget_limiter, - prisma_client, - proxy_logging_obj, - user_api_key_cache, + proxy_logging_obj: ProxyLogging, +) -> None: + team_obj: Final = ( + None + if valid_token.team_id is None + else await _key_team_object( + valid_token, valid_token.team_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj + ) + ) + valid_token.team_object_permission = team_obj.object_permission if team_obj is not None else None + if team_obj is not None: + # Guardrails or other metadata added after the key was cached apply without a cache eviction. + valid_token.team_metadata = team_obj.metadata + if valid_token.project_id is None: + return + project_obj: Final = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if project_obj is not None: + valid_token.project_metadata = project_obj.metadata + valid_token.project_alias = project_obj.project_alias + + +async def _key_team_object( + valid_token: UserAPIKeyAuth, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj: + fetched: Final = ( + None + if team_id == UI_TEAM_ID + else await _team_object_or_none(team_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj) + ) + if fetched is not None: + return fetched + return LiteLLM_TeamTableCachedObj( + team_id=team_id, + max_budget=valid_token.team_max_budget, + soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, + spend=valid_token.team_spend, + tpm_limit=valid_token.team_tpm_limit, + rpm_limit=valid_token.team_rpm_limit, + tpd_limit=valid_token.team_tpd_limit, + blocked=valid_token.team_blocked, + models=_token_team_models(valid_token), + metadata=valid_token.team_metadata, + object_permission_id=valid_token.team_object_permission_id, + object_permission=await _resolve_object_permission_for_unresolvable_team( + object_permission_id=valid_token.team_object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ), ) - if valid_token is not None: - valid_token = _update_key_budget_with_temp_budget_increase(valid_token) - user_obj: LiteLLM_UserTable | None = None - valid_token_dict: dict = {} - if valid_token is not None: - # Got Valid Token from Cache, DB - # Run checks for - # 1. If token can call model - ## 1a. If token can call fallback models (if client-side fallbacks given) - # 2. If user_id for this token is in budget - # 3. If the user spend within their own team is within budget - # 4. If 'user' passed to /chat/completions, /embeddings endpoint is in budget - # 5. If token is expired - # 6. If token spend is under Budget for the token - # 7. If token spend per model is under budget per model - # 8. If token spend is under team budget - # 9. If team spend is under team budget - - ## base case ## key is disabled - if valid_token.blocked is True: - raise ProxyException( - message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - ) - await _prefetch_referenced_auth_objects( - valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client - ) - - # Check 2. If user_id for this token is in budget - done in common_checks() - if valid_token.user_id is not None: - try: - with tracer.trace("litellm.proxy.auth.get_user_object"): - user_obj = await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - verbose_logger.debug( - "litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - %s", - e, - ) - user_obj = None - - if user_obj is not None: - # The joint verification-token view carries the key's columns only, so the - # user's own per-model budget reaches enforcement and the post-call - # increment through the row fetched here. - valid_token.user_model_max_budget = user_obj.model_max_budget - - if ( - user_obj is not None - and isinstance(user_obj.metadata, dict) - and user_obj.metadata.get("scim_active") is False - ): - raise ProxyException( - message=( - f"Authentication Error, User={valid_token.user_id} has been deactivated via SCIM. " - "Keys owned by this user cannot be used." - ), - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - # Check 2a. Check if model has zero cost - if so, skip all budget checks - model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - skip_budget_checks: Final = request_skips_budget_checks( - route=route, model=model, llm_router=llm_router, request_body=request_data - ) - - # Check 3. Check if user is in their team budget - if not skip_budget_checks and valid_token.team_member_spend is not None: - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - if prisma_client is not None and _user_id is not None and _team_id is not None: - _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id) - - team_member_info = await user_api_key_cache.async_get_cache( - key=_cache_key, - model_type=LiteLLM_TeamMembership, - ) - if team_member_info is None: - # read from DB - _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, - include={"litellm_budget_table": True}, - ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership.model_validate(_db_member.model_dump()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) - - if team_member_info is not None and team_member_info.litellm_budget_table is not None: - team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget( - now=datetime.now(timezone.utc), - ) - if team_member_budget is not None and team_member_budget > 0: - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = valid_token.team_member_spend - if valid_token.user_id is not None and valid_token.team_id is not None: - team_member_spend = await get_current_spend( - counter_key=team_member_counter(valid_token.user_id, valid_token.team_id), - fallback_spend=team_member_spend, - max_budget=team_member_budget, - ) - if team_member_spend >= team_member_budget: - # common_checks sends this alert on requests that get past here, so only the - # request rejected here sends it from the builder. - team_member_max_budget_alert_check( - team_id=_team_id, - team_alias=valid_token.team_alias, - team_metadata=valid_token.team_metadata, - organization_id=valid_token.org_id, - user_id=_user_id, - user_email=user_obj.user_email if user_obj is not None else None, - proxy_logging_obj=proxy_logging_obj, - spend=team_member_spend, - max_budget=team_member_budget, - ) - _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - message=( - f"Budget has been exceeded! TeamMember={_entity_id} " - f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" - ), - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=_entity_id, - ) - - # Check 3. If token is expired - if valid_token.expires is not None: - current_time = datetime.now(timezone.utc) - if isinstance(valid_token.expires, datetime): - expiry_time = valid_token.expires - else: - expiry_time = datetime.fromisoformat(valid_token.expires) - if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: - expiry_time = expiry_time.replace(tzinfo=timezone.utc) - verbose_proxy_logger.debug( - "Checking if token expired, expiry time %s and current time %s", expiry_time, current_time - ) - if expiry_time < current_time: - # Token exists but is expired. - raise ProxyException( - message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", - type=ProxyErrorTypes.expired_key, - code=status.HTTP_401_UNAUTHORIZED, - param=abbreviate_api_key(api_key=api_key), - ) - - if not skip_budget_checks: - with tracer.trace("litellm.proxy.auth.budget_checks"): - # Check 4. Max Budget Alert Check (runs before budget enforcement - # so multi-threshold 100% alerts fire on the request that crosses - # max_budget, before BudgetExceededError is raised below) - await virtual_key_max_budget_alert_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 5. Token Spend is under budget - if RouteChecks.is_llm_api_route(route=route): - await virtual_key_max_budget_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 6. Soft Budget Check - await virtual_key_soft_budget_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 5. Token Model Spend is under Model budget - max_budget_per_model: Final = valid_token.model_max_budget - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - if ( - max_budget_per_model is not None - and isinstance(max_budget_per_model, dict) - and len(max_budget_per_model) > 0 - and prisma_client is not None - and current_models - and valid_token.token is not None - ): - ## GET THE SPEND FOR THIS MODEL - for model_name in current_models: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model_name=model_name, - request_data=request_data, - request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, llm_model_list - ), - llm_router=llm_router, - ) - - # Recompute after a potential budget-fallback rewrite so - # the user check below validates the final model - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - # Check 5a. Internal user model_max_budget - if current_models: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=current_models, - ) - - # Check 6: Additional Common Checks across jwt + key auth - if valid_token.team_id is not None: - try: - if valid_token.team_id == UI_TEAM_ID: - raise TeamNotFoundError(team_id=UI_TEAM_ID) - with tracer.trace("litellm.proxy.auth.get_team_object"): - _team_obj = await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except HTTPException: - token_team_models: Final = _token_team_models(valid_token) - _team_obj = LiteLLM_TeamTableCachedObj( - team_id=valid_token.team_id, - max_budget=valid_token.team_max_budget, - soft_budget=valid_token.team_soft_budget, - model_max_budget=valid_token.team_model_max_budget, - spend=valid_token.team_spend, - tpm_limit=valid_token.team_tpm_limit, - rpm_limit=valid_token.team_rpm_limit, - tpd_limit=valid_token.team_tpd_limit, - blocked=valid_token.team_blocked, - models=token_team_models, - metadata=valid_token.team_metadata, - object_permission_id=valid_token.team_object_permission_id, - object_permission=await _resolve_object_permission_for_unresolvable_team( - object_permission_id=valid_token.team_object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ), - ) - else: - _team_obj = None - - if _team_obj is not None: - valid_token.team_object_permission = _team_obj.object_permission - # Keep team_metadata in sync with the freshly fetched team so that - # guardrails (or any other metadata) added after the key was cached - # are picked up on subsequent requests without a cache eviction. - valid_token.team_metadata = _team_obj.metadata - else: - valid_token.team_object_permission = None - - # Fetch project object if key belongs to a project - _project_obj = None - if valid_token.project_id is not None: - _project_obj = await get_project_object( - project_id=valid_token.project_id, +async def _team_object_or_none( + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj | None: + try: + with tracer.trace("litellm.proxy.auth.get_team_object"): + return await get_team_object( + team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - if _project_obj is not None: - valid_token.project_metadata = _project_obj.metadata - valid_token.project_alias = _project_obj.project_alias + except HTTPException: + return None - global_proxy_spend = None - if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget - cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY - with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): - global_proxy_spend = await fetch_global_spend_with_event_coordination( # rebind-ok: pre-existing rebinding on a rename-only line - cache_key=cache_key, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) - if global_proxy_spend is not None: - call_info: Final = CallInfo( - token=valid_token.token, - spend=global_proxy_spend, - max_budget=litellm.max_budget, - user_id=litellm_proxy_admin_name, - team_id=valid_token.team_id, - event_group=Litellm_EntityType.PROXY, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="proxy_budget", - user_info=call_info, - ) - ) - # Token passed all checks - if valid_token is None: - raise HTTPException(401, detail="Invalid API key") - if valid_token.token is None: - raise HTTPException(401, detail="Invalid API key, no token associated") - api_key = valid_token.token - - valid_token_dict = valid_token.model_dump(exclude_none=True) - valid_token_dict.pop("token", None) - # budget_throttle_pct is excluded from model_dump (it must not leak - # into serialized responses), so carry the request-scoped decision - # forward by hand to the auth object the rate limiter receives. - if valid_token.budget_throttle_pct is not None: - valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct - - if _end_user_object is not None or litellm.max_end_user_budget_id is not None: - valid_token_dict.update(end_user_params) - if _end_user_object is not None: - valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission - - # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions - # sso/login, ui/login, /key functions and /user functions - # this will never be allowed to call /chat/completions - - if valid_token is None: - # No token was found when looking up in the DB - raise ProxyException( - message="Authentication Error, Invalid proxy server token passed", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, +async def _master_key_auth( + effects: _BuilderEffects, admin_user_id: str, route: str, start_time: datetime +) -> UserAPIKeyAuth: + # A stable alias replaces the raw master key so neither the key nor its hash reaches spend logs, + # metrics labels, audit trails or rate-limit buckets. + master_key_auth: Final = await _return_user_api_key_auth_obj( + user_obj=None, + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + parent_otel_span=effects.parent_otel_span, + valid_token_dict={**effects.end_user.params, "user_id": admin_user_id}, + route=route, + start_time=start_time, + ) + asyncio.create_task( + cache_key_object( + hashed_token=hash_token(effects.token), + user_api_key_obj=master_key_auth, + user_api_key_cache=effects.user_api_key_cache, + proxy_logging_obj=effects.proxy_logging_obj, ) - if valid_token_dict is not None: - virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj( - user_obj=user_obj, - api_key=api_key, - parent_otel_span=parent_otel_span, - valid_token_dict=valid_token_dict, - route=route, - start_time=start_time, - ) - virtual_key_auth_obj.via_virtual_key = True - return virtual_key_auth_obj + ) + with_end_user: Final = _with_end_user(master_key_auth, effects.end_user, effects.parent_otel_span) + with_end_user.via_virtual_key = True + return with_end_user + + +async def _authenticated( + outcome: Outcome, + effects: _BuilderEffects, + admin_user_id: str, + route: str, + start_time: datetime, + jwt_auth_enabled: bool, +) -> UserAPIKeyAuth: + match outcome: + case Deny(denial=denial): + raise denial_exception( + denial, token=effects.token, route=route, jwt_auth_enabled=jwt_auth_enabled, key=effects.resolved_key + ) + case Allow(grant=grant): + return await _granted(grant, effects, admin_user_id, route, start_time) + + +async def _granted( + grant: Grant, effects: _BuilderEffects, admin_user_id: str, route: str, start_time: datetime +) -> UserAPIKeyAuth: + match grant: + case NoMasterKey(token=token): + return UserAPIKeyAuth( + api_key=token, user_role=LitellmUserRoles.INTERNAL_USER, parent_otel_span=effects.parent_otel_span + ) + case UserAuthRoute(): + return UserAPIKeyAuth() + case MasterKey(): + return await _master_key_auth(effects, admin_user_id, route, start_time) + case CachedAdminKey(): + return _with_end_user(effects.key, effects.end_user, effects.parent_otel_span) + case VirtualKey(budget_throttle_pct=budget_throttle_pct): + return await _virtual_key_auth(effects, budget_throttle_pct, route, start_time) + case WithinModelBudgets(): + raise RuntimeError("the request auth core never ends on a model budget check") + case _: + assert_never(grant) + + +async def _virtual_key_auth( + effects: _BuilderEffects, budget_throttle_pct: float | None, route: str, start_time: datetime +) -> UserAPIKeyAuth: + key: Final = effects.key + token: Final = key.token + if token is None: + raise RuntimeError("the auth core granted a virtual key without a token") + end_user: Final = effects.end_user + # budget_throttle_pct is excluded from model_dump so it never leaks into serialized responses. + throttle: Final = {} if budget_throttle_pct is None else {"budget_throttle_pct": budget_throttle_pct} + end_user_params: Final = ( + end_user.params if end_user.end_user_object is not None or litellm.max_end_user_budget_id is not None else {} + ) + end_user_permission: Final = ( + {} + if end_user.end_user_object is None + else {"end_user_object_permission": end_user.end_user_object.object_permission} + ) + virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj( + user_obj=effects.owner, + api_key=token, + parent_otel_span=effects.parent_otel_span, + valid_token_dict={ + **{k: v for k, v in key.model_dump(exclude_none=True).items() if k != "token"}, + **throttle, + **end_user_params, + **end_user_permission, + }, + route=route, + start_time=start_time, + ) + virtual_key_auth_obj.via_virtual_key = True + return virtual_key_auth_obj async def _safe_fetch(label: str, awaitable): @@ -3427,6 +3152,8 @@ async def _safe_fetch(label: str, awaitable): e, ) raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.debug( "centralized auth: %s fetch swallowed (%s: %s)", @@ -3459,36 +3186,6 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached ) -def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseException) -> bool: - """Whether the token's own team fields may stand in for a team that failed to - resolve, without widening access. - - The UI dashboard mints every session key against the ``UI_TEAM_ID`` sentinel, - which by design never has a team row, so a failed lookup for it is not a - degraded read to be treated with suspicion; it always vouches, exactly as it - always safely has (these keys are restricted elsewhere to UI-only routes). - - For every other team, a team that is provably gone is a definitive answer, - not a degraded read, so nothing may stand in for it and no setting may - override that. - - Otherwise the team's grant is merely unknown. A token carrying one may vouch, - since replaying a recorded grant cannot widen it and denying every team key - while the row is briefly unreadable would trade the widening for an outage. A - token carrying none may not: ``team_models=[]`` reads as every model and - ``team_blocked=False`` as unblocked. ``allow_requests_on_db_unavailable`` opts - back out, and is only consulted here because the failure is known by this - point to be a degraded read. - """ - if valid_token.team_id == UI_TEAM_ID: - return True - if isinstance(lookup_error, TeamNotFoundError): - return False - if valid_token.team_models: - return True - return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - - async def _inherit_org_identity( user_api_key_auth_obj: UserAPIKeyAuth, team_object: LiteLLM_TeamTableCachedObj | None, @@ -3545,68 +3242,10 @@ async def _load_team_member_default_model_budget( skip_budget_checks: bool, proxy_logging_obj: ProxyLogging | None, ) -> None: - team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None - team_metadata: Final[Mapping[str, object] | None] = ( - cast( # cast-ok: team metadata is a string-keyed JSON object - Mapping[str, object], team_metadata_value - ) - if isinstance(team_metadata_value, Mapping) - else None - ) - team_member_budget_id: Final[object | None] = ( - team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None - ) - if ( - not isinstance(user_api_key_auth_obj.user_id, str) - or not isinstance(user_api_key_auth_obj.team_id, str) - or prisma_client is None - ): - return + from litellm.proxy.auth import budget_context_adapter - try: - membership: Final = await get_team_membership( - user_id=user_api_key_auth_obj.user_id, - team_id=user_api_key_auth_obj.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - member_budget_row: Final = membership.litellm_budget_table if membership is not None else None - member_budget_value: Final[object] = ( - getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None - ) - if ( - membership is not None - and membership.budget_id is not None - and membership.budget_id != team_member_budget_id - and isinstance(member_budget_value, Mapping) - and member_budget_value - ): - user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream - member_budget_value - ) - return - if not isinstance(team_member_budget_id, str): - return - team_member_default_budget: Final = await get_team_member_default_budget( - budget_id=team_member_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - raise_on_lookup_error=True, - ) - except Exception: - if not skip_budget_checks: - raise - verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True) - return - - default_budget_value: Final[object] = ( - getattr(team_member_default_budget, "model_max_budget", None) - if team_member_default_budget is not None - else None - ) - user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream - default_budget_value if isinstance(default_budget_value, Mapping) else None + return await budget_context_adapter.load_team_member_default_model_budget( + user_api_key_auth_obj, team_object, prisma_client, user_api_key_cache, skip_budget_checks, proxy_logging_obj ) @@ -3654,41 +3293,34 @@ async def run_centralized_common_checks( user_custom_auth, ) - # Public routes (e.g. /health/liveness) are exempt from - # auth in the builder — the wrapper must not retroactively apply - # authz on top, or k8s readiness probes and other unauthenticated - # callers get 401. - if not force_virtual_key_checks and ( - route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes(current_route=route) - ): - return - - # User-configured pass-through endpoints with ``auth: false`` are - # explicitly unauthenticated — the builder returns an empty - # UserAPIKeyAuth() and the request is forwarded as-is. Running - # common_checks on the empty token would reject the request as - # admin-only. The "auth" flag on the endpoint config is the - # contract; honor it. pass_through_endpoints: Final = general_settings.get("pass_through_endpoints", None) - if not force_virtual_key_checks and pass_through_endpoints is not None: - for endpoint in pass_through_endpoints: - if isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True: - return - - # No-auth dev mode: master_key unset AND no JWT/OAuth2 auth - # configured. The builder returns an INTERNAL_USER token for any - # api_key; the proxy is unauthenticated by configuration. - # Running common_checks would block every admin route on these - # deployments where that was previously not the contract. If any - # authn is enabled (JWT, OAuth2, OAuth2-proxy), authz must run. - if not force_virtual_key_checks and is_no_auth_dev_mode(master_key, general_settings): - return - - if ( - not force_virtual_key_checks - and user_custom_auth is not None - and not general_settings.get("custom_auth_run_common_checks", False) - ): + plan: Final = admission.common_plan( + admission.CommonGate( + force_virtual_key_checks=force_virtual_key_checks, + public_route=route in LiteLLMRoutes.public_routes.value, + additional_public_route=route_in_additonal_public_routes(current_route=route), + unauthenticated_passthrough=any( + isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True + for endpoint in (pass_through_endpoints or ()) + ), + master_key_present=master_key is not None, + jwt_enabled=bool(general_settings.get("enable_jwt_auth", False)), + oauth_enabled=bool(general_settings.get("enable_oauth2_auth", False)), + oauth_proxy_enabled=bool(general_settings.get("enable_oauth2_proxy_auth", False)), + custom_auth=user_custom_auth is not None, + custom_auth_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)), + ) + ) + request.state.auth_ownership_plan = admission.ownership_plan( + admission.Ownership( + virtual_key=user_api_key_auth_obj.via_virtual_key, + custom_auth=user_api_key_auth_obj.authenticated_by_custom_auth, + jwt=bool(user_api_key_auth_obj.jwt_claims), + oauth=bool(general_settings.get("enable_oauth2_auth", False)) + or bool(general_settings.get("enable_oauth2_proxy_auth", False)), + ) + ) + if plan.skip is not None: return parent_otel_span: Final = user_api_key_auth_obj.parent_otel_span @@ -3713,7 +3345,17 @@ async def run_centralized_common_checks( user_api_key_auth_obj.end_user_id = end_user_id fetch_coros: Final = [] - if user_api_key_auth_obj.team_id is not None and user_api_key_auth_obj.team_id != UI_TEAM_ID: + fetch_plan: Final = auth_context.fetch( + auth_context.Fetch( + user_api_key_auth_obj.team_id, + UI_TEAM_ID, + user_api_key_auth_obj.user_id, + user_api_key_auth_obj.project_id, + end_user_id, + ) + ) + if fetch_plan.team: + assert user_api_key_auth_obj.team_id is not None fetch_coros.append( _safe_fetch( "team", @@ -3729,7 +3371,8 @@ async def run_centralized_common_checks( else: fetch_coros.append(_safe_fetch("team", _noop_none())) - if user_api_key_auth_obj.user_id is not None: + if fetch_plan.user: + assert user_api_key_auth_obj.user_id is not None fetch_coros.append( _safe_fetch( "user", @@ -3746,7 +3389,8 @@ async def run_centralized_common_checks( else: fetch_coros.append(_safe_fetch("user", _noop_none())) - if user_api_key_auth_obj.project_id is not None: + if fetch_plan.project: + assert user_api_key_auth_obj.project_id is not None fetch_coros.append( _safe_fetch( "project", @@ -3761,7 +3405,8 @@ async def run_centralized_common_checks( else: fetch_coros.append(_safe_fetch("project", _noop_none())) - if end_user_id: + if fetch_plan.end_user: + assert end_user_id is not None fetch_coros.append( _safe_fetch( "end_user", @@ -3809,38 +3454,46 @@ async def run_centralized_common_checks( global_spend_result, ) = await asyncio.gather(*fetch_coros, return_exceptions=True) - # ProxyException / BudgetExceededError are authorization failures — - # propagate so the wrapper renders them. HTTPException is fallback - # material (404 from get_team_object is the only known producer). - for r in ( - team_result, - user_result, - project_result, - end_user_result, - global_spend_result, - ): - if isinstance(r, (ProxyException, litellm.BudgetExceededError)): - raise r - - # Use BaseException (not HTTPException) in the narrowing checks so - # mypy can narrow ``Any | BaseException`` to the typed object in the - # else branch. After the for-loop above, the only BaseException that - # can still appear here is HTTPException (other listed re-raises were - # propagated; non-listed exceptions were already swallowed to None). - team_object: LiteLLM_TeamTableCachedObj | None - if isinstance(team_result, BaseException): - # Token-derived fallback only valid when a team_id is set; - # _team_obj_from_token asserts that precondition. - if user_api_key_auth_obj.team_id is None: - team_object = None - elif _token_can_vouch_for_team(user_api_key_auth_obj, team_result): - team_object = _team_obj_from_token(user_api_key_auth_obj) - else: - raise team_result - else: - team_object = ( - _team_obj_from_token(user_api_key_auth_obj) if user_api_key_auth_obj.team_id == UI_TEAM_ID else team_result + results: Final = (team_result, user_result, project_result, end_user_result, global_spend_result) + failed: Final = access_policy.decide( + access_policy.ConcurrentFailures( + failures=tuple(isinstance(value, (ProxyException, litellm.BudgetExceededError)) for value in results), + terminal=tuple(isinstance(value, NativeAuthFailure) for value in results), ) + ).failure_index + if failed is not None: + failure: Final = results[failed] + assert isinstance(failure, BaseException) + raise failure + team_facts: Final = auth_context.Team( + user_api_key_auth_obj.team_id, + UI_TEAM_ID, + failed=isinstance(team_result, BaseException), + missing=isinstance(team_result, TeamNotFoundError), + recorded_models=tuple(user_api_key_auth_obj.team_models or ()), + ) + initial_team: Final = auth_context.team(team_facts) + team_plan: Final = ( + auth_context.team( + replace(team_facts, allow_unavailable=PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()) + ) + if initial_team.source == "ReadUnavailablePolicy" + else initial_team + ) + team_object: LiteLLM_TeamTableCachedObj | None + match team_plan.source: + case "None": + team_object = None + case "Token": + team_object = _team_obj_from_token(user_api_key_auth_obj) + case "Loaded": + assert not isinstance(team_result, BaseException) + team_object = team_result + case "Failure": + assert isinstance(team_result, BaseException) + raise team_result + case _: + raise NativeAuthFailure("incomplete team context policy") user_object: LiteLLM_UserTable | None = user_result if isinstance(user_result, LiteLLM_UserTable) else None project_object: Final[LiteLLM_ProjectTableCachedObj | None] = ( @@ -3875,12 +3528,16 @@ async def run_centralized_common_checks( # username/password login) keeps the row's personal budget (max_budget, # and user_email for budget alerts) so common_checks enforces it; every # other field keeps its default. - if user_api_key_auth_obj.user_role == LitellmUserRoles.PROXY_ADMIN: - budget_row: Final = ( - user_object - if isinstance(user_object, LiteLLM_UserTable) and user_api_key_auth_obj.user_id != litellm_proxy_admin_name - else None + admin_plan: Final = auth_context.admin( + auth_context.Admin( + user_api_key_auth_obj.user_role, + user_api_key_auth_obj.user_id, + litellm_proxy_admin_name, + isinstance(user_object, LiteLLM_UserTable), ) + ) + if admin_plan.force_admin: + budget_row: Final = user_object if admin_plan.preserve_budget else None user_object = LiteLLM_UserTable( user_id=user_api_key_auth_obj.user_id or litellm_proxy_admin_name, user_role=LitellmUserRoles.PROXY_ADMIN, @@ -4076,35 +3733,17 @@ async def _apply_key_end_user_default_budget_to_token( parent_otel_span: Span | None, keep_token_limits: bool, ) -> None: - """The builder's end-user pass runs before the key is resolved, so only here can the key's - ``end_user_budget_id`` win over the proxy-wide default on the token that reservation reads. - On the virtual-key path the token's end-user limits are the builder's proxy-wide defaults and - the key budget replaces them wholesale. With ``keep_token_limits`` (custom auth) the token's - limits are caps the custom auth callable set, so the key budget only fills the ones it left - unset.""" - default_budget: Final = ( - end_user_object.litellm_budget_table - if end_user_object is not None - else await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, - ) - ) - if default_budget is None: - return + from litellm.proxy.auth import budget_context_adapter - if not keep_token_limits or valid_token.end_user_max_budget is None: - valid_token.end_user_max_budget = default_budget.max_budget - if not keep_token_limits or valid_token.end_user_tpm_limit is None: - valid_token.end_user_tpm_limit = default_budget.tpm_limit - if not keep_token_limits or valid_token.end_user_rpm_limit is None: - valid_token.end_user_rpm_limit = default_budget.rpm_limit - if not keep_token_limits or valid_token.end_user_tpd_limit is None: - valid_token.end_user_tpd_limit = default_budget.tpd_limit - if not keep_token_limits or valid_token.end_user_model_max_budget is None: - valid_token.end_user_model_max_budget = default_budget.model_max_budget + return await budget_context_adapter.apply_key_end_user_default_budget_to_token( + valid_token, + end_user_object, + key_end_user_budget_id, + prisma_client, + user_api_key_cache, + parent_otel_span, + keep_token_limits, + ) async def _reserve_budget_after_common_checks( @@ -4124,7 +3763,16 @@ async def _reserve_budget_after_common_checks( request: Request | None = None, ) -> None: user_api_key_auth_obj.budget_reservation = None - if not skip_budget_checks and general_settings.get("disable_budget_reservation") is not True: + plan: Final = admission.reservation_plan( + admission.Reservation( + skip_budget_checks=skip_budget_checks, + disabled=general_settings.get("disable_budget_reservation") is True, + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, + fail_closed=general_settings.get("fail_closed_budget_enforcement") is True, + billing_agent=user_api_key_auth_obj.billing_agent_policy is not None, + ) + ) + if plan.reserve: from litellm.proxy.spend_tracking.budget_reservation import ( reserve_budget_for_request, ) @@ -4141,11 +3789,8 @@ async def _reserve_budget_after_common_checks( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, - apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, - fail_closed_budget_enforcement=( - general_settings.get("fail_closed_budget_enforcement") is True - or user_api_key_auth_obj.billing_agent_policy is not None - ), + apply_user_budget_to_team_keys=plan.apply_user_budget_to_team_keys, + fail_closed_budget_enforcement=plan.fail_closed, raw_body=await read_raw_json_body(request=request), ) if request is not None: @@ -4278,6 +3923,8 @@ async def _authorize_authenticated_request( route=route, force_virtual_key_checks=force_virtual_key_checks, ) + except NativeAuthFailure: + raise except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( e=e, @@ -4399,6 +4046,8 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque fan_out_provider(), ) ) + except NativeAuthFailure: + raise except Exception as exc: # noqa: BLE001 # telemetry routing is best-effort and must never break authentication verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc) @@ -4407,7 +4056,7 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque async def user_api_key_auth( request: Request, api_key: str = fastapi.Security(api_key_header), - azure_api_key_header: str = fastapi.Security(azure_api_key_header), + azure_api_key_header: str | None = fastapi.Security(azure_api_key_header), anthropic_api_key_header: str | None = fastapi.Security(anthropic_api_key_header), google_ai_studio_api_key_header: str | None = fastapi.Security(google_ai_studio_api_key_header), azure_apim_header: str | None = fastapi.Security(azure_apim_header), @@ -4448,6 +4097,8 @@ async def user_api_key_auth( request_data=request_data, custom_litellm_key_header=custom_litellm_key_header, ) + except NativeAuthFailure: + raise except Exception: # The body was read first, so a caller who sent both a malformed body and # a rejected key used to get the 400; the response is unchanged, and the @@ -4503,6 +4154,8 @@ async def user_api_key_auth( # any future consumer must treat a missing principal as deny, not allow. try: request.state.principal = _resolve_request_principal(request, user_api_key_auth_obj) + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.warning("Principal projection at auth seam failed (non-fatal): %s", e) @@ -4557,61 +4210,12 @@ async def _return_user_api_key_auth_obj( return UserAPIKeyAuth.model_validate(user_api_key_kwargs) -def get_api_key_from_custom_header(request: Request, custom_litellm_key_header_name: str) -> str: - """ - Get API key from custom header - - Args: - request (Request): Request object - custom_litellm_key_header_name (str): Custom header name - - Returns: - Optional[str]: API key - """ - api_key: str = "" - # use this as the virtual key passed to litellm proxy - custom_litellm_key_header_name = custom_litellm_key_header_name.lower() - _headers: Final = {k.lower(): v for k, v in request.headers.items()} - verbose_proxy_logger.debug( - "searching for custom_litellm_key_header_name= %s, in headers=%s", - custom_litellm_key_header_name, - _headers, - ) - custom_api_key: Final = _headers.get(custom_litellm_key_header_name) - if custom_api_key: - api_key = get_bearer_token(api_key=custom_api_key) # rebind-ok: pre-existing rebinding on a rename-only line - verbose_proxy_logger.debug( - "Found custom API key using header: %s, setting api_key=%s", - custom_litellm_key_header_name, - abbreviate_api_key(api_key), - ) - else: - verbose_proxy_logger.exception( - "No LiteLLM Virtual Key pass. Please set header=%s: Bearer ", custom_litellm_key_header_name - ) - return api_key - - -def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): - valid_token_metadata: Final = valid_token.metadata - if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata: - expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"]) - if expiry.tzinfo is None: - expiry = expiry.replace(tzinfo=timezone.utc) - if expiry > datetime.now(timezone.utc): - return valid_token_metadata["temp_budget_increase"] - return None - - def _update_key_budget_with_temp_budget_increase( valid_token: UserAPIKeyAuth, ) -> UserAPIKeyAuth: - if valid_token.max_budget is None: - return valid_token - temp_budget_increase: Final = _get_temp_budget_increase(valid_token) - if not temp_budget_increase: - return valid_token - return valid_token.model_copy(update={"max_budget": valid_token.max_budget + temp_budget_increase}) + from litellm.proxy.auth import key_budget_context_adapter + + return key_budget_context_adapter.update_key_budget_with_temp_budget_increase(valid_token) async def _lookup_end_user_and_apply_budget( @@ -4622,119 +4226,40 @@ async def _lookup_end_user_and_apply_budget( user_api_key_cache, proxy_logging_obj, ): - """Look up end_user from DB and apply budget limits to valid_token.""" - end_user_object = None - key_end_user_budget_id: Final = get_key_end_user_budget_id(valid_token.metadata) - try: - end_user_object = await get_end_user_object( - end_user_id=valid_token.end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - token_end_user_max_budget=valid_token.end_user_max_budget, - key_end_user_budget_id=key_end_user_budget_id, - ) - if end_user_object is not None: - end_user_params = { - "end_user_id": valid_token.end_user_id, - "allowed_model_region": end_user_object.allowed_model_region, - } - if end_user_object.litellm_budget_table is not None: - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=end_user_object.litellm_budget_table, - end_user_id=valid_token.end_user_id or "", - ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user_params - ) - elif key_end_user_budget_id is not None or litellm.max_end_user_budget_id is not None: - default_budget: Final = await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, - ) - if default_budget is not None: - end_user_params = {"end_user_id": valid_token.end_user_id} - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=default_budget, - end_user_id=valid_token.end_user_id or "", - ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user_params - ) - if valid_token.end_user_max_budget is None: - valid_token.end_user_max_budget = default_budget.max_budget - except Exception as e: - if isinstance(e, litellm.BudgetExceededError): - raise e - verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) - return valid_token, end_user_object + from litellm.proxy.auth import end_user_context_adapter + + return await end_user_context_adapter.lookup_end_user_and_apply_budget( + valid_token, route, parent_otel_span, prisma_client, user_api_key_cache, proxy_logging_obj + ) async def enforce_key_and_fallback_model_access( *, valid_token: UserAPIKeyAuth, - request_data: dict, + request_data: dict[str, object], route: str, request: Request | None, - llm_model_list: list | None, + llm_model_list: Sequence[object] | None, llm_router: litellm.Router | None, ) -> None: - """ - Key-level model allowlist and client fallbacks (same as standard auth). - Not included in common_checks — common_checks enforces team/user/project model access only. - """ - await _normalize_claude_model(request_data, valid_token, request, route) - await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) - _resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) - config: Final = valid_token.config - - if config != {}: - model_list: Final = config.get("model_list", []) - new_model_list: Final = model_list - verbose_proxy_logger.debug("\n new llm router model list %s", new_model_list) - elif isinstance(valid_token.models, list) and "all-team-models" in valid_token.models: - pass - else: - model: Final = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - - if model is not None: - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - - fallback_names: Final = tuple( - name - for target in iter_request_fallback_targets(request_data) - if (name := fallback_target_model_name(target)) is not None - ) - - for _name in dict.fromkeys(fallback_names): # dedupe, preserve order - await can_key_call_model( - model=_name, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - await is_valid_fallback_model( - model=_name, - llm_router=llm_router, - user_model=None, - ) + facts: Final = await read_key_and_fallback_model_access_facts( + valid_token=valid_token, + request_data=request_data, + route=route, + request=request, + llm_model_list=llm_model_list, + llm_router=llm_router, + ) + decision: Final = key_and_fallback_model_access(facts) + match decision.kind: + case "Allow": + return + case "Deny": + if decision.denial is None: + raise NativeAuthFailure("missing native model denial") + raise_model_denial(decision.denial, decision.denial.model or "") + case _: + raise NativeAuthFailure("unanswered native model read") _enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access # pyright: ignore[reportUnknownVariableType] # existing tests use this alias @@ -4758,8 +4283,14 @@ async def _run_post_custom_auth_checks( ) # 1. Look up end_user object from DB if end_user_id is set + custom_facts: Final = common_flow.CustomFacts( + valid_token.end_user_id, + False, + bool(general_settings.get("custom_auth_run_common_checks", False)), + ) + custom_plan: Final = common_flow.custom(custom_facts) end_user_object = None - if valid_token.end_user_id is not None: + if custom_plan.read_end_user: valid_token, end_user_object = await _lookup_end_user_and_apply_budget( valid_token=valid_token, route=route, @@ -4772,19 +4303,15 @@ async def _run_post_custom_auth_checks( # gate skips it for custom-auth deployments unless # custom_auth_run_common_checks is set. Enforce it here on that path # so an over-budget end user can't keep making requests. - if end_user_object is not None and not general_settings.get("custom_auth_run_common_checks", False): + if common_flow.custom(replace(custom_facts, end_user_present=end_user_object is not None)).check_end_user: + assert end_user_object is not None await check_end_user_budget(end_user_obj=end_user_object, route=route) # 2. Check token expiry if valid_token.expires is not None: current_time: Final = datetime.now(timezone.utc) - if isinstance(valid_token.expires, datetime): - expiry_time = valid_token.expires - else: - expiry_time = datetime.fromisoformat(valid_token.expires) - if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: - expiry_time = expiry_time.replace(tzinfo=timezone.utc) - if expiry_time < current_time: + expiry_time: Final = _aware_expiry(valid_token.expires) + if expired(expiry_time, current_time): raise ProxyException( message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, @@ -4792,7 +4319,7 @@ async def _run_post_custom_auth_checks( param=(abbreviate_api_key(api_key=valid_token.token) if valid_token.token else ""), ) - if general_settings.get("custom_auth_run_common_checks", False): + if custom_plan.check_model_access: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, @@ -4802,14 +4329,13 @@ async def _run_post_custom_auth_checks( llm_router=llm_router, ) - current_model = _get_model_from_request_context( + current_model: Final = _get_model_from_request_context( request_data=request_data, route=route, request=request, llm_router=llm_router, team_id=valid_token.team_id, ) - current_models = _get_model_names_for_budget_checks(model=current_model) # A zero-cost model cannot move any counter, so refusing it means refusing on # spend some other model accrued. The JWT and virtual-key paths already skip @@ -4817,45 +4343,12 @@ async def _run_post_custom_auth_checks( # be refused under custom auth and served under the other two. skip_budget_checks: Final = ( is_model_cost_zero(model=current_model, llm_router=llm_router) - if current_model is not None and llm_router is not None + if common_flow.custom( + replace(custom_facts, model_present=current_model is not None, router_present=llm_router is not None) + ).read_pricing else False ) - # 3. Check key-level model_max_budget - max_budget_per_model: Final = valid_token.model_max_budget - if ( - not skip_budget_checks - and max_budget_per_model is not None - and isinstance(max_budget_per_model, dict) - and len(max_budget_per_model) > 0 - and current_models - and valid_token.token is not None - ): - for model_name in current_models: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model_name=model_name, - request_data=request_data, - request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, llm_model_list - ), - llm_router=llm_router, - ) - - # Recompute after a potential budget-fallback rewrite so - # the end-user check below validates the final model - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - # 3b. Attach and check the internal user's model_max_budget. # Custom auth builds its own token, so unlike the main path nothing has # loaded the user row yet. The attach is unconditional because the post-call # spend hook reads this field off the token: gating it on the same condition @@ -4869,30 +4362,45 @@ async def _run_post_custom_auth_checks( proxy_logging_obj=proxy_logging_obj, ) valid_token.user_model_max_budget = user_budget # rebind-ok: the spend hook reads it off this token - if not skip_budget_checks and current_models: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=current_models, + if common_flow.read(common_flow.BudgetSkip(skip_budget_checks, False)) == "CheckBudgets": + await _enforce_model_budgets( + _AUTH_CORE.check_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=valid_token, + route=route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=cast( # cast-ok: router model lists are deployment sequences + Sequence[object] | None, llm_model_list + ), + ), ) - # 4. Check end-user model_max_budget + # Read after the budget checks, which may have rerouted the request to a fallback model. + current_models: Final = _get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + ) + end_user_mmb: Final = valid_token.end_user_model_max_budget - if ( - not skip_budget_checks - and end_user_mmb is not None - and isinstance(end_user_mmb, dict) - and len(end_user_mmb) > 0 - and current_models - and valid_token.end_user_id is not None - ): - for model_name in current_models: - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_mmb, - model=model_name, - team_id=valid_token.team_id, - ) + if common_flow.read(common_flow.BudgetSkip(skip_budget_checks, False)) == "CheckBudgets": + from litellm.proxy.auth import budget_context_adapter + + await budget_context_adapter.check_end_user_model_budget( + valid_token.end_user_id, + end_user_mmb if isinstance(end_user_mmb, dict) else None, + model_max_budget_limiter, + current_models, + valid_token.team_id, + scope="CustomEndUser", + ) # team / user / end_user / project context objects are fetched by # the centralized common_checks gate in user_api_key_auth after @@ -4912,6 +4420,51 @@ async def _run_post_custom_auth_checks( return valid_token +async def _authorize_resolved_key( + identity: UserAPIKeyAuth, + key_hash: str, + request: Request, + request_data: dict[str, object], # mutable-ok: model access checks rewrite the request model in place + route: str, +) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + master_key, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + start_time: Final = datetime.now(timezone.utc) + effects: Final = _BuilderEffects( + request=request, + request_data=request_data, + token=key_hash, + preresolved_key=identity, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + outcome: Final = await drive( + _AUTH_CORE.start_resolved_key( + RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)), + _core_settings(master_key=master_key, allow_user_auth=False, prisma_client=prisma_client), + start_time, + ), + effects, + ) + return await _authenticated( + outcome, + effects, + admin_user_id=litellm_proxy_admin_name, + route=route, + start_time=start_time, + jwt_auth_enabled=False, + ) + + async def authorize_internal_virtual_key( key_hash: str, request: Request, request_data: dict[str, object] ) -> UserAPIKeyAuth: @@ -4925,17 +4478,8 @@ async def authorize_internal_virtual_key( ) route: Final = get_request_route(request=request) await pre_db_read_auth_checks(request_data=request_data, request=request, route=route) - auth: Final = await validate_resolved_virtual_key( - request=request, - request_data=request_data, - valid_token=identity, - api_key=key_hash, - route=route, - start_time=datetime.now(timezone.utc), - parent_otel_span=None, - end_user_id=None, - end_user_params={}, - _end_user_object=None, + auth: Final = await _authorize_resolved_key( + identity=identity, key_hash=key_hash, request=request, request_data=request_data, route=route ) auth.budget_reservation = None recovered: Final = await _authorize_authenticated_request( diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 36650629300..327b4c106de 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,5 +1,6 @@ import asyncio import json +import math from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass from types import MappingProxyType @@ -18,6 +19,8 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # leg pubsub_capable_client, ) from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET +from litellm.redis.stores.connection import connect +from litellm.rust_bridge._native import NativeAuthSubscription class SpendCounterCopies(Protocol): @@ -65,7 +68,7 @@ def _cache_invalidation_message_json(cache_key: str, new_value: float | None = N def _finite_number_or_none(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, (int, float)): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): return None return float(value) @@ -239,7 +242,14 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_spend_counters", "_task", "_user_api_key_cache") + __slots__ = ( + "_additional_in_memory_caches", + "_native", + "_redis_cache", + "_spend_counters", + "_task", + "_user_api_key_cache", + ) def __init__( self, @@ -253,12 +263,31 @@ class AuthCacheInvalidationSubscriber: self._additional_in_memory_caches = tuple(additional_in_memory_caches) self._spend_counters: Final = spend_counters self._task: asyncio.Task[None] | None = None + self._native: NativeAuthSubscription | None = None def start(self) -> None: if self._task is not None: return + shared: Final = self._user_api_key_cache.shared + client: Final = connect(self._redis_cache).client + if shared is not None and client is not None: + self._native = shared.subscribe(client, auth_cache_invalidation_channel(self._redis_cache)) + self._task = asyncio.create_task(self._consume_native(self._native)) + return self._task = asyncio.create_task(self._run()) + async def _consume_native(self, subscription: NativeAuthSubscription) -> None: + try: + while True: + if (data := await subscription.next()) is not None: + self._apply_message({"data": data}) + finally: + await subscription.close() + + @property + def is_ready(self) -> bool: + return self._native is not None and self._native.is_connected() + async def stop(self) -> None: task: Final = self._task if task is None: @@ -269,6 +298,7 @@ class AuthCacheInvalidationSubscriber: await task except asyncio.CancelledError: pass + self._native = None async def _run(self) -> None: backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index d58c58fa1a5..1adbc1b0620 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -19,10 +19,12 @@ from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import TypeAdapter, ValidationError +from litellm._auth.failure import NativeAuthFailure from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.repositories.native_auth_repository import native_auth_repository if TYPE_CHECKING: from prisma import models as prisma_models @@ -106,6 +108,8 @@ class RegistryReadThrough: return False try: found: Final = await self._resync(key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # a failed read-through must surface the original miss error, not a 500 verbose_proxy_logger.warning("registry read-through for %r failed: %s", key, e) return False @@ -280,7 +284,6 @@ async def _resync_agents(agent_id_or_name: str) -> bool: return True prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None - table: Final = agents_table(prisma_client) id_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id_or_name} name_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_name": agent_id_or_name} include_permission: Final[LiteLLM_AgentsTableInclude] = { @@ -290,8 +293,13 @@ async def _resync_agents(agent_id_or_name: str) -> bool: async with AGENT_RECONCILE_LOCK: if _agent_from_registry(agent_id_or_name) is not None: return True - row: Final = await table.find_unique(where=id_filter, include=include_permission) or await table.find_unique( - where=name_filter, include=include_permission + native: Final = native_auth_repository(prisma_client) + row: Final = ( + await native.one("agent", [agent_id_or_name], AgentResponse) + or await native.one("agent_name", [agent_id_or_name], AgentResponse) + if native is not None + else await agents_table(prisma_client).find_unique(where=id_filter, include=include_permission) + or await agents_table(prisma_client).find_unique(where=name_filter, include=include_permission) ) if row is None: return False diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 91529ee36fc..477f0518f41 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -68,9 +68,8 @@ class UserApiKeyCache(DualCache): User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition, ``key_object_cache``, so churn in the other management objects cannot evict them. - Without Redis both partitions are this class's Python memory. ``attach_redis_cache`` moves the whole - keyspace, both memory partitions included, onto the native ``AuthCacheStore`` (``shared``); redis-py never - reads or writes it, and ``redis_cache`` stays None. + Both partitions live in native Moka, including without Redis. ``attach_redis_cache`` adds the shared + Redis tier without creating another local authority; ``redis_cache`` stays None. ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. @@ -103,7 +102,7 @@ class UserApiKeyCache(DualCache): ) self._auth_cache: Final = auth_cache self._redis_connection: RedisCache | None = None - self._shared: AuthCacheStore | None = None + self._shared: AuthCacheStore | None = AuthCacheStore(connect(None).auth_cache(self._settings())) self.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) @property @@ -113,7 +112,7 @@ class UserApiKeyCache(DualCache): @property def has_redis(self) -> bool: - return self._shared is not None + return self._redis_connection is not None @property def redis_connection(self) -> RedisCache | None: @@ -150,12 +149,29 @@ class UserApiKeyCache(DualCache): def attach_redis_cache( self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None ) -> None: - if redis_cache is None or self._shared is not None: + if redis_cache is None or self._redis_connection is not None: return if default_redis_ttl is not None: self.update_cache_ttl(default_in_memory_ttl=None, default_redis_ttl=default_redis_ttl) self._redis_connection = redis_cache - self._shared = self._auth_cache(redis_cache, self._settings()) + if self._auth_cache is _connected_auth_cache and self._shared is not None: + client: Final = connect(redis_cache).client + if client is not None: + self._shared.attach(client) + else: + self._shared = self._auth_cache(redis_cache, self._settings()) + + def begin_fill(self, key: str) -> int: + assert self._shared is not None + return self._shared.begin_fill(key) + + def current_generation(self) -> int: + assert self._shared is not None + return self._shared.current_generation() + + def set_if_generation(self, key: str, value: CacheValue, generation: int, ttl: float | None = None) -> bool: + assert self._shared is not None + return self._shared.set_if_generation(key, CacheCodec.serialize(value, model_type=None), generation, ttl) def get_local(self, key: str) -> object: """This worker's memory copy alone.""" @@ -309,11 +325,24 @@ class UserApiKeyCache(DualCache): return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object): + async def async_set_cache( + self, + key: str | None, + value: CacheValue, + local_only: bool = False, + *, + fill_generation: int | None = None, + **kwargs: object, + ): """Inside a request the Redis SET rides the request's pipeline (memory is written at once); anywhere else, or with options the pipeline does not carry, it goes to Redis directly as before.""" model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + if fill_generation is not None: + if key is not None: + assert self._shared is not None + self._shared.set_if_generation(key, payload, fill_generation, _ttl(kwargs)) + return None ttl: Final = kwargs.get("ttl") pipelined: Final = ( key is not None diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4ed3ac30be7..c44d0e2d47f 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -30,6 +30,12 @@ from litellm.proxy.spend_tracking.spend_counter_batch import ( record_spend_counter_value, ) from litellm.redis.stores.spend_counters import END_USER_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendCounterStore +from litellm.repositories.native_auth_repository import ( + BudgetWindowRow, + NativeAuthDatabaseError, + SpendRow, + native_auth_repository, +) from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( @@ -134,6 +140,24 @@ class SpendCounterReseed: # or tag names that merely contain ":window:". if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + prefix, separator, suffix = counter_key.partition("spend:") + kind, delimiter, entity_id = suffix.partition(":") + if ( + prefix + or not separator + or not delimiter + or kind not in {"key", "team", "team_member", "user", "org", "project"} + ): + return None + if kind == "team_member" and ":" not in entity_id: + return None + try: + return await native.spend("member" if kind == "team_member" else kind, entity_id) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception("SpendCounterReseed.from_db: native read failed for %s", counter_key) + return None try: row: Final = await bounded_db_lookup( SpendCounterReseed._counter_row(prisma_client, counter_key), name="spend_counter" @@ -181,6 +205,15 @@ class SpendCounterReseed: async def end_user_from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: if prisma_client is None or not counter_key.startswith(END_USER_COUNTER_PREFIX): return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + try: + return await native.spend("end_user", counter_key[len(END_USER_COUNTER_PREFIX) :]) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception( + "SpendCounterReseed.end_user_from_db: native read failed for %s", counter_key + ) + return None where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]} try: row: Final = await bounded_db_lookup( @@ -296,6 +329,24 @@ class SpendCounterReseed: if row_entity_type is None: return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + try: + window: Final = await native.one( + "budget_window", [entity_type, entity_id, window_duration], BudgetWindowRow + ) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_table: native read failed for %s=%s window=%s", + entity_type, + entity_id, + window_duration, + ) + return None + if window is None or _as_utc(window.window_start) < _as_utc(expected_window_start): + return None + return float(window.spend or 0.0) + try: row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique( where={ @@ -367,6 +418,18 @@ class SpendCounterReseed: group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) if group_field is None: return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + try: + aggregate: Final = await native.one( + "window_spend", [entity_type, entity_id, window_start.isoformat()], SpendRow + ) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_spend_logs: native read failed for %s=%s", entity_type, entity_id + ) + return None + return float(aggregate.spend or 0.0) if aggregate is not None else 0.0 where: Final = { group_field: entity_id, "startTime": {"gte": window_start}, diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 4daef2592f7..f2317e49c9d 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -15,6 +15,8 @@ from typing import ( ) import litellm +from litellm._auth import model_budget_data +from litellm._auth.model_budget_data import ModelBudgetEntity from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -41,6 +43,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload if TYPE_CHECKING: + from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading from litellm.router import Router VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" @@ -923,21 +926,6 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ) - @with_service_target("model_budgets") - async def get_fallback_model_within_budget( - self, - user_api_key_dict: UserAPIKeyAuth, - model: str, - ) -> str | None: - budget_fallbacks: Final[dict[str, list[str]]] = user_api_key_dict.budget_fallbacks or {} - for fallback_model in budget_fallbacks.get(model, []): - try: - await self.is_key_within_model_budget(user_api_key_dict=user_api_key_dict, model=fallback_model) - return fallback_model - except litellm.BudgetExceededError: - continue - return None - @with_service_target("model_budgets") async def is_user_within_model_budget( self, @@ -1003,9 +991,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): entity_type=Litellm_EntityType.TEAM, entity_id=team_id, team_id=team_id, - model_max_budget=team_model_budget_entries( - model=model, team_model_max_budget=team_model_max_budget, key_model_max_budget=key_model_max_budget - ), + model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, model=model, exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", ) @@ -1029,6 +1016,27 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ) + async def read_model_budget_raw( + self, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> "ModelBudgetReading": + from litellm.proxy.auth.model_budget_adapter import ModelBudgetIO, read_model_budget + + io: Final = ModelBudgetIO( + self.dual_cache, + self._model_access_groups, + self._shared_windows, + _PROCESS_STARTED_AT, + clock=self._clock, + monotonic=time.monotonic, + ) + return await read_model_budget(io, entity, entity_id, team_id, model, budgets, key_budgets) + async def _is_entity_within_model_budget( self, entity_type: Litellm_EntityType, @@ -1037,65 +1045,32 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model_max_budget: Mapping[str, object] | None, model: str, exceeded_message: str, + key_model_max_budget: Mapping[str, object] | None = None, ) -> bool: - for candidate in expand_model_latest_alias(model): - await self._is_single_entity_within_model_budget( - entity_type=entity_type, - entity_id=entity_id, - team_id=team_id, - model_max_budget=model_max_budget, - model=candidate, - exceeded_message=exceeded_message, - ) - return True - - async def _is_single_entity_within_model_budget( - self, - entity_type: Litellm_EntityType, - entity_id: str | None, - team_id: str | None, - model_max_budget: Mapping[str, object] | None, - model: str, - exceeded_message: str, - ) -> bool: - if not model_max_budget: - return True - resolved: Final = resolve_model_budget( - model=model, - model_max_budget=model_max_budget, - model_access_groups=self._model_access_groups, - team_id=team_id, + entities: Final[Mapping[Litellm_EntityType, ModelBudgetEntity]] = { + Litellm_EntityType.KEY: "Key", + Litellm_EntityType.USER: "User", + Litellm_EntityType.TEAM: "Team", + Litellm_EntityType.TEAM_MEMBER: "TeamMember", + Litellm_EntityType.END_USER: "EndUser", + } + reading: Final = await self.read_model_budget_raw( + entities[entity_type], + entity_id, + team_id, + model, + model_max_budget, + key_model_max_budget, ) - if resolved is None: - verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value) - return True - - max_budget: Final = resolved.budget_config.max_budget - if max_budget is None or max_budget < 0: - return True - - local_spend: Final = await self._get_spend_for_model_budget( - entity_type=entity_type, - entity_id=entity_id, - model=model, - resolved=resolved, - ) - current_spend: Final = ( - local_spend - if local_spend >= max_budget - else await self._with_shared_spend( - entity_type=entity_type, entity_id=entity_id, resolved=resolved, local_spend=local_spend - ) - ) - if current_spend >= max_budget: + decision: Final = model_budget_data.check(reading.snapshot) + if decision.denied: + assert reading.snapshot.max_budget is not None raise litellm.BudgetExceededError( - message=( - exceeded_message - if resolved.budget_model in _budget_model_candidates(model) - else f"{exceeded_message}, budget_model={resolved.budget_model}" - ), - current_cost=current_spend, - max_budget=max_budget, + message=exceeded_message + if reading.named_model + else f"{exceeded_message}, budget_model={reading.budget_model}", + current_cost=decision.current_spend, + max_budget=reading.snapshot.max_budget, entity_type=entity_type.value, entity_id=entity_id, ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d54eda45160..e368a282240 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -6586,9 +6586,7 @@ async def regenerate_key_fn( # value in that field. regenerate_target_key: Final = data.key if data and data.key else key is_master_key_regeneration: Final = ( - data is not None - and data.new_master_key is not None - and is_master_key(api_key=regenerate_target_key, _master_key=master_key) + data is not None and data.new_master_key is not None and is_master_key(regenerate_target_key, master_key) ) if ( @@ -6616,7 +6614,7 @@ async def regenerate_key_fn( detail={"error": "DB not connected. prisma_client is None"}, ) - _is_master_key_valid: Final = is_master_key(api_key=key, _master_key=master_key) + _is_master_key_valid: Final = is_master_key(key, master_key) if master_key is not None and data and _is_master_key_valid: if data.new_master_key is None: diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index be24b1630b0..948d2613008 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -19,6 +19,7 @@ from starlette.types import Scope from typing_extensions import assert_never import litellm +from litellm._auth.runtime import litellm_header_token from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.constants import MCP_ALL_TOOLS_WILDCARD, MCP_PEEKED_BODY_SCOPE_KEY @@ -48,7 +49,6 @@ from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_age from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.oauth2_check import OAuth2TokenValidationError from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports - _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export oauth2_introspection_applies, run_centralized_common_checks, @@ -110,7 +110,7 @@ _CACHED_VALUE: Final = TypeAdapter(object) def _normalize_caller_admission_credential(value: str) -> str: - return _get_bearer_token_or_received_api_key(strip_auth_scheme(value, "Bearer")).strip() + return litellm_header_token(strip_auth_scheme(value, "Bearer")).strip() def _admission_request(scope: Scope) -> Request: @@ -797,7 +797,7 @@ class MCPRequestHandler: # for a delegated server, so validate it: identity / spend / rate # limits resolve and any stored upstream token can be forwarded. validated_user_api_key_auth = await user_api_key_auth( - api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}", + api_key=f"Bearer {litellm_header_token(litellm_api_key)}", request=request, ) elif MCPRequestHandler._target_servers_are_true_passthrough( @@ -1278,7 +1278,7 @@ class MCPRequestHandler: permissions, budgets, rate limits) and the sealed upstream token is injected under the server's per-server auth-header key, while the leak-defense chokepoint strips the envelope ``Authorization`` itself from egress.""" - presented_token: Final = _get_bearer_token_or_received_api_key(litellm_api_key) + presented_token: Final = litellm_header_token(litellm_api_key) explicit_auth: Final = await user_api_key_auth(api_key=f"Bearer {presented_token}", request=request) result: Final = await MCPRequestHandler._open_dcr_bridge_envelope( server=server, diff --git a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py index d321f5b5861..74739e67ea6 100644 --- a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py @@ -4,7 +4,6 @@ import hmac import json import re from collections.abc import Iterator, Mapping -from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol @@ -14,7 +13,7 @@ from fastapi.responses import StreamingResponse from pydantic import TypeAdapter, ValidationError from starlette.types import Message -import litellm +from litellm._auth.failure import NativeAuthFailure from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.exceptions import BadRequestError @@ -24,14 +23,12 @@ from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared client factory has untyped params ) -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import check_model_access_helper, check_tools_allowlist -from litellm.proxy.auth.handle_jwt import JWTAuthManager -from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import check_tools_allowlist +from litellm.proxy.auth.codex_turn_adapter import CodexAccess as _LitellmResponsesAccess from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.codex_model_catalog import ( CodexCatalogRow, - bundled_codex_models, codex_catalog_etag, codex_entry_names_default_reasoning_level, merged_codex_catalog, @@ -69,7 +66,6 @@ _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _INPUT_ITEMS: Final = TypeAdapter(tuple[object, ...]) _CACHED_SLUGS: Final = TypeAdapter(frozenset[str]) _CACHED_MODELS_ETAG: Final = TypeAdapter(dict[Literal["etag"], str]) -_MODEL_NAMES: Final = TypeAdapter(list[str]) _HEADER_DERIVED_STATE: Final = frozenset({"_cached_headers"}) _ACCOUNT_CATALOG_HEADERS: Final = frozenset({"authorization", "chatgpt-account-id", "originator", "user-agent"}) _CODEX_USER_AGENT_PRODUCT: Final = re.compile(r"([^/]+)/(\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.+-]+)?)") @@ -108,19 +104,6 @@ _HEADERS: Final = frozenset( ) -def _is_chatgpt(deployment: Deployment) -> bool: - return ( - deployment.litellm_params.custom_llm_provider == _CHATGPT_PROVIDER - or deployment.litellm_params.model.partition("/")[0] == _CHATGPT_PROVIDER - ) - - -def _serves_by_name(deployment: Deployment, names: frozenset[str]) -> bool: - return not _is_chatgpt(deployment) and ( - deployment.model_name in names or deployment.model_info.team_public_model_name in names - ) - - def _candidate_deployments(model: str, llm_router: "Router", team_id: str | None) -> tuple[Deployment, ...]: registry: Final = llm_router.registry() candidates: Final = tuple( @@ -130,37 +113,15 @@ def _candidate_deployments(model: str, llm_router: "Router", team_id: str | None def routes_to_litellm(model: object, llm_router: "Router | None", team_id: str | None) -> bool: - if ( - not isinstance(model, str) - or llm_router is None - or model in bundled_codex_models() - or f"{_CHATGPT_PROVIDER}/{model}" in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed - ): - return False - deployments: Final = _candidate_deployments(model, llm_router, team_id) - names: Final = frozenset({model, llm_router.registry().resolve(model, kind="alias") or model}) - return bool(deployments) and all(_serves_by_name(deployment, names) for deployment in deployments) + from litellm.proxy.auth.codex_model_adapter import routes_to_litellm as native - -def _key_and_team_allow(model: str, llm_router: "Router", user_api_key_dict: UserAPIKeyAuth) -> bool: - allowed_lists: Final = ( - _MODEL_NAMES.validate_python(user_api_key_dict.models), - _MODEL_NAMES.validate_python(user_api_key_dict.team_models), - ) - return all( - check_model_access_helper(model, llm_router, allowed, team_id=user_api_key_dict.team_id) - for allowed in allowed_lists - ) + return native(model, llm_router, team_id) def native_chatgpt_model(model: str, llm_router: "Router", user_api_key_dict: UserAPIKeyAuth) -> str | None: - if model.startswith(f"{_CHATGPT_PROVIDER}/") or _key_and_team_allow(model, llm_router, user_api_key_dict): - return None - chatgpt_model: Final = f"{_CHATGPT_PROVIDER}/{model}" - deployments: Final = _candidate_deployments(chatgpt_model, llm_router, user_api_key_dict.team_id) - if not deployments or not all(_is_chatgpt(deployment) for deployment in deployments): - return None - return chatgpt_model if _key_and_team_allow(chatgpt_model, llm_router, user_api_key_dict) else None + from litellm.proxy.auth.codex_model_adapter import native_chatgpt_model as native + + return native(model, llm_router, user_api_key_dict) class _SharedCache(Protocol): @@ -177,69 +138,10 @@ def codex_shared_cache() -> _SharedCache: return proxy_logging_obj.internal_usage_cache.dual_cache -def _key_may_call_litellm_responses(user_api_key_dict: UserAPIKeyAuth) -> bool: - if not RouteChecks.is_allowed_on_proxy(_LITELLM_RESPONSES_PATH): - return False - try: - _ = RouteChecks.should_call_route(route=_LITELLM_RESPONSES_PATH, valid_token=user_api_key_dict) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - RouteChecks.custom_admin_only_route_check(route=_LITELLM_RESPONSES_PATH) - except HTTPException: - return False - return True - - -async def _jwt_policy_admits_litellm_responses( - gateway_token: str, model: str, request_headers: Mapping[str, str] -) -> bool: - from litellm.proxy.proxy_server import ( - general_settings_view, - jwt_handler, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - try: - _ = await JWTAuthManager.auth_builder( - api_key=gateway_token, - jwt_handler=jwt_handler, - request_data={"model": model}, - general_settings=dict(general_settings_view()), - route=_LITELLM_RESPONSES_PATH, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=proxy_logging_obj, - request_headers=request_headers, - request_method="POST", - ) - except Exception: # noqa: BLE001 # any JWT policy refusal keeps the caller on ChatGPT's models - return False - return True - - -@dataclass(frozen=True, slots=True) -class _LitellmResponsesAccess: - key_admits: bool - jwt_token: str | None - request_headers: Mapping[str, str] - - async def admits(self, model: str) -> bool: - if not self.key_admits: - return False - if self.jwt_token is None: - return True - return await _jwt_policy_admits_litellm_responses(self.jwt_token, model, self.request_headers) - - def _litellm_responses_access(user_api_key_dict: UserAPIKeyAuth, request: Request) -> _LitellmResponsesAccess: - standard_jwt: Final = bool(user_api_key_dict.jwt_claims) and not user_api_key_dict.token - return _LitellmResponsesAccess( - key_admits=_key_may_call_litellm_responses(user_api_key_dict), - jwt_token=_gateway_token(request) if standard_jwt else None, - request_headers=request.headers, - ) + from litellm.proxy.auth.codex_turn_adapter import read_codex_access + + return read_codex_access(user_api_key_dict, request) def _account_slugs_cache_key(headers: Mapping[str, str]) -> str | None: @@ -324,12 +226,9 @@ async def _serves_through_litellm( headers: Mapping[str, str], cache: _SharedCache, ) -> bool: - if not isinstance(model, str) or not routes_to_litellm(model, llm_router, user_api_key_dict.team_id): - return False - if not await access.admits(model): - return False - account_slugs: Final = await _account_slugs(headers, caller, cache) - return account_slugs is not None and model not in account_slugs + from litellm.proxy.auth.codex_turn_adapter import serves_through_litellm + + return await serves_through_litellm(model, llm_router, user_api_key_dict, access, caller, headers, cache) def _gateway_token(request: Request) -> str: @@ -353,19 +252,9 @@ def _forwarded_headers(request: Request) -> Mapping[str, str]: async def native_turn_serves_through_litellm( model: str, llm_router: "Router", request: Request, user_api_key_dict: UserAPIKeyAuth ) -> bool: - chatgpt_bearer: Final = _chatgpt_bearer(request) - gateway_token: Final = _gateway_token(request) - if not chatgpt_bearer or not gateway_token or hmac.compare_digest(chatgpt_bearer.encode(), gateway_token.encode()): - return False - return await _serves_through_litellm( - model, - llm_router, - user_api_key_dict, - _litellm_responses_access(user_api_key_dict, request), - _models_etag_caller(user_api_key_dict, gateway_token), - _forwarded_headers(request), - codex_shared_cache(), - ) + from litellm.proxy.auth.codex_turn_adapter import native_turn_serves_through_litellm as native + + return await native(model, llm_router, request, user_api_key_dict) def _json_object(value: object) -> Mapping[str, object]: @@ -588,6 +477,8 @@ async def _litellm_catalog_rows( return () try: rows: Final = await codex_catalog_rows_for_key(user_api_key_dict) + except NativeAuthFailure: + raise except Exception as error: # noqa: BLE001 # a failed listing must never fail the ChatGPT answer it extends verbose_proxy_logger.warning( "Codex catalog keeps only ChatGPT's models: listing the key's LiteLLM models raised %s", diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ba8153ebb6d..0b12a940f73 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -31,6 +31,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider +from litellm._auth.runtime import scheme_credential from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( @@ -99,7 +100,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul _get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export api_key_header, custom_litellm_key_header, - get_bearer_token, + get_bearer_token, # pyright: ignore[reportUnusedImport] # backwards-compatible package export is_no_auth_dev_mode, user_api_key_auth, user_api_key_auth_accepted_websocket_for_model, @@ -2858,17 +2859,7 @@ _CREDENTIALLESS_VERTEX_MISSING_CREDENTIAL_DETAIL: Final = ( def _normalize_credential_value(value: str) -> str: - """Reduce a header value to the bare token, matching how ``user_api_key_auth`` - reads a caller's key. - - Reuses the auth module's ``_get_bearer_token`` so the caller-key comparison - strips exactly the schemes authentication accepts (``Bearer`` / ``bearer`` / - ``Basic`` / ``AWS4-HMAC-SHA256`` credential), rather than re-deriving a - narrower normalization here. ``_get_bearer_token`` returns ``""`` for a value - with no recognized scheme prefix, so a bare token (or a real Google - credential that carries no scheme) falls back to its own value. - """ - return get_bearer_token(value) or value + return scheme_credential(value) or value _VERTEX_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-goog-api-key"}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d21b6169133..0cdc4ed1575 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -46,6 +46,7 @@ from typing import ( Final, Literal, NamedTuple, + NoReturn, Optional, Protocol, TypeAlias, @@ -3378,7 +3379,7 @@ def _fail_closed_budget_enforcement() -> bool: return general_settings.get("fail_closed_budget_enforcement") is True -def _raise_budget_unverifiable(counter_key: str) -> None: +def _raise_budget_unverifiable(counter_key: str) -> NoReturn: verbose_proxy_logger.warning( "fail_closed_budget_enforcement: rejecting request — spend for %s could " "not be verified against Redis or the database", @@ -3436,46 +3437,28 @@ async def get_current_spend( and cached in-process for a few seconds, so a persistently stale counter drives at most one read per counter per window rather than one per request. """ + from functools import partial + + from litellm._auth.spend_counter_policy import Spend + from litellm.proxy.auth.spend_counter_adapter import SpendIO, resolve_spend + current, verified = await _read_spend_counter_estimate(counter_key=counter_key, fallback_spend=fallback_spend) - if fallback_authoritative: - verified = True - - if max_budget is None or current >= max_budget: - return current - - # Cheap staleness signal for primary counters: the counter reads below the - # spend this caller already knows about. Window counters have no such signal - # (fallback is 0), so they always re-check, bounded by the cache. Strict mode - # (fail_closed_budget_enforcement) always re-checks against the authoritative - # source too, so a counter that is stale-low at the same time as the caller's - # cached spend cannot slip through; the 5s cache keeps that bounded. - is_window: Final = window_start is not None - if fallback_spend > current or is_window or _fail_closed_budget_enforcement(): - authoritative: Final = await _authoritative_floor_spend( - counter_key=counter_key, - window_entity_type=window_entity_type, - window_entity_id=window_entity_id, - window_duration=window_duration, - window_start=window_start, - ) - if authoritative is not None: - verified = True - if authoritative > current: - await _repair_stale_spend_counter(counter_key=counter_key, db_spend=authoritative) - return authoritative - elif fallback_spend > current: - # nothing to read (tag counters, an end user without a row or a DB client, a - # failed read); fallback_spend is the authoritative recorded value loaded in auth. - return fallback_spend - - # Opt-in hard guarantee: when the spend backing this admit decision came - # only from a per-pod cache (Redis and DB both unreadable), reject rather - # than admit on an unverifiable budget. No-op unless the flag is set, so - # default behavior is unchanged. - if not verified and _fail_closed_budget_enforcement(): - _raise_budget_unverifiable(counter_key) - - return current + return await resolve_spend( + Spend(current, verified, fallback_spend, max_budget, window_start is not None, fallback_authoritative), + SpendIO( + read_strict=_fail_closed_budget_enforcement, + read_floor=partial( + _authoritative_floor_spend, + counter_key=counter_key, + window_entity_type=window_entity_type, + window_entity_id=window_entity_id, + window_duration=window_duration, + window_start=window_start, + ), + repair=lambda spend: _repair_stale_spend_counter(counter_key=counter_key, db_spend=spend), + unavailable=partial(_raise_budget_unverifiable, counter_key), + ), + ) @with_service_target(SPEND_COUNTERS_TARGET) @@ -8322,6 +8305,9 @@ class ProxyConfig: Returns: dict: router_settings, or None if no settings found """ + from litellm._auth.failure import NativeAuthFailure + from litellm.repositories.native_auth_repository import NativeAuthDatabaseError + # 1. Try key-level router_settings # user_api_key_dict is already the cached/authenticated key object — # no DB call needed. @@ -8344,6 +8330,10 @@ class ProxyConfig: team_settings: Final = self._parse_router_settings_value(getattr(team_obj, "router_settings", None)) if team_settings is not None: return team_settings + except NativeAuthDatabaseError: + pass + except NativeAuthFailure: + raise except Exception: # If team lookup fails, no team-level settings available pass diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8f0584bd0d0..9855ff77225 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,10 +2,9 @@ from __future__ import annotations import asyncio import json -import math import time from collections.abc import AsyncIterator, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType @@ -15,6 +14,9 @@ from fastapi import HTTPException, status from pydantic import ConfigDict, TypeAdapter import litellm +from litellm._auth import common_budgets, spend_counter_policy +from litellm._auth import reservation as reservation_policy +from litellm._auth.failure import NativeAuthRuntimeError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route @@ -29,6 +31,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded +from litellm.proxy.auth.common_budget_adapter import read_member_limit_facts from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, @@ -182,9 +185,16 @@ def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: requests never reach this branch so their concurrent-overspend protection is untouched. """ - if valid_token is None: + decision: Final = spend_counter_policy.decide( + spend_counter_policy.ThrottleKey( + counter_key, valid_token is not None, valid_token.token if valid_token is not None else None + ) + ) + if decision.kind == "Skip": return False - return counter_key == key_counter(str(valid_token.token)) and should_throttle_budget_exceeded(valid_token) + if decision.kind != "ReadThrottle" or valid_token is None: + raise NativeAuthRuntimeError("Unexpected throttle-counter decision") + return should_throttle_budget_exceeded(valid_token) async def _apply_over_budget_reservation_policy( @@ -204,26 +214,29 @@ async def _apply_over_budget_reservation_policy( down to what is left, unless strict enforcement is on, because the known estimate already does not fit; anything else hard-blocks by raising. """ - if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): - await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) - applied_entries.remove(entry) - return reservation_cost - - remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) - if remaining_before_reservation <= 1e-12: - _raise_counter_budget_exceeded(counter=counter, current_cost=current_spend) - if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: - _raise_counter_budget_exceeded( - counter=counter, - current_cost=current_spend - reservation_cost, - estimated_cost=reservation_cost, + decision: Final = spend_counter_policy.decide( + spend_counter_policy.OverBudget( + _key_reservation_should_release_for_throttle(counter.counter_key, valid_token), + current_spend, + counter.max_budget, + reservation_cost, + fail_closed_budget_enforcement, ) - await _resize_applied_reservation( - entries=applied_entries, - current_reserved_cost=reservation_cost, - new_reserved_cost=remaining_before_reservation, ) - return remaining_before_reservation + match decision: + case spend_counter_policy.Action(kind="Release"): + await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) + applied_entries.remove(entry) + return reservation_cost + case spend_counter_policy.Rejected(current=current, estimate=estimate): + _raise_counter_budget_exceeded(counter=counter, current_cost=current, estimated_cost=estimate) + case spend_counter_policy.Resize(cost=cost): + await _resize_applied_reservation( + entries=applied_entries, current_reserved_cost=reservation_cost, new_reserved_cost=cost + ) + return cost + case _: + raise NativeAuthRuntimeError("Unexpected reservation-overflow decision") def _raise_counter_budget_exceeded( @@ -263,11 +276,20 @@ async def reserve_budget_for_request( fail_closed_budget_enforcement: bool = False, raw_body: bytes | None = None, ) -> dict | None: - if valid_token is None or not RouteChecks.is_llm_api_route(route=route): + start: Final = reservation_policy.Gate(token_present=valid_token is not None) + if reservation_policy.action(start).kind == "Skip": return None - if RouteChecks.is_unbilled_route(route): + assert valid_token is not None + route_facts: Final = replace(start, llm_route=RouteChecks.is_llm_api_route(route=route)) + if reservation_policy.action(route_facts).kind == "Skip": return None - if get_model_from_request(request_body, route, llm_router=llm_router) is None: + billing_facts: Final = replace(route_facts, unbilled=RouteChecks.is_unbilled_route(route)) + if reservation_policy.action(billing_facts).kind == "Skip": + return None + model_facts: Final = replace( + billing_facts, model_present=get_model_from_request(request_body, route, llm_router=llm_router) is not None + ) + if reservation_policy.action(model_facts).kind == "Skip": return None counters: Final = await _get_budget_counters( @@ -282,7 +304,8 @@ async def reserve_budget_for_request( end_user_object=end_user_object, apply_user_budget_to_team_keys=apply_user_budget_to_team_keys, ) - if not counters: + counter_facts: Final = replace(model_facts, counter_count=len(counters)) + if reservation_policy.action(counter_facts).kind == "Skip": return None input_token_counts: Final = await count_request_input_tokens( @@ -302,8 +325,9 @@ async def reserve_budget_for_request( # estimate_request_max_cost still returns None when the model is unknown # to the cost map (no token-priced cost fields, e.g. image/audio routes). # In that case we fall back to read-time enforcement only. - if reservation_cost is None or reservation_cost <= 0: + if reservation_policy.action(replace(counter_facts, cost_read=True, cost=reservation_cost)).kind == "Skip": return None + assert reservation_cost is not None applied_entries: Final[list[dict[str, float | str]]] = [] try: @@ -494,9 +518,21 @@ async def _get_budget_counters( apply_user_budget_to_team_keys: bool = False, ) -> list[_BudgetCounter]: counters: Final[list[_BudgetCounter]] = [] + sources: Final = reservation_policy.sources( + reservation_policy.Sources( + key_present=valid_token.token is not None, + team_present=team_object is not None, + team_id=team_object.team_id if team_object is not None else None, + user_present=user_object is not None, + user_id=user_object.user_id if user_object is not None else None, + apply_user_budget_to_team_keys=apply_user_budget_to_team_keys, + ) + ) - if valid_token.token is not None: - if valid_token.max_budget is not None and valid_token.max_budget > 0: + if sources.key: + assert valid_token.token is not None + if reservation_policy.counter("Key", valid_token.max_budget): + assert valid_token.max_budget is not None counters.append( _BudgetCounter( counter_key=key_counter(valid_token.token), @@ -517,9 +553,11 @@ async def _get_budget_counters( ) ) - if team_object is not None and team_object.team_id is not None: + if sources.team: + assert team_object is not None and team_object.team_id is not None team_id: Final = team_object.team_id - if team_object.max_budget is not None and team_object.max_budget > 0: + if reservation_policy.counter("Team", team_object.max_budget): + assert team_object.max_budget is not None counters.append( _BudgetCounter( counter_key=team_counter(team_id), @@ -540,14 +578,8 @@ async def _get_budget_counters( ) ) - is_team_key: Final = team_object is not None and team_object.team_id is not None - if ( - (not is_team_key or apply_user_budget_to_team_keys) - and user_object is not None - and user_object.user_id is not None - and user_object.max_budget is not None - and user_object.max_budget > 0 - ): + if sources.user and reservation_policy.counter("User", user_object.max_budget if user_object is not None else None): + assert user_object is not None and user_object.user_id is not None and user_object.max_budget is not None counters.append( _BudgetCounter( counter_key=user_counter(user_object.user_id), @@ -616,7 +648,9 @@ async def _get_end_user_budget_counter( end_user_id: str | None, end_user_object: object, ) -> _BudgetCounter | None: - end_user_id = end_user_id or valid_token.end_user_id + end_user_id = reservation_policy.identity( + reservation_policy.Identity("EndUser", end_user_id, valid_token.end_user_id) + ) if end_user_id is None: return None @@ -629,8 +663,9 @@ async def _get_end_user_budget_counter( budget_table: Final = _get_value(end_user_object, "litellm_budget_table") max_budget = _to_float(_get_value(budget_table, "max_budget")) - if max_budget is None or max_budget <= 0: + if not reservation_policy.counter("EndUser", max_budget): return None + assert max_budget is not None return _BudgetCounter( counter_key=end_user_counter(end_user_id), @@ -669,8 +704,9 @@ async def _get_tag_budget_counters( continue budget_table = _get_value(tag_object, "litellm_budget_table") max_budget = _to_float(_get_value(budget_table, "max_budget")) - if max_budget is None or max_budget <= 0: + if not reservation_policy.counter("Tag", max_budget): continue + assert max_budget is not None counters.append( _BudgetCounter( counter_key=tag_counter(tag_name), @@ -712,8 +748,9 @@ async def _get_model_access_group_budget_counters( def _model_access_group_counter(group: str, budget: ModelAccessGroupBudget | None) -> _BudgetCounter | None: """A counter for one group, or nothing when the group carries no budget to reserve against.""" - if budget is None or budget.max_budget is None or budget.max_budget <= 0: + if not reservation_policy.counter("AccessGroup", budget.max_budget if budget is not None else None): return None + assert budget is not None and budget.max_budget is not None return _BudgetCounter( counter_key=model_access_group_counter(group), source_cache_key=model_access_group_cache_key(group), @@ -742,8 +779,19 @@ async def _get_team_member_budget_counter( user_object: LiteLLM_UserTable | None, user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: - if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: + if ( + reservation_policy.action( + reservation_policy.MemberGate( + team_object is not None, + team_object.team_id if team_object is not None else None, + user_object is not None, + valid_token.user_id, + ) + ).kind + == "Skip" + ): return None + assert team_object is not None and valid_token.user_id is not None membership_cache_key: Final = team_membership_reservation_cache_key( user_id=valid_token.user_id, team_id=team_object.team_id @@ -759,25 +807,33 @@ async def _get_team_member_budget_counter( member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None now: Final = datetime.now(timezone.utc) - team_member_budget: float | None = None - if member_budget_row is not None and member_budget_row.max_budget is not None: - team_member_budget = member_budget_row.effective_max_budget(now=now) - else: + default_cap: float | None = None + if ( + reservation_policy.action( + reservation_policy.MemberCap(member_budget_row.max_budget if member_budget_row is not None else None) + ).kind + == "ReadDefault" + ): default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): default_budget: Final = _CACHED_VALUE.validate_python( await user_api_key_cache.async_get_cache(key=f"team_member_default_budget:{default_budget_id}") ) - default_cap: Final = _to_float(_get_value(default_budget, "max_budget")) - if default_cap is not None and default_cap > 0: - team_member_budget = default_cap + ( - member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 - ) + default_cap = _to_float(_get_value(default_budget, "max_budget")) - if team_member_budget is None or team_member_budget <= 0: + team_member_budget: Final = common_budgets.member_limit( + read_member_limit_facts( + member_budget_row, + default_cap if reservation_policy.counter("Key", default_cap) else None, + now, + ) + ) + + if not reservation_policy.counter("Member", team_member_budget): return None + assert team_member_budget is not None - team_member_spend = cast(LiteLLM_TeamMembership, team_membership).spend if team_membership is not None else 0.0 + team_member_spend = team_membership.spend if team_membership is not None else 0.0 return _BudgetCounter( counter_key=team_member_counter(valid_token.user_id, team_object.team_id), source_cache_key=membership_cache_key, @@ -794,11 +850,11 @@ async def _get_org_budget_counter( team_object: LiteLLM_TeamTable | None, user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: - org_id: str | None = None - if valid_token.org_id is not None: - org_id = valid_token.org_id - elif team_object is not None and team_object.organization_id is not None: - org_id = team_object.organization_id + org_id: Final = reservation_policy.identity( + reservation_policy.Identity( + "Organization", valid_token.org_id, team_object.organization_id if team_object is not None else None + ) + ) if org_id is None: return None @@ -813,8 +869,9 @@ async def _get_org_budget_counter( return None org_max_budget: Final = _to_float(_get_value(org_budget_table, "max_budget")) - if org_max_budget is None or org_max_budget <= 0: + if not reservation_policy.counter("Organization", org_max_budget): return None + assert org_max_budget is not None org_spend: Final = _to_float(_get_value(org_table, "spend")) or 0.0 return _BudgetCounter( @@ -847,8 +904,9 @@ async def _get_project_budget_counter( return None project_max_budget: Final = _to_float(_get_value(project_budget_table, "max_budget")) - if project_max_budget is None or project_max_budget <= 0 or not math.isfinite(project_max_budget): + if not reservation_policy.counter("Project", project_max_budget): return None + assert project_max_budget is not None return _BudgetCounter( counter_key=project_counter(valid_token.project_id), @@ -875,10 +933,11 @@ def _get_budget_limit_counters( window_dict = _coerce_window(window) budget_duration = window_dict.get("budget_duration") max_budget = _to_float(window_dict.get("max_budget")) - if not budget_duration or max_budget is None or max_budget <= 0: + facts = reservation_policy.Window(bool(budget_duration), max_budget) + if reservation_policy.action(facts).kind == "Skip": continue window_start = get_budget_window_start(window_dict) - if window_start is None: + if reservation_policy.action(replace(facts, start_present=window_start is not None)).kind == "WarnWindow": verbose_proxy_logger.warning( "Skipping budget window with invalid duration for %s=%s: %s", entity_type, @@ -886,6 +945,7 @@ def _get_budget_limit_counters( budget_duration, ) continue + assert max_budget is not None and window_start is not None counters.append( _BudgetCounter( counter_key=f"{entity_prefix}:window:{budget_duration}", @@ -937,7 +997,10 @@ async def _reservation_counter_loaded(counter: _BudgetCounter, fail_closed_budge try: await _initialize_reservation_counter(counter=counter) except _CounterReservationUnavailable: - if fail_closed_budget_enforcement: + if ( + spend_counter_policy.decide(spend_counter_policy.ReservationFailure(fail_closed_budget_enforcement)).kind + == "Unavailable" + ): _raise_reservation_unavailable(counter_key=counter.counter_key) return False return True @@ -1013,7 +1076,12 @@ async def _reserve_reservable_counters( if reserved_values is None: for entry in entries: applied_entries.remove(entry) - if fail_closed_budget_enforcement: + if ( + spend_counter_policy.decide( + spend_counter_policy.ReservationFailure(fail_closed_budget_enforcement) + ).kind + == "Unavailable" + ): _raise_reservation_unavailable(counter_key=group[0].counter_key) continue for counter, entry, reserved_value in zip(group, entries, reserved_values): @@ -1023,7 +1091,10 @@ async def _reserve_reservable_counters( current_spend = reserved_value - (charged_cost - reservation_cost) else: current_spend = current_spend_by_counter_key[counter.counter_key] + reservation_cost - if current_spend > counter.max_budget: + if ( + spend_counter_policy.decide(spend_counter_policy.Overflow(current_spend, counter.max_budget)).kind + == "CheckOverflow" + ): reservation_cost = await _apply_over_budget_reservation_policy( counter=counter, valid_token=valid_token, @@ -1044,14 +1115,16 @@ def _reservation_groups( """Every counter the batch read says still has room for the estimate is charged in one pipeline. As soon as one does not, the counters are charged one at a time so the over-budget policy settles each before the next is touched, and a rejection charges nothing after it.""" - if not counters: - return () - if all( - current_spend_by_counter_key[counter.counter_key] + reservation_cost <= counter.max_budget - for counter in counters - ): - return (tuple(counters),) - return tuple((counter,) for counter in counters) + decision: Final = spend_counter_policy.decide( + spend_counter_policy.ReservationGroups( + tuple(current_spend_by_counter_key[counter.counter_key] for counter in counters), + tuple(counter.max_budget for counter in counters), + reservation_cost, + ) + ) + if not isinstance(decision, spend_counter_policy.Groups): + raise NativeAuthRuntimeError("Unexpected reservation grouping decision") + return tuple(tuple(counters[index] for index in group) for group in decision.groups) async def _reserve_counters( @@ -1163,7 +1236,12 @@ async def _set_reserved_entries_actual_cost( ) ) inconsistent: Final = tuple(item for item, ok in zip(adjustments, consistent) if not ok) - if inconsistent and not reseed_on_inconsistent: + if ( + spend_counter_policy.decide( + spend_counter_policy.ResizeConsistency(bool(inconsistent), reseed_on_inconsistent) + ).kind + == "RejectResize" + ): # Pre-call admission resize: the in-flight reservation cost is not yet # persisted, so the DB floor would discard it. Keep the original # fail-closed behavior (raise -> reserve_budget_for_request releases and @@ -1206,10 +1284,7 @@ async def _counter_can_apply_adjustment( current_value, _ = await read_spend_counter_cache_value(counter_key=counter_key) except (TypeError, ValueError): return False - if current_value is None: - return False - - return not (adjustment < 0 and current_value + adjustment < -1e-12) + return spend_counter_policy.decide(spend_counter_policy.Adjustment(current_value, adjustment)).kind == "Compatible" async def _release_applied_entries_best_effort( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 15b2d18eeae..a50067e00f2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -2,7 +2,6 @@ import ipaddress import json import os import re -import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt @@ -14,6 +13,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm +from litellm._auth.key import is_master_key as is_master_key # noqa: PLC0414 # public re-export from litellm._logging import verbose_proxy_logger from litellm.constants import ( BILLED_RETRY_ATTEMPT_METADATA_KEY, @@ -135,16 +135,6 @@ def _get_max_string_length_prompt_in_db() -> int: return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB -def is_master_key(api_key: str | None, _master_key: str | None) -> bool: - """ - Raw-only constant-time master-key comparison. The hashed form is never - considered equivalent — only the raw master-key string matches. - """ - if _master_key is None or api_key is None: - return False - return secrets.compare_digest(api_key, _master_key) - - _is_master_key: Final = is_master_key diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b9189f143f..e807749d7da 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -141,6 +141,7 @@ from litellm import ( ModelResponseStream, Router, ) +from litellm._auth.key import hash_token as hash_token # noqa: PLC0414 # public re-export from litellm._internal_context import service_target from litellm._logging import redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes @@ -331,7 +332,9 @@ if TYPE_CHECKING: from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord from litellm.proxy.db.model_usage_rollup import ModelUsageTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction + from litellm.repositories.database_settings import NativeDatabaseSettings from litellm.repositories.prisma_protocols import TableActions + from litellm.rust_bridge._native import NativeDatabase from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep @@ -5674,6 +5677,8 @@ class PrismaClient: http_client: "HttpConfig | None" = None, engine_exit_wait_seconds: float = _SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS, ): + self.native_auth_database: NativeDatabase | None = None + self._native_database_url: Final = database_url ## init logging object self.baseline_accounting_transactions: list[BaselineAccountingRecord] = [] self.baseline_accounting_lock: Final = asyncio.Lock() @@ -5773,6 +5778,7 @@ class PrismaClient: self.db = writer_wrapper else: self.db = writer_wrapper # Client to connect to Prisma db + self._native_reader_url: Final = read_replica_url if isinstance(self.db, RoutingPrismaWrapper) else None self._db_reconnect_lock = asyncio.Lock() self._closing: bool = False self._engine_exit_wait_seconds: Final = engine_exit_wait_seconds @@ -5841,6 +5847,23 @@ class PrismaClient: return self.db.read_target return self.db + @property + def native_writer_settings(self) -> "NativeDatabaseSettings": + from litellm.repositories.database_settings import NativeDatabaseSettings + + return NativeDatabaseSettings.from_url(self._native_database_url, self.writer_db.token_auth) + + def set_native_auth_database(self, native: "NativeDatabase | None") -> None: + self.native_auth_database = native + + @property + def native_reader_settings(self) -> "NativeDatabaseSettings | None": + from litellm.repositories.database_settings import NativeDatabaseSettings + + if self._native_reader_url is None or not isinstance(self.db, RoutingPrismaWrapper): + return None + return NativeDatabaseSettings.from_url(self._native_reader_url, self.db.reader.token_auth) + def tx( self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT, max_wait: timedelta = _PRISMA_DEFAULT_TX_MAX_WAIT ) -> "TransactionManager": @@ -5878,11 +5901,8 @@ class PrismaClient: # Default to success if metadata parsing fails return "success" - def hash_token(self, token: str): - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token + def hash_token(self, token: str) -> str: + return hash_token(token) def jsonify_object(self, data: Mapping[str, object]) -> dict[str, object]: db_data: Final[dict[str, object]] = copy.deepcopy(dict(data)) @@ -8457,15 +8477,6 @@ async def send_email( return True -def hash_token(token: str): - import hashlib - - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token - - def _is_pbkdf2_row(stored: str) -> bool: return re.fullmatch(r"pbkdf2:[A-Za-z0-9+/]{64}", stored) is not None @@ -8473,7 +8484,6 @@ def _is_pbkdf2_row(stored: str) -> bool: def hash_password(password: str) -> str: """Hash a password using scrypt with a random salt, or PBKDF2 in FIPS mode.""" import base64 - import hashlib import os if fips_mode_active(): @@ -8488,7 +8498,6 @@ def hash_password(password: str) -> str: def verify_password(password: str, stored: str) -> bool: """Verify a password against a stored hash. Supports scrypt and SHA256, or only PBKDF2 in FIPS mode.""" import base64 - import hashlib import secrets if fips_mode_active(): diff --git a/litellm/redis/stores/auth_cache/__init__.py b/litellm/redis/stores/auth_cache/__init__.py index 7891344022f..7e9f388be9c 100644 --- a/litellm/redis/stores/auth_cache/__init__.py +++ b/litellm/redis/stores/auth_cache/__init__.py @@ -10,7 +10,7 @@ from litellm.caching.redis_batch import Batch, active_request_redis_batches from litellm.redis.stores._route import route from litellm.redis.stores.auth_cache.handle import AuthCacheHandle from litellm.redis.stores.auth_cache.types import AUTH_CACHE_TARGET, AuthCacheSettings, WriteBack -from litellm.rust_bridge._native import LocalSizes +from litellm.rust_bridge._native import LocalSizes, NativeAuthSubscription, NativeRedisClient __all__ = [ "AUTH_CACHE_TARGET", @@ -36,6 +36,22 @@ class AuthCacheStore: def __init__(self, handle: AuthCacheHandle) -> None: self._handle: Final = handle + def attach(self, client: NativeRedisClient) -> bool: + return self._handle.attach(client) + + def begin_fill(self, key: str) -> int: + return self._handle.current_generation() + + def current_generation(self) -> int: + return self._handle.current_generation() + + def set_if_generation(self, key: str, value: object, generation: int, ttl: float | None = None) -> bool: + self._forget_miss(key) + return self._handle.set_if_generation(key, value, generation, _redis_ttl(ttl)) + + def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: + return self._handle.subscribe(client, channel) + def reconfigure(self, settings: AuthCacheSettings) -> None: """Both memory partitions keep what they hold, as ``update_cache_ttl`` and ``update_in_memory_max_size`` do.""" self._handle.reconfigure(memory_entries=settings.memory_entries, memory_ttl_seconds=settings.memory_ttl) diff --git a/litellm/redis/stores/auth_cache/handle.py b/litellm/redis/stores/auth_cache/handle.py index 9a2d70079d1..fc91ae2ef35 100644 --- a/litellm/redis/stores/auth_cache/handle.py +++ b/litellm/redis/stores/auth_cache/handle.py @@ -5,10 +5,14 @@ from typing import Protocol from litellm.caching.redis_batch import NativePending from litellm.redis.stores.auth_cache.types import WriteBack -from litellm.rust_bridge._native import LocalSizes, NativeSend +from litellm.rust_bridge._native import LocalSizes, NativeAuthSubscription, NativeRedisClient, NativeSend class AuthCacheHandle(Protocol): + def attach(self, client: NativeRedisClient) -> bool: ... + def current_generation(self) -> int: ... + def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: ... + def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: ... def get(self, key: str) -> Awaitable[object]: ... def get_sync(self, key: str) -> object: ... def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Awaitable[None]: ... diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index d60a2099477..ab6e677bcd4 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -1,9 +1,13 @@ import asyncio from collections.abc import AsyncGenerator, Coroutine from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace +from datetime import datetime, timezone from typing import TYPE_CHECKING, Final +from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.token_auth import parse_database_token_expiration +from litellm.repositories.database_settings import NativeDatabaseSettings from litellm.rust_bridge._native import NativeDatabase if TYPE_CHECKING: @@ -14,23 +18,89 @@ if TYPE_CHECKING: class Database: prisma: "PrismaClient" native: NativeDatabase + password_tasks: tuple[asyncio.Task[None], ...] = () @classmethod async def open(cls, prisma: "PrismaClient") -> "Database": try: - native: Final = await NativeDatabase.connect() + configured: Final = getattr(prisma, "native_writer_settings", None) + writer: Final = ( + await _initial_password(configured) if isinstance(configured, NativeDatabaseSettings) else None + ) + reader_configured: Final = getattr(prisma, "native_reader_settings", None) + reader: Final = ( + await _initial_password(reader_configured) + if isinstance(reader_configured, NativeDatabaseSettings) + else None + ) + native: Final = await NativeDatabase.connect(writer, reader) except BaseException: await prisma.aclose() raise - return cls(prisma=prisma, native=native) + prisma.set_native_auth_database(native if writer is not None else None) + return cls( + prisma=prisma, + native=native, + password_tasks=tuple( + asyncio.create_task(_refresh_password(native, settings, is_reader)) + for settings, is_reader in ((writer, False), (reader, True)) + if settings is not None and settings.password_provider is not None + ), + ) async def aclose(self) -> None: + self.prisma.set_native_auth_database(None) + for task in self.password_tasks: + task.cancel() + await asyncio.gather(*self.password_tasks, return_exceptions=True) try: await self.prisma.aclose() finally: await self.native.close() +async def _initial_password(settings: NativeDatabaseSettings) -> NativeDatabaseSettings: + if settings.password_provider is None: + return settings + return replace(settings, password=await asyncio.to_thread(settings.password_provider)) + + +def _refresh_delay(settings: NativeDatabaseSettings, password: str | None) -> float: + expires: Final = ( + parse_database_token_expiration(settings.token_auth, password) + if settings.token_auth is not None and password is not None + else None + ) + return ( + 600.0 + if expires is None + else max(10.0, (expires - datetime.now(timezone.utc).replace(tzinfo=None)).total_seconds() - 180.0) + ) + + +async def _refresh_password(native: NativeDatabase, settings: NativeDatabaseSettings, reader: bool) -> None: + if settings.password_provider is None: + return + await asyncio.sleep(_refresh_delay(settings, settings.password)) + while True: + await asyncio.sleep(await _rotate_password(native, settings, reader)) + + +async def _rotate_password(native: NativeDatabase, settings: NativeDatabaseSettings, reader: bool) -> float: + if settings.password_provider is None: + return 600.0 + try: + password: Final = await asyncio.to_thread(settings.password_provider) + if reader: + native.set_passwords(reader_password=password) + else: + native.set_passwords(writer_password=password) + except Exception: + verbose_proxy_logger.warning("Native database credential refresh failed") + return 30.0 + return _refresh_delay(settings, password) + + @asynccontextmanager async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]: task: Final = asyncio.create_task(work) diff --git a/litellm/repositories/database_settings.py b/litellm/repositories/database_settings.py new file mode 100644 index 00000000000..99f9a98ef07 --- /dev/null +++ b/litellm/repositories/database_settings.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass, field +from functools import partial +from typing import Final +from urllib.parse import parse_qsl, unquote, urlsplit + +from litellm.proxy.db.token_auth import DatabaseTokenAuth, mint_database_token, parse_iam_endpoint_from_url + + +class NativeDatabaseSettingsError(ValueError): + def __init__(self, detail: str) -> None: + super().__init__(f"Invalid native database settings. {detail}") + + +def _password(auth: DatabaseTokenAuth, url: str) -> str: + return unquote(mint_database_token(auth, parse_iam_endpoint_from_url(url))) + + +@dataclass(frozen=True, slots=True) +class NativeDatabaseSettings: + host: str + port: int + username: str + password: str | None = field(repr=False) + database: str + schema: str + ssl_mode: str + ssl_root_cert: str | None + max_connections: int + acquire_timeout: float + idle_timeout: float | None + statement_cache_capacity: int + connect_timeout: float + password_provider: Callable[[], str] | None = field(default=None, repr=False) + token_auth: DatabaseTokenAuth | None = field(default=None, repr=False) + + @classmethod + def from_url(cls, url: str, token_auth: DatabaseTokenAuth | None = None) -> NativeDatabaseSettings: + parsed: Final = urlsplit(url) + if parsed.scheme not in ("postgres", "postgresql"): + raise NativeDatabaseSettingsError("PostgreSQL is required") + params: Final = dict(parse_qsl(parsed.query, keep_blank_values=True)) + host: Final = params.get("host", parsed.hostname or "localhost") + ssl_mode: Final = params.get("sslmode", "prefer") + return cls( + host=host, + port=parsed.port or 5432, + username=unquote(parsed.username or "postgres"), + password=None if parsed.password is None else unquote(parsed.password), + database=unquote(parsed.path.lstrip("/") or "postgres"), + schema=params.get("schema", "public"), + ssl_mode="verify-full" if params.get("sslaccept") == "strict" and ssl_mode != "disable" else ssl_mode, + ssl_root_cert=params.get("sslrootcert", params.get("sslcert")), + max_connections=int(params.get("connection_limit", str((os.cpu_count() or 1) * 2 + 1))), + acquire_timeout=float(params.get("pool_timeout", "10")), + idle_timeout=float(params.get("max_idle_connection_lifetime", "300")) or None, + statement_cache_capacity=0 if params.get("pgbouncer", "false").lower() == "true" else 100, + connect_timeout=float(params.get("connect_timeout", "5")), + password_provider=None if token_auth is None else partial(_password, token_auth, url), + token_auth=token_auth, + ) diff --git a/litellm/repositories/native_auth_repository.py b/litellm/repositories/native_auth_repository.py new file mode 100644 index 00000000000..e3d9a2f4752 --- /dev/null +++ b/litellm/repositories/native_auth_repository.py @@ -0,0 +1,194 @@ +import time +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from inspect import getattr_static +from itertools import chain +from typing import Final, Protocol, TypeVar + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm.litellm_core_utils.core_helpers import require_attributes +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.team import Member +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded, bounded_db_lookup, db_lookup_gate +from litellm.types.proxy.auth.user_api_key_auth import LiteLLM_VerificationTokenView + +_Model = TypeVar("_Model", bound=BaseModel) # rebind-ok: TypeVar declarations cannot be annotated Final +_ROWS: Final = TypeAdapter(list[dict[str, object]]) +_SPEND: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) +_BATCH_ID_KINDS: Final = frozenset( + { + "key", + "user", + "team", + "org", + "project", + "end_user", + "budget", + "object_permission", + "access_group", + "tag", + "model_access_group", + "agent", + "agent_name", + "retired_agent", + "managed_vector_store", + "agent_access_group", + } +) +_BATCH_SIZE: Final = 1000 + + +class ModelAccessGroupRow(BaseModel): + model_config = ConfigDict(frozen=True) + access_group_name: str + spend: float | None = 0.0 + litellm_budget_table: LiteLLM_BudgetTable | None = None + + +class EndUserRegistryRow(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str + + +class TagRegistryRow(BaseModel): + model_config = ConfigDict(frozen=True) + tag_name: str + + +class AgentIdRow(BaseModel): + model_config = ConfigDict(frozen=True) + agent_id: str + + +class SpendRow(BaseModel): + model_config = ConfigDict(frozen=True) + spend: float | None = None + + +class BudgetWindowRow(SpendRow): + window_start: datetime + + +class NativeAuthDatabaseError(NativeAuthFailure): + pass + + +class NativeAuthDeadlineExceeded(NativeAuthDatabaseError, DBLookupDeadlineExceeded): + pass + + +class AuthReadHandle(Protocol): + async def auth_read( + self, + kind: str, + ids: list[str], # mutable-ok: PyO3 requires a Python list at this boundary + writer: bool = False, + ) -> object: ... + + async def auth_spend(self, kind: str, id: str) -> object: ... + + +class _NativeAuthClient(Protocol): + @property + def native_auth_database(self) -> AuthReadHandle | None: ... + + +@dataclass(frozen=True, slots=True) +class NativeAuthRepository: + database: AuthReadHandle + + async def key(self, token: str, *, writer: bool = False) -> LiteLLM_VerificationTokenView | None: + rows: Final = await self.read("key", [token], writer=writer) + if not rows: + return None + try: + row: Final = rows[0] + members: Final = _ROWS.validate_python(row.get("team_members_with_roles") or []) + matches: Final = tuple( + Member.model_validate(member) + for member in members + if member.get("user_id") is not None and member.get("user_id") == row.get("user_id") + ) + projection: Final = LiteLLM_VerificationTokenView.model_validate( + { + **row, + "team_models": row.get("team_models") if row.get("team_models") is not None else [], + "team_blocked": row.get("team_blocked") if row.get("team_blocked") is not None else False, + "team_member": matches[-1] if matches else None, + "last_refreshed_at": time.time(), + } + ) + return ( + projection.model_copy(update={"expires": projection.expires.isoformat()}) + if isinstance(projection.expires, datetime) + else projection + ) + except Exception as error: # noqa: BLE001 # preserve terminal native projection errors + raise NativeAuthRuntimeError("Invalid native auth key projection") from error + + async def read(self, kind: str, ids: Sequence[str], writer: bool = False) -> Sequence[Mapping[str, object]]: + try: + value: Final = await bounded_db_lookup(self._read(kind, ids, writer), name=kind) + return _ROWS.validate_python(value) + except NativeAuthFailure: + raise + except DBLookupDeadlineExceeded as error: + raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error + except OSError as error: + raise NativeAuthDatabaseError(f"Native auth {kind} database read failed") from error + except Exception as error: # noqa: BLE001 # native errors must not enter legacy fallback paths + raise NativeAuthRuntimeError(f"Native auth {kind} read failed") from error + + async def _read(self, kind: str, ids: Sequence[str], writer: bool) -> object: + async with db_lookup_gate.current(): + return await self.database.auth_read(kind, list(ids), writer=writer) + + async def one(self, kind: str, ids: Sequence[str], model: type[_Model], *, writer: bool = False) -> _Model | None: + rows: Final = await self.many(kind, ids, model, writer=writer) + return rows[0] if rows else None + + async def many( + self, kind: str, ids: Sequence[str], model: type[_Model], *, writer: bool = False + ) -> Sequence[_Model]: + rows: Final = await self._many_rows(kind, ids, writer=writer) + try: + return tuple(model.model_validate(row) for row in rows) + except Exception as error: # noqa: BLE001 # invalid native projections are terminal + raise NativeAuthRuntimeError(f"Invalid native auth {kind} projection") from error + + async def _many_rows(self, kind: str, ids: Sequence[str], *, writer: bool) -> Sequence[Mapping[str, object]]: + if kind not in _BATCH_ID_KINDS or len(ids) <= _BATCH_SIZE: + return await self.read(kind, ids, writer=writer) + unique_ids: Final = tuple(dict.fromkeys(ids)) + chunks: Final = [ + await self.read(kind, unique_ids[start : start + _BATCH_SIZE], writer=writer) + for start in range(0, len(unique_ids), _BATCH_SIZE) + ] + return tuple(chain.from_iterable(chunks)) + + async def spend(self, kind: str, entity_id: str) -> float | None: + try: + value: Final = await bounded_db_lookup(self._spend(kind, entity_id), name="spend_counter") + return _SPEND.validate_python(value) + except NativeAuthFailure: + raise + except DBLookupDeadlineExceeded as error: + raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error + except OSError as error: + raise NativeAuthDatabaseError(f"Native auth {kind} database spend read failed") from error + except Exception as error: # noqa: BLE001 # a failed native read is not an authoritative miss + raise NativeAuthRuntimeError(f"Native auth {kind} spend read failed") from error + + async def _spend(self, kind: str, entity_id: str) -> object: + async with db_lookup_gate.current(): + return await self.database.auth_spend(kind, entity_id) + + +def native_auth_repository(client: object) -> NativeAuthRepository | None: + if getattr_static(client, "native_auth_database", None) is None: + return None + database: Final = require_attributes(client, _NativeAuthClient, "native_auth_database").native_auth_database + return NativeAuthRepository(database) if database is not None else None diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 77af87f5eb3..4e7c6449b18 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -7,6 +7,7 @@ from pydantic import JsonValue from typing_extensions import disjoint_base from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings +from litellm.repositories.database_settings import NativeDatabaseSettings from litellm.rust_bridge.execution import Execution from litellm.rust_bridge.public_call import NativeCall @@ -22,6 +23,80 @@ def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... +@final +class NativeAuthContinuation: + def resume(self, answer: object) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... + +def request_auth_start( + facts: object, settings: object, now: object, resolved: bool +) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... +def request_auth_budgets(entry: str) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... +def request_auth_access(facts: object) -> dict[str, object]: ... +def request_auth_agent_access(facts: object) -> str: ... +def request_auth_agent_grants(facts: object) -> object: ... +def request_auth_context(facts: object) -> object: ... +def request_auth_reservation(facts: object) -> object: ... +def request_auth_spend_counter_policy(facts: object) -> object: ... +def request_auth_common_flow(facts: object) -> object: ... +def request_auth_common_read(facts: object) -> object: ... +def request_auth_common_budget_scopes(facts: object) -> object: ... +def request_auth_custom_flow(facts: object) -> object: ... +def request_auth_expired(expires: object, now: object) -> bool: ... +def request_auth_agent_invocation(facts: object) -> dict[str, object]: ... +def request_auth_admission(facts: object) -> dict[str, object]: ... +def request_auth_managed(facts: object) -> dict[str, object]: ... +def request_auth_common_budget(facts: object) -> dict[str, object]: ... +def request_auth_budget_alert(facts: object) -> bool: ... +def request_auth_budget_window_start(facts: object) -> float | None: ... +def request_auth_member_limit(facts: object) -> float | None: ... +def request_auth_budget_price(facts: object) -> dict[str, object]: ... +def request_auth_budget_route(facts: object) -> str: ... +def request_auth_model_budget_resolve(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_query(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_check(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_aliases(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_scope(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_team_exclusions(facts: object) -> object: ... +def request_auth_select_api_key(keys: object) -> dict[str, str] | None: ... +def request_auth_select_websocket_api_key(keys: object) -> dict[str, str]: ... +def request_auth_scheme_credential(value: str) -> str | None: ... +def request_auth_litellm_header_token(value: str) -> str: ... +def request_auth_route_access(request: object) -> object: ... +def request_auth_org_route(facts: object) -> object: ... +def request_auth_org_admin(memberships: object, organizations: tuple[str, ...]) -> bool: ... +def request_auth_model_access(request: object) -> object: ... +def request_auth_model_groups(facts: object) -> object: ... +def request_auth_fallback_access(facts: object) -> object: ... +def request_auth_key_model_access(facts: object) -> object: ... +def request_auth_unpriced_model(facts: object) -> object: ... +def request_auth_model_attribution(facts: object) -> object: ... +def request_auth_fallback_candidate(facts: object) -> object: ... +def request_auth_personal_budget(facts: object) -> object: ... +def request_auth_personal_models(facts: object) -> object: ... +def request_auth_owner_teams(facts: object) -> object: ... +def request_auth_team_override(facts: object) -> object: ... +def request_auth_model_helpers(facts: object) -> object: ... +def request_auth_codex_model(facts: object) -> object: ... +def request_auth_codex_normalization(facts: object) -> object: ... +def request_auth_claude_membership(facts: object) -> object: ... +def request_auth_claude_normalization(facts: object) -> object: ... +def request_auth_codex_turn(facts: object) -> object: ... +def request_auth_route_scope(facts: object) -> object: ... +def request_auth_personal_owner(facts: object) -> object: ... +def request_auth_personal_model_request(facts: object) -> object: ... +def request_auth_member_models(facts: object) -> object: ... +def request_auth_resolved_models(facts: object) -> object: ... +def request_auth_personal_model_listing(facts: object) -> object: ... +def request_auth_agent_ceiling(facts: object) -> object: ... +def request_auth_agent_group_load(facts: object) -> object: ... +def request_auth_team_aliases(facts: object) -> object: ... +def request_auth_managed_models(facts: object) -> object: ... +def request_auth_agent_ceiling_source(facts: object) -> object: ... +def request_auth_managed_model_grants(facts: object) -> object: ... +def request_auth_model_projection(facts: object, reads: object) -> object: ... +def request_auth_route_enable(facts: object) -> object: ... +def request_auth_llm_route(context: object, configured: bool | None = None) -> object: ... + @final class NativeClickHouseSpendConfig: def __new__(cls, database: str, url: str, retention_days: int) -> NativeClickHouseSpendConfig: ... @@ -133,9 +208,7 @@ class NativeRouterUsage: def attach(self, client: NativeRedisClient) -> bool: ... @property def has_redis(self) -> bool: ... - def increment( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> Future[float]: ... + def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float]: ... def increment_deployment( self, deployment_id: str, @@ -144,12 +217,8 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> Future[None]: ... - def get_many( - self, counters: Sequence[_UsageCounter] - ) -> Future[list[float | None]]: ... - def increment_sync( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> float: ... + def get_many(self, counters: Sequence[_UsageCounter]) -> Future[list[float | None]]: ... + def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float: ... def increment_deployment_sync( self, deployment_id: str, @@ -158,9 +227,7 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> None: ... - def get_many_sync( - self, counters: Sequence[_UsageCounter] - ) -> list[float | None]: ... + def get_many_sync(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... @@ -331,12 +398,18 @@ class NativeAuthCache: ) -> NativeAuthCache: ... @property def has_redis(self) -> bool: ... + def attach(self, client: NativeRedisClient) -> bool: ... + def current_generation(self) -> int: ... + def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: ... + def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: ... def get(self, key: str, request: NativeRequestRedis | None = None) -> Future[object]: ... def reconfigure(self, *, memory_entries: int, memory_ttl_seconds: float | None = None) -> None: ... def get_sync(self, key: str) -> object: ... def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ... def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... - def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ... + def declare_set( + self, key: str, value: object, ttl_seconds: float | None, send: NativeSend + ) -> NativePending[None]: ... def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ... def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ... def write_back(self, entries: Sequence[tuple[str, object, float | None]], send: NativeSend) -> Future[None]: ... @@ -350,6 +423,12 @@ class NativeAuthCache: def flush_local(self) -> None: ... def len_local(self) -> LocalSizes: ... +@final +class NativeAuthSubscription: + def is_connected(self) -> bool: ... + def next(self) -> Future[bytes | None]: ... + def close(self) -> Future[None]: ... + @final class SpendRead: def __new__(cls, value: float | None, from_redis: bool) -> SpendRead: ... @@ -390,7 +469,15 @@ class NativeSpendCounters: @final class NativeDatabase: @staticmethod - def connect() -> Future[NativeDatabase]: ... + def connect( + writer_settings: NativeDatabaseSettings | None = None, + reader_settings: NativeDatabaseSettings | None = None, + ) -> Future[NativeDatabase]: ... + def auth_read(self, kind: str, ids: list[str], writer: bool = False) -> Future[list[dict[str, JsonValue]]]: ... + def auth_spend(self, kind: str, id: str) -> Future[float | None]: ... + @property + def reader_available(self) -> bool: ... + def set_passwords(self, writer_password: str | None = None, reader_password: str | None = None) -> None: ... def close(self) -> Future[None]: ... @final @@ -1074,7 +1161,9 @@ class RouterCoreHandle: callbacks: Mapping[str, bool] | None = None, ) -> RouterCall[_SuccessReportSteps]: ... # Count a success against any id, driven like `report_failure`; `done` carries `key`. - def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ... + def count_success( + self, id: str, callbacks: Mapping[str, bool] | None = None + ) -> RouterCall[_SuccessReportSteps]: ... # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... # Every deployment as a `model_list` entry, in config order. @@ -1171,3 +1260,31 @@ class RouterCall(Generic[_S]): # candidate ids, pick first; pre_route: see _PreRouteStep; cache: see _CacheStep; secrets, build and model_info: # see those steps; first call: None. def next(self, outcome: object) -> _S: ... + +def request_auth_budget_exemption(facts: object) -> dict[str, object]: ... +def request_auth_budget_alert_plan(facts: object) -> dict[str, object]: ... +def request_auth_context_models(facts: object) -> object: ... +def request_auth_context_end_user(facts: object) -> object: ... +def request_auth_context_token(facts: object) -> object: ... +def request_auth_context_member(facts: object) -> object: ... +def request_auth_context_temporary(facts: object) -> object: ... +def request_auth_context_throttle(facts: object) -> object: ... +def request_auth_context_percentage(facts: object) -> float | None: ... +def request_auth_end_user_normalize_id(facts: object) -> object: ... +def request_auth_end_user_default_id(facts: object) -> object: ... +def request_auth_end_user_apply(facts: object) -> object: ... +def request_auth_end_user_registry(facts: object) -> object: ... +def request_auth_end_user_assignments(facts: object) -> object: ... +def request_auth_end_user_lookup(facts: object) -> object: ... +def request_auth_end_user_validation(facts: object) -> object: ... +def request_auth_budget_lookup_organization(facts: object) -> object: ... +def request_auth_budget_lookup_member(facts: object) -> object: ... +def request_auth_budget_lookup_emails(facts: object) -> object: ... +def request_auth_budget_lookup_soft(facts: object) -> bool: ... +def request_auth_route_role(facts: object) -> object: ... +def request_auth_route_role_api(facts: object) -> object: ... +def request_auth_route_role_jwt(facts: object) -> object: ... +def request_auth_route_role_inside(facts: object) -> object: ... +def request_auth_route_role_expand(facts: object) -> object: ... +def request_auth_route_role_listed(facts: object) -> object: ... +def request_auth_org_admin_raw(facts: object) -> object: ... diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 1183c36eda7..92067be067a 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -9,6 +9,7 @@ from typing import ( from pydantic import ConfigDict, Field, PrivateAttr, model_validator from typing_extensions import TypeIs # noqa: TID251 # narrows the before-validator input without copying it +from litellm._auth.key import hash_token as hash_token # noqa: PLC0414 # public re-export from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.team import Member from litellm.models.verification_token import LiteLLM_VerificationToken @@ -37,15 +38,6 @@ def is_jwt(token: str | None) -> bool: return len(parts) == 3 -def hash_token(token: str) -> str: - import hashlib - - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token - - _UntypedDict: TypeAlias = dict[Any, Any] diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 0ac1e1d2521..2f27bd31b71 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -85,12 +85,14 @@ class HeldStatementRelay: *, after: bytes | None = None, armed: bool = True, + response: bool = False, ) -> None: self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._trigger: Final = trigger self._after: Final = after + self._response: Final = response self._marker_length: Final = max(len(trigger), len(after or b"")) self._server: Final = LoopbackServer(self.port, self._serve) self._loop: Final = self._server.loop @@ -144,8 +146,8 @@ class HeldStatementRelay: await close_writer(writer) await asyncio.gather( - forward(client_reader, server_writer, True), - forward(server_reader, client_writer, False), + forward(client_reader, server_writer, not self._response), + forward(server_reader, client_writer, self._response), ) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 262114ee32b..327264ba2fa 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -1,18 +1,18 @@ import json -import threading import os import re import signal import subprocess import sys +import threading import time import uuid from collections.abc import Callable, Generator, Iterable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor -from contextlib import AbstractContextManager, contextmanager, suppress +from contextlib import AbstractContextManager, ExitStack, contextmanager, suppress from dataclasses import dataclass -from enum import Enum from datetime import UTC, datetime, timedelta +from enum import Enum from ipaddress import IPv4Address from itertools import accumulate from pathlib import Path @@ -417,7 +417,7 @@ def kill_stray_process(pid: int) -> None: @contextmanager -def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str]) -> Iterator[None]: +def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str] | psutil.Process) -> Iterator[None]: target: Final = psutil.Process(process.pid) _suspend(target) try: @@ -427,6 +427,20 @@ def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str]) -> Iterator target.resume() +@contextmanager +def paused_prisma_engines(worker_pid: int) -> Iterator[None]: + engines: Final = tuple( + child + for child in psutil.Process(worker_pid).children(recursive=True) + if "query-engine" in child.name() + ) + assert engines, f"No Prisma query engine under worker {worker_pid}" + with ExitStack() as resources: + for engine in engines: + resources.enter_context(paused(engine)) + yield + + @contextmanager def owned_proxy( gateway: Gateway, diff --git a/tests/integration/authorization/test_native_auth_io.py b/tests/integration/authorization/test_native_auth_io.py new file mode 100644 index 00000000000..ac69fbcb14c --- /dev/null +++ b/tests/integration/authorization/test_native_auth_io.py @@ -0,0 +1,149 @@ +import os +import uuid +from hashlib import sha256 +from typing import Final + +import httpx +import psycopg +import pytest +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.process import paused_prisma_engines + +_OBSERVATIONS: Final = TypeAdapter(dict[str, list[dict[str, JsonValue]]]) + + +def _worker_pid(gateway: Gateway) -> int: + response: Final = gateway.request("GET", "/health/integration-readiness") + assert response.status_code == 200, response.text + return int(response.headers["x-integration-worker-pid"]) + + +def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"native auth storage {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _observations(gateway: Gateway) -> list[dict[str, JsonValue]]: + response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=5, trust_env=False) + response.raise_for_status() + return _OBSERVATIONS.validate_json(response.content)["requests"] + + +def test_cold_and_warm_context_reads_do_not_need_prisma(gateway: Gateway, peer: Gateway) -> None: + primary_pid: Final = _worker_pid(gateway) + peer_pid: Final = _worker_pid(peer) + assert primary_pid != peer_pid + with gateway.scenario() as scenario: + model: Final = scenario.model() + team: Final = scenario.team(models=[model]) + owner: Final = scenario.member(team) + project: Final = scenario.project(team, models=[model]) + key: Final = scenario.key(team_id=team, user_id=owner, project_id=project, models=[model]) + with paused_prisma_engines(peer_pid): + cold: Final = _chat(peer, model, key) + assert cold.status_code == 200, cold.text + warm: Final = _chat(peer, model, key) + assert warm.status_code == 200, warm.text + assert len(_observations(gateway)) == 2 + + +def test_warm_key_read_does_not_wait_for_the_key_table(gateway: Gateway, peer: Gateway) -> None: + database_url: Final = gateway.database_url or os.environ["DATABASE_URL"] + with gateway.scenario() as scenario: + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + assert _chat(peer, model, key).status_code == 200 + with psycopg.connect(database_url) as connection: + connection.execute('LOCK TABLE "LiteLLM_VerificationToken" IN ACCESS EXCLUSIVE MODE') + warm: Final = _chat(peer, model, key) + assert warm.status_code == 200, warm.text + assert len(_observations(gateway)) == 2 + + +def test_revocation_reaches_a_warmed_peer_without_prisma(gateway: Gateway, peer: Gateway) -> None: + peer_pid: Final = _worker_pid(peer) + assert _worker_pid(gateway) != peer_pid + with gateway.scenario() as scenario: + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + for worker in (gateway, peer): + response: Final = _chat(worker, model, key) + assert response.status_code == 200, response.text + with paused_prisma_engines(peer_pid): + gateway.post("/key/block", {"key": key}) + denied: Final = eventually( + lambda: _chat(peer, model, key), lambda response: response.status_code == 401, seconds=10 + ) + before: Final = _observations(gateway) + assert len(before) >= 2 + assert _chat(peer, model, key).status_code == 401 + assert _observations(gateway) == before + assert denied.json()["error"]["type"] == "auth_error", denied.text + rows: Final = gateway.rows( + 'SELECT blocked FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ) + assert rows == [{"blocked": True}] + + +@pytest.mark.parametrize("context", ("team", "project", "user")) +def test_context_changes_reach_a_warmed_peer(gateway: Gateway, peer: Gateway, context: str) -> None: + peer_pid: Final = _worker_pid(peer) + with gateway.scenario() as scenario: + model: Final = scenario.model() + team: Final = scenario.team(models=[model]) + owner: Final = scenario.member(team) + project: Final = scenario.project(team, models=[model]) + key: Final = scenario.key(team_id=team, user_id=owner, project_id=project, models=[model]) + warm: Final = _chat(peer, model, key) + assert warm.status_code == 200, warm.text + with paused_prisma_engines(peer_pid): + if context == "team": + gateway.post("/team/update", {"team_id": team, "blocked": True}) + elif context == "project": + gateway.post("/project/update", {"project_id": project, "blocked": True}) + else: + gateway.post("/user/update", {"user_id": owner, "metadata": {"scim_active": False}}) + denied: Final = eventually( + lambda: _chat(peer, model, key), lambda response: response.status_code in (401, 422), seconds=10 + ) + before: Final = _observations(gateway) + assert before + repeated: Final = _chat(peer, model, key) + assert repeated.status_code == denied.status_code, repeated.text + assert _observations(gateway) == before + + +@pytest.mark.parametrize("source", ("redis", "database")) +def test_spend_reads_share_native_counters_and_reseed_without_prisma( + gateway: Gateway, peer: Gateway, source: str +) -> None: + database_url: Final = gateway.database_url or os.environ["DATABASE_URL"] + peer_pid: Final = _worker_pid(peer) + with gateway.scenario() as scenario: + model: Final = scenario.model() + control: Final = scenario.key(models=[model]) + allowed: Final = _chat(peer, model, control) + assert allowed.status_code == 200, allowed.text + key: Final = scenario.key(models=[model], max_budget=1) + hashed: Final = sha256(key.encode()).hexdigest() + counter: Final = f"spend:key:{hashed}" + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), socket_timeout=2) as cache: + if source == "redis": + cache.set(counter, 2, ex=60) + else: + with psycopg.connect(database_url) as connection: + connection.execute('UPDATE "LiteLLM_VerificationToken" SET spend=2 WHERE token=%s', (hashed,)) + cache.delete(counter) + with paused_prisma_engines(peer_pid): + denied: Final = _chat(peer, model, key) + assert denied.status_code == 422, denied.text + value: Final = cache.get(counter) + assert value is not None and float(value) == 2 + assert len(_observations(gateway)) == 1 diff --git a/tests/integration/caching/test_auth_cache_native.py b/tests/integration/caching/test_auth_cache_native.py new file mode 100644 index 00000000000..73f5a7a4bf5 --- /dev/null +++ b/tests/integration/caching/test_auth_cache_native.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +import uuid +from collections.abc import Callable +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from typing import Final + +import psycopg +import pytest +from redis.asyncio import Redis + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.auth.auth_checks import get_team_object +from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects, prefetch_identity_keys +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AuthCacheInvalidationSubscriber, + auth_cache_invalidation_channel, + evict_local, + evict_shared, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.redis.stores.connection import connect +from litellm.redis.stores.spend_counters import SpendCounterStore +from litellm.repositories.database_settings import NativeDatabaseSettings +from litellm.rust_bridge._native import NativeDatabase +from tests.integration._support.database_relay import HeldStatementRelay +from tests.integration._support.redis_process import owned_redis + +KEY: Final = hashlib.sha256(b"native-cache-test-key").hexdigest() + + +async def until(predicate: Callable[[], bool], seconds: float = 12) -> None: + async with asyncio.timeout(seconds): + while not predicate(): + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace", (None, "phase3")) +@pytest.mark.parametrize("shared", (False, True)) +async def test_native_subscriber_invalidates_and_recovers_without_serving_outage_fills( + tmp_path: Path, namespace: str | None, shared: bool +) -> None: + with owned_redis(tmp_path) as server: + redis_cache: Final = RedisCache( + host=server.host, + port=server.port, + namespace=namespace, + socket_timeout=0.3, + socket_connect_timeout=0.3, + ) + cache: Final = UserApiKeyCache(redis_cache=redis_cache if shared else None) + additional: Final = InMemoryCache() + spend: Final = SpendCounterStore(connect(None).spend_counters()) + subscriber: Final = AuthCacheInvalidationSubscriber(redis_cache, cache, (additional,), spend) + publisher: Final = Redis(host=server.host, port=server.port) + channel: Final = auth_cache_invalidation_channel(redis_cache) + subscriber.start() + started: Final = cache.current_generation() + try: + assert not subscriber.is_ready + await until(lambda: subscriber.is_ready) + assert cache.current_generation() > started + fill: Final = cache.begin_fill(KEY) + assert cache.set_if_generation(KEY, {"token": KEY, "models": ["allowed"]}, fill) + assert cache.get_local(KEY) == {"token": KEY, "models": ["allowed"]} + spend.set_local("spend:key:k-1", 99.0) + assert ( + await publisher.publish( + channel, json.dumps({"cache_key": "spend:key:k-1", "new_value": 0.0, "ttl": 60}) + ) + == 1 + ) + await until(lambda: spend.get_local("spend:key:k-1") == 0.0) + assert cache.get_local(KEY) == {"token": KEY, "models": ["allowed"]} + assert ( + await publisher.publish( + channel, json.dumps({"cache_key": "spend_db_floor:spend:key:k-1", "new_value": 3.0}) + ) + == 1 + ) + await until(lambda: additional.get_cache("spend_db_floor:spend:key:k-1") == 3.0) + assert await publisher.publish(channel, json.dumps({"cache_key": "spend:key:k-1"})) == 1 + await until(lambda: spend.get_local("spend:key:k-1") is None) + additional.set_cache("model-cache", "stale") + assert await publisher.publish(channel, json.dumps({"cache_key": KEY})) == 1 + await until(lambda: cache.get_local(KEY) is None) + assert not cache.set_if_generation(KEY, {"token": KEY}, fill) + assert await publisher.publish(channel, json.dumps({"cache_key": "model-cache"})) == 1 + await until(lambda: additional.get_cache("model-cache") is None) + + cache.set_local(KEY, {"token": KEY}, None) + before_stop: Final = cache.current_generation() + server.stop() + await until(lambda: cache.current_generation() > before_stop) + assert not subscriber.is_ready + assert cache.get_local(KEY) is None + outage: Final = cache.begin_fill(KEY) + assert not cache.set_if_generation(KEY, {"token": KEY}, outage) + cache.set_local(KEY, {"token": KEY}, None) + assert cache.get_local(KEY) is None + server.start() + await until(lambda: subscriber.is_ready) + assert cache.current_generation() > outage + assert not cache.set_if_generation(KEY, {"token": KEY}, outage) + assert cache.set_if_generation(KEY, {"token": KEY}, cache.begin_fill(KEY)) + assert await publisher.publish(channel, json.dumps({"cache_key": KEY})) == 1 + await until(lambda: cache.get_local(KEY) is None) + finally: + await subscriber.stop() + await publisher.aclose() + assert cache.get_local(KEY) is None + assert not cache.set_if_generation(KEY, {"token": KEY}, cache.begin_fill(KEY)) + + +@pytest.mark.asyncio +async def test_management_eviction_reaches_native_memory_and_shared_redis(tmp_path: Path) -> None: + with owned_redis(tmp_path) as server: + redis_cache: Final = RedisCache(host=server.host, port=server.port, namespace="management") + first: Final = UserApiKeyCache(redis_cache=redis_cache) + second: Final = UserApiKeyCache(redis_cache=redis_cache) + await first.async_set_cache(KEY, {"token": KEY}) + assert await second.async_get_cache(KEY) == {"token": KEY} + evict_local((KEY,), first) + assert first.get_local(KEY) is None + await evict_shared((KEY,), first) + second.delete_local(KEY) + assert await second.async_get_cache(KEY) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", (True, False)) +async def test_prefetch_does_not_restore_an_invalidated_redis_reply(tmp_path: Path, identity: bool) -> None: + key: Final = KEY if identity else "team_id:prefetch-race" + value: Final = {"token": KEY, "blocked": False, "marker": "held-redis-prefetch"} + with owned_redis(tmp_path) as server: + relay: Final = HeldStatementRelay(server.host, server.port, b"held-redis-prefetch", response=True) + relay.start() + coordination: Final = RedisCache(host=server.host, port=server.port) + cache: Final = UserApiKeyCache(redis_cache=RedisCache(host=server.host, port=relay.port, socket_timeout=30)) + writer: Final = UserApiKeyCache(redis_cache=coordination) + subscriber: Final = AuthCacheInvalidationSubscriber(coordination, cache) + redis: Final = Redis(host=server.host, port=server.port) + subscriber.start() + try: + await until(lambda: subscriber.is_ready) + await writer.async_set_cache(key, value) + fill: Final = asyncio.create_task( + prefetch_identity_keys((key,), cache) + if identity + else prefetch_auth_objects(AuthObjectRefs(team_id="prefetch-race"), cache, None) + ) + await until(relay.held.is_set) + generation: Final = cache.current_generation() + await evict_shared((key,), writer) + assert ( + await redis.publish(auth_cache_invalidation_channel(coordination), json.dumps({"cache_key": key})) == 1 + ) + await until(lambda: cache.current_generation() > generation) + assert await redis.get(key) is None + assert cache.get_local(key) is None + relay.release() + await fill + assert cache.get_local(key) is None + assert await cache.async_get_cache(key) is None + finally: + relay.release() + await subscriber.stop() + await redis.aclose() + relay.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefetch", (True, False)) +async def test_sqlx_prefetch_does_not_restore_local_or_shared_context_after_invalidation( + tmp_path: Path, prefetch: bool +) -> None: + team_id: Final = f"prefetch-{uuid.uuid4().hex}" + key: Final = f"team_id:{team_id}" + settings: Final = NativeDatabaseSettings.from_url(os.environ["DATABASE_URL"]) + relay: Final = HeldStatementRelay(settings.host, settings.port, team_id.encode(), response=True) + relay.start() + database: Final = await NativeDatabase.connect( + replace(settings, host="127.0.0.1", port=relay.port, ssl_mode="disable") + ) + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute('INSERT INTO "LiteLLM_TeamTable" (team_id, blocked) VALUES (%s, false)', (team_id,)) + with owned_redis(tmp_path) as server: + coordination: Final = RedisCache(host=server.host, port=server.port) + cache: Final = UserApiKeyCache(redis_cache=coordination) + writer: Final = UserApiKeyCache(redis_cache=coordination) + subscriber: Final = AuthCacheInvalidationSubscriber(coordination, cache) + redis: Final = Redis(host=server.host, port=server.port) + subscriber.start() + try: + await until(lambda: subscriber.is_ready) + fill: Final = asyncio.create_task( + prefetch_auth_objects( + AuthObjectRefs(team_id=team_id), cache, SimpleNamespace(native_auth_database=database) + ) + if prefetch + else get_team_object(team_id, SimpleNamespace(native_auth_database=database), cache) + ) + await until(relay.held.is_set) + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute('UPDATE "LiteLLM_TeamTable" SET blocked=true WHERE team_id=%s', (team_id,)) + generation: Final = cache.current_generation() + await evict_shared((key,), writer) + assert ( + await redis.publish(auth_cache_invalidation_channel(coordination), json.dumps({"cache_key": key})) == 1 + ) + await until(lambda: cache.current_generation() > generation) + assert await redis.get(key) is None + assert cache.get_local(key) is None + relay.release() + await fill + assert await redis.get(key) is None + assert cache.get_local(key) is None + await prefetch_auth_objects( + AuthObjectRefs(team_id=team_id), cache, SimpleNamespace(native_auth_database=database) + ) + assert cache.get_local(key)["blocked"] is True + finally: + relay.release() + await subscriber.stop() + await redis.aclose() + await database.close() + relay.stop() + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute('DELETE FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team_id,)) diff --git a/tests/unit/_auth/__init__.py b/tests/unit/_auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/_auth/budget_reference.py b/tests/unit/_auth/budget_reference.py new file mode 100644 index 00000000000..a0219ea8259 --- /dev/null +++ b/tests/unit/_auth/budget_reference.py @@ -0,0 +1,117 @@ +import math +from dataclasses import replace +from typing import Final + +from litellm._auth.common_budgets import BudgetDecision, BudgetFacts +from litellm._auth.budget_pricing import ExemptionFacts, ModelPriceFacts +from litellm._auth.model_budget_data import ModelBudgetDecision, ModelBudgetSnapshot + + +def price(f: ModelPriceFacts) -> bool: + if f.cached is not None: + return f.cached + if f.lookup_failed or not f.group_found or f.input_price is None or f.output_price is None: + return False + if f.input_price > 0 or f.output_price > 0: + return False + deployments: Final = f.deployments or () + return any(d.has_input_price or d.has_output_price for d in deployments) and not any( + (d.input_price is not None and d.input_price > 0) + or (d.output_price is not None and d.output_price > 0) + or (d.query_price is not None and d.query_price > 0) + or (d.ptu_count is not None and d.cost_per_ptu_per_hour is not None) + for d in deployments + ) + + +def exemption(f: ExemptionFacts) -> bool: + route: Final = f.route + if not route.side_effect_route and ( + route.discovery_route or (not route.llm_route and not route.registered_pass_through) + ): + return True + if route.body_present and ( + route.route in ("/mcp-rest/tools/list", "/v1/mcp/tools") + or route.route not in ("/mcp/tools/call", "/mcp-rest/tools/call") + and (route.route == "/mcp" or route.route.startswith("/mcp/")) + and route.method in ("initialize", "notifications/initialized", "ping", "tools/list") + ): + return True + return f.model_input_present and f.router_present and all(price(p) for p in f.prices) + + +def common(f: BudgetFacts) -> BudgetDecision: + limit: Final = f.limit + if limit is None: + return BudgetDecision() + if f.scope == "User" and f.team_key and not f.apply_user_budget_to_team_keys: + return BudgetDecision() + if f.scope == "EndUser" and f.info_route: + return BudgetDecision() + if f.scope == "Project" and not math.isfinite(limit): + return BudgetDecision() + if f.scope == "AccessGroup" and limit <= 0: + return BudgetDecision() + if f.scope == "Proxy" and (not limit > 0 or f.skip or not f.llm_route or f.model_list_route): + return BudgetDecision() + if f.spend is None: + return BudgetDecision(read_spend=f.scope not in ("KeySoft", "TeamSoft", "ProjectSoft", "Proxy")) + if f.scope.endswith("Soft"): + enabled: Final = { + "KeySoft": bool(limit), + "TeamSoft": f.token_present and bool(f.alert_recipients), + "ProjectSoft": f.token_present, + }[f.scope] + return BudgetDecision(alert="soft_budget" if enabled and f.spend >= limit else None) + denied: Final = ( + not f.spend <= limit + if f.scope == "Tag" + else not f.spend < limit + if f.scope == "AccessGroup" + else f.spend > limit + if f.scope == "EndUser" + else math.isfinite(limit) and (f.spend > limit if f.scope in ("Team", "Proxy") else f.spend >= limit) + ) + result: Final = BudgetDecision( + denied=denied, + alert=( + "token_budget" + if f.scope == "Key" + else "user_budget" + if f.scope == "User" + else "team_budget" + if f.scope == "Team" and denied and f.token_present + else "project_budget" + if f.scope == "Project" and denied and f.token_present + else "organization_budget" + if f.scope == "Organization" and denied + else "team_member_threshold" + if f.scope == "TeamMember" and math.isfinite(limit) + else None + ), + ) + if ( + f.scope == "Key" + and denied + and f.throttle_opt_in + and (f.tpm_limit is not None or f.rpm_limit is not None) + and f.throttle_percentage is not None + and 0 < f.throttle_percentage <= 1 + ): + return replace(result, denied=False, throttle_percentage=f.throttle_percentage) + return result + + +def model(f: ModelBudgetSnapshot) -> ModelBudgetDecision: + local: Final = (f.local_spend or 0.0) + (f.legacy_spend or 0.0) + current: Final = ( + max(local, f.shared_spend) + if f.shared_spend is not None + and f.shared_started_at is not None + and f.shared_window_seconds is not None + and f.now < f.shared_started_at + f.shared_window_seconds + else local + ) + return ModelBudgetDecision( + None, f.max_budget is not None and not f.max_budget < 0 and current >= f.max_budget, current + ) diff --git a/tests/unit/_auth/conftest.py b/tests/unit/_auth/conftest.py new file mode 100644 index 00000000000..e114732e771 --- /dev/null +++ b/tests/unit/_auth/conftest.py @@ -0,0 +1,28 @@ +from collections.abc import Callable +from typing import Final, Protocol, cast + +import pytest + +from litellm._auth import credential, rust_core +from litellm._auth.decision import AuthCore +from tests.unit._auth.reference_core import PythonCore +from tests.unit._auth.shadow import ShadowCore +from litellm._auth.credential import PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys + + +class CredentialSelector(Protocol): + def select_api_key(self, keys: PresentedKeys) -> SelectedKey | None: ... + def select_websocket_api_key(self, keys: WebSocketKeys) -> str | WebSocketDenial: ... + def scheme_credential(self, value: str) -> str | None: ... + def litellm_header_token(self, value: str) -> str: ... + + +@pytest.fixture(params=(PythonCore, rust_core.RustCore, ShadowCore), ids=("python", "rust", "shadow")) +def core(request: pytest.FixtureRequest) -> AuthCore: + factory: Final = cast(Callable[[], AuthCore], request.param) + return factory() + + +@pytest.fixture(params=("python", "rust")) +def credentials(request: pytest.FixtureRequest) -> CredentialSelector: + return credential if cast(str, request.param) == "python" else rust_core diff --git a/tests/unit/_auth/model_alias_reference.py b/tests/unit/_auth/model_alias_reference.py new file mode 100644 index 00000000000..b656699e894 --- /dev/null +++ b/tests/unit/_auth/model_alias_reference.py @@ -0,0 +1,38 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry + + +def _entry(key: str, catalog: Mapping[str, object]) -> AliasedCostEntry | None: + value: Final = catalog.get(key) + return value if isinstance(value, AliasedCostEntry) and value.names[:1] not in ((), (key,)) else None + + +def _provider(key: str, value: AliasedCostEntry, providers: Sequence[str]) -> str | None: + prefix, separator, _ = key.partition("/") + if separator and prefix in providers: + return prefix + provider: Final = value.get("litellm_provider") + return provider if isinstance(provider, str) else None + + +def _target(key: str, provider: str, catalog: Mapping[str, object], providers: Sequence[str]) -> str | None: + value: Final = _entry(key, catalog) + if value is None or _provider(key, value, providers) != provider: + return None + return value.names[0].removeprefix(f"{provider}/") + + +def expand(model: str, catalog: Mapping[str, object], providers: Sequence[str]) -> tuple[str, ...]: + """Alias expansion frozen from main 2f430873aead937b9112e683703956eb1ab179af.""" + provider, separator, name = model.partition("/") + if not separator: + value: Final = _entry(model, catalog) + return (model,) if value is None else (model, value.names[0]) + target: Final = ( + _target(f"{provider}/{name}", provider, catalog, providers) + or _target(name, provider, catalog, providers) + or name + ) + return (model, f"{provider}/{target}") if target != name else (model,) diff --git a/tests/unit/_auth/raw_answers.py b/tests/unit/_auth/raw_answers.py new file mode 100644 index 00000000000..3ef95c7c294 --- /dev/null +++ b/tests/unit/_auth/raw_answers.py @@ -0,0 +1,40 @@ +from litellm._auth.decision import ModelBudgetExceeded, ModelBudgetRead +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.model_budget_data import ModelBudgetSnapshot + + +def fallback_read(model: str, authorized: bool = True) -> FallbackAuthorizationFacts: + return FallbackAuthorizationFacts( + key=GroupedModelScope( + ModelAccessFacts((model,), (model,) if authorized else ("not-allowed",), object_type="key") + ) + ) + + +def budget_read(exceeded: ModelBudgetExceeded | None = None) -> ModelBudgetRead: + if exceeded is None: + return ModelBudgetRead(ModelBudgetSnapshot()) + return ModelBudgetRead( + ModelBudgetSnapshot(max_budget=exceeded.max_budget, local_spend=exceeded.current_cost), + exceeded.message, + exceeded.entity_type, + exceeded.entity_id, + ) + + +def candidate_read(model: str, usable: bool = True) -> FallbackCandidateFacts: + return FallbackCandidateFacts( + router_present=True, + user_present=True, + team_present=True, + membership="Loaded", + price=ModelPriceFacts(cached=True), + authorization=fallback_read(model, authorized=usable), + attribution_reset=True, + attribution=ModelGroupAttributionFacts(), + attribution_stamped=True, + ) diff --git a/tests/unit/_auth/raw_model_reference.py b/tests/unit/_auth/raw_model_reference.py new file mode 100644 index 00000000000..bc484145f0e --- /dev/null +++ b/tests/unit/_auth/raw_model_reference.py @@ -0,0 +1,52 @@ +from dataclasses import replace +from typing import Final + +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts +from litellm._auth.model_access import ModelAccessDecision +from tests.unit._auth.support.model_reference import model_access + + +def grouped(scope: GroupedModelScope) -> ModelAccessDecision: + direct: Final = model_access(scope.facts) + if direct.kind == "Allow" or not scope.group_ids: + return direct + if scope.group_models is None: + raise AssertionError("unread resource groups") + return model_access(replace(scope.facts, grants=scope.group_models)) if scope.group_models else direct + + +def initial_access(facts: KeyAndFallbackModelAccessFacts) -> ModelAccessDecision | None: + if facts.config_nonempty or "all-team-models" in facts.key_models: + return None + if not facts.primary_loaded: + raise AssertionError("unread primary") + if facts.primary is not None: + primary: Final = grouped(facts.primary) + if primary.kind != "Allow": + return primary + if facts.fallback_names is None: + raise AssertionError("unread fallbacks") + for index, _ in enumerate(facts.fallback_names): + if (result := grouped(facts.fallback_scopes[index])).kind != "Allow": + return result + if facts.routing[index] != "Routed": + raise AssertionError("routing failure should be projected by the effect") + return None + + +def fallback_access(facts: FallbackAuthorizationFacts) -> bool: + if facts.refuse_agent_keys and facts.agent_id is not None: + return False + if facts.key is None: + raise AssertionError("unread key scope") + if grouped(facts.key).kind != "Allow": + return False + if facts.team is not None and grouped(facts.team).kind != "Allow": + assigned: Final = replace(facts.team.facts, grants=facts.assigned_models or ()) + if not facts.assigned_models or model_access(assigned).kind != "Allow": + return False + return all( + scope is None or model_access(scope).kind == "Allow" + for scope in (facts.embedded_team, facts.member, facts.project) + ) diff --git a/tests/unit/_auth/reference_core.py b/tests/unit/_auth/reference_core.py new file mode 100644 index 00000000000..6cb7a030233 --- /dev/null +++ b/tests/unit/_auth/reference_core.py @@ -0,0 +1,266 @@ +import math +from dataclasses import dataclass +from datetime import datetime +from typing import Final + +from litellm._auth.decision import ( + AdminKey, + AlertKeyBudget, + AlertKeySoftBudget, + AlertProxyBudget, + AlertTeamMemberBudget, + Allow, + CachedAdminKey, + CheckBudgetExemption, + Deny, + EnforceModelAccess, + EvictPresentedKey, + ExpiredKey, + InspectKey, + KeyBlocked, + KeyBudgetExceeded, + KeyFacts, + KeyWithoutToken, + LoadKey, + LoadKeyContext, + LoadKeyOwner, + MalformedKey, + MasterKey, + MasterKeyOnlyRoute, + MissingKey, + ModelAccessDenied, + NoDatabase, + NoMasterKey, + NonAdminKey, + NotAVirtualKey, + OwnerDeactivated, + OwnerFacts, + PrefetchAuthObjects, + PrepareKey, + PresentedKey, + ReadKeySpend, + ReadProxySpend, + ReadTeamMemberBudget, + RequestFacts, + ResolveEndUser, + ResolvePresentedKey, + SendKeyBudgetAlerts, + Settings, + Step, + TeamMemberBudget, + TeamMemberBudgetExceeded, + TeamMemberStep, + UserAuthDisabled, + UserAuthRoute, + VirtualKey, + WithinModelBudgets, +) +from litellm._auth.key import hash_token, is_master_key +from tests.unit._auth.raw_model_reference import initial_access +from tests.unit._auth.reference_model_budgets import ( + check_model_budgets, + check_team_member_model_budgets, + check_user_model_budgets, +) + +_USER_AUTH_ROUTE: Final = "/user/auth" +_VIRTUAL_KEY_PREFIX: Final = "sk-" + + +@dataclass(frozen=True, slots=True) +class _KeyCheck: + settings: Settings + now: datetime + is_llm_api_route: bool + key_label: str + + +class PythonCore: + def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: + token: Final = facts.token + if settings.master_key is None: + return Allow(NoMasterKey(token=token)) + if token is None: + return Deny(MissingKey()) + if token == "": + return Deny(MalformedKey()) + if facts.route == _USER_AUTH_ROUTE: + return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(UserAuthDisabled()) + if is_master_key(token, settings.master_key): + return ResolveEndUser(prefetch_key=False, resume=lambda: Allow(MasterKey())) + hashed_token: Final = hash_token(token) + return ResolveEndUser( + prefetch_key=True, + resume=lambda: ResolvePresentedKey( + hashed_token=hashed_token, + resume=lambda presented: _after_presented_key(presented, token, hashed_token, facts, settings, now), + ), + ) + + def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: + return _check_key(_KeyCheck(settings, now, facts.is_llm_api_route, key_label=facts.token or "")) + + def check_model_budgets(self) -> Step: + return check_model_budgets(then=lambda: Allow(WithinModelBudgets())) + + def check_user_model_budgets(self) -> Step: + return check_user_model_budgets(then=lambda: Allow(WithinModelBudgets())) + + def check_team_member_model_budgets(self) -> TeamMemberStep: + return check_team_member_model_budgets() + + +def _after_presented_key( + presented: PresentedKey | None, + token: str, + hashed_token: str, + facts: RequestFacts, + settings: Settings, + now: datetime, +) -> Step: + match presented: + case AdminKey(expires=expires) if expires is not None and expires < now: + return EvictPresentedKey( + hashed_token=hashed_token, + resume=lambda: Deny(ExpiredKey(expires=expires, now=now, key_label=token)), + ) + case AdminKey(): + return Allow(CachedAdminKey()) + case NonAdminKey() | None: + pass + if facts.route in settings.master_key_only_routes: + return Deny(MasterKeyOnlyRoute()) + if not settings.has_db: + return Deny(NoDatabase()) + if presented is not None: + cached: Final = _KeyCheck(settings, now, facts.is_llm_api_route, key_label=token) + return PrepareKey(loaded_from_db=False, resume=lambda: _check_key(cached)) + if not token.startswith(_VIRTUAL_KEY_PREFIX): + return Deny(NotAVirtualKey()) + loaded: Final = _KeyCheck(settings, now, facts.is_llm_api_route, key_label=hashed_token) + return LoadKey( + hashed_token=hashed_token, + resume=lambda: PrepareKey(loaded_from_db=True, resume=lambda: _check_key(loaded)), + ) + + +def _check_key(check: _KeyCheck) -> Step: + return InspectKey(resume=lambda key: _after_inspect(check, key)) + + +def _after_inspect(check: _KeyCheck, key: KeyFacts) -> Step: + if key.blocked: + return Deny(KeyBlocked()) + return EnforceModelAccess( + resume=lambda facts: ( + Deny(ModelAccessDenied(denial)) + if (denial := initial_access(facts)) is not None + else PrefetchAuthObjects(resume=lambda: _check_owner(check, key)) + ) + ) + + +def _check_owner(check: _KeyCheck, key: KeyFacts) -> Step: + user_id: Final = key.user_id + if user_id is None: + return _check_budget_exemption(check, key) + return LoadKeyOwner(user_id=user_id, resume=lambda owner: _after_owner(check, key, user_id, owner)) + + +def _after_owner(check: _KeyCheck, key: KeyFacts, user_id: str, owner: OwnerFacts | None) -> Step: + if owner is not None and owner.deactivated: + return Deny(OwnerDeactivated(user_id=user_id)) + return _check_budget_exemption(check, key) + + +def _check_budget_exemption(check: _KeyCheck, key: KeyFacts) -> Step: + from tests.unit._auth.budget_reference import exemption + + return CheckBudgetExemption(resume=lambda facts: _check_team_member_budget(check, key, exemption(facts))) + + +def _check_team_member_budget(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step: + user_id: Final = key.user_id + team_id: Final = key.team_id + if exempt or key.team_member_spend is None or not check.settings.has_db or user_id is None or team_id is None: + return _check_expiry(check, key, exempt) + return ReadTeamMemberBudget( + user_id=user_id, + team_id=team_id, + resume=lambda budget: _after_team_member_budget(check, key, user_id, team_id, budget), + ) + + +def _after_team_member_budget( + check: _KeyCheck, key: KeyFacts, user_id: str, team_id: str, budget: TeamMemberBudget | None +) -> Step: + if budget is None or budget.spend < budget.max_budget: + return _check_expiry(check, key, exempt=False) + return AlertTeamMemberBudget( + user_id=user_id, + team_id=team_id, + budget=budget, + resume=lambda: Deny( + TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=budget.spend, max_budget=budget.max_budget) + ), + ) + + +def _check_expiry(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step: + if key.expires is not None and key.expires < check.now: + return Deny(ExpiredKey(expires=key.expires, now=check.now, key_label=check.key_label)) + if exempt: + return _load_context(check, key, budget_throttle_pct=None) + return SendKeyBudgetAlerts(resume=lambda: _check_key_budget(check, key)) + + +def _check_key_budget(check: _KeyCheck, key: KeyFacts) -> Step: + max_budget: Final = key.max_budget + if not check.is_llm_api_route or max_budget is None: + return _check_soft_budget(check, key, budget_throttle_pct=None) + return ReadKeySpend( + max_budget=max_budget, + resume=lambda spend: AlertKeyBudget( + spend=spend, resume=lambda: _after_key_spend(check, key, spend, max_budget) + ), + ) + + +def _after_key_spend(check: _KeyCheck, key: KeyFacts, spend: float, max_budget: float) -> Step: + # A NaN max_budget compares false against any spend, so a non-finite budget means no limit, never a bypass. + if not math.isfinite(max_budget) or spend < max_budget: + return _check_soft_budget(check, key, budget_throttle_pct=None) + if key.budget_throttle_pct is not None: + return _check_soft_budget(check, key, budget_throttle_pct=key.budget_throttle_pct) + return Deny(KeyBudgetExceeded(spend=spend, max_budget=max_budget)) + + +def _check_soft_budget(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + def model_budgets() -> Step: + return check_model_budgets(then=lambda: _load_context(check, key, budget_throttle_pct)) + + if key.soft_budget and key.spend >= key.soft_budget: + return AlertKeySoftBudget(resume=model_budgets) + return model_budgets() + + +def _load_context(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + return LoadKeyContext(resume=lambda: _check_proxy_budget(check, key, budget_throttle_pct)) + + +def _check_proxy_budget(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + if not (check.settings.proxy_max_budget > 0 and check.settings.has_db): + return _grant(key, budget_throttle_pct) + return ReadProxySpend( + resume=lambda spend: ( + _grant(key, budget_throttle_pct) + if spend is None + else AlertProxyBudget(spend=spend, resume=lambda: _grant(key, budget_throttle_pct)) + ) + ) + + +def _grant(key: KeyFacts, budget_throttle_pct: float | None) -> Step: + if not key.has_token: + return Deny(KeyWithoutToken()) + return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) diff --git a/tests/unit/_auth/reference_model_budgets.py b/tests/unit/_auth/reference_model_budgets.py new file mode 100644 index 00000000000..07bc4ae2315 --- /dev/null +++ b/tests/unit/_auth/reference_model_budgets.py @@ -0,0 +1,175 @@ +from tests.unit._auth.test_fallback_candidate import reference as candidate_reference +from collections.abc import Callable +from typing import Final + +from litellm._auth.decision import ( + Allow, + AuthorizeBudgetFallback, + CheckKeyModelBudget, + CheckTeamMemberModelBudget, + CheckUserModelBudget, + Deny, + FallbackTarget, + ModelBudgetExceeded, + ModelBudgetRead, + ModelBudgetScope, + ReadModelBudgetScope, + ReadTeamMemberFallbacks, + ReadTeamMemberModelBudgetScope, + RerouteTeamMemberRequest, + RewriteRequestModel, + Step, + TeamMemberModelBudgetScope, + TeamMemberStep, + VetTeamMemberFallback, + WithinModelBudgets, +) +from tests.unit._auth.budget_reference import model +from tests.unit._auth.raw_model_reference import fallback_access + + +def exceeded(reading: ModelBudgetRead) -> ModelBudgetExceeded | None: + result: Final = model(reading.snapshot) + if not result.denied: + return None + assert reading.snapshot.max_budget is not None + return ModelBudgetExceeded( + reading.message, result.current_spend, reading.snapshot.max_budget, reading.entity_type, reading.entity_id + ) + + +def check_model_budgets(then: Callable[[], Step]) -> Step: + return ReadModelBudgetScope(resume=lambda scope: _check_key_budgets(scope, then)) + + +def check_user_model_budgets(then: Callable[[], Step]) -> Step: + return ReadModelBudgetScope(resume=lambda scope: _check_user_budgets(scope, scope.models, then)) + + +def _check_key_budgets(scope: ModelBudgetScope, then: Callable[[], Step]) -> Step: + if not scope.key_has_model_budgets or not scope.models: + return _check_user_budgets(scope, scope.models, then) + return _check_key_budget_of( + scope, + scope.models, + scope.rewritable_model, + then=lambda: ReadModelBudgetScope(resume=lambda rescoped: _check_user_budgets(rescoped, rescoped.models, then)), + ) + + +def _check_key_budget_of( + scope: ModelBudgetScope, models: tuple[str, ...], rewritable_model: str | None, then: Callable[[], Step] +) -> Step: + if not models: + return then() + model: Final = models[0] + rest: Final = models[1:] + + def after(exceeded: ModelBudgetExceeded | None) -> Step: + if exceeded is None: + return _check_key_budget_of(scope, rest, rewritable_model, then) + if model != rewritable_model: + return Deny(exceeded) + return _first_fallback_within_budget( + scope.key_budget_fallbacks.get(model, ()), + found=lambda fallback: _reroute( + fallback, exceeded, then=lambda: _check_key_budget_of(scope, rest, fallback, then) + ), + exhausted=lambda: Deny(exceeded), + ) + + return CheckKeyModelBudget(model=model, resume=lambda reading: after(exceeded(reading))) + + +def _first_fallback_within_budget( + candidates: tuple[str, ...], found: Callable[[str], Step], exhausted: Callable[[], Step] +) -> Step: + if not candidates: + return exhausted() + candidate: Final = candidates[0] + return CheckKeyModelBudget( + model=candidate, + resume=lambda reading: ( + found(candidate) + if exceeded(reading) is None + else _first_fallback_within_budget(candidates[1:], found, exhausted) + ), + ) + + +def _reroute(fallback: str, exceeded: ModelBudgetExceeded, then: Callable[[], Step]) -> Step: + return AuthorizeBudgetFallback( + model=fallback, + resume=lambda facts: ( + RewriteRequestModel(model=fallback, resume=then) if fallback_access(facts) else Deny(exceeded) + ), + ) + + +def _check_user_budgets(scope: ModelBudgetScope, models: tuple[str, ...], then: Callable[[], Step]) -> Step: + if not scope.user_has_model_budgets or not models: + return then() + return CheckUserModelBudget( + model=models[0], + resume=lambda reading: ( + _check_user_budgets(scope, models[1:], then) if (failure := exceeded(reading)) is None else Deny(failure) + ), + ) + + +def check_team_member_model_budgets() -> TeamMemberStep: + return ReadTeamMemberModelBudgetScope( + resume=lambda scope: ( + _check_team_member_budget_of(scope, scope.models, scope.rewritable_model) + if scope.has_model_budgets + else Allow(WithinModelBudgets()) + ) + ) + + +def _check_team_member_budget_of( + scope: TeamMemberModelBudgetScope, models: tuple[str, ...], rewritable_model: str | None +) -> TeamMemberStep: + if not models: + return Allow(WithinModelBudgets()) + model: Final = models[0] + rest: Final = models[1:] + + def reroute(target: FallbackTarget) -> TeamMemberStep: + return RerouteTeamMemberRequest( + model=target.model, resume=lambda: _check_team_member_budget_of(scope, rest, target.model) + ) + + def after(exceeded: ModelBudgetExceeded | None) -> TeamMemberStep: + if exceeded is None: + return _check_team_member_budget_of(scope, rest, rewritable_model) + if model != rewritable_model or not scope.can_fall_back: + return Deny(exceeded) + return ReadTeamMemberFallbacks( + model=model, + resume=lambda targets: _first_usable_fallback( + tuple(target for target in targets if target.model != model), + found=reroute, + exhausted=lambda: Deny(exceeded), + ), + ) + + return CheckTeamMemberModelBudget(model=model, resume=lambda reading: after(exceeded(reading))) + + +def _first_usable_fallback( + targets: tuple[FallbackTarget, ...], + found: Callable[[FallbackTarget], TeamMemberStep], + exhausted: Callable[[], TeamMemberStep], +) -> TeamMemberStep: + if not targets: + return exhausted() + target: Final = targets[0] + return VetTeamMemberFallback( + target=target, + resume=lambda facts: ( + found(target) + if candidate_reference(facts).usable + else _first_usable_fallback(targets[1:], found, exhausted) + ), + ) diff --git a/tests/unit/_auth/reference_policy.py b/tests/unit/_auth/reference_policy.py new file mode 100644 index 00000000000..27a0951b551 --- /dev/null +++ b/tests/unit/_auth/reference_policy.py @@ -0,0 +1,186 @@ +from collections.abc import Mapping +from typing import Final + +from litellm._auth import access, admission, managed + + +def access_decision(facts: access.AccessFacts) -> access.AccessDecision: + match facts: + case access.ConcurrentFailures(failures=failures, terminal=terminal): + return access.AccessDecision( + failure_index=next( + (index for index, failed in enumerate(terminal) if failed), + next((index for index, failed in enumerate(failures) if failed), None), + ) + ) + case access.UserParamRequest(enabled=enabled): + return access.AccessDecision(read="UserParam" if enabled else None) + case access.ClientTagRequest(enabled=enabled, llm_route=llm_route): + return access.AccessDecision( + read=None + if not enabled or llm_route is False + else "ClientTagRoute" + if llm_route is None + else "ClientTags" + ) + case access.GuardrailRequest(requested_keys=keys): + modifies: Final = any( + key in keys + for key in ( + "guardrails", + "disable_global_guardrails", + "disable_global_guardrail", + "opted_out_global_guardrails", + ) + ) + return access.AccessDecision(read="GuardrailPermission" if modifies else None) + case access.TraceRequest(required=required): + return access.AccessDecision(next_trace="Gate" if required else None) + case access.Search(name=name, allowed=allowed, strict=strict): + return access.AccessDecision(denial="Search" if (strict or allowed) and name not in allowed else None) + case access.Blocked(blocked=blocked): + return access.AccessDecision(denial="Blocked" if blocked else None) + case access.UserParam( + enabled=enabled, method=method, llm_route=llm_route, mcp_route=mcp_route, has_user=has_user + ): + return access.AccessDecision( + denial="UserParam" + if enabled and method and method.upper() == "POST" and llm_route and not mcp_route and not has_user + else None + ) + case access.ClientTags(enabled=enabled, llm_route=llm_route, has_tags=has_tags): + return access.AccessDecision(denial="ClientTags" if enabled and llm_route and has_tags else None) + case access.Guardrails(requested_keys=keys, modify_guardrails=modify): + return access.AccessDecision( + denial="Guardrails" + if modify is False + and any( + key in keys + for key in ( + "guardrails", + "disable_global_guardrails", + "disable_global_guardrail", + "opted_out_global_guardrails", + ) + ) + else None + ) + case access.AgentTrace(required=required, header=header): + return access.AccessDecision(denial="AgentTrace" if required and not header else None) + case access.Tools(requested=requested, key_allowed=key, team_allowed=team): + tool_allowed: Final = key or team + rejected: Final = tuple(tool for tool in requested if tool_allowed and tool not in tool_allowed) + return access.AccessDecision(denial="Tools" if rejected else None, rejected=rejected) + case access.GrantLayers(): + if not facts.deny_by_default or not facts.identity_present or facts.master_key or facts.dashboard: + return access.AccessDecision() + layers: Final[tuple[access.GrantLayer, ...]] = ( + *(("key",) if facts.virtual_key and not facts.session_token else ()), + *(("team",) if facts.team_present else ()), + ) + return access.AccessDecision(layers=layers or ("user",)) + case access.VectorStores(requested=requested, allowed=allowed, strict=strict): + if not allowed: + return access.AccessDecision(denial="MissingVectorGrant" if strict else None) + rejected_index: Final = next((index for index, store in enumerate(requested) if store not in allowed), None) + return access.AccessDecision( + denial="VectorStores" if rejected_index is not None else None, rejected_index=rejected_index + ) + case access.TeamLookup(dashboard=dashboard, missing=missing, recorded_models=models, allow_unavailable=allow): + if not dashboard and not missing and not models and allow is None: + return access.AccessDecision(read="UnavailablePolicy") + return access.AccessDecision( + denial="TeamLookup" if not dashboard and (missing or (not models and not allow)) else None + ) + case access.TraceGate( + required=required, method=method, llm_route=llm, passthrough=passthrough, telemetry=telemetry + ): + return access.AccessDecision( + next_trace="Header" + if required and method not in ("GET", "HEAD", "OPTIONS") and (llm or passthrough) and not telemetry + else None + ) + case access.TraceStep(source=source, value=value, present=present, headers_only=headers_only): + if value: + return access.AccessDecision() + if source == "Header" and not headers_only: + return access.AccessDecision(next_trace="Body") + if source == "Body": + return access.AccessDecision(next_trace="Selected" if present else "Traceparent") + if source == "Traceparent": + return access.AccessDecision(next_trace="Selected") + if source == "Selected" and not present: + return access.AccessDecision(next_trace="Other") + return access.AccessDecision(denial="TeamTrace") + case access.TeamTrace(): + gate: Final = access_decision( + access.TraceGate(facts.required, facts.method, facts.llm_route, facts.passthrough, facts.telemetry) + ) + trace_present: Final = facts.header or ( + not facts.headers_only + and ( + facts.body_trace + or (not facts.body_trace_present and facts.traceparent_id) + or (facts.selected_trace if facts.selected_trace_present else facts.other_trace) + ) + ) + return access.AccessDecision(denial="TeamTrace" if gate.next_trace and not trace_present else None) + + +def common_plan(facts: admission.CommonGate) -> admission.CommonPlan: + if facts.force_virtual_key_checks: + return admission.CommonPlan(None) + if not facts.force_virtual_key_checks and (facts.public_route or facts.additional_public_route): + return admission.CommonPlan("PublicRoute") + if facts.unauthenticated_passthrough: + return admission.CommonPlan("UnauthenticatedPassthrough") + if ( + not facts.master_key_present + and not facts.jwt_enabled + and not facts.oauth_enabled + and not facts.oauth_proxy_enabled + ): + return admission.CommonPlan("NoAuthDevelopment") + if facts.custom_auth and not facts.custom_auth_common_checks: + return admission.CommonPlan("CustomAuth") + return admission.CommonPlan(None) + + +def reservation_plan(facts: admission.Reservation) -> admission.ReservationPlan: + return admission.ReservationPlan( + not facts.skip_budget_checks and not facts.disabled, + facts.apply_user_budget_to_team_keys, + facts.fail_closed or facts.billing_agent, + ) + + +def actor(facts: managed.Actor) -> managed.ManagedDecision: + agent, context = facts.agent, facts.context + if not agent.enabled or agent.identity is None or not agent.identity.active: + return managed.ManagedDecision("Deny", "Disabled", None) + if context is None: + return managed.ManagedDecision("Deny", "MissingContext", None) + if context.agent_id != agent.agent_id or context.binding_revision != agent.identity.revision: + return managed.ManagedDecision("Deny", "ChangedIdentity", None) + if agent.execution_mode not in (context.mode, "both"): + return managed.ManagedDecision("Deny", "ExecutionMode", None) + if context.mode == "delegated" and not context.user_id: + return managed.ManagedDecision("Deny", "MissingHuman", None) + return managed.ManagedDecision("Allow", None, None) + + +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 + + +def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool: + if entry.get("tiered_pricing") is not None: + return True + for key, value in entry.items(): + if "cost_per" not in key: + continue + if _is_positive_cost(value): + return True + if isinstance(value, dict) and any(_is_positive_cost(nested) for nested in value.values()): + return True + return False diff --git a/tests/unit/_auth/shadow.py b/tests/unit/_auth/shadow.py new file mode 100644 index 00000000000..af753d85178 --- /dev/null +++ b/tests/unit/_auth/shadow.py @@ -0,0 +1,115 @@ +import math +from collections.abc import Callable, Mapping +from dataclasses import dataclass, fields, is_dataclass, replace +from datetime import datetime +from functools import reduce +from typing import Final + +from litellm._auth import decision +from tests.unit._auth.reference_core import PythonCore +from litellm._auth.rust_core import ( + AUTH_STEP, + MEMBER_STEP, + OBJECT_MAPPING, + OBJECT_TUPLE, + CoreStep, + RustCore, + field_value, +) + + +class ShadowMismatch(RuntimeError): + pass + + +def snapshot(value: object) -> object: + if isinstance(value, float) and math.isnan(value): + return (float, "nan") + if is_dataclass(value) and not isinstance(value, type): + return ( + type(value), + tuple( + (field.name, snapshot(field_value(value, field.name))) + for field in fields(value) + if field.name != "resume" + ), + ) + if isinstance(value, Mapping): + return tuple((key, snapshot(item)) for key, item in OBJECT_MAPPING.validate_python(value).items()) + if isinstance(value, (tuple, list)): + return tuple(snapshot(item) for item in OBJECT_TUPLE.validate_python(value)) + return value + + +def compare(primary: object, reference: object) -> None: + if snapshot(primary) != snapshot(reference): + raise ShadowMismatch(f"auth shadow mismatch at {type(primary).__name__}") + + +@dataclass(frozen=True, slots=True) +class _RecordedEffect: + signature: object + answers: tuple[object, ...] + + +def _replay(reference: CoreStep, record: _RecordedEffect) -> CoreStep: + if snapshot(reference) != record.signature or isinstance(reference, (decision.Allow, decision.Deny)): + raise ShadowMismatch(f"auth shadow effect mismatch at {type(reference).__name__}") + return _resume(reference.resume, record.answers) + + +def _resume(continuation: Callable[..., CoreStep], answers: tuple[object, ...]) -> CoreStep: + return continuation(*answers) + + +def _shadow( + primary: CoreStep, reference: Callable[[], CoreStep], records: tuple[_RecordedEffect, ...] = () +) -> CoreStep: + if isinstance(primary, (decision.Allow, decision.Deny)): + compare(primary, reduce(_replay, records, reference())) + return primary + primary_resume: Final[Callable[..., CoreStep]] = primary.resume + signature: Final = snapshot(primary) + + def resume(*recorded_answers: object) -> CoreStep: + recorded: Final = _RecordedEffect(signature, recorded_answers) + next_primary: Final = _resume(primary_resume, recorded_answers) + return _shadow(next_primary, reference, (*records, recorded)) + + return replace(primary, resume=resume) + + +class ShadowCore: + def __init__(self, primary: decision.AuthCore | None = None, reference: decision.AuthCore | None = None) -> None: + self.primary: Final = primary if primary is not None else RustCore() + self.reference: Final = reference if reference is not None else PythonCore() + + def start(self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow(self.primary.start(facts, settings, now), lambda: self.reference.start(facts, settings, now)), + ) + + def start_resolved_key( + self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime + ) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow( + self.primary.start_resolved_key(facts, settings, now), + lambda: self.reference.start_resolved_key(facts, settings, now), + ), + ) + + def check_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow(self.primary.check_model_budgets(), self.reference.check_model_budgets) + ) + + def check_user_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow(self.primary.check_user_model_budgets(), self.reference.check_user_model_budgets) + ) + + def check_team_member_model_budgets(self) -> decision.TeamMemberStep: + return MEMBER_STEP.validate_python( + _shadow(self.primary.check_team_member_model_budgets(), self.reference.check_team_member_model_budgets), + ) diff --git a/tests/unit/_auth/support/claude_reference.py b/tests/unit/_auth/support/claude_reference.py new file mode 100644 index 00000000000..a72792a4abb --- /dev/null +++ b/tests/unit/_auth/support/claude_reference.py @@ -0,0 +1,103 @@ +"""Request-auth Claude policy frozen from main 2f430873aead937b9112e683703956eb1ab179af.""" + +from collections.abc import Container, Mapping +from dataclasses import dataclass +from typing import Final + +from starlette.requests import Request + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.http_parsing_utils import rewrite_request_model +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class RoutingNames: + router: RegistryReader | None + team_id: str | None + alias_maps: tuple[object, ...] + + def __contains__(self, name: object) -> bool: + if not isinstance(name, str): + return False + if name in litellm.model_alias_map or any( + isinstance(aliases, Mapping) and name in aliases for aliases in self.alias_maps + ): + return True + if self.router is None: + return False + registry: Final = self.router.registry() + return ( + name in registry.aliases() + or registry.has_deployment(name) + or bool(registry.names(name, kind="route_candidates", team_id=self.team_id)) + ) + + +def unmarked(name: str) -> str: + return name[:-4] if name.lower().endswith("[1m]") else name + + +def decoded(view_id: str) -> str | None: + encoded: Final = unmarked(view_id).removeprefix("claude-router-") + if encoded == unmarked(view_id): + return None + try: + model_id: Final = bytes.fromhex(encoded).decode() + except (ValueError, UnicodeDecodeError): + return None + return model_id if f"claude-router-{model_id.encode().hex()}" == unmarked(view_id) else None + + +def group_name(view_id: str, routing_names: Container[str]) -> str | None: + if view_id in routing_names: + return None + plain: Final = unmarked(view_id) + if plain != view_id and plain in routing_names: + return plain + model_id: Final = decoded(view_id) + return model_id if model_id and model_id in routing_names else None + + +async def normalize( + data: dict[str, object], # mutable-ok: native and reference both apply the request-local rewrite + token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + from litellm.proxy.auth.user_api_key_auth import ( + _get_model_from_request_context, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # raw request extraction is not the frozen policy + ) + from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + + if ( + route + not in frozenset( # comprehension-ok: frozen main route set construction + f"/{prefix}{endpoint}" + for prefix in ("", "v1/") + for endpoint in ("messages", "chat/completions", "responses") + ) + or llm_router is None + ): + return + if request is not None and request.scope.get("litellm.claude_model_normalized") is True: + return + requested: Final = _get_model_from_request_context(data, route, request, llm_router, token.team_id) + if not isinstance(requested, str) or requested != data.get("model"): + return + if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"): + return + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + aliases: Final[object] = settings.get("model_group_alias") if isinstance(settings, Mapping) else None + source: Final = group_name( + requested, + RoutingNames(llm_router, token.team_id, (token.aliases, token.team_model_aliases, aliases)), # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen raw alias maps retain legacy unparameterized mappings + ) + if request is not None: + request.scope["litellm.claude_model_normalized"] = True # rebind-ok: frozen request-local effect + if source is None: + return + rewrite_request_model(data, request, source) diff --git a/tests/unit/_auth/support/codex_reference.py b/tests/unit/_auth/support/codex_reference.py new file mode 100644 index 00000000000..66d57b362ac --- /dev/null +++ b/tests/unit/_auth/support/codex_reference.py @@ -0,0 +1,68 @@ +"""Codex authorization frozen from main 2f430873aead937b9112e683703956eb1ab179af.""" + +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_api import RegistryReader +from litellm.types.router import Deployment +from tests.unit._auth.support.model_registry_reference import ( + _helper, # pyright: ignore[reportPrivateUsage] # independently frozen model helper, not production policy +) + +_MODEL_NAMES: Final = TypeAdapter(list[str]) + + +def _is_chatgpt(deployment: Deployment) -> bool: + return ( + deployment.litellm_params.custom_llm_provider == "chatgpt" + or deployment.litellm_params.model.partition("/")[0] == "chatgpt" + ) + + +def _serves_by_name(deployment: Deployment, names: frozenset[str]) -> bool: + return not _is_chatgpt(deployment) and ( + deployment.model_name in names or deployment.model_info.team_public_model_name in names + ) + + +def _candidates(model: str, router: RegistryReader, team_id: str | None) -> tuple[Deployment, ...]: + registry: Final = router.registry() + candidates: Final = tuple( + registry.deployment(model_id) for model_id in registry.names(model, kind="route_candidates", team_id=team_id) + ) + return tuple(deployment for deployment in candidates if deployment is not None) + + +def routes_to_litellm( + model: object, + router: RegistryReader | None, + team_id: str | None, + bundled: tuple[str, ...], + prices: Mapping[str, object], +) -> bool: + if not isinstance(model, str) or router is None or model in bundled or f"chatgpt/{model}" in prices: + return False + deployments: Final = _candidates(model, router, team_id) + names: Final = frozenset({model, router.registry().resolve(model, kind="alias") or model}) + return bool(deployments) and all(_serves_by_name(deployment, names) for deployment in deployments) + + +def _key_and_team_allow(model: str, router: RegistryReader, token: UserAPIKeyAuth) -> bool: + allowed: Final = ( + _MODEL_NAMES.validate_python(token.models), # pyright: ignore[reportUnknownMemberType] # frozen list projection from the legacy token model + _MODEL_NAMES.validate_python(token.team_models), + ) + return all(_helper(model, router, grants, token.team_id, None) for grants in allowed) + + +def native_chatgpt_model(model: str, router: RegistryReader, token: UserAPIKeyAuth) -> str | None: + if model.startswith("chatgpt/") or _key_and_team_allow(model, router, token): + return None + prefixed: Final = f"chatgpt/{model}" + deployments: Final = _candidates(prefixed, router, token.team_id) + if not deployments or not all(_is_chatgpt(deployment) for deployment in deployments): + return None + return prefixed if _key_and_team_allow(prefixed, router, token) else None diff --git a/tests/unit/_auth/support/fallback_reference.py b/tests/unit/_auth/support/fallback_reference.py new file mode 100644 index 00000000000..d3c47e288a7 --- /dev/null +++ b/tests/unit/_auth/support/fallback_reference.py @@ -0,0 +1,68 @@ +from dataclasses import replace +from typing import Final + +from litellm._auth.fallback_access import FallbackAuthorizationDecision as Decision +from litellm._auth.fallback_access import FallbackAuthorizationFacts as Facts +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessDecision +from tests.unit._auth.support.model_reference import model_access + + +def grouped(scope: GroupedModelScope, *, merge: bool = False) -> ModelAccessDecision: + primary: Final = model_access(scope.facts) + if primary.kind != "Deny" or not scope.group_ids: + return primary + if scope.group_models is None: + return ModelAccessDecision("ReadGroups") + if not scope.group_models: + return primary + grants: Final = tuple(dict.fromkeys((*scope.facts.grants, *scope.group_models))) if merge else scope.group_models + return model_access(replace(scope.facts, grants=grants)) + + +def authorization(facts: Facts) -> Decision: + if facts.refuse_agent_keys and facts.agent_id is not None: + return Decision("Deny") + if facts.key is None: + return Decision("ReadKey") + key: Final = grouped(facts.key) + if key.kind == "ReadGroups": + return Decision("ReadKeyGroups") + if key.kind != "Allow": + return Decision("Deny", key) + if facts.team_present and facts.team is None: + return Decision("ReadTeam") + if facts.team is not None: + team: Final = grouped(facts.team, merge=True) + if team.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if team.kind != "Allow": + if team.kind != "Deny": + return Decision("Deny", team) + if facts.assigned_models is None: + return Decision("ReadAssignedGroups") + if not facts.assigned_models: + return Decision("Deny", team) + if facts.assigned_scope is None: + return Decision("ReadAssignedScope") + if model_access(facts.assigned_scope).kind != "Allow": + return Decision("Deny", team) + elif facts.embedded_team is None and facts.embedded_team_grants: + return Decision("ReadEmbeddedTeam") + elif facts.embedded_team is not None: + embedded: Final = model_access(facts.embedded_team) + if embedded.kind != "Allow": + return Decision("Deny", embedded) + if facts.member is None and facts.member_grants: + return Decision("ReadMember") + if facts.member is not None: + member: Final = model_access(facts.member) + if member.kind != "Allow": + return Decision("Deny", member) + if facts.project is None and facts.project_grants: + return Decision("ReadProject") + if facts.project is not None: + project: Final = model_access(facts.project) + if project.kind != "Allow": + return Decision("Deny", project) + return Decision("Allow") diff --git a/tests/unit/_auth/support/model_reference.py b/tests/unit/_auth/support/model_reference.py new file mode 100644 index 00000000000..5409f2a5389 --- /dev/null +++ b/tests/unit/_auth/support/model_reference.py @@ -0,0 +1,91 @@ +import re +from dataclasses import replace +from typing import Final + +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts + + +def _pattern(model: str, pattern: str) -> bool: + return "*" in pattern and re.match("^" + pattern.replace("*", ".*") + "$", model) is not None + + +def helper(model: str, facts: ModelAccessFacts, *, team_aliases: bool = True) -> bool: + ordinary: Final = tuple(grant for grant in facts.grants if grant != "all-team-models") + expanded: Final = tuple(dict.fromkeys((*ordinary, *facts.registry.names))) + grants: Final = ( + (expanded or facts.grants) + if "all-team-models" in facts.grants and facts.team_id is not None and facts.registry.present + else facts.grants + ) + record: Final = facts.registry.records.get(model) + groups: Final = record.access_groups if record is not None else () + if facts.registry.present and any(grant in groups for grant in grants): + return True + filtered: Final = tuple(grant for grant in grants if grant not in groups) + if team_aliases and model in facts.aliases.team: + return True + if any(_pattern(model, grant) for grant in filtered): + return True + provider: Final = record.provider if record is not None else None + if ( + provider is not None + and (provider.stripped_model != model or "/" not in model) + and any(_pattern(f"{provider.provider}/{provider.stripped_model}", grant) for grant in filtered) + ): + return True + return not grants or "*" in filtered or "all-proxy-models" in filtered or model in filtered + + +def _single(model: str, facts: ModelAccessFacts) -> ModelAccessDecision: + after_team: Final = facts.aliases.team.get(model, model) + after_key: Final = facts.aliases.key.get(after_team, after_team) + after_global: Final = facts.aliases.global_.get(after_key, after_key) + dispatched: Final = facts.aliases.key.get(after_global, after_global) + key_applied: Final = after_key != after_team or dispatched != after_global + alias: Final = ( + facts.aliases.global_.get(model) if model in facts.aliases.global_ else facts.registry.aliases.get(model) + ) + parent: Final = facts.registry.compaction_parents.get(model) + potential: Final = ( + (dispatched,) + if key_applied + else (model, *((parent,) if parent is not None else ()), *((alias,) if alias else ())) + ) + if any(helper(candidate, facts) for candidate in potential): + return ModelAccessDecision("Allow") + record: Final = facts.registry.records.get(model) + public: Final = ( + record.deployment_public_name if record is not None and record.deployment_team == facts.team_id else None + ) + owning: Final = ( + ( + *((public,) if public else ()), + *(alias for alias, target in facts.aliases.team.items() if target == model), + ) + if not key_applied + and facts.registry.present + and facts.team_id is not None + and model.startswith(f"model_name_{facts.team_id}_") + and record is not None + and record.deployment_team == facts.team_id + else () + ) + if any(helper(candidate, facts, team_aliases=False) for candidate in owning): + return ModelAccessDecision("Allow") + return ModelAccessDecision("Deny", model, facts.grants, facts.object_type) + + +def _list_item(model: str, facts: ModelAccessFacts) -> ModelAccessDecision: + decision: Final = model_access( + replace(facts, requested=(model,), fallback_depth=facts.fallback_depth + 1, is_list=False) + ) + return replace(decision, model=model) if decision.kind == "Depth" else decision + + +def model_access(facts: ModelAccessFacts) -> ModelAccessDecision: + if facts.fallback_depth >= facts.max_depth: + return ModelAccessDecision("Depth") + decisions: Final = ( + _list_item(model, facts) if facts.is_list else _single(model, facts) for model in facts.requested + ) + return next((decision for decision in decisions if decision.kind != "Allow"), ModelAccessDecision("Allow")) diff --git a/tests/unit/_auth/support/model_registry_reference.py b/tests/unit/_auth/support/model_registry_reference.py new file mode 100644 index 00000000000..996d157c179 --- /dev/null +++ b/tests/unit/_auth/support/model_registry_reference.py @@ -0,0 +1,189 @@ +"""Phase 1 model policy, with ownership guard updated from main 2f430873aead.""" + +import re +from collections.abc import Mapping, Sequence +from typing import Final, Literal + +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.proxy._types import LiteLLM_UserTable, ModelAccessDeniedProxyException, ProxyErrorTypes +from litellm.proxy.auth.model_access_denied import ( + customer_model_access_denied_client_message, + model_access_denied_client_message, +) +from litellm.router_api import RegistryReader +from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent + + +def _pattern(model: str, pattern: str) -> bool: + return "*" in pattern and re.match("^" + pattern.replace("*", ".*") + "$", model) is not None + + +def _provider_pattern(model: str, pattern: str) -> bool: + try: + stripped, provider, _, _ = litellm.get_llm_provider(model=model) + except Exception: # noqa: BLE001 # the Phase 1 provider inference helper swallowed lookup failures + return False + if stripped == model and "/" in model: + return False + return _pattern(f"{provider}/{stripped}", pattern) + + +def _helper( + model: str, + router: RegistryReader | None, + grants: Sequence[str], + team: str | None, + aliases: Mapping[str, str] | None, +) -> bool: + groups: Final = router.registry().access_groups(model_name=model, team_id=team) if router else {} + names: Final = ( + router.registry().names() if "all-team-models" in grants and team is not None and router is not None else None + ) + ordinary: Final = tuple(grant for grant in grants if grant != "all-team-models") + expanded: Final = ( + (tuple(dict.fromkeys((*ordinary, *names))) or tuple(grants)) if names is not None else tuple(grants) + ) + if groups and any(grant in groups for grant in expanded): + return True + filtered: Final = tuple(grant for grant in expanded if grant not in groups) + if aliases and model in aliases: + return True + if any(_pattern(model, grant) for grant in filtered): + return True + if any("*" in grant and _provider_pattern(model, grant) for grant in filtered): + return True + return (not filtered and not expanded) or "*" in filtered or "all-proxy-models" in filtered or model in filtered + + +def _owning( + model: str, router: RegistryReader | None, team: str | None, aliases: Mapping[str, str] | None +) -> tuple[str, ...]: + if router is None or team is None or not model.startswith(f"model_name_{team}_"): + return () + deployment: Final = router.registry().deployments(model, source="first_in_group") + if deployment is None or deployment.model_info.team_id != team: + return () + public: Final = ( + deployment.model_info.team_public_model_name + if deployment is not None and deployment.model_info.team_id == team + else None + ) + return (*((public,) if public else ()), *(alias for alias, target in (aliases or {}).items() if target == model)) + + +def can_object_call_model( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"] = "user", + fallback_depth: int = 0, +) -> Literal[True]: + if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: + raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}") # noqa: TRY002 # frozen Phase 1 exception contract + if not isinstance(model, str): + for item in model: + can_object_call_model( + item, + llm_router, + models, + team_model_aliases, + team_id, + key_model_aliases, + object_type, + fallback_depth + 1, + ) + return True + return _single(model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type) + + +def _single( + model: str, + router: RegistryReader | None, + grants: Sequence[str], + team_aliases: Mapping[str, str] | None, + team: str | None, + key_aliases: Mapping[str, str] | None, + object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"], +) -> Literal[True]: + parent: Final = native_compaction_parent(model) + target: Final = ( + litellm.model_alias_map[model] + if model in litellm.model_alias_map + else router.registry().resolve(model, kind="alias") + if router is not None and model in router.registry().aliases() + else None + ) + after_team: Final = (team_aliases or {}).get(model, model) + after_key: Final = (key_aliases or {}).get(after_team, after_team) + after_global: Final = litellm.model_alias_map.get(after_key, after_key) + dispatched: Final = (key_aliases or {}).get(after_global, after_global) + key_applied: Final = after_key != after_team or dispatched != after_global + candidates: Final = ( + (dispatched,) + if key_applied + else (model, *((parent,) if parent is not None else ()), *((target,) if target else ())) + ) + owning: Final = () if key_applied else _owning(model, router, team, team_aliases) + if any(_helper(candidate, router, grants, team, team_aliases) for candidate in candidates): + return True + if any(_helper(candidate, router, grants, team, None) for candidate in owning): + return True + raise ModelAccessDeniedProxyException( + message=customer_model_access_denied_client_message(model) + if object_type == "customer" + else model_access_denied_client_message(model), + internal_message=f"{object_type} not allowed to access model. This {object_type} can only access models={list(grants)}. Tried to access {model}", + type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type), + param="model", + code=403, + ) + + +class _Grants(BaseModel, frozen=True): + models: tuple[str, ...] = () + + +async def can_user_call_model( + model: str, llm_router: RegistryReader | None, user_object: LiteLLM_UserTable | None +) -> Literal[True]: + if user_object is None: + return True + grants: Final = _Grants.model_validate(user_object, from_attributes=True).models + if "no-default-models" in grants: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model), + internal_message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + if not grants: + return True + return can_object_call_model(model, llm_router, grants) + + +def entry_has_priced_metric(entry: Mapping[str, object]) -> bool: + if "tiered_pricing" in entry and entry["tiered_pricing"] is not None: + return True + return any(_priced(key, value) for key, value in entry.items()) + + +def _priced(key: str, value: object) -> bool: + if "cost_per" not in key: + return False + if isinstance(value, bool): + return False + if isinstance(value, (int, float)): + return value > 0 + if isinstance(value, dict): + return any( + isinstance(number, (int, float)) and not isinstance(number, bool) and number > 0 + for number in TypeAdapter(Mapping[object, object]).validate_python(value).values() + ) + return False diff --git a/tests/unit/_auth/support/route_enable_reference.py b/tests/unit/_auth/support/route_enable_reference.py new file mode 100644 index 00000000000..63c15a2a4f6 --- /dev/null +++ b/tests/unit/_auth/support/route_enable_reference.py @@ -0,0 +1,39 @@ +from typing import Final + +from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteFeature, RouteSetting +from tests.unit._auth.support.route_reference import RouteChecks + + +def _setting(facts: RouteSetting | None, feature: RouteFeature) -> RouteEnableDecision | None: + if facts is None: + return RouteEnableDecision("ReadSettings", feature) + if facts.environment_present and not facts.premium: + return RouteEnableDecision("Deny", feature, "License") + if facts.secret is None: + return RouteEnableDecision("ReadSecret", feature) + if facts.secret: + return RouteEnableDecision("Deny", feature, "Disabled") + return None + + +def route_enable(facts: RouteEnableFacts) -> RouteEnableDecision: + if facts.failed is not None: + return RouteEnableDecision("PropagateHttp" if facts.failed == "Http" else "Allow") + if not facts.imported: + return RouteEnableDecision("ReadImport") + if RouteChecks.check_route_access(facts.context.route, facts.context.catalog.groups.get("management_routes", ())): + admin: Final = _setting(facts.admin, "Admin") + if admin is not None: + return admin + if not RouteChecks.is_llm_api_route(facts.context.route, configured_eu_jev=True): + if facts.context.route.rstrip("/") != "/eu_jev/v1/systemone": + return RouteEnableDecision("Allow") + if facts.eu_jev_configured is None: + return RouteEnableDecision("ReadEuJev") + if facts.eu_jev_configured: + return RouteEnableDecision("Allow") + if facts.exemptions is None: + return RouteEnableDecision("ReadExemptions") + if facts.context.route in facts.exemptions: + return RouteEnableDecision("Allow") + return _setting(facts.llm, "Llm") or RouteEnableDecision("Allow") diff --git a/tests/unit/_auth/support/route_reference.py b/tests/unit/_auth/support/route_reference.py new file mode 100644 index 00000000000..3ce40db3218 --- /dev/null +++ b/tests/unit/_auth/support/route_reference.py @@ -0,0 +1,1233 @@ +import itertools +import re +from collections.abc import Collection, Iterable, Mapping, Sequence +from typing import Final, Protocol, cast + +from fastapi import HTTPException, Request, status +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + MODEL_LIST_ROUTES, + UI_TEAM_ID, + CommonProxyErrors, + KeyManagementRoutes, + LiteLLM_UserTable, + LiteLLMRoutes, + LitellmUserRoles, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, +) + + +def user_is_org_admin(request_data: Mapping[str, object], user_object: LiteLLM_UserTable | None = None) -> bool: + if user_object is None or user_object.organization_memberships is None: + return False + singular: Final = request_data.get("organization_id") + multiple: Final = request_data.get("organizations") + candidates: Final = ( + *((singular,) if singular is not None else ()), + *(TypeAdapter(tuple[object, ...]).validate_python(multiple) if isinstance(multiple, list) else ()), + ) + admins: Final = frozenset( + membership.organization_id + for membership in user_object.organization_memberships + if membership.user_role == LitellmUserRoles.ORG_ADMIN.value and membership.organization_id is not None # pyright: ignore[reportUnnecessaryComparison] # Phase 1 rejected absent legacy membership IDs + ) + return bool(candidates) and all(candidate in admins for candidate in candidates) + + +def managed_agent_route_allowed(route: str, method: str | None) -> bool: + if route in ("/agents", "/v1/agents"): + return method in (None, "GET", "HEAD") + if route in ( + "/realtime", + "/v1/realtime", + "/openai/v1/realtime", + "/realtime/translations", + "/v1/realtime/translations", + "/live/sessions", + "/v1/live/sessions", + "/openai/v1/live/sessions", + ): + return method in (None, "GET") + model_routes: Final = frozenset( + f"{prefix}/{operation}" + for prefix, operation in itertools.product( + ("", "/v1"), + ( + "chat/completions", + "completions", + "embeddings", + "responses", + "messages", + "messages/count_tokens", + "images/generations", + "images/edits", + "audio/transcriptions", + "audio/speech", + "moderations", + "rerank", + "decisions", + "systemone", + "extraction", + "ocr", + ), + ) + ) | frozenset( + ( + "/openai/v1/responses", + "/v2/rerank", + "/claude_code_gateway/v1/messages", + "/claude_code_gateway/v1/messages/count_tokens", + "/cursor/chat/completions", + ) + ) + paths: Final = ( + "/engines/{model:path}/chat/completions", + "/engines/{model:path}/completions", + "/engines/{model:path}/embeddings", + "/openai/deployments/{model:path}/chat/completions", + "/openai/deployments/{model:path}/completions", + "/openai/deployments/{model:path}/embeddings", + "/openai/deployments/{model:path}/images/generations", + "/openai/deployments/{model:path}/images/edits", + "/v1beta/models/{model_name:path}:countTokens", + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/{model_name:path}:streamGenerateContent", + "/models/{model_name:path}:countTokens", + "/models/{model_name:path}:generateContent", + "/models/{model_name:path}:streamGenerateContent", + ) + if route in model_routes or RouteChecks.check_route_access(route, paths): + return method in (None, "POST") + mcp: Final = tuple( + candidate + for candidate in LiteLLMRoutes.mcp_inference_routes.value + if candidate not in ("/token", "/introspect") + ) + return RouteChecks.check_route_access(route, mcp) or RouteChecks.check_route_access( + route, LiteLLMRoutes.agent_inference_routes.value + ) + + +class _ReferenceToken(Protocol): + @property + def allowed_routes(self) -> Sequence[str] | None: ... + + @property + def metadata(self) -> Mapping[str, object]: ... + + @property + def team_metadata(self) -> Mapping[str, object] | None: ... + + +def _token(value: UserAPIKeyAuth) -> _ReferenceToken: + return cast(_ReferenceToken, value) + + +# Management write routes denied to PROXY_ADMIN_VIEW_ONLY. Adding a new write +# endpoint to a management router REQUIRES adding it here too — the surrounding +# check falls through to "allow" if the route is not matched, which previously +# let view-only admins call /team/block, /team/unblock, /key/bulk_update, etc. +_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset( + [ + # user + "/user/new", + "/management/v1/users/bulk", + "/user/delete", + "/management/v1/users/bulk_delete", + "/user/bulk_update", + # team + "/team/new", + "/management/v1/teams/{team_id}/members/bulk_delete", + "/management/v1/teams/{team_id}/members/bulk_update", + "/team/update", + "/team/delete", + "/team/block", + "/team/unblock", + "/team/permissions_update", + "/team/permissions_bulk_update", + # model + "/model/new", + "/model/update", + "/model/delete", + # JWT key mapping + "/jwt/key/mapping/new", + "/jwt/key/mapping/update", + "/jwt/key/mapping/delete", + # key management — keep in sync with KeyManagementRoutes write entries + KeyManagementRoutes.KEY_GENERATE.value, + KeyManagementRoutes.KEY_UPDATE.value, + KeyManagementRoutes.KEY_DELETE.value, + KeyManagementRoutes.KEY_REGENERATE.value, + KeyManagementRoutes.KEY_GENERATE_SERVICE_ACCOUNT.value, + KeyManagementRoutes.KEY_BLOCK.value, + KeyManagementRoutes.KEY_UNBLOCK.value, + KeyManagementRoutes.KEY_BULK_UPDATE.value, + KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value, + ] +) + +# Suffixes for `/key/{key_id}/...` path-parameterized write routes that the +# enum templates with `{key_id}`. The blocklist above can't match templated +# paths directly because the request route carries the resolved key id. +_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spend") + +_LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp")) +_EU_JEV_DECISION_ROUTE: Final = "/eu_jev/v1/systemone" +_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS)) + +# Writes that only touch the caller's own credential row (BYOK connections and +# per-user MCP credentials). Method-scoped because the generic credential CRUD +# handlers share these paths: the route check cannot admit PATCH/DELETE on a +# name like "user_connections" without also opening the admin CRUD handlers. +_SELF_SCOPED_CREDENTIAL_ROUTES: Final = ( + ("GET", "/credentials/user_connections"), + ("POST", "/credentials/{credential_name:path}/user_connection/start"), + ("POST", "/credentials/{credential_name:path}/user_connection/poll"), + ("DELETE", "/credentials/{credential_name:path}/user_connection"), + ("POST", "/v1/mcp/server/oauth/{server_id}/register"), + ("POST", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-env-vars"), + ("DELETE", "/v1/mcp/server/{server_id}/user-env-vars"), +) + + +_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( + { + *MODEL_LIST_ROUTES, + "/utils/token_counter", + "/utils/transform_request", + "/responses/input_tokens", + "/v1/responses/input_tokens", + "/openai/v1/responses/input_tokens", + } +) +_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"}) +_TOKEN_COUNTING_ACTION: Final = "countTokens" + + +class RouteChecks: + @staticmethod + def check_route_enabled(route: str) -> None: + try: + from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + + EnterpriseRouteChecks.should_call_route(route=route) + except HTTPException as e: + raise e + except Exception: + pass + + @staticmethod + def should_call_route( + route: str, + valid_token: UserAPIKeyAuth, + request: Request | None = None, + ): + """ + Check if management route is disabled and raise exception + """ + RouteChecks.check_route_enabled(route) + # Check if Virtual Key is allowed to call the route - Applies to all Roles + RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token, request=request) + return True + + @staticmethod + def is_virtual_key_allowed_to_call_route( + route: str, + valid_token: UserAPIKeyAuth, + request: Request | None = None, + ) -> bool: + """ + Raises Exception if Virtual Key is not allowed to call the route + """ + token: Final = _token(valid_token) + + # Only check if token.allowed_routes is set and is a list with at least one item + if token.allowed_routes is None: + return True + if not isinstance(token.allowed_routes, list): + return True + if len(token.allowed_routes) == 0: + return True + + denied_auth_enforced_pass_through_route = False + + # explicit check for allowed routes (exact match or prefix match) + for allowed_route in token.allowed_routes: + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): + return True + + ## check if 'allowed_route' is a field name in LiteLLMRoutes + if any(allowed_route in LiteLLMRoutes._member_names_ for allowed_route in token.allowed_routes): + for allowed_route in token.allowed_routes: + if allowed_route in LiteLLMRoutes._member_names_: + if RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes[allowed_route].value, + ): + if ( + allowed_route in _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS + and RouteChecks.is_auth_enforced_pass_through_route( + route=route, + method=RouteChecks.get_request_method(request=request), + ) + ): + if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): + return True + denied_auth_enforced_pass_through_route = True + else: + return True + + ################################################ + # For llm_api_routes, also check registered pass-through endpoints + ################################################ + if allowed_route in _LLM_API_ROUTE_GROUPS: + if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET": + return True + + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): + if RouteChecks.is_auth_enforced_pass_through_route( + route=route, + method=RouteChecks.get_request_method(request=request), + ): + if RouteChecks.check_passthrough_route_access( + route=route, user_api_key_dict=valid_token + ): + return True + denied_auth_enforced_pass_through_route = True + else: + return True + + # Method-aware carve-out: allow GET on the two + # read-only MCP-server discovery endpoints + # (`/v1/mcp/server` and `/v1/mcp/server/{server_id}`) + # so virtual keys with allowed_routes=["llm_api_routes"] + # can list/inspect MCP servers. The GET handlers in + # mcp_management_endpoints.py sanitize the response + # for restricted virtual keys (stripping url, + # headers, env, credentials). POST/PUT/DELETE on + # these paths are admin-only management writes and + # are intentionally not covered. + if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route( + route=route, request=request + ): + return True + + # Agent registry CRUD moved from llm_api_routes into + # management_routes so DISABLE_LLM_API_ENDPOINTS stops + # blocking it. Keys configured with + # allowed_routes=["llm_api_routes"] before that split + # could reach these paths, so keep them reachable here; + # the handlers in agent_endpoints/endpoints.py still + # enforce proxy-admin on writes and scope reads by role. + if allowed_route == "llm_api_routes" and RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes.agent_management_routes.value, + ): + return True + + # check if wildcard pattern is allowed + for allowed_route in token.allowed_routes: + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): + return True + + if denied_auth_enforced_pass_through_route: + raise RouteChecks._auth_pass_through_denied_exception(route=route) + + if token.metadata.get("password_reset_required") is True: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + "This account's password must be changed before the session can be used: " + "it was either found in a known data breach or set by an admin. " + "Change it via POST /user/password/change (UI: /ui/change-password), then log in again." + ), + ) + + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {token.allowed_routes}. Tried to call route: {route}", + ) + + @staticmethod + def _mask_user_id(user_id: str) -> str: + """ + Mask user_id to prevent leaking sensitive information in error messages + + Args: + user_id (str): The user_id to mask + + Returns: + str: Masked user_id showing only first 2 and last 2 characters + """ + from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + + if not user_id or len(user_id) <= 4: + return "***" + + # Use SensitiveDataMasker with custom configuration for user_id + masker: Final = SensitiveDataMasker(visible_prefix=6, visible_suffix=2, mask_char="*") + + return masker.mask_value(user_id) + + @staticmethod + def _raise_admin_only_route_exception( + user_obj: LiteLLM_UserTable | None, + route: str, + ) -> None: + """ + Raise exception for routes that require proxy admin access + + Args: + user_obj (Optional[LiteLLM_UserTable]): The user object + route (str): The route being accessed + + Raises: + ProxyException: 401 with user role and masked user_id information + """ + user_role = "unknown" + user_id = "unknown" + if user_obj is not None: + user_role = user_obj.user_role or "unknown" + user_id = user_obj.user_id or "unknown" + + masked_user_id: Final = RouteChecks._mask_user_id(user_id) + raise ProxyException( + message=( + "Authentication Error, Only proxy admin can be used to generate, delete, update info for new " + f"keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" + ), + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + + @staticmethod + def non_proxy_admin_allowed_routes_check( + user_obj: LiteLLM_UserTable | None, + _user_role: LitellmUserRoles | None, + route: str, + request: Request, + valid_token: UserAPIKeyAuth, + request_data: Mapping[str, object] | None, + ): + """ + Checks if Non Proxy Admin User is allowed to access the route + """ + token: Final = _token(valid_token) + + # Check user has defined custom admin routes + RouteChecks.custom_admin_only_route_check( + route=route, + ) + + if RouteChecks.is_auth_enforced_pass_through_route( + route=route, + method=RouteChecks.get_request_method(request=request), + ): + RouteChecks._require_auth_pass_through_access( + route=route, + valid_token=valid_token, + jwt_team_allowed_routes=RouteChecks._jwt_team_allowed_routes(valid_token=valid_token), + ) + elif RouteChecks.is_llm_api_route(route=route): + pass + elif RouteChecks.is_info_route(route=route): + # check if user allowed to call an info route + if route == "/key/info": + # handled by function itself + pass + elif route == "/user/info": + # check if user can access this route + query_params: Final = request.query_params + user_id: Final = query_params.get("user_id") + verbose_proxy_logger.debug("user_id: %s & valid_token.user_id: %s", user_id, valid_token.user_id) + if ( + user_id + and user_id != valid_token.user_id + and _user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"key not allowed to access this user's info. user_id={user_id}, key's user_id={valid_token.user_id}", + ) + elif route == "/v2/user/info": + # handled by the endpoint itself (full RBAC in handler) + pass + elif route == "/model/info": + # /model/info just shows models user has access to + pass + elif route == "/team/info": + pass # handled by function itself + elif ( + route in LiteLLMRoutes.global_spend_tracking_routes.value + and getattr(valid_token, "permissions", None) is not None + and "get_spend_routes" in getattr(valid_token, "permissions", []) + ): + pass + elif route == "/lens" or route.startswith("/lens/"): + pass # Lens authorizes the delegated role, including its POST-based read endpoints. + elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: + RouteChecks._check_proxy_admin_viewer_access( + route=route, + _user_role=_user_role, + request_data=request_data, + request=request, + ) + elif ( + _user_role == LitellmUserRoles.INTERNAL_USER.value + and ( + RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value) + or RouteChecks.is_self_scoped_credential_route(route=route, request=request) + ) + or user_is_org_admin(request_data=request_data, user_object=user_obj) # pyright: ignore[reportArgumentType] # preserve the legacy request dictionary contract in this reference + and RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value) + or _user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + and RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes.internal_user_view_only_routes.value, + ) + or RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.self_managed_routes.value) + ): + pass + elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"): + pass # authN/authZ handled by api itself + elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token) or ( + valid_token.is_team_service_account + and RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.team_service_account_key_routes.value + ) + ): + pass + elif token.allowed_routes is not None: + # check if route is in allowed_routes (exact match or prefix match) + route_allowed = False + for allowed_route in token.allowed_routes: + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): + route_allowed = True + break + + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): + route_allowed = True + break + + if not route_allowed: + RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) + else: + RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) + + @staticmethod + def custom_admin_only_route_check(route: str): + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # preserve the legacy settings import in this reference + premium_user, + ) + + if "admin_only_routes" in general_settings: + if premium_user is not True: + verbose_proxy_logger.error( + "Trying to use 'admin_only_routes' this is an Enterprise only feature. %s", + CommonProxyErrors.not_premium_user.value, + ) + return + if route in general_settings["admin_only_routes"]: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route. Route={route} is an admin only route", + ) + + @staticmethod + def is_unbilled_route(route: str) -> bool: + resource, _, action = route.rsplit("/", 1)[-1].partition(":") + return route in _UNBILLED_ROUTES or resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION + + @staticmethod + def is_llm_api_route(route: str, configured_eu_jev: bool | None = None) -> bool: + """ + Helper to checks if provided route is an OpenAI route + + + Returns: + - True: if route is an OpenAI route + - False: if route is not an OpenAI route + """ + # Ensure route is a string before performing checks + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # reference retains the legacy runtime validation + return False + + if route in LiteLLMRoutes.openai_routes.value: + return True + + if route in LiteLLMRoutes.anthropic_routes.value: + return True + + if route in LiteLLMRoutes.google_routes.value: + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.litellm_native_routes.value): + return True + + # fuzzy match routes like "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + # Check for routes with placeholders or wildcard patterns + for openai_route in LiteLLMRoutes.openai_routes.value: + # Replace placeholders with regex pattern + # placeholders are written as "/threads/{thread_id}" + if "{" in openai_route: + if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): + return True + # Check for wildcard patterns like "/containers/*" + if RouteChecks.is_wildcard_pattern(pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): + return True + + # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" + for google_route in LiteLLMRoutes.google_routes.value: + if "{" in google_route: + if RouteChecks._route_matches_pattern(route=route, pattern=google_route): + return True + + # Check for Anthropic routes with placeholders + for anthropic_route in LiteLLMRoutes.anthropic_routes.value: + if "{" in anthropic_route: + if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): + return True + + if RouteChecks._is_azure_openai_route(route=route): + return True + + for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: + if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): + return True + return RouteChecks.is_builtin_eu_jev_decision_route(route=route, configured_eu_jev=configured_eu_jev) + + @staticmethod + def is_builtin_eu_jev_decision_route(route: str, configured_eu_jev: bool | None = None) -> bool: + if route.rstrip("/") != _EU_JEV_DECISION_ROUTE: + return False + if configured_eu_jev is not None: + return not configured_eu_jev + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass_through_endpoints imports the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash( + route=_EU_JEV_DECISION_ROUTE + ) + + @staticmethod + def _is_get_mcp_server_discovery_route(route: str, request: Request | None) -> bool: + """ + Returns True if `request` is a GET against one of the two read-only + MCP-server discovery paths: + + - GET `/v1/mcp/server` (list) + - GET `/v1/mcp/server/{server_id}` (single server, single segment) + + Multi-segment paths (`/v1/mcp/server/{id}/approve`, etc.) and any + non-GET method return False, so admin-only management writes on the + same path prefix are not reachable through this carve-out. + """ + if request is None or request.method.upper() != "GET": + return False + if route == "/v1/mcp/server": + return True + prefix: Final = "/v1/mcp/server/" + if not route.startswith(prefix): + return False + remainder: Final = route[len(prefix) :] + return bool(remainder) and "/" not in remainder + + @staticmethod + def is_management_route(route: str) -> bool: + """ + Check if route is a management route + """ + return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value) + + @staticmethod + def is_info_route(route: str) -> bool: + """ + Check if route is an info route + """ + return route in LiteLLMRoutes.info_routes.value + + @staticmethod + def is_dashboard_session_info_route(route: str, team_id: str | None) -> bool: + """ + Check if an Admin UI session key is calling an info route the dashboard polls in the background + """ + return team_id == UI_TEAM_ID and RouteChecks.is_info_route(route) + + @staticmethod + def _is_azure_openai_route(route: str) -> bool: + """ + Check if route is a route from AzureOpenAI SDK client + + eg. + route='/openai/deployments/vertex_ai/gemini-1.5-flash/chat/completions' + """ + # Ensure route is a string before attempting regex matching + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # reference retains the legacy runtime validation + return False + # Add support for deployment and engine model paths + deployment_pattern: Final = r"^/openai/deployments/[^/]+/[^/]+/chat/completions$" + engine_pattern: Final = r"^/engines/[^/]+/chat/completions$" + + if re.match(deployment_pattern, route) or re.match(engine_pattern, route): + return True + return False + + @staticmethod + def _route_matches_pattern(route: str, pattern: str) -> bool: + """ + Check if route matches the pattern placed in proxy/_types.py + + Example: + - pattern: "/threads/{thread_id}" + - route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + - returns: True + + + - pattern: "/key/{token_id}/regenerate" + - route: "/key/regenerate/82akk800000000jjsk" + - returns: False, pattern is "/key/{token_id}/regenerate" + """ + # Ensure route is a string before attempting regex matching + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # reference retains the legacy runtime validation + return False + + def _placeholder_to_regex(match: re.Match[str]) -> str: + placeholder: Final = match.group(0).strip("{}") + if not placeholder.endswith(":path"): + return r"[^/]+" + # A ":path" placeholder takes whatever the router's own path + # converter takes, slashes and colons alike, so an id spelled with + # either (or both) still matches the template it was mounted under. + # + # Unless the template puts a ":" literal of its own after the + # placeholder: the Google routes end in ":generateContent" and + # friends, and there the value has to stop before that suffix + # rather than swallow it and match a different verb. + # + # "[\s\S]" rather than ".", because "." stops at a newline and the + # path converter does not: a %0A anywhere in the value would leave + # the route unmatched here while still reaching the handler, which + # turns this gate into a bypass for the lists built on it. + return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" + + pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) + # Anchor the pattern to match the entire string + pattern = f"^{pattern}$" + if re.match(pattern, route): + return True + return False + + @staticmethod + def is_wildcard_pattern(pattern: str) -> bool: + """ + Check if pattern is a wildcard pattern + """ + return pattern.endswith("*") + + _is_wildcard_pattern = is_wildcard_pattern + + @staticmethod + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + """ + Check if route matches the wildcard pattern + + eg. + + pattern: "/scim/v2/*" + route: "/scim/v2/Users" + - returns: True + + pattern: "/scim/v2/*" + route: "/chat/completions" + - returns: False + + + pattern: "/scim/v2/*" + route: "/scim/v2/Users/123" + - returns: True + + """ + if pattern.endswith("*"): + # Get the prefix (everything before the wildcard) + prefix: Final = pattern[:-1] + return route.startswith(prefix) + else: + # If there's no wildcard, the pattern and route should match exactly + return route == pattern + + @staticmethod + def _route_matches_allowed_route(route: str, allowed_route: str) -> bool: + """ + Check if route matches the allowed_route pattern. + Supports both exact match and prefix match. + + Examples: + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6" -> True (exact match) + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6/v1/chat/completions" -> True (prefix match) + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-600" -> False (not a valid prefix) + + Args: + route: The actual route being accessed + allowed_route: The allowed route pattern + + Returns: + bool: True if route matches (exact or prefix), False otherwise + """ + # Exact match + if route == allowed_route: + return True + # Prefix match - ensure we add "/" to prevent false matches like /fake-openai-proxy-600 + if route.startswith(allowed_route + "/"): + return True + return False + + @staticmethod + def check_route_access(route: str, allowed_routes: Collection[str]) -> bool: + """ + Check if a route has access by checking both exact matches and patterns + + Args: + route (str): The route to check + allowed_routes (Sequence): Allowed routes/patterns + + Returns: + bool: True if route is allowed, False otherwise + """ + ######################################################### + # exact match route is in allowed_routes + ######################################################### + if route in allowed_routes: + return True + + ######################################################### + # wildcard match route is in allowed_routes + # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* + ######################################################### + if any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in allowed_routes + if RouteChecks.is_wildcard_pattern(pattern=allowed_route) + ): + return True + + ######################################################### + # pattern match route is in allowed_routes + # pattern: "/threads/{thread_id}" + # route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + # returns: True + ######################################################### + if any( # Check pattern match + RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes + ): + return True + + return False + + @staticmethod + def get_request_method(request: Request | None) -> str | None: + if request is None: + return None + + try: + method: Final = cast( # cast-ok: request.method is str at runtime; tests hand a MagicMock + object, request.method + ) + except (AttributeError, KeyError): + return None + if not isinstance(method, str): + return None + + return method.upper() + + _get_request_method = get_request_method + + @staticmethod + def is_self_scoped_credential_route(route: str, request: Request | None) -> bool: + method: Final = RouteChecks.get_request_method(request) + if method is None: + return False + return any( + method == expected_method and RouteChecks.check_route_access(route=route, allowed_routes=(pattern,)) + for expected_method, pattern in _SELF_SCOPED_CREDENTIAL_ROUTES + ) + + @staticmethod + def is_auth_enforced_pass_through_route(route: str, method: str | None = None) -> bool: + """ + True for config/DB pass-through endpoints registered with auth=true. + + These routes are injected into ``openai_routes`` for spend/budget hooks but + must not inherit blanket ``openai_routes`` RBAC; access is gated by + ``allowed_passthrough_routes`` on the key or team. + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + route_info: Final = InitPassThroughEndpointHelpers.get_registered_pass_through_route(route=route, method=method) + if route_info is None: + return False + return route_info.get("auth") is True + + @staticmethod + def _auth_pass_through_denied_exception(route: str) -> HTTPException: + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"Key/team not allowed to access passthrough route {route}. " + "Configure `allowed_passthrough_routes` on the team or key." + ), + ) + + @staticmethod + def _route_matches_denied_route(route: str, denied_route: str) -> bool: + """A `/` entry denies every route, since every route sits under the root.""" + normalized_denied_route: Final = denied_route.rstrip("/") or "/" + return ( + normalized_denied_route == "/" + or RouteChecks._route_matches_allowed_route(route=route, allowed_route=normalized_denied_route) + or RouteChecks.route_matches_wildcard_pattern(route=route, pattern=denied_route) + ) + + @staticmethod + def matching_denied_passthrough_route( + route: str, metadata_sources: Iterable[Mapping[str, object] | None] + ) -> str | None: + """ + First ``denied_passthrough_routes`` entry across ``metadata_sources`` that matches ``route``. + Unlike the allowlist (key list, else team list), every source's deny list applies. + """ + denied_routes: Final = tuple( + itertools.chain.from_iterable( + cast( # cast-ok: management endpoints validate this metadata key as a list of route strings on write + "list[str]", (metadata or {}).get("denied_passthrough_routes") or [] + ) + for metadata in metadata_sources + ) + ) + if not denied_routes: + return None + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + forwarded_routes: Final = InitPassThroughEndpointHelpers.forwarded_routes(route) + return next( + ( + denied_route + for denied_route in denied_routes + if any( + RouteChecks._route_matches_denied_route(route=candidate, denied_route=denied_route) + for candidate in forwarded_routes + ) + ), + None, + ) + + @staticmethod + def passthrough_route_denied_exception(route: str, denied_route: str) -> HTTPException: + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"Key/team denied access to passthrough route {route}. " + f"Matched `{denied_route}` in `denied_passthrough_routes`." + ), + ) + + @staticmethod + def _raise_if_passthrough_route_denied(route: str, valid_token: UserAPIKeyAuth) -> None: + token: Final = _token(valid_token) + denied_route: Final = RouteChecks.matching_denied_passthrough_route( + route=route, + metadata_sources=(token.metadata, token.team_metadata), + ) + if denied_route is not None: + raise RouteChecks.passthrough_route_denied_exception(route=route, denied_route=denied_route) + + @staticmethod + def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool: + """ + Explicit paths and trailing-wildcard prefixes grant auth=true pass-through. Blanket grants never do: + a named route group like ``openai_routes`` is only ever compared as a path, and an entry that names + no path segment (``*``, ``/*``) is skipped. + """ + return any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in team_allowed_routes + if allowed_route.rstrip("*").strip("/") + ) + + @staticmethod + def _jwt_team_allowed_routes(valid_token: UserAPIKeyAuth) -> Collection[str]: + """``team_allowed_routes`` for team tokens built by JWT auth; JWT-mapped virtual keys stay key-scoped.""" + if valid_token.jwt_claims is None or valid_token.token is not None or valid_token.team_id is None: + return () + + from litellm.proxy.proxy_server import jwt_handler + + return jwt_handler.litellm_jwtauth.team_allowed_routes + + @staticmethod + def _require_auth_pass_through_access( + route: str, + valid_token: UserAPIKeyAuth, + jwt_team_allowed_routes: Collection[str] = (), + ) -> None: + """ + Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the + key or team, or an explicit JWT ``team_allowed_routes`` entry. A key or team + ``denied_passthrough_routes`` match blocks the route even when one of those grants it. + """ + RouteChecks._raise_if_passthrough_route_denied(route=route, valid_token=valid_token) + if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): + return + if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes): + return + raise RouteChecks._auth_pass_through_denied_exception(route=route) + + @staticmethod + def check_passthrough_route_access(route: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Check if route is a passthrough route. + Supports both exact match and prefix match. + """ + token: Final = _token(user_api_key_dict) + metadata: Final = token.metadata + team_metadata: Final = token.team_metadata or {} + if metadata is None and team_metadata is None: # pyright: ignore[reportUnnecessaryComparison] # reference retains the legacy runtime validation + return False + if "allowed_passthrough_routes" not in metadata and "allowed_passthrough_routes" not in team_metadata: + return False + if ( + metadata.get("allowed_passthrough_routes") is None + and team_metadata.get("allowed_passthrough_routes") is None + ): + return False + + allowed_passthrough_routes: Final = cast( + Sequence[str], + metadata.get("allowed_passthrough_routes") or team_metadata.get("allowed_passthrough_routes") or [], + ) + + # Check if route matches any allowed passthrough route (exact or prefix match) + for allowed_route in allowed_passthrough_routes: + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): + return True + + return False + + @staticmethod + def is_assistants_api_request(request: Request) -> bool: + """ + Returns True if `thread` or `assistant` is in the request path + + Args: + request (Request): The request object + + Returns: + bool: True if `thread` or `assistant` is in the request path, False otherwise + """ + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route + + route: Final = get_request_route(request) + if "thread" in route or "assistant" in route: + return True + return False + + _is_assistants_api_request = is_assistants_api_request + + @staticmethod + def is_generate_content_route(route: str) -> bool: + """ + Returns True if this is a google generateContent or streamGenerateContent route + + These routes from google allow passing key=api_key in the query params + """ + if "generateContent" in route: + return True + if "streamGenerateContent" in route: + return True + return False + + # HTTP methods that are intrinsically read-only and therefore safe to + # default-allow for PROXY_ADMIN_VIEW_ONLY. Anything else (POST/PUT/PATCH/ + # DELETE) is treated as a write attempt and goes through the explicit + # write-allowlist below. + _SAFE_HTTP_METHODS = frozenset({"GET", "HEAD", "OPTIONS"}) + + # Explicit write routes that PROXY_ADMIN_VIEW_ONLY must NEVER call. The + # role-principle is "no writes, ever" — the management_routes list is the + # authoritative source for which non-llm routes are writes; we just need + # to filter out the read endpoints (info / list) that share the prefix. + # A cleaner approach is to denylist by HTTP verb (POST/PUT/PATCH/DELETE); + # this block stays as a backstop in case a write is implemented as GET. + _ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset( + [ + "/user/new", + "/management/v1/users/bulk", + "/user/delete", + "/management/v1/users/bulk_delete", + "/user/bulk_update", + "/team/new", + "/management/v1/teams/{team_id}/members/bulk_delete", + "/management/v1/teams/{team_id}/members/bulk_update", + "/team/update", + "/team/delete", + "/model/new", + "/model/update", + "/model/delete", + "/key/generate", + "/key/delete", + "/key/update", + "/key/regenerate", + "/key/service-account/generate", + "/key/block", + "/key/unblock", + "/team/key/bulk_update", + ] + ) + + @staticmethod + def _check_proxy_admin_viewer_access( + route: str, + _user_role: str, + request_data: Mapping[str, object] | None, + request: Request | None = None, + ) -> None: + """ + Check access for PROXY_ADMIN_VIEW_ONLY role. + + Admin Viewer follows a read-parity-with-Proxy-Admin rule: anything Proxy + Admin can read/list/get, Admin Viewer can read/list/get. The only + exclusions are cost-incurring inference routes (Playground, /chat/ + completions, etc.) and any state-mutating request. + + Implementation: + 1. LLM/inference routes → 403 (cost-incurring). + 2. Safe HTTP method (GET/HEAD/OPTIONS) → allow by default. This is + the read-parity guarantee — every new GET endpoint added anywhere + in the codebase is automatically readable by Admin Viewer + without needing to remember to add it to an allowlist. + 3. Unsafe HTTP method (POST/PUT/PATCH/DELETE): + - Allow `/user/update` only when restricted to user_email. + - Allow `/user/password/change` (endpoint only writes the caller's own row). + - Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. + - Otherwise allow only if the route is in admin_viewer_routes / + global_spend_tracking_routes (legacy explicit-allow set). + - Else 403. + """ + if RouteChecks.is_llm_api_route(route=route): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", + ) + + # Check if this is a write operation on management routes + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value): + # For management routes, only allow read operations or specific allowed updates + if route == "/user/update": + # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY + if request_data is not None and isinstance(request_data, dict): + _params_updated: Final = request_data.keys() + for param in _params_updated: + if param != "user_email": + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email can be updated", + ) + elif RouteChecks.check_route_access(route=route, allowed_routes=_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES) or ( + route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) + ): + # Block write operations for PROXY_ADMIN_VIEW_ONLY + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) + # Allow read operations on management routes (like /user/info, /team/info, /model/info) + method: Final = request.method.upper() if request is not None else "GET" + is_safe_method: Final = method in RouteChecks._SAFE_HTTP_METHODS + + # ── Safe HTTP method: default-allow ────────────────────────────── + if is_safe_method: + return + + # ── Unsafe HTTP method: explicit checks ────────────────────────── + # Allow `/user/update` for self-service email change. + if route == "/user/update": + if request_data is not None and isinstance(request_data, dict): + for param in request_data: + if param != "user_email": + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"user not allowed to access this route, role= {_user_role}. " + f"Trying to access: {route} and updating invalid param: {param}. " + "only user_email can be updated" + ), + ) + return + + # Self-service password change; the endpoint only writes the caller's own row. + if route == "/user/password/change": + return + + # Self-service logout; the endpoint only revokes the caller's own session key. + if route == "/session/logout": + return + + if RouteChecks.is_self_scoped_credential_route(route=route, request=request): + return + + # Hard-block known write routes regardless of HTTP method (defensive + # — these are POSTs in practice, but pinning them here protects + # against future GET-shaped writes). + if RouteChecks.check_route_access( + route=route, allowed_routes=RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES + ) or (route.startswith("/key/") and route.endswith("/regenerate")): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) + + # Legacy explicit-allow sets (kept for routes that are POST but + # semantically read-only, e.g. /spend/calculate). Both admin_viewer_routes + # and global_spend_tracking_routes are reads/listings. + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value): + return + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value): + return + + # NOTE: We intentionally do NOT fall back to allowing all + # `management_routes`. That set is a mix of reads (info/list — handled + # via the safe-method branch above) and writes (`/team/block`, + # `/team/permissions_update`, `/jwt/key/mapping/{new,update,delete}`, + # `/key/bulk_update`, `/key/{id}/reset_spend`). A blanket allow would + # let Admin Viewer POST these write endpoints — violating the + # "no writes, ever" rule. Default-deny instead. + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) diff --git a/tests/unit/_auth/test_access.py b/tests/unit/_auth/test_access.py new file mode 100644 index 00000000000..17d161dec1e --- /dev/null +++ b/tests/unit/_auth/test_access.py @@ -0,0 +1,110 @@ +from collections.abc import Callable +from itertools import product +from typing import Final, cast + +import pytest + +from litellm._auth import access +from tests.unit._auth.reference_policy import access_decision + + +@pytest.fixture(params=("reference", "native")) +def decide(request: pytest.FixtureRequest) -> Callable[[access.AccessFacts], access.AccessDecision]: + return access_decision if cast(str, request.param) == "reference" else access.decide + + +@pytest.mark.parametrize( + "facts", + tuple( + access.UserParam(*values) + for values in product((False, True), (None, "POST", "post", "GET"), (False, True), (False, True), (False, True)) + ), +) +def test_user_param_gate_preserves_method_route_and_presence( + decide: Callable[[access.AccessFacts], access.AccessDecision], facts: access.UserParam +) -> None: + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize("facts", tuple(access.GrantLayers(*values) for values in product((False, True), repeat=7))) +def test_default_deny_layers_preserve_master_dashboard_and_session_exceptions( + decide: Callable[[access.AccessFacts], access.AccessDecision], facts: access.GrantLayers +) -> None: + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize( + "key,team,expected", + ((None, None, ()), ((), ("a",), ("b", "b")), (("b",), ("a",), ("a",)), (("a", "b"), (), ()), ((), (), ())), +) +def test_key_tools_override_team_and_keep_duplicate_denials( + decide: Callable[[access.AccessFacts], access.AccessDecision], + key: tuple[str, ...] | None, + team: tuple[str, ...] | None, + expected: tuple[str, ...], +) -> None: + assert decide(access.Tools(("a", "b", "b"), key, team)) == access.AccessDecision( + denial="Tools" if expected else None, rejected=expected + ) + + +@pytest.mark.parametrize( + "strict,allowed,expected", + ( + (False, None, None), + (False, (), None), + (True, None, "MissingVectorGrant"), + (True, (), "MissingVectorGrant"), + (True, ("a",), "VectorStores"), + (False, ("a",), "VectorStores"), + ), +) +def test_vector_grants_distinguish_missing_from_denied_and_keep_first_index( + decide: Callable[[access.AccessFacts], access.AccessDecision], + strict: bool, + allowed: tuple[str, ...] | None, + expected: access.AccessDenial | None, +) -> None: + assert decide(access.VectorStores(("a", "b", "a"), allowed, strict)) == access.AccessDecision( + denial=expected, rejected_index=1 if expected == "VectorStores" else None + ) + + +@pytest.mark.parametrize( + "facts", + tuple( + access.TraceStep(source, value, present, headers_only) + for source, value, present, headers_only in product( + ("Header", "Body", "Traceparent", "Selected", "Other"), + (None, "", " ", "trace"), + (False, True), + (False, True), + ) + ), +) +def test_trace_source_precedence_and_read_short_circuits( + decide: Callable[[access.AccessFacts], access.AccessDecision], facts: access.TraceStep +) -> None: + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize( + "key", + ("guardrails", "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", "unrelated"), +) +@pytest.mark.parametrize("modify", (None, False, True)) +def test_guardrail_permission_requires_literal_false( + decide: Callable[[access.AccessFacts], access.AccessDecision], key: str, modify: bool | None +) -> None: + facts: Final = access.Guardrails((key,), modify) + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize( + "enabled,route,read", + ((False, None, None), (True, None, "ClientTagRoute"), (True, False, None), (True, True, "ClientTags")), +) +def test_client_tags_only_read_route_and_request_state_when_enabled( + decide: Callable[[access.AccessFacts], access.AccessDecision], enabled: bool, route: bool | None, read: str | None +) -> None: + assert decide(access.ClientTagRequest(enabled, route)).read == read diff --git a/tests/unit/_auth/test_admission.py b/tests/unit/_auth/test_admission.py new file mode 100644 index 00000000000..3a7c77dd8db --- /dev/null +++ b/tests/unit/_auth/test_admission.py @@ -0,0 +1,21 @@ +from itertools import product + +import pytest + +from litellm._auth import admission +from tests.unit._auth import reference_policy + + +@pytest.mark.parametrize("facts", tuple(admission.CommonGate(*values) for values in product((False, True), repeat=10))) +def test_common_gate_has_the_reference_skip_priority(facts: admission.CommonGate) -> None: + assert admission.common_plan(facts) == reference_policy.common_plan(facts) + + +@pytest.mark.parametrize("facts", tuple(admission.Reservation(*values) for values in product((False, True), repeat=5))) +def test_reservation_plan_preserves_billing_agent_fail_closed(facts: admission.Reservation) -> None: + assert admission.reservation_plan(facts) == reference_policy.reservation_plan(facts) + + +@pytest.mark.parametrize("facts", tuple(admission.Ownership(*values) for values in product((False, True), repeat=4))) +def test_observation_never_transfers_request_ownership(facts: admission.Ownership) -> None: + assert admission.ownership_plan(facts).owner == "Python" diff --git a/tests/unit/_auth/test_agent_ceilings.py b/tests/unit/_auth/test_agent_ceilings.py new file mode 100644 index 00000000000..a32fb1f6488 --- /dev/null +++ b/tests/unit/_auth/test_agent_ceilings.py @@ -0,0 +1,91 @@ +from collections.abc import Callable +from itertools import chain +from typing import Final + +import pytest + +from litellm._auth.agent_ceilings import AgentCeilingDecision as Decision +from litellm._auth.agent_ceilings import AgentCeilingFacts as Facts +from litellm._auth.agent_ceilings import AgentCeilingGroup as Group +from litellm._auth.agent_ceilings import AgentGroupLoadDecision as LoadDecision +from litellm._auth.agent_ceilings import AgentGroupLoadFacts as LoadFacts +from litellm._auth.agent_ceilings import agent_ceiling, agent_group_load +from litellm._auth.failure import NativeAuthFailure + + +def reference(facts: Facts) -> Decision: + if facts.ids is None: + return Decision("ReadIds") + if not facts.ids: + return Decision("Complete") + if facts.groups is None: + return Decision("ReadGroups", authoritative=facts.check_db_only) + indices: Final = tuple(index for index, group in enumerate(facts.groups) if group.status == "Found") + return Decision( + "Complete", + present=True, + ids=tuple(name for name, group in zip(facts.ids, facts.groups, strict=True) if group.status != "Missing"), + models=frozenset(chain.from_iterable(facts.groups[index].models for index in indices)), + mcp_server_ids=frozenset(chain.from_iterable(facts.groups[index].mcp_server_ids for index in indices)), + agent_ids=frozenset(chain.from_iterable(facts.groups[index].agent_ids for index in indices)), + indices=indices, + ) + + +def load_reference(facts: LoadFacts) -> LoadDecision: + if not facts.database_present: + return LoadDecision("DatabaseUnavailable") + if facts.status is None: + return LoadDecision("Read") + if facts.status == "Success": + return LoadDecision("Return") + return LoadDecision("Raise" if facts.status == "NotFound" or facts.check_db_only else "HttpUnavailable") + + +@pytest.mark.parametrize("policy", (reference, agent_ceiling)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (Facts(), Decision("ReadIds")), + (Facts(()), Decision("Complete")), + (Facts(("group",)), Decision("ReadGroups")), + (Facts(("group",), check_db_only=True), Decision("ReadGroups", authoritative=True)), + (Facts(("group",), (Group("Missing"),)), Decision("Complete", present=True)), + (Facts(("group",), (Group("Unavailable"),)), Decision("Complete", present=True, ids=("group",))), + ( + Facts( + ("one", "two", "one"), + (Group("Found", ("model",), ("server",), ("agent",)), Group("Missing"), Group("Found", ("model",))), + ), + Decision( + "Complete", + present=True, + ids=("one", "one"), + models=frozenset(("model",)), + mcp_server_ids=frozenset(("server",)), + agent_ids=frozenset(("agent",)), + indices=(0, 2), + ), + ), + ), +) +def test_ceiling_union_and_missing_vs_unavailable( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize("policy", (load_reference, agent_group_load)) +@pytest.mark.parametrize("database", (False, True)) +@pytest.mark.parametrize("strict", (False, True)) +@pytest.mark.parametrize("status", (None, "Success", "NotFound", "HttpError")) +def test_group_read_error_policy( + policy: Callable[[LoadFacts], LoadDecision], database: bool, strict: bool, status: str | None +) -> None: + facts: Final = LoadFacts(database, strict, status) # pyright: ignore[reportArgumentType] # parameter table covers the complete typed status domain + assert policy(facts) == load_reference(facts) + + +def test_mismatched_group_answers_are_terminal() -> None: + with pytest.raises(NativeAuthFailure): + agent_ceiling(Facts(("group",), ())) diff --git a/tests/unit/_auth/test_agent_grants.py b/tests/unit/_auth/test_agent_grants.py new file mode 100644 index 00000000000..946d1a351e3 --- /dev/null +++ b/tests/unit/_auth/test_agent_grants.py @@ -0,0 +1,193 @@ +from dataclasses import replace +from itertools import product +from typing import Final, assert_never + +import pytest + +from litellm._auth import agent_grants as policy + + +def _decision( + action: policy.Action = "Complete", grants: tuple[str, ...] | None = None, team_id: str | None = None +) -> policy.Decision: + return policy.Decision(action, grants, team_id) + + +def reference(facts: policy.Facts) -> policy.Decision: + match facts: + case policy.Resolve(): + if facts.managed: + return _decision("ReadManaged") + if not facts.key_team_loaded: + return _decision("ReadKeyTeam") + if not facts.strict and facts.role != "proxy_admin" and facts.default_deny is None: + return _decision("ReadDefault") + if (facts.strict or (facts.role != "proxy_admin" and facts.default_deny)) and facts.key_team is None: + return _decision(grants=()) + if not facts.caller_loaded: + return _decision("ReadCaller") + if not facts.ceiling_loaded: + return _decision("ReadCeiling") + return _decision( + grants=facts.ceiling + if facts.own is None + else facts.own + if facts.ceiling is None + else tuple(item for item in facts.own if item in facts.ceiling) + ) + case policy.Failure(): + return _decision("Unavailable") if facts.strict else _decision(grants=() if facts.closed else None) + case policy.OwnerTeam(): + if facts.permission_id_present and not facts.permission_present: + return _decision("Unavailable") if facts.strict else _decision(grants=()) + return _decision("ReadDeclared") + case policy.OwnerUnion(): + if any(source.unavailable for source in facts.teams): + return _decision("Unavailable") + restrictions: Final = ( + *(source.grants for source in facts.teams), + *(() for source in facts.withheld if source.unavailable or source.grants is not None), + *(((),) if facts.unreadable else ()), + ) + return _decision( + grants=None + if restrictions and all(grants is None for grants in restrictions) + else tuple(sorted({item for grants in restrictions for item in (grants or ())})) + ) + case policy.BillingTeam(): + restricted: Final = tuple( + source for source in facts.sources if source.grants is not None and facts.agent_id in source.grants + ) + selected: Final = restricted or ( + facts.sources + if all(source.grants is None and not source.unavailable for source in facts.sources) + else () + ) + return _decision( + team_id=min((source.team_id for source in selected if source.team_id is not None), default=None) + ) + case policy.Human(): + return _decision( + grants=tuple( + sorted( + { + item + for source in facts.sources + if source.team_id is None or source.team_id == facts.team_id + for item in (source.grants or ()) + } + ) + ) + ) + case policy.Delegated(): + if any(source.team_id is None and facts.agent_id in (source.grants or ()) for source in facts.sources): + return _decision(team_id=facts.team_id) + teams: Final = frozenset( + source.team_id + for source in facts.sources + if source.team_id is not None + and facts.agent_id in (source.grants or ()) + and (facts.allowed is None or source.team_id in facts.allowed) + ) + if facts.team_id in teams: + return _decision(team_id=facts.team_id) + if not facts.explicit and teams: + return _decision(team_id=min(teams)) + return _decision("SelectTeam") + case policy.TeamLookup(): + if not facts.subject: + return _decision() + if facts.team_id is None: + return _decision( + "ReadPersonalEligibility" + if facts.personal is None + else "ReadOwnerTeams" + if facts.personal + else "Complete" + ) + if not facts.database or (facts.loaded and not facts.row_present): + return _decision(grants=() if facts.strict else None) + return _decision("ReadDeclared" if facts.loaded else "ReadTeam") + case policy.GroupMatches(): + return _decision( + grants=tuple( + agent.agent_id for agent in facts.agents if any(group in agent.groups for group in facts.groups) + ) + ) + case _: + assert_never(facts) + + +_SOURCES: Final = ( + (), + (policy.Source("b", None),), + (policy.Source("b", ("agent",)), policy.Source("a", ("agent", "other"))), + (policy.Source("b", None, True),), + (policy.Source(None, ("agent",)),), + (policy.Source("b", ()),), +) +_RESOLVE: Final = tuple( + policy.Resolve( + managed, + strict, + role, + loaded, + key, + default, + caller_loaded=caller, + own=("agent", "other"), + ceiling_loaded=ceiling_loaded, + ceiling=ceiling, + ) + for managed, strict, role, loaded, key, default, caller, ceiling_loaded, ceiling in product( + (False, True), + (False, True), + (None, "proxy_admin", "internal_user"), + (False, True), + (None, (), ("agent",)), + (None, False, True), + (False, True), + (False, True), + (None, (), ("agent",)), + ) +) +_CASES: Final[tuple[policy.Facts, ...]] = ( + *_RESOLVE, + *( + policy.OwnerUnion(teams, withheld, unreadable) + for teams, withheld, unreadable in product(_SOURCES, _SOURCES, (False, True)) + ), + *(policy.Failure(strict, closed) for strict, closed in product((False, True), repeat=2)), + *(policy.OwnerTeam(strict, declared, loaded) for strict, declared, loaded in product((False, True), repeat=3)), + *(policy.BillingTeam(sources, "agent") for sources in _SOURCES), + *(policy.Human(sources, team) for sources, team in product(_SOURCES, (None, "a", "b"))), + *( + policy.Delegated(sources, "agent", team, explicit, allowed) + for sources, team, explicit, allowed in product(_SOURCES, (None, "a", "b"), (False, True), (None, (), ("b",))) + ), + *( + policy.TeamLookup(subject, team, database, strict, personal, loaded, row) + for subject, team, database, strict, personal, loaded, row in product( + (False, True), + (None, "", "team"), + (False, True), + (False, True), + (None, False, True), + (False, True), + (False, True), + ) + ), + policy.GroupMatches((policy.ConfiguredAgent("agent", ("group",)), policy.ConfiguredAgent("other", ())), ("group",)), +) + + +@pytest.mark.parametrize("facts", _CASES) +def test_agent_grants_match_reference(facts: policy.Facts) -> None: + assert policy.decide(facts) == reference(facts) + + +def test_default_deny_precedes_caller_and_ceiling_reads() -> None: + facts: Final = policy.Resolve(False, False, "internal_user", key_team_loaded=True) + assert policy.decide(facts).action == "ReadDefault" + assert policy.decide(replace(facts, default_deny=True)) == _decision(grants=()) + assert policy.decide(replace(facts, default_deny=False)).action == "ReadCaller" diff --git a/tests/unit/_auth/test_auth_context.py b/tests/unit/_auth/test_auth_context.py new file mode 100644 index 00000000000..8585e627b18 --- /dev/null +++ b/tests/unit/_auth/test_auth_context.py @@ -0,0 +1,64 @@ +from itertools import product +from typing import Final + +import pytest + +from litellm._auth import auth_context + +_FETCH: Final = tuple( + auth_context.Fetch(team, "ui", user, project, end_user) + for team, user, project, end_user in product( + (None, "ui", "", "team"), (None, "", "user"), (None, "", "project"), (None, "", "customer") + ) +) +_TEAM: Final = tuple( + auth_context.Team(team, "ui", failed, missing, models, allow) + for team, failed, missing, models, allow in product( + (None, "ui", "team"), (False, True), (False, True), ((), ("model",)), (None, False, True) + ) +) +_ADMIN: Final = tuple( + auth_context.Admin(role, user, "admin", row) + for role, user, row in product((None, "proxy_admin", "internal_user"), (None, "admin", "user"), (False, True)) +) + + +@pytest.mark.parametrize("facts", _FETCH) +def test_context_fetches_preserve_missing_and_empty_ids(facts: auth_context.Fetch) -> None: + assert auth_context.fetch(facts) == auth_context.FetchPlan( + facts.team_id is not None and facts.team_id != facts.ui_team_id, + facts.user_id is not None, + facts.project_id is not None, + bool(facts.end_user_id), + ) + + +@pytest.mark.parametrize("facts", _TEAM) +def test_team_context_preserves_fallback_rules(facts: auth_context.Team) -> None: + expected: Final = ( + ("Token" if facts.team_id == facts.ui_team_id else "Loaded") + if not facts.failed + else ( + "None" + if facts.team_id is None + else "Token" + if facts.team_id == facts.ui_team_id + else "Failure" + if facts.missing + else "Token" + if facts.recorded_models + else "ReadUnavailablePolicy" + if facts.allow_unavailable is None + else "Token" + if facts.allow_unavailable + else "Failure" + ) + ) + assert auth_context.team(facts).source == expected + + +@pytest.mark.parametrize("facts", _ADMIN) +def test_admin_projection_preserves_personal_budgets(facts: auth_context.Admin) -> None: + assert auth_context.admin(facts) == auth_context.AdminPlan( + facts.role == "proxy_admin", facts.user_row and facts.user_id != facts.admin_id + ) diff --git a/tests/unit/_auth/test_budget_pricing.py b/tests/unit/_auth/test_budget_pricing.py new file mode 100644 index 00000000000..b193bfe90c5 --- /dev/null +++ b/tests/unit/_auth/test_budget_pricing.py @@ -0,0 +1,78 @@ +from math import inf, nan +from typing import Final + +import pytest + +from litellm._auth import budget_pricing as native +from litellm._auth.budget_pricing import DeploymentPrice, ModelPriceFacts, PriceDecision, RouteBudgetFacts + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (ModelPriceFacts(cached=True, lookup_failed=True), PriceDecision(True, None)), + (ModelPriceFacts(cached=False), PriceDecision(False, None)), + (ModelPriceFacts(lookup_failed=True), PriceDecision(False, None)), + (ModelPriceFacts(), PriceDecision(False, False)), + (ModelPriceFacts(group_found=True), PriceDecision(False, False)), + (ModelPriceFacts(group_found=True, input_price=0, output_price=0), PriceDecision(False, None, True)), + (ModelPriceFacts(group_found=True, input_price=0, output_price=0, deployments=()), PriceDecision(False, False)), + ), +) +def test_pricing_fail_closed_and_read_plan(facts: ModelPriceFacts, expected: PriceDecision) -> None: + assert native.price(facts) == expected + + +@pytest.mark.parametrize( + ("price", "expected"), ((None, False), (-1, True), (0, True), (1, False), (inf, False), (nan, True)) +) +def test_zero_price_keeps_reference_nonfinite_semantics(price: float | None, expected: bool) -> None: + facts: Final = ModelPriceFacts( + group_found=True, input_price=price, output_price=0, deployments=(DeploymentPrice(has_input_price=True),) + ) + assert native.price(facts) == PriceDecision(expected, expected) + + +@pytest.mark.parametrize( + ("deployment", "expected"), + ( + (DeploymentPrice(), False), + (DeploymentPrice(has_input_price=True), True), + (DeploymentPrice(has_output_price=True), True), + (DeploymentPrice(has_input_price=True, input_price=1), False), + (DeploymentPrice(has_input_price=True, output_price=1), False), + (DeploymentPrice(has_input_price=True, query_price=1), False), + (DeploymentPrice(has_input_price=True, ptu_count=0, cost_per_ptu_per_hour=0), False), + (DeploymentPrice(has_input_price=True, ptu_count=0), True), + (DeploymentPrice(has_input_price=True, input_price=nan), True), + ), +) +def test_explicit_deployment_cost_and_ptu_override_zero_group_price( + deployment: DeploymentPrice, expected: bool +) -> None: + facts: Final = ModelPriceFacts(group_found=True, input_price=0, output_price=0, deployments=(deployment,)) + assert native.price(facts) == PriceDecision(expected, expected) + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (RouteBudgetFacts(side_effect_route=True, discovery_route=True), "ReadPricing"), + (RouteBudgetFacts(discovery_route=True, llm_route=True), "Skip"), + (RouteBudgetFacts(llm_route=True), "ReadPricing"), + (RouteBudgetFacts(), "ReadPassThrough"), + (RouteBudgetFacts(registered_pass_through=False), "Skip"), + (RouteBudgetFacts(registered_pass_through=True), "ReadPricing"), + (RouteBudgetFacts(llm_route=True, route="/mcp", method="tools/list", body_present=True), "Skip"), + (RouteBudgetFacts(llm_route=True, route="/mcp", method="tools/list", body_present=False), "ReadPricing"), + ( + RouteBudgetFacts(llm_route=True, route="/mcp/tools/call", method="tools/list", body_present=True), + "ReadPricing", + ), + (RouteBudgetFacts(llm_route=True, route="/mcp-rest/tools/list", body_present=True), "Skip"), + (RouteBudgetFacts(llm_route=True, route="/mcp/team", method="initialize", body_present=True), "Skip"), + (RouteBudgetFacts(llm_route=True, route="/mcp", method="tools/call", body_present=True), "ReadPricing"), + ), +) +def test_route_and_mcp_exemption_policy(facts: RouteBudgetFacts, expected: str) -> None: + assert native.route(facts) == expected diff --git a/tests/unit/_auth/test_claude_normalization.py b/tests/unit/_auth/test_claude_normalization.py new file mode 100644 index 00000000000..57ca8e71ce3 --- /dev/null +++ b/tests/unit/_auth/test_claude_normalization.py @@ -0,0 +1,264 @@ +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from typing import Final, cast + +import pytest +from starlette.requests import Request + +import litellm +from litellm._auth.claude_normalization import ClaudeNormalizationFacts, claude_membership, claude_normalization +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.claude_normalization_adapter import normalize_claude_model, read_claude_membership_facts +from litellm.router_api import RouterRegistry +from litellm.rust_bridge import _native +from tests.unit._auth.support import claude_reference as reference + + +@dataclass(frozen=True, slots=True) +class Trace: + fail_at: str | None = None + calls: list[tuple[object, ...]] = field(default_factory=list) # mutable-ok: ordered raw IO observations + + def read(self, operation: str, *values: object) -> None: + self.calls.append((operation, *values)) + if operation == self.fail_at: + raise ValueError(operation) + + +@dataclass(frozen=True, slots=True) +class Registry: + trace: Trace + alias_names: tuple[str, ...] = () + deployments: tuple[str, ...] = () + candidates: tuple[str, ...] = () + + def aliases(self) -> Mapping[str, object]: + self.trace.read("aliases") + return dict.fromkeys(self.alias_names) + + def has_deployment(self, name: str) -> bool: + self.trace.read("deployment", name) + return name in self.deployments + + def names(self, name: str, *, kind: str, team_id: str | None) -> tuple[str, ...]: + self.trace.read("names", name, kind, team_id) + return ("deployment-id",) if name in self.candidates else () + + +@dataclass(frozen=True, slots=True) +class Router: + reader: Registry + + def registry(self) -> RouterRegistry: + self.reader.trace.read("registry") + return cast(RouterRegistry, self.reader) # cast-ok: raw registry fixture supplies the selected IO + + +@dataclass(frozen=True, slots=True) +class Settings: + trace: Trace + aliases: object = None + + async def get_hierarchical_router_settings( + self, *, user_api_key_dict: UserAPIKeyAuth, prisma_client: object, proxy_logging_obj: object + ) -> Mapping[str, object]: + self.trace.read("settings") + return {"model_group_alias": self.aliases} + + +@dataclass(frozen=True, slots=True) +class LookupOnly(Mapping[str, object]): + entries: Mapping[str, object] + trace: Trace + + def __getitem__(self, key: str) -> object: + self.trace.read("lookup", key) + return self.entries[key] + + def __len__(self) -> int: + return len(self.entries) + + def __iter__(self) -> Iterator[str]: + raise RuntimeError("unrequested whole-map read") + + +@pytest.mark.parametrize("present", (False, True)) +def test_membership_only_reads_the_native_selected_alias_name(monkeypatch: pytest.MonkeyPatch, present: bool) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + trace: Final = Trace() + aliases: Final = LookupOnly({"picker": None} if present else {}, trace) + facts: Final = read_claude_membership_facts("picker", None, None, (aliases,)) + assert claude_membership(facts).present is present + assert trace.calls == [("lookup", "picker")] + + +_PICKER: Final = "claude-router-" + "model".encode().hex() +_MODEL_NAMES: Final = ( + "model", + "model[1m]", + "model[1M]", + _PICKER, + _PICKER + "[1m]", + "claude-router-6D6F64656C", + "claude-router-6d 6f64656c", + "claude-router-ff", + "claude-router-1", + "claude-router-", + "claude-router-zz[1m]", + "claude-router-" + "模型".encode().hex(), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", _MODEL_NAMES) +@pytest.mark.parametrize("source", ("global", "key", "team", "settings", "registry", "deployment", "candidates")) +@pytest.mark.parametrize("claim", ("requested", "target", "missing")) +async def test_normalization_matches_frozen_main_rewrite_scope_and_read_order( + monkeypatch: pytest.MonkeyPatch, model: str, source: str, claim: str +) -> None: + name: Final = model if claim == "requested" else (reference.decoded(model) or reference.unmarked(model)) + names: Final = () if claim == "missing" else (name,) + monkeypatch.setattr(litellm, "model_alias_map", dict.fromkeys(names, "target") if source == "global" else {}) + token: Final = UserAPIKeyAuth( + team_id="team", + aliases=dict.fromkeys(names, "target") if source == "key" else {}, + team_model_aliases=dict.fromkeys(names, "target") if source == "team" else {}, + ) + + async def run(native: bool) -> tuple[Mapping[str, object], object, tuple[tuple[object, ...], ...]]: + trace: Final = Trace() + router: Final = Router( + Registry( + trace, + names if source == "registry" else (), + names if source == "deployment" else (), + names if source == "candidates" else (), + ) + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr( + proxy_server, "proxy_config", Settings(trace, dict.fromkeys(names) if source == "settings" else None) + ) + request: Final = Request({"type": "http", "path": "/v1/messages", "method": "POST", "headers": []}) + data: Final[dict[str, object]] = {"model": model, "messages": ()} # mutable-ok: real request rewrite target + await (normalize_claude_model if native else reference.normalize)(data, token, request, "/v1/messages") + return data, request.scope.get("litellm.claude_model_normalized"), tuple(trace.calls) + + assert await run(True) == await run(False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + ( + "/messages", + "/v1/messages", + "/chat/completions", + "/v1/chat/completions", + "/responses", + "/v1/responses", + "/models", + "/backend-api/codex/responses", + ), +) +@pytest.mark.parametrize("router_present", (False, True)) +@pytest.mark.parametrize("normalized", (False, True, 1)) +@pytest.mark.parametrize("request_present", (False, True)) +async def test_normalization_early_exits_and_idempotency_match_main( + monkeypatch: pytest.MonkeyPatch, route: str, router_present: bool, normalized: object, request_present: bool +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + + async def run(native: bool) -> tuple[Mapping[str, object], object, tuple[tuple[object, ...], ...]]: + trace: Final = Trace() + monkeypatch.setattr( + proxy_server, "llm_router", Router(Registry(trace, deployments=("model",))) if router_present else None + ) + monkeypatch.setattr(proxy_server, "proxy_config", Settings(trace)) + request: Final = ( + Request( + { + "type": "http", + "path": route, + "method": "POST", + "headers": [], + "litellm.claude_model_normalized": normalized, + } + ) + if request_present + else None + ) + data: Final[dict[str, object]] = {"model": _PICKER} # mutable-ok: real request rewrite target + await (normalize_claude_model if native else reference.normalize)(data, UserAPIKeyAuth(), request, route) + return ( + data, + request.scope.get("litellm.claude_model_normalized") if request is not None else None, + tuple(trace.calls), + ) + + assert await run(True) == await run(False) + + +@pytest.mark.parametrize("aliases", (None, [], {"picker": None}, {"other": "picker"}, {"picker": 42})) +@pytest.mark.parametrize("position", (0, 1, 2)) +def test_raw_alias_membership_preserves_keys_not_targets( + monkeypatch: pytest.MonkeyPatch, aliases: object, position: int +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + maps: Final = tuple(aliases if index == position else None for index in range(3)) + first: Final = Trace() + second: Final = Trace() + facts: Final = read_claude_membership_facts("picker", Router(Registry(first)), "team", maps) + expected: Final = "picker" in reference.RoutingNames(Router(Registry(second)), "team", maps) + assert claude_membership(facts).present == expected + assert first.calls == second.calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_at", ("settings", "registry", "aliases", "deployment", "names")) +async def test_raw_read_errors_preserve_order_and_do_not_mark_request( + monkeypatch: pytest.MonkeyPatch, fail_at: str +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + + async def run(native: bool) -> tuple[tuple[object, ...], ...]: + trace: Final = Trace(fail_at) + monkeypatch.setattr(proxy_server, "llm_router", Router(Registry(trace))) + monkeypatch.setattr(proxy_server, "proxy_config", Settings(trace)) + request: Final = Request({"type": "http", "path": "/messages", "method": "POST", "headers": []}) + data: Final[dict[str, object]] = {"model": _PICKER} # mutable-ok: real request rewrite target + with pytest.raises(ValueError, match=fail_at): + await (normalize_claude_model if native else reference.normalize)( + data, UserAPIKeyAuth(), request, "/messages" + ) + assert data == {"model": _PICKER} + assert "litellm.claude_model_normalized" not in request.scope + return tuple(trace.calls) + + assert await run(True) == await run(False) + + +def test_non_string_request_never_reads_body_model_or_settings() -> None: + facts: Final = ClaudeNormalizationFacts("/messages", False, router_present=True, model_loaded=True) + assert claude_normalization(facts).kind == "Complete" + + +@pytest.mark.asyncio +async def test_native_membership_failure_is_terminal_before_registry_io(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = NativeAuthRuntimeError("native membership failed") + trace: Final = Trace() + monkeypatch.setattr(proxy_server, "llm_router", Router(Registry(trace))) + monkeypatch.setattr(proxy_server, "proxy_config", Settings(trace)) + + def fail(_facts: object) -> object: + raise failure + + monkeypatch.setattr(_native, "request_auth_claude_membership", fail) + data: Final[dict[str, object]] = {"model": _PICKER} # mutable-ok: real request rewrite target + with pytest.raises(NativeAuthRuntimeError) as raised: + await normalize_claude_model(data, UserAPIKeyAuth(), None, "/messages") + assert raised.value is failure + assert data == {"model": _PICKER} + assert trace.calls == [("settings",)] diff --git a/tests/unit/_auth/test_codex_models.py b/tests/unit/_auth/test_codex_models.py new file mode 100644 index 00000000000..3dba18e0894 --- /dev/null +++ b/tests/unit/_auth/test_codex_models.py @@ -0,0 +1,140 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Final, cast + +import pytest + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth import codex_model_adapter as native +from litellm.router_api import RouterRegistry +from litellm.rust_bridge import _native +from litellm.types.router import Deployment +from tests.unit._auth.support import codex_reference as reference + + +@dataclass(frozen=True, slots=True) +class Registry: + deployments: tuple[Deployment | None, ...] + alias: str | None = None + groups: Mapping[str, tuple[str, ...]] = field(default_factory=dict[str, tuple[str, ...]]) + calls: list[tuple[object, ...]] = field(default_factory=list) # mutable-ok: ordered raw registry IO trace + + def names(self, model: str | None = None, *, kind: str = "all", team_id: str | None = None) -> tuple[str, ...]: + self.calls.append(("names", model, kind, team_id)) + return ( + tuple(str(index) for index in range(len(self.deployments))) + if kind == "route_candidates" + else ("model", "chatgpt/model") + ) + + def deployment(self, model_id: str) -> Deployment | None: + self.calls.append(("deployment", model_id)) + return self.deployments[int(model_id)] + + def resolve(self, model: str, *, kind: str) -> str | None: + self.calls.append(("resolve", model, kind)) + return self.alias + + def access_groups(self, model_name: str, team_id: str | None = None) -> tuple[str, ...]: + self.calls.append(("groups", model_name, team_id)) + return self.groups.get(model_name, ()) + + +@dataclass(frozen=True, slots=True) +class Router: + reader: Registry + + def registry(self) -> RouterRegistry: + self.reader.calls.append(("registry",)) + return cast(RouterRegistry, self.reader) # cast-ok: raw registry double supplies the selected IO methods + + +def deployment( + name: str, backend: str = "openai/backend", provider: str | None = None, public: str | None = None +) -> Deployment: + return Deployment.model_validate( + { + "model_name": name, + "litellm_params": {"model": backend, "custom_llm_provider": provider}, + "model_info": {"team_public_model_name": public}, + } + ) + + +_DEPLOYMENTS: Final = ( + (), + (None,), + (deployment("model"),), + (deployment("target"),), + (deployment("internal", public="model"),), + (deployment("other"),), + (deployment("model", "chatgpt/backend"),), + (deployment("model", provider="chatgpt"),), + (deployment("model"), deployment("model", "chatgpt/backend")), + (None, deployment("model")), +) + + +@pytest.mark.parametrize("model", (None, [], 5, "", "model", "bundled", "priced", "chatgpt/model")) +@pytest.mark.parametrize("alias", (None, "target", "")) +@pytest.mark.parametrize("deployments", _DEPLOYMENTS) +@pytest.mark.parametrize("router_present", (False, True)) +def test_route_candidates_match_frozen_main_and_read_order( + monkeypatch: pytest.MonkeyPatch, + model: object, + alias: str | None, + deployments: tuple[Deployment | None, ...], + router_present: bool, +) -> None: + monkeypatch.setattr(native, "bundled_codex_models", lambda: {"bundled": None}) + prices: Final[Mapping[str, object]] = {"chatgpt/priced": {}} + monkeypatch.setattr(litellm, "model_cost", prices) + old: Final = Router(Registry(deployments, alias)) + new: Final = Router(Registry(deployments, alias)) + expected: Final = reference.routes_to_litellm(model, old if router_present else None, "team", ("bundled",), prices) + assert native.routes_to_litellm(model, new if router_present else None, "team") == expected + assert new.reader.calls == old.reader.calls + + +@pytest.mark.parametrize("model", ("model", "bundled", "chatgpt/model", "")) +@pytest.mark.parametrize( + "key", ((), ("model",), ("chatgpt/model",), ("other",), ("*",), ("all-team-models",), ("group",)) +) +@pytest.mark.parametrize("team", ((), ("model",), ("chatgpt/model",), ("other",))) +@pytest.mark.parametrize("deployments", _DEPLOYMENTS) +def test_chatgpt_prefix_matches_frozen_main_and_grant_read_order( + model: str, key: tuple[str, ...], team: tuple[str, ...], deployments: tuple[Deployment | None, ...] +) -> None: + groups: Final = {"model": ("group",), "chatgpt/model": ("group",)} + old: Final = Router(Registry(deployments, groups=groups)) + new: Final = Router(Registry(deployments, groups=groups)) + token: Final = UserAPIKeyAuth(models=list(key), team_models=list(team), team_id="team") + assert native.native_chatgpt_model(model, new, token) == reference.native_chatgpt_model(model, old, token) + assert new.reader.calls == old.reader.calls + + +def test_native_model_failure_is_terminal_before_registry_reads(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = NativeAuthRuntimeError("injected native failure") + + def fail(_facts: object) -> object: + raise failure + + monkeypatch.setattr(_native, "request_auth_codex_model", fail) + router: Final = Router(Registry((deployment("model"),))) + with pytest.raises(NativeAuthRuntimeError) as error: + native.routes_to_litellm("model", router, "team") + assert error.value is failure + assert router.reader.calls == [] + + +@pytest.mark.parametrize("route", ("/backend-api/codex/responses", "/v1/responses", "/messages")) +def test_native_route_guard_precedes_all_model_reads(route: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(native, "bundled_codex_models", lambda: dict[str, object]()) + monkeypatch.setattr(litellm, "model_cost", {}) + router: Final = Router(Registry((deployment("model"),))) + assert native.is_codex_native_litellm_turn(route, "model", router, "team") == ( + route == "/backend-api/codex/responses" + ) + assert bool(router.reader.calls) == (route == "/backend-api/codex/responses") diff --git a/tests/unit/_auth/test_codex_normalization.py b/tests/unit/_auth/test_codex_normalization.py new file mode 100644 index 00000000000..3108271f83b --- /dev/null +++ b/tests/unit/_auth/test_codex_normalization.py @@ -0,0 +1,168 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest +from fastapi import Request + +from litellm._auth.codex_normalization import CodexNormalizationFacts, codex_normalization +from litellm._auth.codex_turn import CodexAccessFacts, CodexTurnFacts +from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_normalization_adapter import resolve_codex_native_chatgpt_model +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_reference import RouteChecks as FrozenRoutes +from tests.unit._auth.test_codex_models import Registry, Router, deployment +from tests.unit._auth.test_codex_turn import ( + _KEY, # pyright: ignore[reportPrivateUsage] # immutable raw key facts shared by the composed planner + _ROUTING, # pyright: ignore[reportPrivateUsage] # immutable raw routing facts shared by the composed planner +) +from tests.unit._auth.test_codex_turn import reference as reference_turn + +_NATIVE_ROUTE: Final = "/backend-api/codex/responses" +_TURN: Final = CodexTurnFacts( + "model", access=CodexAccessFacts(_KEY, False), routing=_ROUTING, catalog_loaded=True, slugs=() +) +_ALIAS: Final = CodexNormalizationFacts( + "Alias", + "/chat/completions", + True, + True, + context=route_context("/chat/completions"), + resolved=False, + marked=True, + model_loaded=True, + model="model", + explicit=False, + provider=False, + routing=_ROUTING, + body_present=True, + settings_loaded=True, + target="target", + provider_after=False, + turn=_TURN, +) + + +def reference(facts: CodexNormalizationFacts) -> str | None: + """Completed alias reads in _resolve_router_settings_model_group_alias on main 2f430873aead.""" + if not facts.request_present or not facts.router_present or not FrozenRoutes.is_llm_api_route(facts.route): + return None + if facts.resolved: + return None + if facts.explicit or facts.provider and facts.route != _NATIVE_ROUTE: + return None + if facts.provider and facts.routing is not None and facts.model in (facts.routing.bundled or ()): + return None + if facts.model is None or not facts.body_present: + return None + if not facts.target or facts.target == facts.model: + return None + if facts.provider_after and (facts.turn is None or not reference_turn(facts.turn)): + return None + return facts.target + + +def native(facts: CodexNormalizationFacts) -> str | None: + decision: Final = codex_normalization(facts) + assert decision.kind in ("Complete", "Rewrite") + return decision.target + + +@pytest.mark.parametrize("policy", (reference, native)) +@pytest.mark.parametrize("route", ("/key/info", "/chat/completions", _NATIVE_ROUTE)) +@pytest.mark.parametrize("provider", (False, True)) +@pytest.mark.parametrize("target", (None, "", "model", "target")) +@pytest.mark.parametrize("body", (False, True)) +def test_alias_mutation_only_after_main_admission( + policy: Callable[[CodexNormalizationFacts], str | None], route: str, provider: bool, target: str | None, body: bool +) -> None: + facts: Final = replace( + _ALIAS, + route=route, + context=route_context(route), + provider=provider, + provider_after=provider, + target=target, + body_present=body, + ) + expected: Final = ( + "target" + if route != "/key/info" and (not provider or route == _NATIVE_ROUTE) and target == "target" and body + else None + ) + assert policy(facts) == expected + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (replace(_ALIAS, request_present=False, context=None), "Complete"), + (replace(_ALIAS, router_present=False, context=None), "Complete"), + (replace(_ALIAS, context=None), "ReadContext"), + (replace(_ALIAS, resolved=None), "ReadResolved"), + (replace(_ALIAS, resolved=True), "Complete"), + (replace(_ALIAS, marked=False, explicit=None), "MarkResolved"), + (replace(_ALIAS, model_loaded=False, explicit=None), "ReadModel"), + (replace(_ALIAS, explicit=True, body_present=None), "Complete"), + (replace(_ALIAS, explicit=None), "ReadExplicit"), + (replace(_ALIAS, provider=None), "ReadProvider"), + (replace(_ALIAS, body_present=None), "ReadBody"), + (replace(_ALIAS, settings_loaded=False), "ReadSettings"), + (replace(_ALIAS, provider_after=None), "ReadProviderAfter"), + ( + replace(_ALIAS, route=_NATIVE_ROUTE, context=route_context(_NATIVE_ROUTE), provider=True, routing=None), + "ReadRouting", + ), + ( + replace( + _ALIAS, + route=_NATIVE_ROUTE, + context=route_context(_NATIVE_ROUTE), + provider=True, + provider_after=True, + turn=None, + ), + "ReadTurn", + ), + ( + replace( + _ALIAS, + route=_NATIVE_ROUTE, + context=route_context(_NATIVE_ROUTE), + provider=True, + provider_after=True, + turn=replace(_TURN, slugs=("model",)), + ), + "Complete", + ), + (CodexNormalizationFacts("Prefix", _NATIVE_ROUTE, True, model="model", model_loaded=True), "ReadRouter"), + (CodexNormalizationFacts("Prefix", _NATIVE_ROUTE, True, True, model="model", model_loaded=True), "ReadPrefix"), + (CodexNormalizationFacts("Prefix", "/v1/responses", True, model="model", model_loaded=True), "Complete"), + ), +) +def test_normalization_read_and_mutation_order(facts: CodexNormalizationFacts, expected: str) -> None: + assert codex_normalization(facts).kind == expected + + +@pytest.mark.parametrize("requested_scope", (None, "client-alias")) +def test_prefix_adapter_preserves_client_model_and_rewrites_request( + monkeypatch: pytest.MonkeyPatch, requested_scope: str | None +) -> None: + router: Final = Router(Registry((deployment("chatgpt/model", "chatgpt/backend"),))) + monkeypatch.setattr(proxy_server, "llm_router", router) + token: Final = UserAPIKeyAuth(models=["chatgpt/model"], team_models=["chatgpt/model"]) + request: Final = Request({"type": "http", "path": _NATIVE_ROUTE, "headers": [], "method": "POST"}) + if requested_scope is not None: + request.scope[CLIENT_REQUESTED_MODEL_SCOPE_KEY] = requested_scope + data: Final[dict[str, object]] = { # mutable-ok: production mutates request-local body + "model": "model", + "input": "hello", + } + resolve_codex_native_chatgpt_model(data, token, request, _NATIVE_ROUTE) + assert data == {"model": "chatgpt/model", "input": "hello"} + assert request.scope[CLIENT_REQUESTED_MODEL_SCOPE_KEY] == (requested_scope or "model") + before: Final = tuple(router.reader.calls) + resolve_codex_native_chatgpt_model(data, token, request, _NATIVE_ROUTE) + assert tuple(router.reader.calls) == before diff --git a/tests/unit/_auth/test_codex_turn.py b/tests/unit/_auth/test_codex_turn.py new file mode 100644 index 00000000000..132f4377edb --- /dev/null +++ b/tests/unit/_auth/test_codex_turn.py @@ -0,0 +1,226 @@ +from collections.abc import Callable, Mapping +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.codex_models import CodexCandidate, CodexModelFacts +from litellm._auth.codex_turn import ( + CodexAccessFacts, + CodexCredentials, + CodexKeyFacts, + CodexTurnFacts, + JwtRead, + codex_access, + codex_jwt, + codex_jwt_token, + codex_key, + codex_turn, +) +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.route_enable import RouteEnableFacts +from litellm._auth.route_scope import RouteScopeFacts, RouteValues +from litellm._auth.routes import VirtualRouteFacts +from litellm.proxy.auth import handle_jwt +from litellm.proxy.auth.codex_turn_adapter import read_jwt_result +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_reference import RouteChecks as FrozenRoutes + +_CONTEXT: Final = route_context("/v1/responses") +_KEY: Final = CodexKeyFacts( + RouteScopeFacts("Proxy", _CONTEXT.route, False, False), + RouteEnableFacts(_CONTEXT, failed="Other"), + VirtualRouteFacts(_CONTEXT, None), + True, + "internal_user", + RouteScopeFacts("Admin", _CONTEXT.route, False, False), +) +_ROUTING: Final = CodexModelFacts( + "Routes", + "model", + True, + bundled=(), + priced=False, + candidates=(CodexCandidate(True, "openai", "openai/backend", "model"),), + alias_loaded=True, +) +_KEYS: Final = ( + _KEY, + replace(_KEY, proxy=RouteScopeFacts("Proxy", _CONTEXT.route, True, True, values=RouteValues())), + replace(_KEY, enabled=RouteEnableFacts(_CONTEXT, failed="Http")), + replace(_KEY, virtual_route=VirtualRouteFacts(_CONTEXT, ("/other",))), + replace( + _KEY, admin=RouteScopeFacts("Admin", _CONTEXT.route, True, True, values=RouteValues(names=(_CONTEXT.route,))) + ), + replace( + _KEY, + role="proxy_admin", + admin=RouteScopeFacts("Admin", _CONTEXT.route, True, True, values=RouteValues(names=(_CONTEXT.route,))), + ), +) + + +def _scope(facts: RouteScopeFacts | None) -> bool: + assert facts is not None + if not facts.present or facts.kind == "Admin" and not facts.premium: + return True + assert facts.values is not None + values: Final = facts.values.text if facts.values.text is not None else facts.values.names + return (facts.route in values) == (facts.kind == "Proxy") + + +def reference(facts: CodexTurnFacts) -> bool: + """Completed raw facts evaluated in main 2f430873aead order, without native policy.""" + credentials: Final = facts.credentials + if facts.native and ( + credentials is None or not credentials.bearer_present or not credentials.gateway_present or credentials.equal + ): + return False + routing: Final = facts.routing + if ( + facts.model is None + or routing is None + or not routing.router_present + or routing.priced + or facts.model in (routing.bundled or ()) + ): + return False + candidates: Final = routing.candidates + names: Final = (facts.model, routing.alias or facts.model) + if not candidates or any( + candidate.provider == "chatgpt" + or (candidate.model or "").partition("/")[0] == "chatgpt" + or (candidate.name not in names and candidate.public_name not in names) + for candidate in candidates + ): + return False + access: Final = facts.access + assert access is not None + key: Final = access.key + if not _scope(key.proxy): + return False + assert key.enabled is not None and key.virtual_route is not None + if key.enabled.failed == "Http": + return False + if key.virtual_route.allowed_routes and not FrozenRoutes.check_route_access( + _CONTEXT.route, key.virtual_route.allowed_routes + ): + return False + if key.role != "proxy_admin" and not _scope(key.admin): + return False + if access.jwt_token_present and access.jwt == "Exception": + return False + return facts.slugs is not None and facts.model not in facts.slugs + + +def native(facts: CodexTurnFacts) -> bool: + decision: Final = codex_turn(facts) + assert decision.kind == "Complete" + return decision.allowed + + +@pytest.mark.parametrize("policy", (reference, native)) +@pytest.mark.parametrize("key", _KEYS) +@pytest.mark.parametrize("jwt", (None, "Returned", "Exception")) +@pytest.mark.parametrize( + "credentials", + ( + CodexCredentials(True, True, False), + CodexCredentials(False, True, False), + CodexCredentials(True, False, False), + CodexCredentials(True, True, True), + ), +) +@pytest.mark.parametrize("native_turn", (False, True)) +@pytest.mark.parametrize("routed", (False, True)) +@pytest.mark.parametrize("slugs", (None, (), ("other",), ("model",))) +def test_turn_destination_matches_main( + policy: Callable[[CodexTurnFacts], bool], + key: CodexKeyFacts, + jwt: JwtRead | None, + credentials: CodexCredentials, + native_turn: bool, + routed: bool, + slugs: tuple[str, ...] | None, +) -> None: + facts: Final = CodexTurnFacts( + "model", + native_turn, + credentials, + CodexAccessFacts(key, jwt is not None, jwt), + _ROUTING if routed else replace(_ROUTING, bundled=("model",)), + True, + slugs, + ) + expected: Final = ( + (not native_turn or credentials == CodexCredentials(True, True, False)) + and routed + and key in (_KEYS[0], _KEYS[5]) + and jwt != "Exception" + and slugs is not None + and "model" not in slugs + ) + assert policy(facts) == expected + + +@pytest.mark.parametrize( + ("facts", "stage"), + ( + (CodexTurnFacts("model", native=True), "ReadCredentials"), + (CodexTurnFacts("model", native=True, credentials=CodexCredentials(True, True, False)), "ReadAccess"), + (CodexTurnFacts("model", access=CodexAccessFacts(_KEYS[1], False)), "ReadRouting"), + (CodexTurnFacts("model", access=CodexAccessFacts(_KEY, True), routing=_ROUTING), "ReadJwt"), + (CodexTurnFacts("model", access=CodexAccessFacts(_KEY, False), routing=_ROUTING), "ReadCatalog"), + ), +) +def test_turn_retains_main_read_order(facts: CodexTurnFacts, stage: str) -> None: + assert codex_turn(facts).kind == stage + + +@pytest.mark.parametrize( + ("facts", "stage"), + ( + (CodexKeyFacts(), "ReadProxy"), + (CodexKeyFacts(proxy=_KEY.proxy), "ReadEnabled"), + (CodexKeyFacts(proxy=_KEY.proxy, enabled=_KEY.enabled), "ReadVirtual"), + (replace(_KEY, role_loaded=False), "ReadRole"), + (replace(_KEY, admin=None), "ReadAdmin"), + ), +) +def test_key_route_reads_are_selected_in_order(facts: CodexKeyFacts, stage: str) -> None: + assert codex_key(facts).kind == stage + + +@pytest.mark.parametrize("claims", (False, True)) +@pytest.mark.parametrize("token", (False, True)) +def test_standard_jwt_selection(claims: bool, token: bool) -> None: + assert codex_jwt_token(claims_present=claims, token_present=token).allowed == (claims and not token) + + +@pytest.mark.parametrize("result", ("Returned", "Exception")) +def test_jwt_validator_return_is_not_reinterpreted(result: JwtRead) -> None: + assert codex_jwt(result).allowed == (result == "Returned") + assert codex_access(CodexAccessFacts(_KEYS[1], True, result)).allowed is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", (None, ValueError("validator rejected"), NativeAuthRuntimeError("native failed"))) +async def test_jwt_wrapper_never_turns_a_native_fault_into_subscription_fallback( + monkeypatch: pytest.MonkeyPatch, failure: Exception | None +) -> None: + class Validator: + @staticmethod + async def auth_builder( + *, api_key: str, request_data: Mapping[str, object], general_settings: Mapping[str, object], route: str + ) -> bool: + if failure is not None: + raise failure + return False + + monkeypatch.setattr(handle_jwt, "JWTAuthManager", Validator) + if isinstance(failure, NativeAuthRuntimeError): + with pytest.raises(NativeAuthRuntimeError) as raised: + await read_jwt_result("test-credential", "model", {}) + assert raised.value is failure + else: + assert await read_jwt_result("test-credential", "model", {}) == ("Returned" if failure is None else "Exception") diff --git a/tests/unit/_auth/test_common_budgets.py b/tests/unit/_auth/test_common_budgets.py new file mode 100644 index 00000000000..3d3efb4fdcd --- /dev/null +++ b/tests/unit/_auth/test_common_budgets.py @@ -0,0 +1,134 @@ +from collections.abc import Callable +from dataclasses import replace +from math import inf, nan +from typing import Final, get_args + +import pytest + +from litellm._auth import common_budgets as native +from litellm._auth.common_budgets import ( + AlertFacts, + BudgetDecision, + BudgetFacts, + BudgetScope, + MemberLimitFacts, + WindowFacts, +) + +from .budget_reference import common as reference + + +@pytest.mark.parametrize("scope", get_args(BudgetScope)) +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, nan, inf, -inf)) +@pytest.mark.parametrize("spend", (None, -1.0, 0.0, 1.0, 2.0, nan, inf, -inf)) +def test_numeric_budget_case_table_matches_reference( + scope: BudgetScope, limit: float | None, spend: float | None +) -> None: + facts: Final = BudgetFacts(scope, limit, spend, token_present=True, alert_recipients=("owner",)) + assert native.check(facts) == reference(facts) + + +@pytest.mark.parametrize("policy", (reference, native.check), ids=("reference", "native")) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (BudgetFacts("Team", 1, 1), BudgetDecision()), + (BudgetFacts("Key", 1, 1), BudgetDecision(denied=True, alert="token_budget")), + (BudgetFacts("Team", 1, 2, token_present=True), BudgetDecision(denied=True, alert="team_budget")), + (BudgetFacts("User", 1, 2, team_key=True), BudgetDecision()), + ( + BudgetFacts("User", 1, 2, team_key=True, apply_user_budget_to_team_keys=True), + BudgetDecision(denied=True, alert="user_budget"), + ), + (BudgetFacts("EndUser", 1, 2, info_route=True), BudgetDecision()), + (BudgetFacts("EndUser", 1, 1), BudgetDecision()), + (BudgetFacts("Project", inf), BudgetDecision()), + (BudgetFacts("Team", inf), BudgetDecision(read_spend=True)), + (BudgetFacts("Proxy", 1, 2, skip=True), BudgetDecision()), + (BudgetFacts("Proxy", 1, 2, llm_route=False), BudgetDecision()), + (BudgetFacts("Proxy", 1, 2, model_list_route=True), BudgetDecision()), + (BudgetFacts("Proxy", 1, 2), BudgetDecision(denied=True)), + ( + BudgetFacts("TeamSoft", 0, 0, token_present=True, alert_recipients=("a",)), + BudgetDecision(alert="soft_budget"), + ), + (BudgetFacts("TeamSoft", 1, 2, token_present=True), BudgetDecision()), + (BudgetFacts("ProjectSoft", 1, 2, token_present=True), BudgetDecision(alert="soft_budget")), + (BudgetFacts("KeySoft", 0, 1), BudgetDecision()), + (BudgetFacts("KeySoft", -1, 0), BudgetDecision(alert="soft_budget")), + (BudgetFacts("Tag", 0, nan), BudgetDecision(denied=True)), + (BudgetFacts("AccessGroup", nan, 0), BudgetDecision(denied=True)), + ), +) +def test_reference_and_native_obey_the_same_policy_table( + policy: Callable[[BudgetFacts], BudgetDecision], + facts: BudgetFacts, + expected: BudgetDecision, +) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize("percentage", (None, -1.0, 0.0, 0.5, 1.0, 1.1, nan, inf)) +@pytest.mark.parametrize("opt_in", (False, True)) +@pytest.mark.parametrize(("tpm", "rpm"), ((None, None), (0, None), (None, 0), (100, 100))) +def test_throttles_only_an_exceeded_opted_in_key( + percentage: float | None, opt_in: bool, tpm: int | None, rpm: int | None +) -> None: + facts: Final = BudgetFacts( + "Key", 1, 2, throttle_percentage=percentage, throttle_opt_in=opt_in, tpm_limit=tpm, rpm_limit=rpm + ) + assert native.check(facts) == reference(facts) + assert native.check(replace(facts, spend=0)) == BudgetDecision(alert="token_budget") + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (AlertFacts(8, 10), True), + (AlertFacts(10, 10), False), + (AlertFacts(0, 0), False), + (AlertFacts(nan, 10), False), + (AlertFacts(nan, 10, True, 80, team_member=True), True), + (AlertFacts(0, 0, True, 0, team_member=True), False), + (AlertFacts(10, 10, True, 80), True), + (AlertFacts(7, 10, True, 80), False), + (AlertFacts(10, 10, True, None), False), + (AlertFacts(10, 10, False, team_member=True), False), + ), +) +def test_threshold_alert_policy(facts: AlertFacts, expected: bool) -> None: + assert native.alert(facts) is expected + + +@pytest.mark.parametrize( + ("member", "default", "increase", "expiry", "expected"), + ( + (1, 5, 2, 101, 3), + (1, 5, 2, 100, 1), + (1, 5, 2, 99, 1), + (0, 5, 2, 101, 2), + (None, 5, 2, 101, 7), + (None, 0, 2, 101, None), + (None, -1, 2, 101, None), + (None, None, 2, 101, None), + (1, 5, None, 101, 1), + (1, 5, 2, None, 1), + ), +) +def test_temporary_member_budget_expiry_is_exclusive( + member: float | None, default: float | None, increase: float | None, expiry: float | None, expected: float | None +) -> None: + assert native.member_limit(MemberLimitFacts(member, default, increase, expiry, 100)) == expected + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (WindowFacts(60, 100, 200), 40), + (WindowFacts(60, None, 200), 140), + (WindowFacts(None, 100, 200), None), + (WindowFacts(0, 100, 200), 100), + ), +) +def test_counter_window_start(facts: WindowFacts, expected: float | None) -> None: + assert native.window_start(facts) == expected diff --git a/tests/unit/_auth/test_common_flow.py b/tests/unit/_auth/test_common_flow.py new file mode 100644 index 00000000000..365a10e654c --- /dev/null +++ b/tests/unit/_auth/test_common_flow.py @@ -0,0 +1,147 @@ +from itertools import product +from typing import Final + +import pytest + +from litellm._auth import common_flow + +_FACTS: Final = tuple( + common_flow.Facts(model, skip, team, user, token, user_id, agent_id, customer, models) + for model, skip, team, user, token, user_id, agent_id, customer, models in product( + (False, True), + (False, True), + (False, True), + (False, True), + (False, True), + (None, "", "user"), + (None, "", "agent"), + (False, True), + ((), ("model",)), + ) +) +_BUDGETS: Final = tuple( + common_flow.BudgetFacts(token, groups, customer, table) + for token, groups, customer, table in product((False, True), (0, 1), (False, True), (False, True)) +) +_CUSTOM: Final = tuple( + common_flow.CustomFacts(end_user_id, end_user, common, model, router) + for end_user_id, end_user, common, model, router in product( + (None, "", "customer"), (False, True), (False, True), (False, True), (False, True) + ) +) + + +@pytest.mark.parametrize("facts", _FACTS) +def test_common_model_plan_matches_reference(facts: common_flow.Facts) -> None: + assert common_flow.plan(facts) == common_flow.Plan( + membership=facts.token_present + and (facts.model_present or not facts.skip_budgets) + and facts.team_present + and facts.user_id is not None, + team_models=facts.model_present and facts.team_present, + member_models=facts.model_present and facts.team_present and facts.token_present and bool(facts.user_id), + user_models=facts.model_present and not facts.team_present and facts.user_present, + customer_models=facts.customer_present and bool(facts.customer_models), + agent_trace=facts.token_present and bool(facts.agent_id), + ) + + +@pytest.mark.parametrize("entity", (False, True)) +@pytest.mark.parametrize("budget", (False, True)) +@pytest.mark.parametrize("mapping", (False, True)) +@pytest.mark.parametrize("team", (False, True)) +def test_model_budget_scope_matches_reference(entity: bool, budget: bool, mapping: bool, team: bool) -> None: + expected: Final = "CheckModelBudgets" if entity and budget and mapping and team else "Continue" + assert common_flow.read(common_flow.ModelBudgets("Member", entity, budget, mapping, team)) == expected + + +@pytest.mark.parametrize("router", (False, True)) +@pytest.mark.parametrize("disabled", (False, True)) +def test_fallback_eligibility_matches_reference(router: bool, disabled: bool) -> None: + expected: Final = "AllowFallback" if router and not disabled else "Continue" + assert common_flow.read(common_flow.FallbackAllowed(router, disabled)) == expected + + +@pytest.mark.parametrize( + "facts, expected", + ( + (common_flow.ModelBudgets("Key"), "ReadModelBudgetPresence"), + (common_flow.ModelBudgets("Key", budget_present=False), "Continue"), + (common_flow.ModelBudgets("Key", budget_present=True), "ReadModelBudgetToken"), + (common_flow.ModelBudgets("User"), "ReadModelBudgetUser"), + (common_flow.ModelBudgets("User", entity_present=False), "Continue"), + (common_flow.ModelBudgets("User", entity_present=True), "ReadModelBudgetShape"), + (common_flow.ModelBudgets("User", entity_present=True, mapping=False), "Continue"), + (common_flow.ModelBudgets("Member", entity_present=True), "ReadModelBudgetTeam"), + (common_flow.ModelBudgets("Member", entity_present=True, team_present=False), "Continue"), + (common_flow.FallbackAllowed(False), "Continue"), + (common_flow.FallbackAllowed(True), "ReadFallbackDisabled"), + ), +) +def test_scope_reads_preserve_reference_short_circuit( + facts: common_flow.ModelBudgets | common_flow.FallbackAllowed, expected: common_flow.ReadPlan +) -> None: + assert common_flow.read(facts) == expected + + +@pytest.mark.parametrize("facts", _BUDGETS) +def test_budget_scopes_preserve_reference_priority(facts: common_flow.BudgetFacts) -> None: + assert common_flow.budgets(facts) == ( + "Team", + "TeamWindows", + *(("KeyWindows",) if facts.token_present else ()), + "TeamSoft", + "Organization", + "Tags", + *(("Groups",) if facts.group_count else ()), + "User", + "Member", + *(("Customer",) if facts.customer_present and facts.customer_budget_present else ()), + ) + + +@pytest.mark.parametrize("facts", _CUSTOM) +def test_custom_flow_matches_reference(facts: common_flow.CustomFacts) -> None: + assert common_flow.custom(facts) == common_flow.CustomPlan( + facts.end_user_id is not None, + facts.end_user_present and not facts.common_checks, + facts.common_checks, + facts.model_present and facts.router_present, + ) + + +@pytest.mark.parametrize("skip", (False, True)) +@pytest.mark.parametrize("exempt", (None, False, True)) +def test_budget_reads_preserve_short_circuit(skip: bool, exempt: bool | None) -> None: + expected: Final = "SkipBudgets" if skip or exempt else "ReadExemption" if exempt is None else "CheckBudgets" + assert common_flow.read(common_flow.BudgetSkip(skip, exempt)) == expected + + +@pytest.mark.parametrize("enabled", (False, True)) +@pytest.mark.parametrize("route", (None, False, True)) +def test_pricing_read_order(enabled: bool, route: bool | None) -> None: + expected: Final = "Continue" if not enabled or route is False else "ReadRoute" if route is None else "ReadUnpriced" + assert common_flow.read(common_flow.Pricing(enabled, route)) == expected + + +@pytest.mark.parametrize("count", (0, 1, 2)) +def test_unpriced_denial(count: int) -> None: + assert common_flow.read(common_flow.Unpriced(count)) == ("RejectUnpriced" if count else "Continue") + + +@pytest.mark.parametrize("router", (False, True)) +@pytest.mark.parametrize("multipart", (None, False, True)) +@pytest.mark.parametrize("resource", (None, False, True)) +def test_fallback_reads_preserve_body_binding_guards( + router: bool, multipart: bool | None, resource: bool | None +) -> None: + expected: Final = ( + "Continue" + if not router or multipart or resource + else "ReadContentType" + if multipart is None + else "ReadResourceBinding" + if resource is None + else "ReadFallbackSettings" + ) + assert common_flow.read(common_flow.MemberFallback(router, multipart, resource)) == expected diff --git a/tests/unit/_auth/test_credential.py b/tests/unit/_auth/test_credential.py new file mode 100644 index 00000000000..5c9c21a1d10 --- /dev/null +++ b/tests/unit/_auth/test_credential.py @@ -0,0 +1,161 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from tests.unit._auth.conftest import CredentialSelector + +from litellm._auth.credential import ( + KeySource, + PresentedKeys, + SelectedKey, + WebSocketDenial, + WebSocketKeys, +) + +AWS4_BEARER: Final = ( + "AWS4-HMAC-SHA256 Credential=Bearer token/date/region/service/aws4_request, SignedHeaders=host, Signature=sig" +) +AWS4_RAW: Final = ( + "AWS4-HMAC-SHA256 Credential=token/date/region/service/aws4_request, SignedHeaders=host, Signature=sig" +) + +PRECEDENCE: Final = ( + ("custom_header", KeySource.CUSTOM_HEADER, "Bearer {}"), + ("litellm_header", KeySource.LITELLM_HEADER, "Bearer {}"), + ("authorization", KeySource.AUTHORIZATION, "Bearer {}"), + ("api_key", KeySource.API_KEY, "{}"), + ("x_api_key", KeySource.X_API_KEY, "{}"), + ("google", KeySource.GOOGLE, "{}"), + ("apim", KeySource.APIM, "{}"), + ("query", KeySource.QUERY, "{}"), + ("pass_through", KeySource.PASS_THROUGH, "{}"), +) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("Bearer token", "token"), + ("bearer token", "token"), + ("Basic dXNlcjpwYXNz", "dXNlcjpwYXNz"), + (AWS4_BEARER, "token"), + (AWS4_RAW, "token"), + ("Bearer ", ""), + ("Bearer token", " token"), + ("BEARER token", "token"), + ("BASIC dXNlcjpwYXNz", "dXNlcjpwYXNz"), + ("AWS4-HMAC-SHA256 SignedHeaders=host", None), + ("token", None), + ("Bearer", None), + ("Token token", None), + ("Bearertoken", None), + ], +) +def test_scheme_credential_strips_a_known_scheme_in_any_case( + credentials: CredentialSelector, value: str, expected: str | None +): + assert credentials.scheme_credential(value) == expected + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("token", "token"), + ("Bearer token", "token"), + (AWS4_BEARER, "token"), + ("AWS4-HMAC-SHA256 malformed", "AWS4-HMAC-SHA256 malformed"), + ("Bearer ", ""), + ], +) +def test_litellm_header_strips_a_known_scheme_or_keeps_the_raw_value( + credentials: CredentialSelector, value: str, expected: str +): + assert credentials.litellm_header_token(value) == expected + + +@pytest.mark.parametrize("rank", range(len(PRECEDENCE))) +def test_higher_priority_source_wins_over_every_lower_one(credentials: CredentialSelector, rank: int): + higher, source, template = PRECEDENCE[rank] + winner: Final = template.format(higher) + losers: Final = {field: lower_template.format("loser") for field, _, lower_template in PRECEDENCE[rank + 1 :]} + assert credentials.select_api_key(PresentedKeys(**losers, **{higher: winner})) == SelectedKey( + source=source, presented=winner, token=higher + ) + + +@pytest.mark.parametrize( + ("keys", "expected"), + [ + (PresentedKeys(custom_header="custom", x_api_key="lower"), ""), + (PresentedKeys(custom_header="", x_api_key="lower"), ""), + (PresentedKeys(authorization="token", x_api_key="lower"), ""), + (PresentedKeys(authorization="Token x", x_api_key="lower"), ""), + ], +) +def test_selected_source_without_a_known_scheme_yields_an_empty_token( + credentials: CredentialSelector, keys: PresentedKeys, expected: str +): + selected: Final = credentials.select_api_key(keys) + assert selected is not None + assert selected.token == expected + + +@pytest.mark.parametrize( + ("keys", "expected_source"), + [ + (PresentedKeys(authorization="", x_api_key="lower"), KeySource.X_API_KEY), + (PresentedKeys(query="", apim="lower"), KeySource.APIM), + (PresentedKeys(litellm_header=""), KeySource.LITELLM_HEADER), + (PresentedKeys(api_key=""), KeySource.API_KEY), + ], +) +def test_only_empty_authorization_and_query_fall_through( + credentials: CredentialSelector, keys: PresentedKeys, expected_source: KeySource +): + selected: Final = credentials.select_api_key(keys) + assert selected is not None + assert selected.source is expected_source + + +@pytest.mark.parametrize("keys", [PresentedKeys(), PresentedKeys(authorization=""), PresentedKeys(query="")]) +def test_nothing_presented_selects_nothing(credentials: CredentialSelector, keys: PresentedKeys): + assert credentials.select_api_key(keys) is None + + +def test_presented_keeps_the_raw_value_of_the_selected_source( + credentials: CredentialSelector, +): + keys: Final = PresentedKeys(authorization=AWS4_BEARER) + assert credentials.select_api_key(keys) == SelectedKey(KeySource.AUTHORIZATION, AWS4_BEARER, "token") + assert credentials.select_api_key(replace(keys, custom_header="Basic abc")) == SelectedKey( + KeySource.CUSTOM_HEADER, "Basic abc", "abc" + ) + + +@pytest.mark.parametrize( + ("keys", "expected"), + [ + (WebSocketKeys(authorization="Bearer token"), "token"), + (WebSocketKeys(authorization="bearer token "), "token"), + ( + WebSocketKeys(authorization="Bearer token", api_key="other", subprotocols="openai-insecure-api-key.x"), + "token", + ), + (WebSocketKeys(authorization="Basic token"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(authorization="token"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(authorization="Basic token", api_key="other"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(api_key="token"), "token"), + (WebSocketKeys(authorization="", api_key="token"), "token"), + (WebSocketKeys(api_key="token", subprotocols="openai-insecure-api-key.other"), "token"), + (WebSocketKeys(subprotocols="realtime, openai-insecure-api-key.token, openai-beta.realtime-v1"), "token"), + (WebSocketKeys(api_key="", subprotocols="openai-insecure-api-key.token"), "token"), + (WebSocketKeys(subprotocols="realtime"), WebSocketDenial.NO_API_KEY), + (WebSocketKeys(subprotocols="openai-insecure-api-key."), WebSocketDenial.NO_API_KEY), + (WebSocketKeys(), WebSocketDenial.NO_API_KEY), + ], +) +def test_websocket_takes_a_bearer_authorization_then_api_key_then_subprotocol( + credentials: CredentialSelector, keys: WebSocketKeys, expected: str | WebSocketDenial +): + assert credentials.select_websocket_api_key(keys) == expected diff --git a/tests/unit/_auth/test_failure.py b/tests/unit/_auth/test_failure.py new file mode 100644 index 00000000000..d2e42ac3455 --- /dev/null +++ b/tests/unit/_auth/test_failure.py @@ -0,0 +1,36 @@ +import pytest + +from litellm._auth.failure import NativeAuthFailure, native_boundary +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentAccess, AgentRequestHandler + + +@pytest.mark.parametrize("error", (TypeError("bad fact"), ValueError("bad wire"), RuntimeError("native error"))) +def test_native_errors_keep_public_exception_type_and_are_marked_terminal(error: Exception) -> None: + @native_boundary + def entry(value: int) -> int: + raise error + + with pytest.raises(type(error), match=str(error)) as raised: + entry(1) + assert isinstance(raised.value, NativeAuthFailure) + assert raised.value.__cause__ is error + + +@pytest.mark.asyncio +async def test_native_failure_never_becomes_unrestricted_agent_access() -> None: + @native_boundary + def broken_entry() -> None: + raise ValueError("invalid native policy answer") + + async def key_lookup(auth: UserAPIKeyAuth | None, *, strict: bool) -> AgentAccess: + broken_entry() + raise AssertionError("native failure returned") + + async def team_lookup(auth: UserAPIKeyAuth | None, *, strict: bool) -> AgentAccess: + raise AssertionError("team lookup followed native failure") + + with pytest.raises(NativeAuthFailure, match="invalid native policy answer"): + await AgentRequestHandler.resolve_key_team_agent_access( + UserAPIKeyAuth(), lookup_key_access=key_lookup, lookup_team_access=team_lookup + ) diff --git a/tests/unit/_auth/test_fallback_access.py b/tests/unit/_auth/test_fallback_access.py new file mode 100644 index 00000000000..ba7c7036c02 --- /dev/null +++ b/tests/unit/_auth/test_fallback_access.py @@ -0,0 +1,106 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import ( + FallbackAuthorizationDecision, + FallbackAuthorizationFacts, + GroupedModelScope, + fallback_authorization, +) +from litellm._auth.model_access import ModelAccessFacts +from tests.unit._auth.support.fallback_reference import authorization + +_ALLOW: Final = ModelAccessFacts(("model",), ("model",), object_type="key") +_DENY: Final = replace(_ALLOW, grants=("other",)) +_BASE: Final = FallbackAuthorizationFacts(GroupedModelScope(_ALLOW)) + + +@pytest.mark.parametrize("policy", (authorization, fallback_authorization)) +@pytest.mark.parametrize( + ("facts", "kind"), + ( + pytest.param(_BASE, "Allow", id="no-extra-ceilings"), + pytest.param(replace(_BASE, refuse_agent_keys=True, agent_id=""), "Deny", id="agent-presence-not-truthiness"), + pytest.param( + replace(_BASE, key=GroupedModelScope(_DENY, ("group",))), "ReadKeyGroups", id="key-groups-before-team" + ), + pytest.param( + replace(_BASE, key=GroupedModelScope(_DENY, ("group",), ())), "Deny", id="empty-key-group-retains-denial" + ), + pytest.param( + replace(_BASE, key=GroupedModelScope(_DENY, ("group",), ("model",))), "Allow", id="key-group-fallback" + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY, ("group",))), "ReadTeamGroups", id="team-group-fallback" + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY)), "ReadAssignedGroups", id="team-denial-asks-owned-key-grants" + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY), assigned_models=()), + "Deny", + id="empty-assigned-does-not-widen", + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY), assigned_models=("model",)), + "ReadAssignedScope", + id="assigned-reads-scope-before-override", + ), + pytest.param( + replace( + _BASE, + team=GroupedModelScope(_DENY), + assigned_models=("model",), + assigned_scope=replace(_DENY, grants=("model",)), + ), + "Allow", + id="assigned-overrides-team", + ), + pytest.param( + replace(_BASE, embedded_team=_DENY, assigned_models=("model",)), + "Deny", + id="embedded-team-no-assigned-override", + ), + pytest.param(replace(_BASE, member=_DENY), "Deny", id="member-ceiling"), + pytest.param(replace(_BASE, project=_DENY), "Deny", id="project-ceiling"), + ), +) +def test_fallback_authorization_stage_order( + policy: Callable[[FallbackAuthorizationFacts], FallbackAuthorizationDecision], + facts: FallbackAuthorizationFacts, + kind: str, +) -> None: + assert policy(facts).kind == kind + + +def test_key_denial_wins_before_later_scopes() -> None: + decision: Final = fallback_authorization( + replace(_BASE, key=GroupedModelScope(_DENY), project=replace(_DENY, object_type="project")) + ) + assert decision.denial is not None + assert decision.denial.object_type == "key" + + +@pytest.mark.parametrize("policy", (authorization, fallback_authorization)) +@pytest.mark.parametrize( + ("facts", "kind"), + ( + (FallbackAuthorizationFacts(), "ReadKey"), + (FallbackAuthorizationFacts(refuse_agent_keys=True, agent_id="agent"), "Deny"), + (replace(_BASE, team_present=True), "ReadTeam"), + (replace(_BASE, embedded_team_grants=("model",)), "ReadEmbeddedTeam"), + (replace(_BASE, member_grants=("model",)), "ReadMember"), + (replace(_BASE, project_grants=("model",)), "ReadProject"), + (replace(_BASE, key=GroupedModelScope(_DENY), team_present=True, member_grants=("model",)), "Deny"), + (replace(_BASE, member=_DENY, project_grants=("model",)), "Deny"), + ), +) +def test_later_scope_projection_is_requested_only_after_prior_allow( + policy: Callable[[FallbackAuthorizationFacts], FallbackAuthorizationDecision], + facts: FallbackAuthorizationFacts, + kind: str, +) -> None: + assert policy(facts).kind == kind diff --git a/tests/unit/_auth/test_fallback_candidate.py b/tests/unit/_auth/test_fallback_candidate.py new file mode 100644 index 00000000000..f7ccd135b12 --- /dev/null +++ b/tests/unit/_auth/test_fallback_candidate.py @@ -0,0 +1,211 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.common_budgets import BudgetFacts +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope +from litellm._auth.fallback_candidate import CandidateBudgetRead, CandidateBudgetScope, fallback_candidate +from litellm._auth.fallback_candidate import FallbackCandidateDecision as Decision +from litellm._auth.fallback_candidate import FallbackCandidateFacts as Facts +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.model_budget_data import ModelBudgetEntity, ModelBudgetSnapshot +from litellm._auth.routes import RouteCatalog, RouteContext +from litellm._auth.unpriced_models import UnpricedModelFacts +from tests.unit._auth.support.fallback_reference import authorization +from tests.unit._auth.support.model_reference import model_access +from tests.unit._auth.test_model_attribution import reference as attribution + +_ENTITIES: Final[tuple[ModelBudgetEntity, ...]] = ("Key", "User", "EndUser", "Team", "TeamMember") + + +def reference(facts: Facts) -> Decision: + if not all((facts.router_present, facts.user_present, facts.team_present)) or facts.membership == "Failed": + return Decision("Complete") + if facts.membership == "Unread": + return Decision("ReadMembership") + if facts.io_failed: + return Decision("Complete") + if facts.block_unpriced: + if facts.route is None: + return Decision("ReadRoute") + if facts.route.route in facts.route.catalog.groups.get("openai_routes", ()): + if facts.unpriced is None: + return Decision("ReadUnpriced") + if ( + facts.unpriced.group_found + and facts.unpriced.deployments == () + and facts.unpriced.cost_deployments == () + ): + return Decision("Complete") + if facts.customer_present: + if facts.customer_grants is None: + return Decision("ReadCustomerGrants") + if facts.customer_grants: + if facts.customer is None: + return Decision("ReadCustomer") + if model_access(facts.customer).kind != "Allow": + return Decision("Complete", log_failure=True) + if facts.price is None: + return Decision("ReadPrice") + if not facts.price.cached: + if facts.budget_models is None: + return Decision("ReadBudgetModels") + if facts.budget_scopes is None: + return Decision("ReadBudgetScopes") + for index, model in enumerate(facts.budget_models): + if model in facts.budget_models[:index]: + continue + for entity in _ENTITIES: + if not any( + scope.entity == entity and scope.identity_present and scope.configured + for scope in facts.budget_scopes + ): + continue + reading = next( + (reading for reading in facts.budgets if reading.model_index == index and reading.entity == entity), + None, + ) + if reading is None: + return Decision("ReadBudget", model_index=index, entity=entity) + snapshot = reading.snapshot + spend = (snapshot.local_spend or 0.0) + (snapshot.legacy_spend or 0.0) + if snapshot.max_budget is not None and snapshot.max_budget >= 0 and spend >= snapshot.max_budget: + return Decision("Complete") + if facts.authorization is None: + return Decision("ReadAuthorization") + if authorization(facts.authorization).kind == "Deny": + return Decision("Complete") + if not facts.attribution_reset: + return Decision("ResetAttribution") + if facts.attribution is None: + return Decision("ReadAttribution") + groups: Final = attribution(facts.attribution).groups + if not facts.attribution_stamped: + return Decision("StampAttribution", groups=groups) + if groups: + if not facts.group_budgets_read: + return Decision("ReadGroupBudgets", groups=groups) + if any( + budget.limit is not None and budget.limit > 0 and budget.spend is not None and budget.spend >= budget.limit + for budget in facts.group_budgets + ): + return Decision("Complete", log_failure=True) + return Decision("Complete", usable=True) + + +_INITIAL: Final = Facts(True, True, True) +_MEMBER: Final = replace(_INITIAL, membership="Loaded") +_FREE: Final = replace(_MEMBER, price=ModelPriceFacts(cached=True)) +_PAID: Final = replace(_MEMBER, price=ModelPriceFacts(cached=False)) +_SCOPE: Final = ModelAccessFacts(("model",), ("model",), object_type="key") +_AUTH: Final = replace(_FREE, authorization=FallbackAuthorizationFacts(key=GroupedModelScope(_SCOPE))) +_RESET: Final = replace(_AUTH, attribution_reset=True) +_GROUPS: Final = ModelGroupAttributionFacts( + model_present=True, + token_present=True, + router_present=True, + database_present=True, + registry_loaded=True, + registry=("group",), + covering=("group",), + grants=(("group",), (), (), (), ()), +) +_ATTRIBUTED: Final = replace(_RESET, attribution=_GROUPS) +_STAMPED: Final = replace(_ATTRIBUTED, attribution_stamped=True) +_SCOPES: Final = tuple(CandidateBudgetScope(entity, True, True) for entity in _ENTITIES) +_BUDGET: Final = replace(_PAID, budget_models=("model",), budget_scopes=_SCOPES) +_ROUTE: Final = RouteContext( + "/chat/completions", catalog=RouteCatalog(groups={"openai_routes": ("/chat/completions",)}) +) + + +@pytest.mark.parametrize("policy", (reference, fallback_candidate)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_INITIAL, Decision("ReadMembership")), + (replace(_INITIAL, router_present=False), Decision("Complete")), + (replace(_INITIAL, user_present=False), Decision("Complete")), + (replace(_INITIAL, team_present=False), Decision("Complete")), + (replace(_INITIAL, membership="Failed"), Decision("Complete")), + (replace(_MEMBER, io_failed=True), Decision("Complete")), + (replace(_MEMBER, block_unpriced=True), Decision("ReadRoute")), + (replace(_MEMBER, block_unpriced=True, route=_ROUTE), Decision("ReadUnpriced")), + ( + replace( + _MEMBER, + block_unpriced=True, + route=_ROUTE, + unpriced=UnpricedModelFacts("model", True, group_found=True, deployments=(), cost_deployments=()), + customer_present=True, + ), + Decision("Complete"), + ), + (replace(_MEMBER, block_unpriced=True, route=RouteContext("/health")), Decision("ReadPrice")), + (replace(_MEMBER, customer_present=True), Decision("ReadCustomerGrants")), + (replace(_MEMBER, customer_present=True, customer_grants=()), Decision("ReadPrice")), + (replace(_MEMBER, customer_present=True, customer_grants=("model",)), Decision("ReadCustomer")), + ( + replace( + _MEMBER, customer_present=True, customer_grants=("other",), customer=replace(_SCOPE, grants=("other",)) + ), + Decision("Complete", log_failure=True), + ), + (_MEMBER, Decision("ReadPrice")), + (_FREE, Decision("ReadAuthorization")), + (_PAID, Decision("ReadBudgetModels")), + (replace(_PAID, budget_models=("model",)), Decision("ReadBudgetScopes")), + (_BUDGET, Decision("ReadBudget", entity="Key")), + ( + replace( + _BUDGET, budgets=(CandidateBudgetRead(0, "Key", ModelBudgetSnapshot(max_budget=1, local_spend=1)),) + ), + Decision("Complete"), + ), + (replace(_BUDGET, budget_scopes=()), Decision("ReadAuthorization")), + ( + replace(_FREE, authorization=FallbackAuthorizationFacts(refuse_agent_keys=True, agent_id="")), + Decision("Complete"), + ), + (_AUTH, Decision("ResetAttribution")), + (_RESET, Decision("ReadAttribution")), + (_ATTRIBUTED, Decision("StampAttribution", groups=("group",))), + (_STAMPED, Decision("ReadGroupBudgets", groups=("group",))), + ( + replace(_STAMPED, group_budgets_read=True, group_budgets=(BudgetFacts("AccessGroup", 1, 1),)), + Decision("Complete", log_failure=True), + ), + ( + replace(_STAMPED, group_budgets_read=True, group_budgets=(BudgetFacts("AccessGroup", 1, 0),)), + Decision("Complete", usable=True), + ), + (replace(_STAMPED, attribution=ModelGroupAttributionFacts(io_failed=True)), Decision("Complete", usable=True)), + ), +) +def test_candidate_staged_case_table(policy: Callable[[Facts], Decision], facts: Facts, expected: Decision) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize("policy", (reference, fallback_candidate)) +@pytest.mark.parametrize("index", range(5)) +def test_candidate_model_budget_order(policy: Callable[[Facts], Decision], index: int) -> None: + reads: Final = tuple(CandidateBudgetRead(0, entity, ModelBudgetSnapshot()) for entity in _ENTITIES[:index]) + assert policy(replace(_BUDGET, budgets=reads)) == Decision("ReadBudget", entity=_ENTITIES[index]) + + +@pytest.mark.parametrize("policy", (reference, fallback_candidate)) +def test_candidate_deduplicates_model_names_without_reordering(policy: Callable[[Facts], Decision]) -> None: + reads: Final = tuple(CandidateBudgetRead(0, entity, ModelBudgetSnapshot()) for entity in _ENTITIES) + assert policy(replace(_BUDGET, budget_models=("model", "model", "other"), budgets=reads)) == Decision( + "ReadBudget", model_index=2, entity="Key" + ) + + +def test_candidate_native_decode_failure_is_not_a_failed_candidate() -> None: + with pytest.raises(NativeAuthFailure): + fallback_candidate(replace(_BUDGET, budgets=(CandidateBudgetRead(-1, "Key", ModelBudgetSnapshot()),))) diff --git a/tests/unit/_auth/test_key.py b/tests/unit/_auth/test_key.py new file mode 100644 index 00000000000..1d41920e4dd --- /dev/null +++ b/tests/unit/_auth/test_key.py @@ -0,0 +1,35 @@ +import hashlib + +import pytest + +from litellm._auth.key import hash_token, is_master_key + + +@pytest.mark.parametrize("token", ["sk-1234", "", "sk-ünïcode"]) +def test_hash_token_is_sha256_hex_of_utf8(token: str): + assert hash_token(token) == hashlib.sha256(token.encode()).hexdigest() + + +def test_hash_token_differs_per_token(): + assert hash_token("sk-a") != hash_token("sk-b") + + +@pytest.mark.parametrize( + ("token", "master_key", "expected"), + [ + ("sk-master", "sk-master", True), + ("sk-master", "sk-other", False), + ("sk-mast", "sk-master", False), + (None, "sk-master", False), + ("sk-master", None, False), + (None, None, False), + ("sk-mästér", "sk-mästér", True), + ("sk-mästér", "sk-master", False), + ], +) +def test_is_master_key_matches_only_the_exact_raw_key(token: str | None, master_key: str | None, expected: bool): + assert is_master_key(token, master_key) is expected + + +def test_hashed_master_key_is_not_the_master_key(): + assert is_master_key(hash_token("sk-master"), "sk-master") is False diff --git a/tests/unit/_auth/test_key_model_access.py b/tests/unit/_auth/test_key_model_access.py new file mode 100644 index 00000000000..9456f307518 --- /dev/null +++ b/tests/unit/_auth/test_key_model_access.py @@ -0,0 +1,80 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts, key_and_fallback_model_access +from litellm._auth.model_access import ModelAccessFacts + +_ALLOW: Final = GroupedModelScope(ModelAccessFacts(("model",), ("model",), object_type="key")) +_DENY: Final = GroupedModelScope(replace(_ALLOW.facts, grants=("other",))) +_BASE: Final = KeyAndFallbackModelAccessFacts(False, ("model",)) + + +@pytest.mark.parametrize( + ("facts", "kind", "index"), + ( + pytest.param(_BASE, "ReadPrimary", 0, id="first-read"), + pytest.param(replace(_BASE, config_nonempty=True), "Allow", 0, id="config-before-primary"), + pytest.param(replace(_BASE, key_models=("all-team-models",)), "Allow", 0, id="team-sentinel-before-primary"), + pytest.param( + replace(_BASE, primary_loaded=True, primary=_DENY), "Deny", 0, id="primary-denial-before-fallback-parsing" + ), + pytest.param( + replace(_BASE, primary_loaded=True, primary=replace(_DENY, group_ids=("group",))), + "ReadPrimaryGroups", + 0, + id="key-group-before-fallback-parsing", + ), + pytest.param( + replace(_BASE, primary_loaded=True), "ReadFallbackNames", 0, id="no-primary-still-checks-fallbacks" + ), + pytest.param( + replace(_BASE, primary_loaded=True, fallback_names=("model",)), + "ReadFallbackScope", + 0, + id="fallback-needs-scope", + ), + pytest.param( + replace(_BASE, primary_loaded=True, fallback_names=("model",), fallback_scopes=(_ALLOW,)), + "RouteFallback", + 0, + id="allowed-fallback-then-probe", + ), + pytest.param( + replace(_BASE, primary_loaded=True, fallback_names=("model",), fallback_scopes=(_DENY,)), + "Deny", + 0, + id="denied-fallback-before-probe", + ), + pytest.param( + replace( + _BASE, + primary_loaded=True, + fallback_names=("first", "second"), + fallback_scopes=(_ALLOW,), + routing=("Failed",), + ), + "RoutingFailure", + 0, + id="routing-failure-before-next-scope", + ), + pytest.param( + replace( + _BASE, + primary_loaded=True, + fallback_names=("first", "second"), + fallback_scopes=(_ALLOW,), + routing=("Routed",), + ), + "ReadFallbackScope", + 1, + id="next-scope-after-routing", + ), + pytest.param(replace(_BASE, primary_loaded=True, fallback_names=()), "Allow", 0, id="empty-fallbacks"), + ), +) +def test_key_and_fallback_access_order(facts: KeyAndFallbackModelAccessFacts, kind: str, index: int) -> None: + decision: Final = key_and_fallback_model_access(facts) + assert (decision.kind, decision.index) == (kind, index) diff --git a/tests/unit/_auth/test_managed.py b/tests/unit/_auth/test_managed.py new file mode 100644 index 00000000000..9f33aab692e --- /dev/null +++ b/tests/unit/_auth/test_managed.py @@ -0,0 +1,84 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth import managed +from tests.unit._auth import reference_policy + +AGENT: Final = managed.Agent("agent", True, True, managed.BoundIdentity(True, "revision"), "both") +CONTEXT: Final = managed.Context("agent", "revision", "delegated", "user") + + +@pytest.mark.parametrize( + "agent,context", + ( + (AGENT, CONTEXT), + (replace(AGENT, enabled=False), None), + (replace(AGENT, identity=None), CONTEXT), + (replace(AGENT, identity=managed.BoundIdentity(False, "old")), CONTEXT), + (AGENT, None), + (AGENT, replace(CONTEXT, agent_id="other")), + (AGENT, replace(CONTEXT, binding_revision=None)), + (AGENT, replace(CONTEXT, binding_revision="old")), + (replace(AGENT, execution_mode="autonomous"), CONTEXT), + (AGENT, replace(CONTEXT, user_id=None)), + (AGENT, replace(CONTEXT, user_id="")), + ), +) +def test_actor_denial_order_preserves_verified_context_contract( + agent: managed.Agent, context: managed.Context | None +) -> None: + facts: Final = managed.Actor(agent, context) + assert managed.decide(facts) == reference_policy.actor(facts) + + +def test_actor_read_plan_assigns_policy_before_delegation_grant_read() -> None: + facts: Final = managed.ActorLookup("agent", True, context=CONTEXT) + assert managed.decide(facts).action == "ReadAgent" + loaded: Final = replace(facts, agent_loaded=True, agent=AGENT) + assert managed.decide(loaded).action == "AssignPolicy" + assigned: Final = replace(loaded, policy_assigned=True) + assert managed.decide(assigned).action == "ReadHumanGrants" + assert managed.decide(replace(assigned, human_grants=())).denial == "HumanNotPermitted" + assert managed.decide(replace(assigned, human_grants=("agent",))).action == "Allow" + assert managed.decide(replace(assigned, delegation_verified=True)).action == "Allow" + + +def test_missing_actor_still_reads_retirement_before_denial() -> None: + facts: Final = managed.ActorLookup("agent", True, agent_loaded=True, context=CONTEXT) + assert managed.decide(facts).action == "ReadRetired" + assert managed.decide(replace(facts, retired=False)).denial == "Removed" + assert managed.decide(replace(facts, retired=False, context=None)).action == "Allow" + + +def test_declared_grants_read_groups_in_order_and_keep_empty_declarations_restricted() -> None: + facts: Final = managed.DeclaredGrants((), ("group",), ("unified",)) + assert managed.decide(facts).action == "ReadDeclaredGroups" + grouped: Final = replace(facts, group_agents=()) + assert managed.decide(grouped).action == "ReadUnifiedGroups" + loaded: Final = replace(grouped, unified_agents=()) + assert managed.decide(loaded).grants == () + assert managed.decide(managed.DeclaredGrants((), (), ())).grants is None + + +@pytest.mark.parametrize("mode", (None, "autonomous", "delegated")) +@pytest.mark.parametrize("caller", (None, (), ("a",), ("a", "b"))) +@pytest.mark.parametrize("human", ((), ("b",), ("a", "b"))) +def test_managed_grants_intersect_each_authority_without_inheriting_unrestricted( + mode: str | None, caller: tuple[str, ...] | None, human: tuple[str, ...] +) -> None: + facts: Final = managed.ManagedAgentGrants( + True, + mode, + True, + own=("a", "b", "c"), + ceilings=(("a", "b"),), + caller_loaded=True, + caller=caller, + human=human, + ) + caller_ids: Final = frozenset(("a", "b", "c")) if caller is None else frozenset(caller) + human_ids: Final = frozenset(human) if mode == "delegated" else frozenset(("a", "b", "c")) + expected: Final = frozenset(("a", "b")) & caller_ids & human_ids + assert frozenset(managed.decide(facts).grants or ()) == expected diff --git a/tests/unit/_auth/test_managed_models.py b/tests/unit/_auth/test_managed_models.py new file mode 100644 index 00000000000..167ab0da702 --- /dev/null +++ b/tests/unit/_auth/test_managed_models.py @@ -0,0 +1,116 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.managed_models import ManagedModelDecision as Decision +from litellm._auth.managed_models import ManagedModelFacts as Facts +from litellm._auth.managed_models import managed_models +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access +from tests.unit._auth.test_member_models import reference as member_reference + + +def reference(facts: Facts) -> Decision: + if facts.purpose == "All" and not facts.managed_loaded: + return Decision("ReadManaged") + if not facts.model_present or not facts.token_present: + return Decision("Allow") + if facts.purpose == "All" and facts.managed_present: + if not facts.managed_models_valid or not facts.managed_models_count: + return Decision("Deny", ModelAccessDecision("NoAgentModels")) + if facts.managed is None: + return Decision("ReadManagedScope") + managed: Final = model_access(facts.managed) + if managed.kind != "Allow": + return Decision("Deny", managed) + if facts.purpose != "Caller": + if facts.ceilings is None: + return Decision("ReadCeilings") + if not facts.aliases_loaded: + return Decision("ReadAliases") + for index, ceiling in enumerate(facts.ceilings): + if not ceiling: + return Decision("Deny", ModelAccessDecision("EmptyAgentCeiling")) + if index >= len(facts.scopes): + return Decision("ReadCeilingScope", index=index) + if (scope := model_access(facts.scopes[index])).kind != "Allow": + return Decision("Deny", scope) + if facts.purpose == "Ceilings": + return Decision("Allow") + if not facts.caller_loaded: + return Decision("ReadCaller") + if not facts.caller_present: + return Decision("Allow") + if not facts.team_loaded: + return Decision("ReadTeam") + if facts.team_present: + if facts.team is None: + return Decision("ReadTeamScope") + team: Final = grouped(facts.team, merge=True) + if team.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if team.kind != "Allow": + return Decision("Deny", team) + if facts.member is None: + return Decision("ReadMember") + member: Final = member_reference(facts.member) + return Decision("Allow") if member.kind == "Allow" else Decision("Deny", member.denial, member.member_models) + if not facts.user_loaded: + return Decision("ReadUser") + if not facts.user_present or not facts.user_grants: + return Decision("Allow") + if "no-default-models" in facts.user_grants: + return Decision("Deny", ModelAccessDecision("NoDefaultModels")) + if facts.user is None: + return Decision("ReadUserScope") + user: Final = model_access(facts.user) + return Decision("Allow") if user.kind == "Allow" else Decision("Deny", user) + + +_ALL: Final = Facts("All", True, True) +_GROUPS: Final = Facts("Ceilings", True, True) +_CALLER: Final = Facts("Caller", True, True, caller_loaded=True, caller_present=True, team_loaded=True) +_ALLOW: Final = ModelAccessFacts(("model",), ("model",), object_type="agent") +_DENY: Final = replace(_ALLOW, grants=("other",)) +_TEAM: Final = GroupedModelScope(replace(_ALLOW, object_type="team")) + + +@pytest.mark.parametrize("policy", (reference, managed_models)) +@pytest.mark.parametrize( + "facts", + ( + _ALL, + replace(_ALL, model_present=False), + replace(_ALL, managed_loaded=True, model_present=False), + replace(_ALL, managed_loaded=True, managed_present=True), + replace(_ALL, managed_loaded=True, managed_present=True, managed_models_valid=True, managed_models_count=1), + replace( + _ALL, + managed_loaded=True, + managed_present=True, + managed_models_valid=True, + managed_models_count=1, + managed=_DENY, + ), + _GROUPS, + replace(_GROUPS, ceilings=()), + replace(_GROUPS, ceilings=(), aliases_loaded=True), + replace(_GROUPS, ceilings=((),), aliases_loaded=True), + replace(_GROUPS, ceilings=(("model",), ("other",)), aliases_loaded=True, scopes=(_ALLOW,)), + replace(_GROUPS, ceilings=(("model",), ("other",)), aliases_loaded=True, scopes=(_ALLOW, _DENY)), + _CALLER, + replace(_CALLER, team_present=True), + replace(_CALLER, team_present=True, team=replace(_TEAM, group_ids=("g",))), + replace(_CALLER, team_present=True, team=_TEAM), + replace(_CALLER, team_present=True, team=_TEAM, member=MemberModelFacts(True, True, True, ())), + replace(_CALLER, user_loaded=True, user_present=True, user_grants=("no-default-models",)), + replace(_CALLER, user_loaded=True, user_present=True, user_grants=("model",)), + ), +) +def test_managed_caller_and_ceiling_read_denial_order(policy: Callable[[Facts], Decision], facts: Facts) -> None: + assert policy(facts) == reference(facts) diff --git a/tests/unit/_auth/test_member_models.py b/tests/unit/_auth/test_member_models.py new file mode 100644 index 00000000000..b9901fd8706 --- /dev/null +++ b/tests/unit/_auth/test_member_models.py @@ -0,0 +1,56 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.member_models import MemberModelDecision as Decision +from litellm._auth.member_models import MemberModelFacts as Facts +from litellm._auth.member_models import member_models +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Facts) -> Decision: + if not facts.user_present or not facts.team_present: + return Decision("Allow") + if not facts.membership_loaded: + return Decision("ReadMembership") + if facts.grants is None: + return Decision("ReadGrants") + if not facts.grants: + return Decision("Allow") + if facts.scope is None: + return Decision("ReadScope") + denied: Final = model_access(facts.scope) + if denied.kind == "Allow": + return Decision("Allow") + return Decision("Deny", denied, facts.grants if denied.kind == "Deny" else None) + + +_BASE: Final = Facts(True, True) +_GRANTS: Final = replace(_BASE, membership_loaded=True, grants=("allowed",)) +_SCOPE: Final = ModelAccessFacts(("requested",), ("allowed",), object_type="team") + + +@pytest.mark.parametrize("policy", (reference, member_models)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadMembership")), + (replace(_BASE, user_present=False), Decision("Allow")), + (replace(_BASE, team_present=False), Decision("Allow")), + (replace(_BASE, membership_loaded=True), Decision("ReadGrants")), + (replace(_GRANTS, grants=()), Decision("Allow")), + (_GRANTS, Decision("ReadScope")), + (replace(_GRANTS, scope=replace(_SCOPE, grants=())), Decision("Allow")), + ( + replace(_GRANTS, scope=_SCOPE), + Decision("Deny", ModelAccessDecision("Deny", "requested", ("allowed",), "team"), ("allowed",)), + ), + ), +) +def test_member_scope_load_order_and_denial_envelope( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_model_access.py b/tests/unit/_auth/test_model_access.py new file mode 100644 index 00000000000..aee39808184 --- /dev/null +++ b/tests/unit/_auth/test_model_access.py @@ -0,0 +1,285 @@ +from collections.abc import Awaitable, Callable +from dataclasses import replace +from typing import Final, Literal +from unittest.mock import MagicMock + +import pytest + +from litellm._auth import model_access as native +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.model_access import ( + ModelAccessFacts as Facts, +) +from litellm._auth.model_access import ( + ModelAliases as Aliases, +) +from litellm._auth.model_access import ( + ModelProvider as Provider, +) +from litellm._auth.model_access import ( + ModelRecord as Record, +) +from litellm._auth.model_access import ( + ModelRegistry as Registry, +) +from litellm.proxy._types import LiteLLM_UserTable, ModelAccessDeniedProxyException +from litellm.proxy.auth import model_access_adapter +from tests.unit._auth.support import model_reference as reference +from tests.unit._auth.support import model_registry_reference + +_BASE: Final = Facts(("model",), object_type="key") +_CASES: Final = ( + pytest.param(_BASE, "Allow", id="empty-is-unrestricted"), + pytest.param(replace(_BASE, grants=("model",)), "Allow", id="exact-grant"), + pytest.param(replace(_BASE, grants=("other",)), "Deny", id="denied-model"), + pytest.param(replace(_BASE, grants=("*",)), "Allow", id="global-wildcard"), + pytest.param(replace(_BASE, grants=("all-proxy-models",)), "Allow", id="proxy-sentinel"), + pytest.param(replace(_BASE, grants=("mo*",)), "Allow", id="prefix-wildcard"), + pytest.param(replace(_BASE, requested=("model\n",), grants=("mo*",)), "Allow", id="python-final-newline"), + pytest.param(replace(_BASE, requested=("mod\nel",), grants=("mo*l",)), "Deny", id="embedded-newline"), + pytest.param(replace(_BASE, requested=("\u0301a",), grants=(r"\w*",)), "Deny", id="word-combining-mark"), + pytest.param(replace(_BASE, requested=("\u200da",), grants=(r"\w*",)), "Deny", id="word-join-control"), + pytest.param(replace(_BASE, requested=("\u00b2a",), grants=(r"\w*",)), "Allow", id="word-other-number"), + pytest.param( + replace(_BASE, grants=("p/*",), registry=Registry(records={"model": Record(provider=Provider("model", "p"))})), + "Allow", + id="inferred-provider", + ), + pytest.param( + replace( + _BASE, + requested=("p/model",), + grants=("q/*",), + registry=Registry(records={"p/model": Record(provider=Provider("model", "p"))}), + ), + "Deny", + id="different-provider", + ), + pytest.param( + replace(_BASE, grants=("other",), aliases=Aliases(team={"model": "target"})), "Allow", id="team-alias-exists" + ), + pytest.param( + replace(_BASE, grants=("model",), aliases=Aliases(key={"model": "target"})), + "Deny", + id="key-alias-replaces-source", + ), + pytest.param( + replace(_BASE, grants=("target",), aliases=Aliases(key={"model": "target"})), "Allow", id="key-target" + ), + pytest.param( + replace(_BASE, grants=("target",), aliases=Aliases(global_={"model": "target"})), "Allow", id="global-target" + ), + pytest.param( + replace( + _BASE, + grants=("target",), + aliases=Aliases( + team={"model": "middle"}, key={"middle": "third", "global": "target"}, global_={"third": "global"} + ), + ), + "Allow", + id="alias-order", + ), + pytest.param( + replace(_BASE, grants=("target",), registry=Registry(present=True, aliases={"model": "target"})), + "Allow", + id="router-alias", + ), + pytest.param( + replace( + _BASE, + grants=("target",), + registry=Registry(present=True, aliases={"model": "target"}), + aliases=Aliases(global_={"model": "wrong"}), + ), + "Deny", + id="global-over-router", + ), + pytest.param( + replace(_BASE, grants=("parent",), registry=Registry(compaction_parents={"model": "parent"})), + "Allow", + id="compaction-parent", + ), + pytest.param( + replace( + _BASE, + grants=("group",), + registry=Registry(present=True, records={"model": Record(access_groups=("group",))}), + ), + "Allow", + id="registry-group", + ), + pytest.param( + replace( + _BASE, + grants=("other-group",), + registry=Registry(present=True, records={"model": Record(access_groups=("group",))}), + ), + "Deny", + id="different-group", + ), + pytest.param( + replace(_BASE, grants=("all-team-models",), team_id="team", registry=Registry(present=True, names=("model",))), + "Allow", + id="team-sentinel-expands", + ), + pytest.param( + replace(_BASE, grants=("all-team-models",), team_id="team", registry=Registry(present=True)), + "Deny", + id="empty-registry-sentinel", + ), + pytest.param( + replace(_BASE, grants=("all-team-models",), registry=Registry(present=True, names=("model",))), + "Deny", + id="team-sentinel-no-team", + ), + pytest.param( + replace( + _BASE, + requested=("model_name_team_id",), + grants=("public",), + team_id="team", + registry=Registry( + present=True, + records={"model_name_team_id": Record(deployment_team="team", deployment_public_name="public")}, + ), + ), + "Allow", + id="owning-team-public-name", + ), + pytest.param( + replace( + _BASE, + requested=("model_name_team_id",), + grants=("public",), + team_id="team", + registry=Registry( + present=True, + records={"model_name_team_id": Record(deployment_team="other", deployment_public_name="public")}, + ), + ), + "Deny", + id="public-name-wrong-team", + ), + pytest.param( + replace( + _BASE, + requested=("model_name_team_id",), + grants=("other",), + team_id="team", + aliases=Aliases(team={"public": "model_name_team_id"}), + registry=Registry(present=True), + ), + "Deny", + id="owning-alias-not-unconditional", + ), + pytest.param( + replace(_BASE, requested=("model", "other"), grants=("model",), is_list=True), + "Deny", + id="all-list-items-required", + ), + pytest.param(replace(_BASE, requested=(), grants=("other",), is_list=True), "Allow", id="empty-list"), + pytest.param(replace(_BASE, fallback_depth=5), "Depth", id="depth"), + pytest.param(replace(_BASE, fallback_depth=4, is_list=True), "Depth", id="list-depth"), +) + + +@pytest.mark.parametrize("policy", (reference.model_access, native.model_access), ids=("reference", "native")) +@pytest.mark.parametrize(("facts", "kind"), _CASES) +def test_model_policy_case_table( + policy: Callable[[Facts], native.ModelAccessDecision], facts: Facts, kind: str +) -> None: + decision: Final = policy(facts) + assert decision.kind == kind + assert decision == reference.model_access(facts) + + +@pytest.mark.parametrize( + "grant", + ( + r"\w*", r"\W*", r"[\w]*", r"[\W]*", r"[^\w]*", r"[^\W]*", + r"[\w\W]*", r"[^\w\W]*", r"[^\w]+*", r"[\w^]*", r"[\w-]*", + r"[\-\w]*", r"[]\w]*", r"[^]\w]*", r"[[\w]*", r"[\b\w]*", + r"\\w*", r"[\\w]*", r"\b\w*", r"\B\w*", r"\w\b*", r"\w\B*", + r"\W\b*", r"\W\B*", r"(?u:\w)*", r"(?i:\w)*", r"(?i:[\wΙ])*", + r"(?i:[^\wΙ])*", r"[a-zA-Z0-9_]*", r"(?# [\b)\w*", + "(?x:# [\\b)\n\\w)*", "(?x:# [\n)(?-x:#)\\w*", r"(\w)\1*", + ), +) +@pytest.mark.parametrize( + "model", + ( + "\u0301a", "\u200da", "\u203fa", "\u0345a", "²a", "¼a", "Ⅷa", "a²", "a\u0301", + "_", "-", "^", "a", "A", "é", "한글", "αα", "\n", "\b", r"\word", "]a", "[a", "#²", " a", + ), +) +def test_unicode_regex_model_access_matches_python(grant: str, model: str) -> None: + facts: Final = replace(_BASE, requested=(model,), grants=(grant,)) + assert native.model_access(facts) == reference.model_access(facts) + + +@pytest.mark.parametrize( + ("groups", "models", "merge", "kind"), + ( + ((), None, False, "Deny"), + (("g",), None, False, "ReadGroups"), + (("g",), (), False, "Deny"), + (("g",), ("model",), False, "Allow"), + (("g",), ("model",), True, "Allow"), + (("g",), ("other",), False, "Deny"), + ), +) +def test_group_fallback_waits_for_data( + groups: tuple[str, ...], models: tuple[str, ...] | None, merge: bool, kind: str +) -> None: + assert native.grouped_model_access(replace(_BASE, grants=("no",)), groups, models, merge=merge).kind == kind + + +@pytest.mark.parametrize( + ("present", "grants", "kind"), + ( + (False, ("no-default-models",), "Allow"), + (True, ("no-default-models", "model"), "NoDefaultModels"), + (True, (), "Allow"), + (True, ("other",), "Deny"), + ), +) +def test_user_default_scope(present: bool, grants: tuple[str, ...], kind: str) -> None: + assert native.user_model_access(replace(_BASE, grants=grants), present=present).kind == kind + + +@pytest.mark.parametrize( + "policy", (model_registry_reference.can_user_call_model, model_access_adapter.can_user_call_model) +) +@pytest.mark.asyncio +async def test_user_precheck_does_not_read_registry( + policy: Callable[[str, MagicMock, LiteLLM_UserTable | None], Awaitable[Literal[True]]], +) -> None: + router: Final = MagicMock() + user: Final = LiteLLM_UserTable(user_id="user", models=["no-default-models"]) + with pytest.raises(ModelAccessDeniedProxyException, match="model"): + await policy("model", router, user) + assert not router.mock_calls + assert await policy("model", router, None) is True + assert not router.mock_calls + + +def test_native_model_decode_failure_is_terminal() -> None: + with pytest.raises(NativeAuthFailure): + native.model_access(replace(_BASE, fallback_depth=-1)) + + +@pytest.mark.parametrize( + ("managed", "ceilings", "kind"), + ( + (None, (), "Allow"), + ((), (), "NoAgentModels"), + (("model",), ((),), "EmptyAgentCeiling"), + (("model",), (("model",), ("other",)), "Deny"), + (("model",), (("model",),), "Allow"), + ), +) +def test_agent_grants_are_intersected( + managed: tuple[str, ...] | None, ceilings: tuple[tuple[str, ...], ...], kind: str +) -> None: + assert native.agent_model_access(_BASE, managed, ceilings).kind == kind diff --git a/tests/unit/_auth/test_model_attribution.py b/tests/unit/_auth/test_model_attribution.py new file mode 100644 index 00000000000..09f85e0c9e2 --- /dev/null +++ b/tests/unit/_auth/test_model_attribution.py @@ -0,0 +1,87 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.model_attribution import ModelGroupAttributionDecision as Decision +from litellm._auth.model_attribution import ModelGroupAttributionFacts as Facts +from litellm._auth.model_attribution import model_group_attribution + + +def reference(facts: Facts) -> Decision: + if facts.io_failed or not all( + (facts.model_present, facts.token_present, facts.router_present, facts.database_present) + ): + return Decision("Complete") + if not facts.registry_loaded: + return Decision("ReadRegistry") + if facts.registry is not None and not facts.registry: + return Decision("Complete") + if facts.covering is None: + return Decision("ReadCovering") + budgeted: Final = set(facts.covering) if facts.registry is None else set(facts.covering) & set(facts.registry) + if not budgeted: + return Decision("Complete") + match len(facts.grants): + case 0: + return Decision("ReadKey") + case 1: + return Decision("ReadTeam") + case 2: + return Decision( + "ReadMember", + include=facts.team_present and facts.user_present, + lookup=facts.team_present and facts.user_present and not facts.membership_loaded, + ) + case 3: + return Decision("ReadProject") + case 4: + return Decision("ReadOrg", id=facts.token_org or facts.team_org) + case _: + scopes: Final = facts.grants[:4] if facts.org_failed else facts.grants + granted: Final = set[str]().union(*scopes) + groups: Final = tuple(sorted(budgeted & granted)) + return Decision("Complete", groups, stamp=bool(groups)) + + +_BASE: Final = Facts(model_present=True, token_present=True, router_present=True, database_present=True) +_COVERED: Final = replace(_BASE, registry_loaded=True, registry=("a", "b"), covering=("b", "a", "c", "b")) +_FINAL: Final = replace(_COVERED, grants=(("a",), (), ("b",), (), ("c",))) + + +@pytest.mark.parametrize("policy", (reference, model_group_attribution)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadRegistry")), + (replace(_BASE, model_present=False), Decision("Complete")), + (replace(_BASE, token_present=False), Decision("Complete")), + (replace(_BASE, router_present=False), Decision("Complete")), + (replace(_BASE, database_present=False), Decision("Complete")), + (replace(_BASE, io_failed=True), Decision("Complete")), + (replace(_BASE, registry_loaded=True, registry=()), Decision("Complete")), + (replace(_BASE, registry_loaded=True), Decision("ReadCovering")), + (replace(_COVERED, covering=()), Decision("Complete")), + (_COVERED, Decision("ReadKey")), + (replace(_COVERED, grants=((),)), Decision("ReadTeam")), + (replace(_COVERED, grants=((), ())), Decision("ReadMember")), + ( + replace(_COVERED, grants=((), ()), team_present=True, user_present=True), + Decision("ReadMember", include=True, lookup=True), + ), + ( + replace(_COVERED, grants=((), ()), team_present=True, user_present=True, membership_loaded=True), + Decision("ReadMember", include=True), + ), + (replace(_COVERED, grants=((), (), ())), Decision("ReadProject")), + (replace(_COVERED, grants=((), (), (), ()), token_org="key", team_org="team"), Decision("ReadOrg", id="key")), + (replace(_COVERED, grants=((), (), (), ()), token_org="", team_org="team"), Decision("ReadOrg", id="team")), + (_FINAL, Decision("Complete", ("a", "b"), stamp=True)), + (replace(_FINAL, registry=None), Decision("Complete", ("a", "b", "c"), stamp=True)), + (replace(_FINAL, registry=None, org_failed=True), Decision("Complete", ("a", "b"), stamp=True)), + (replace(_FINAL, grants=(("*",), (), (), (), ())), Decision("Complete")), + ), +) +def test_attribution_case_table(policy: Callable[[Facts], Decision], facts: Facts, expected: Decision) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_model_budget_aliases.py b/tests/unit/_auth/test_model_budget_aliases.py new file mode 100644 index 00000000000..0d1e1b5bd60 --- /dev/null +++ b/tests/unit/_auth/test_model_budget_aliases.py @@ -0,0 +1,66 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import Mock, call + +import pytest + +import litellm +from litellm._auth import model_budget_aliases as native +from litellm._auth.model_budget_data import ModelBudgetSnapshot +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry +from litellm.proxy.auth.model_budget_adapter import ModelBudgetIO, read_alias_models +from litellm.proxy.hooks.model_budget_shared_windows import SharedWindows + +from .model_alias_reference import expand + + +@pytest.mark.parametrize("model", ("latest", "openai/latest", "unknown/latest", "bedrock/latest")) +@pytest.mark.parametrize("names", ((), ("latest",), ("openai/latest",), ("target",), ("",), ("openai/target",))) +@pytest.mark.parametrize("provider", (None, "openai", "unknown", "bedrock_converse", 42)) +def test_alias_lookup_matches_frozen_main(model: str, names: tuple[str, ...], provider: object) -> None: + row: Final = AliasedCostEntry({"litellm_provider": provider}, names) + catalog: Final[Mapping[str, object]] = {model: row, "latest": row} + expected: Final = expand(model, catalog, litellm.provider_list) + lookup: Final = Mock(side_effect=catalog.get) + io: Final = ModelBudgetIO( + DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lookup + ) + assert read_alias_models(model, io) == expected + assert lookup.call_args_list[0].args == (model,) + + +@pytest.mark.parametrize("raw", (None, {}, {"alias_of": "target"}, {"names": ["target"]})) +def test_plain_catalog_rows_do_not_become_aliases(raw: object) -> None: + io: Final = ModelBudgetIO( + DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lambda _: raw + ) + assert read_alias_models("latest", io) == ("latest",) + + +@pytest.mark.parametrize("direct", (True, False)) +def test_catalog_reads_bare_name_only_when_qualified_alias_misses(direct: bool) -> None: + row: Final = AliasedCostEntry({"litellm_provider": "openai"}, ("target", "latest")) + catalog: Final = {"openai/latest": row if direct else None, "latest": row} + lookup: Final = Mock(side_effect=catalog.get) + io: Final = ModelBudgetIO( + DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lookup + ) + assert read_alias_models("openai/latest", io) == ("openai/latest", "openai/target") + assert lookup.call_args_list == ([call("openai/latest")] if direct else [call("openai/latest"), call("latest")]) + + +@pytest.mark.parametrize( + ("snapshots", "expected"), + ( + ((), native.ScopeDecision(0, None)), + ((ModelBudgetSnapshot(1, 0),), native.ScopeDecision(1, None)), + ((ModelBudgetSnapshot(1, 1),), native.ScopeDecision(None, 0)), + ((ModelBudgetSnapshot(1, 0), ModelBudgetSnapshot(2, 2)), native.ScopeDecision(None, 1)), + ((ModelBudgetSnapshot(1, 0), ModelBudgetSnapshot(2, 0)), native.ScopeDecision(None, 1)), + ), +) +def test_scope_reads_only_until_first_denial( + snapshots: tuple[ModelBudgetSnapshot, ...], expected: native.ScopeDecision +) -> None: + assert native.scope(native.ScopeFacts(("alias", "target"), snapshots)) == expected diff --git a/tests/unit/_auth/test_model_budget_data.py b/tests/unit/_auth/test_model_budget_data.py new file mode 100644 index 00000000000..314a287681c --- /dev/null +++ b/tests/unit/_auth/test_model_budget_data.py @@ -0,0 +1,111 @@ +from math import inf, isnan, nan +from typing import Final + +import pytest + +from litellm._auth import model_budget_data as native +from litellm._auth.model_budget_data import ( + ModelBudgetEntry, + ModelBudgetQueryFacts, + ModelBudgetResolveFacts, + ModelBudgetSnapshot, +) + +from .budget_reference import model as reference + + +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("local", (None, -1.0, 0.0, 1.0, 2.0, inf, nan)) +@pytest.mark.parametrize("shared", (None, 0.0, 2.0, inf, nan)) +@pytest.mark.parametrize("now", (99, 100)) +def test_raw_counter_and_shared_window_case_table( + limit: float | None, local: float | None, shared: float | None, now: float +) -> None: + facts: Final = ModelBudgetSnapshot(limit, local, 0, shared, 40, 60, now) + result: Final = native.check(facts) + expected: Final = reference(facts) + assert result.denied == expected.denied + assert result.read is None + assert ( + isnan(result.current_spend) and isnan(expected.current_spend) + ) or result.current_spend == expected.current_spend + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1)), "Local"), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(-1)), None), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(None)), None), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1, 0), 60, 59), "Legacy"), + (ModelBudgetQueryFacts("EndUser", ModelBudgetSnapshot(1, 0), 60, 59), "Legacy"), + (ModelBudgetQueryFacts("EndUser", ModelBudgetSnapshot(1, 0), 60, 60), None), + (ModelBudgetQueryFacts("User", ModelBudgetSnapshot(1, 0), 60, 1), "Shared"), + (ModelBudgetQueryFacts("Team", ModelBudgetSnapshot(1, 0), 60, 1, redis_present=True), None), + (ModelBudgetQueryFacts("TeamMember", ModelBudgetSnapshot(1, 1), 60), None), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1, 0), 60, same_counter=True), "Shared"), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1, 0, 0), 60, shared_read=True), None), + ), +) +def test_native_selects_the_next_io_read(facts: ModelBudgetQueryFacts, expected: str | None) -> None: + assert native.query(facts).read == expected + + +def entry(name: str, limit: float | None = 1, duration: str | None = "1h") -> ModelBudgetEntry: + return ModelBudgetEntry(name, limit, duration, 3600 if duration is not None else None) + + +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("overrides", (("alias",), ("target",), ("alias", "target"))) +def test_team_exact_entries_are_removed_across_all_alias_names(limit: float | None, overrides: tuple[str, ...]) -> None: + key: Final = tuple(entry(name, limit) for name in overrides) + facts: Final = tuple(ModelBudgetResolveFacts(model, model, "", (), key, True) for model in ("alias", "target")) + excluded: Final = native.team_exclusions(facts) + assert excluded == (tuple(sorted(overrides)) if limit is not None and limit >= 0 else ()) + resolution: Final = native.resolve( + ModelBudgetResolveFacts( + "alias", + "alias", + "", + (entry("alias"), entry("target"), entry("*")), + access_groups=("alias", "target"), + excluded_entries=excluded, + ) + ) + expected: Final = "alias" if "alias" not in excluded else "target" if "target" not in excluded else "*" + assert resolution.entry is not None and resolution.entry.name == expected + + +@pytest.mark.parametrize( + ("entries", "groups", "expected"), + ( + ((entry("*"), entry("m*"), entry("model")), ("group",), "model"), + ((entry("*"), entry("m*"), entry("group")), ("group",), "group"), + ((entry("*"), entry("m*"), entry("model", duration=None)), (), "m*"), + ((entry("provider/model"), entry("model")), (), "provider/model"), + ((entry("b"), entry("a"), entry("*")), ("b", "a"), "a"), + ((entry("model", None), entry("*")), (), "model"), + ((entry("m?del"),), (), None), + ((entry("*a*"), entry("mo*"), entry("mod*l")), (), "mod*l"), + ), +) +def test_resolution_prioritizes_exact_then_groups_then_specific_globs( + entries: tuple[ModelBudgetEntry, ...], + groups: tuple[str, ...], + expected: str | None, +) -> None: + result: Final = native.resolve(ModelBudgetResolveFacts("provider/model", "", "", entries, access_groups=groups)) + assert (result.entry.name if result.entry is not None else None) == expected + assert result.read_groups is False + + +@pytest.mark.parametrize( + ("key_limit", "expected"), ((0, "*"), (1, "*"), (-1, "model"), (None, "model"), (nan, "model")) +) +def test_only_enforced_named_key_entry_replaces_the_team_entry(key_limit: float | None, expected: str) -> None: + facts: Final = ModelBudgetResolveFacts( + "model", "", "", (entry("model"), entry("*")), (entry("model", key_limit),), team=True, access_groups=() + ) + result: Final = native.resolve(facts) + assert result.entry is not None + assert result.entry.name == expected diff --git a/tests/unit/_auth/test_model_budgets.py b/tests/unit/_auth/test_model_budgets.py new file mode 100644 index 00000000000..d091e3d8c66 --- /dev/null +++ b/tests/unit/_auth/test_model_budgets.py @@ -0,0 +1,332 @@ +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from tests.unit._auth.raw_answers import candidate_read +from dataclasses import dataclass, replace +from typing import Final + +import pytest + +from litellm._auth.decision import ( + Allow, + AuthCore, + Deny, + FallbackTarget, + ModelBudgetExceeded, + ModelBudgetRead, + ModelBudgetScope, + Outcome, + Step, + TeamMemberModelBudgetScope, + WithinModelBudgets, +) +from litellm._auth.driver import drive_model_budgets, drive_team_member_model_budgets +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from tests.unit._auth.raw_answers import budget_read, fallback_read + +Call = tuple[str, str | None] + +SCOPE: Final = ModelBudgetScope( + models=("gpt-4o",), + rewritable_model="gpt-4o", + key_has_model_budgets=True, + key_budget_fallbacks={"gpt-4o": ("gpt-4o-mini", "haiku")}, + user_has_model_budgets=True, +) +ALLOWED: Final = Allow(WithinModelBudgets()) + + +def _exceeded(entity_type: str, model: str) -> ModelBudgetExceeded: + return ModelBudgetExceeded( + message=f"{entity_type} over budget for {model}", + current_cost=2.0, + max_budget=1.0, + entity_type=entity_type, + entity_id="id", + ) + + +@dataclass +class FakeRequest: + scope: ModelBudgetScope + key_over: frozenset[str] = frozenset() + user_over: frozenset[str] = frozenset() + unauthorized: frozenset[str] = frozenset() + calls: tuple[Call, ...] = () + + def _record(self, name: str, model: str | None = None) -> None: + self.calls = (*self.calls, (name, model)) + + async def read_model_budget_scope(self) -> ModelBudgetScope: + self._record("read") + return self.scope + + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: + self._record("key", model) + return budget_read(_exceeded("key", model) if model in self.key_over else None) + + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: + self._record("authorize", model) + return fallback_read(model, model not in self.unauthorized) + + async def rewrite_request_model(self, model: str) -> None: + self._record("rewrite", model) + self.scope = replace( + self.scope, + models=tuple(model if m == self.scope.rewritable_model else m for m in self.scope.models), + rewritable_model=model, + ) + + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: + self._record("user", model) + return budget_read(_exceeded("user", model) if model in self.user_over else None) + + +async def _run(step: Step, request: FakeRequest) -> tuple[Outcome, tuple[Call, ...]]: + outcome: Final = await drive_model_budgets(step, request) + return outcome, request.calls + + +@pytest.mark.asyncio +async def test_within_budget_checks_the_key_then_the_user_without_rerouting( + core: AuthCore, +): + assert await _run(core.check_model_budgets(), FakeRequest(SCOPE)) == ( + ALLOWED, + (("read", None), ("key", "gpt-4o"), ("read", None), ("user", "gpt-4o")), + ) + + +@pytest.mark.asyncio +async def test_over_budget_reroutes_to_the_first_fallback_within_budget_and_checks_the_user_on_it( + core: AuthCore, +): + request: Final = FakeRequest(SCOPE, key_over=frozenset({"gpt-4o", "gpt-4o-mini"})) + assert await _run(core.check_model_budgets(), request) == ( + ALLOWED, + ( + ("read", None), + ("key", "gpt-4o"), + ("key", "gpt-4o-mini"), + ("key", "haiku"), + ("authorize", "haiku"), + ("rewrite", "haiku"), + ("read", None), + ("user", "haiku"), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "key_over", "unauthorized"), + [ + (replace(SCOPE, rewritable_model=None), {"gpt-4o"}, set[str]()), + (replace(SCOPE, rewritable_model="other"), {"gpt-4o"}, set[str]()), + (replace(SCOPE, key_budget_fallbacks={}), {"gpt-4o"}, set[str]()), + (SCOPE, {"gpt-4o", "gpt-4o-mini", "haiku"}, set[str]()), + (SCOPE, {"gpt-4o"}, {"gpt-4o-mini"}), + ], + ids=["nested_model", "different_top_level_model", "no_fallbacks", "every_fallback_over", "fallback_unauthorized"], +) +async def test_an_unreroutable_over_budget_model_is_denied_with_its_own_budget_error( + core: AuthCore, scope: ModelBudgetScope, key_over: set[str], unauthorized: set[str] +): + request: Final = FakeRequest(scope, key_over=frozenset(key_over), unauthorized=frozenset(unauthorized)) + outcome, calls = await _run(core.check_model_budgets(), request) + assert outcome == Deny(_exceeded("key", "gpt-4o")) + assert not any(name in ("rewrite", "user") for name, _ in calls) + + +@pytest.mark.asyncio +async def test_a_later_model_is_checked_after_an_earlier_one_was_rerouted( + core: AuthCore, +): + scope: Final = replace(SCOPE, models=("gpt-4o", "embed")) + request: Final = FakeRequest(scope, key_over=frozenset({"gpt-4o", "embed"})) + outcome, calls = await _run(core.check_model_budgets(), request) + assert outcome == Deny(_exceeded("key", "embed")) + assert calls[-2:] == (("rewrite", "gpt-4o-mini"), ("key", "embed")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "expected_calls"), + [ + ( + replace(SCOPE, key_has_model_budgets=False), + (("read", None), ("user", "gpt-4o")), + ), + ( + replace(SCOPE, user_has_model_budgets=False), + (("read", None), ("key", "gpt-4o"), ("read", None)), + ), + (replace(SCOPE, models=()), (("read", None),)), + ], + ids=["no_key_budgets", "no_user_budgets", "no_models"], +) +async def test_only_budgets_that_exist_are_checked( + core: AuthCore, scope: ModelBudgetScope, expected_calls: tuple[Call, ...] +): + assert await _run(core.check_model_budgets(), FakeRequest(scope)) == (ALLOWED, expected_calls) + + +@pytest.mark.asyncio +async def test_a_user_over_budget_is_denied_on_every_path( + core: AuthCore, +): + over: Final = frozenset({"gpt-4o"}) + for step in (core.check_model_budgets(), core.check_user_model_budgets()): + outcome, _ = await _run(step, FakeRequest(SCOPE, user_over=over)) + assert outcome == Deny(_exceeded("user", "gpt-4o")) + + +@pytest.mark.asyncio +async def test_the_user_only_check_never_reads_key_budgets( + core: AuthCore, +): + request: Final = FakeRequest(SCOPE, key_over=frozenset({"gpt-4o"})) + assert await _run(core.check_user_model_budgets(), request) == ( + ALLOWED, + (("read", None), ("user", "gpt-4o")), + ) + + +MEMBER_SCOPE: Final = TeamMemberModelBudgetScope( + models=("gpt-4o",), rewritable_model="gpt-4o", has_model_budgets=True, can_fall_back=True +) + + +def _target(model: str) -> FallbackTarget: + return FallbackTarget(requested=model, model=model, served_model=model) + + +@dataclass +class FakeTeamMemberRequest: + scope: TeamMemberModelBudgetScope + over: frozenset[str] = frozenset() + fallbacks: tuple[str, ...] = ("gpt-4o", "unusable", "gpt-4o-mini", "haiku") + unusable: frozenset[str] = frozenset({"unusable"}) + calls: tuple[Call, ...] = () + targets: tuple[FallbackTarget, ...] | None = None + vetted: tuple[FallbackTarget, ...] = () + + def _record(self, name: str, model: str | None = None) -> None: + self.calls = (*self.calls, (name, model)) + + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: + self._record("read") + return self.scope + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetRead: + self._record("member", model) + return budget_read(_exceeded("member", model) if model in self.over else None) + + async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: + self._record("fallbacks", model) + return self.targets if self.targets is not None else tuple(_target(fallback) for fallback in self.fallbacks) + + async def vet_team_member_fallback(self, target: FallbackTarget) -> FallbackCandidateFacts: + self._record("vet", target.model) + self.vetted = (*self.vetted, target) + return candidate_read(target.model, target.model not in self.unusable) + + async def reroute_team_member_request(self, model: str) -> None: + self._record("reroute", model) + + +async def _run_member(core: AuthCore, request: FakeTeamMemberRequest) -> tuple[Outcome, tuple[Call, ...]]: + outcome: Final = await drive_team_member_model_budgets(core.check_team_member_model_budgets(), request) + return outcome, request.calls + + +@pytest.mark.asyncio +async def test_large_user_budget_scopes_do_not_exhaust_the_driver_stack(core: AuthCore): + count: Final = 600 + scope: Final = replace(SCOPE, models=("model",) * count, key_has_model_budgets=False) + request: Final = FakeRequest(scope) + assert await drive_model_budgets(core.check_user_model_budgets(), request) == ALLOWED + assert request.calls == (("read", None),) + (("user", "model"),) * count + + +@pytest.mark.asyncio +async def test_member_fallback_preserves_requested_resolved_and_served_aliases(core: AuthCore): + target: Final = FallbackTarget(requested="public-alias", model="resolved-model", served_model="deployment") + request: Final = FakeTeamMemberRequest(MEMBER_SCOPE, over=frozenset({"gpt-4o"}), targets=(target,)) + assert await _run_member(core, request) == ( + ALLOWED, + ( + ("read", None), + ("member", "gpt-4o"), + ("fallbacks", "gpt-4o"), + ("vet", "resolved-model"), + ("reroute", "resolved-model"), + ), + ) + assert request.vetted == (target,) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "expected_calls"), + [ + (MEMBER_SCOPE, (("read", None), ("member", "gpt-4o"))), + (replace(MEMBER_SCOPE, has_model_budgets=False), (("read", None),)), + (replace(MEMBER_SCOPE, models=()), (("read", None),)), + ], + ids=["within_budget", "no_member_budgets", "no_models"], +) +async def test_a_team_member_within_budget_is_allowed( + core: AuthCore, scope: TeamMemberModelBudgetScope, expected_calls: tuple[Call, ...] +): + assert await _run_member(core, FakeTeamMemberRequest(scope)) == (ALLOWED, expected_calls) + + +@pytest.mark.asyncio +async def test_an_over_budget_team_member_reroutes_to_the_first_usable_fallback_other_than_itself( + core: AuthCore, +): + request: Final = FakeTeamMemberRequest( + replace(MEMBER_SCOPE, models=("gpt-4o", "embed")), over=frozenset({"gpt-4o"}) + ) + assert await _run_member(core, request) == ( + ALLOWED, + ( + ("read", None), + ("member", "gpt-4o"), + ("fallbacks", "gpt-4o"), + ("vet", "unusable"), + ("vet", "gpt-4o-mini"), + ("reroute", "gpt-4o-mini"), + ("member", "embed"), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "unusable"), + [ + (replace(MEMBER_SCOPE, rewritable_model=None), frozenset[str]()), + (replace(MEMBER_SCOPE, can_fall_back=False), frozenset[str]()), + (MEMBER_SCOPE, frozenset({"unusable", "gpt-4o-mini", "haiku"})), + ], + ids=["nested_model", "fallbacks_unavailable", "no_usable_fallback"], +) +async def test_an_over_budget_team_member_without_a_usable_fallback_is_denied_with_its_own_error( + core: AuthCore, scope: TeamMemberModelBudgetScope, unusable: frozenset[str] +): + outcome, calls = await _run_member( + core, FakeTeamMemberRequest(scope, over=frozenset({"gpt-4o"}), unusable=unusable) + ) + assert outcome == Deny(_exceeded("member", "gpt-4o")) + assert not any(name == "reroute" for name, _ in calls) + + +@pytest.mark.asyncio +async def test_a_later_model_is_not_rerouted_once_the_request_model_changed( + core: AuthCore, +): + scope: Final = replace(MEMBER_SCOPE, models=("gpt-4o", "gpt-4o")) + request: Final = FakeTeamMemberRequest(scope, over=frozenset({"gpt-4o"})) + outcome, calls = await _run_member(core, request) + assert outcome == Deny(_exceeded("member", "gpt-4o")) + assert [name for name, _ in calls].count("reroute") == 1 diff --git a/tests/unit/_auth/test_model_groups.py b/tests/unit/_auth/test_model_groups.py new file mode 100644 index 00000000000..49afbf26491 --- /dev/null +++ b/tests/unit/_auth/test_model_groups.py @@ -0,0 +1,51 @@ +from dataclasses import replace +from typing import Final, Literal + +import pytest + +from litellm._auth.model_groups import GroupLookup, GroupResources, ResourceGroupFacts, resource_group_decision + +_BASE: Final = ResourceGroupFacts(True, True, True, ("group",), "team", "key") + + +@pytest.mark.parametrize( + ("team_ids", "key_ids", "expected"), + ( + (("team",), (), ("model",)), + ((), ("key",), ("model",)), + (("other",), (), ()), + ((), ("other",), ()), + ((), (), ()), + ), +) +def test_group_ownership_boundary( + team_ids: tuple[str, ...], key_ids: tuple[str, ...], expected: tuple[str, ...] +) -> None: + decision: Final = resource_group_decision( + replace(_BASE, loaded=(GroupLookup("Found", GroupResources(team_ids, key_ids, ("model", "model"))),)) + ) + assert (decision.kind, decision.resources) == ("Complete", expected) + + +@pytest.mark.parametrize("kind", ("Missing", "Failed")) +@pytest.mark.parametrize("check_db_only", (False, True)) +@pytest.mark.parametrize("strict_errors", (False, True)) +def test_strict_lookup_errors_stop_before_next_group( + kind: Literal["Missing", "Failed"], check_db_only: bool, strict_errors: bool +) -> None: + decision: Final = resource_group_decision( + replace( + _BASE, + group_ids=("first", "second"), + check_db_only=check_db_only, + strict_errors=strict_errors, + loaded=(GroupLookup(kind),), + ) + ) + assert (decision.kind, decision.index) == (("Raise", 0) if check_db_only or strict_errors else ("Read", 1)) + + +@pytest.mark.parametrize("field", ("subject_present", "database_present", "cache_present")) +def test_unavailable_dependency_grants_nothing(field: str) -> None: + decision: Final = resource_group_decision(replace(_BASE, **{field: False})) + assert (decision.kind, decision.resources) == ("Complete", ()) diff --git a/tests/unit/_auth/test_model_helpers.py b/tests/unit/_auth/test_model_helpers.py new file mode 100644 index 00000000000..c038776f218 --- /dev/null +++ b/tests/unit/_auth/test_model_helpers.py @@ -0,0 +1,105 @@ +import re +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.model_access import ModelProvider +from litellm._auth.model_helpers import ModelHelperDecision as Decision +from litellm._auth.model_helpers import ModelHelperFacts as Facts +from litellm._auth.model_helpers import model_helper + + +def reference(facts: Facts) -> Decision: + if facts.kind == "Group" and (facts.grants is None or facts.model in facts.grants): + return Decision("Complete", True) + if facts.router_present and facts.groups is None: + return Decision("ReadGroups") + groups: Final = facts.groups or () + grants: Final = facts.grants or () + if facts.kind == "Group": + return Decision("Complete", any(grant in groups for grant in grants)) + expand: Final = "all-team-models" in grants and facts.team_id is not None and facts.router_present + if expand and facts.names is None: + return Decision("ReadNames") + expanded: Final = ( + ( + tuple( + dict.fromkeys((*tuple(grant for grant in grants if grant != "all-team-models"), *(facts.names or ()))) + ) + or grants + ) + if expand + else grants + ) + if any(grant in groups for grant in expanded) or facts.model in facts.team_aliases: + return Decision("Complete", True) + filtered: Final = tuple(grant for grant in expanded if grant not in groups) + patterns: Final = tuple(grant.replace("*", ".*") + "$" for grant in filtered if "*" in grant) + if any(re.match(pattern, facts.model) is not None for pattern in patterns): + return Decision("Complete", True) + if patterns: + if not facts.provider_loaded: + return Decision("ReadProvider") + provider: Final = facts.provider + if ( + provider is not None + and (provider.stripped_model != facts.model or "/" not in facts.model) + and any( + re.match(pattern, f"{provider.provider}/{provider.stripped_model}") is not None for pattern in patterns + ) + ): + return Decision("Complete", True) + return Decision( + "Complete", + not filtered and not expanded or "*" in filtered or "all-proxy-models" in filtered or facts.model in filtered, + ) + + +_GROUP: Final = Facts("Group", "model", ("model",), router_present=True) +_HELPER: Final = replace(_GROUP, kind="Helper") +_READ: Final = replace(_HELPER, groups=()) + + +@pytest.mark.parametrize("policy", (reference, model_helper)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_GROUP, Decision("Complete", True)), + (replace(_GROUP, grants=None), Decision("Complete", True)), + (replace(_GROUP, grants=()), Decision("ReadGroups")), + (replace(_GROUP, grants=(), groups=()), Decision("Complete")), + (replace(_GROUP, grants=("group",), groups=("group",)), Decision("Complete", True)), + (_HELPER, Decision("ReadGroups")), + (_READ, Decision("Complete", True)), + (replace(_READ, grants=()), Decision("Complete", True)), + (replace(_READ, grants=("other",)), Decision("Complete")), + (replace(_READ, grants=("other",), team_aliases=("model",)), Decision("Complete", True)), + (replace(_READ, grants=("group",), groups=("group",)), Decision("Complete", True)), + (replace(_READ, grants=("all-team-models",), team_id="team"), Decision("ReadNames")), + (replace(_READ, grants=("all-team-models",), team_id="team", names=()), Decision("Complete")), + (replace(_READ, grants=("all-team-models",), team_id="team", names=("model",)), Decision("Complete", True)), + (replace(_READ, grants=("mod*",)), Decision("Complete", True)), + (replace(_READ, grants=("provider/*",)), Decision("ReadProvider")), + (replace(_READ, grants=("provider/*",), provider_loaded=True), Decision("Complete")), + ( + replace(_READ, grants=("provider/*",), provider_loaded=True, provider=ModelProvider("model", "provider")), + Decision("Complete", True), + ), + ( + replace( + _READ, + model="unknown/model", + grants=("provider/*",), + provider_loaded=True, + provider=ModelProvider("unknown/model", "provider"), + ), + Decision("Complete"), + ), + ), +) +def test_model_helper_read_order_and_grants( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_owner_teams.py b/tests/unit/_auth/test_owner_teams.py new file mode 100644 index 00000000000..d8d40c79d7a --- /dev/null +++ b/tests/unit/_auth/test_owner_teams.py @@ -0,0 +1,91 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.owner_teams import OwnerTeamBudget, OwnerTeamLoad, owner_teams +from litellm._auth.owner_teams import OwnerTeamDecision as Decision +from litellm._auth.owner_teams import OwnerTeamFacts as Facts +from tests.unit._auth.test_personal_budget import ALLOW, INITIAL +from tests.unit._auth.test_personal_budget import reference as budget + + +def reference(facts: Facts) -> Decision: + if facts.purpose == "Grants" and not (facts.user_present and facts.database_present): + return Decision("Complete") + if not facts.loaded: + ids: Final = tuple(dict.fromkeys(facts.ids)) + return Decision( + "ReadTeams", + ids=tuple(team for team in ids if facts.purpose == "Models" or team and team != facts.ui_team_id), + ) + eligible: Final = tuple( + index + for index, row in enumerate(facts.loads) + if row.present + and not row.blocked + and (facts.purpose == "Models" or facts.org is None or row.org in (None, facts.org)) + ) + unreadable: Final = any(row.failed for row in facts.loads) + if facts.purpose == "Models": + return Decision("Complete", selected=eligible, unreadable=unreadable) + missing: Final = tuple(index for index in eligible if not any(read.index == index for read in facts.budgets)) + if missing: + return Decision("ReadBudgets", indices=missing) + outcomes: Final = tuple((read.index, budget(read.facts).allowed) for read in facts.budgets) + blocked: Final = tuple( + index + for index, row in enumerate(facts.loads) + if row.present and row.blocked and (facts.org is None or row.org in (None, facts.org)) + ) + return Decision( + "Complete", + selected=tuple(index for index in eligible if (index, True) in outcomes), + withheld=(*blocked, *(index for index in eligible if (index, False) in outcomes)), + unreadable=unreadable, + ) + + +_BASE: Final = Facts("Grants", ("", "ui", "a", "a", "b"), "ui", org="org") +_ROWS: Final = replace( + _BASE, + loaded=True, + loads=( + OwnerTeamLoad(True, blocked=True, org="org"), + OwnerTeamLoad(True, org="org"), + OwnerTeamLoad(True, org="other"), + OwnerTeamLoad(False, failed=True), + OwnerTeamLoad(True), + ), +) +_COMPLETE: Final = replace( + _ROWS, budgets=(OwnerTeamBudget(1, ALLOW), OwnerTeamBudget(4, replace(ALLOW, io_failed=True))) +) + + +@pytest.mark.parametrize("policy", (reference, owner_teams)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadTeams", ids=("a", "b"))), + (replace(_BASE, user_present=False), Decision("Complete")), + (replace(_BASE, database_present=False), Decision("Complete")), + (replace(_BASE, purpose="Models"), Decision("ReadTeams", ids=("", "ui", "a", "b"))), + (_ROWS, Decision("ReadBudgets", indices=(1, 4))), + (replace(_ROWS, purpose="Models"), Decision("Complete", selected=(1, 2, 4), unreadable=True)), + (_COMPLETE, Decision("Complete", selected=(1,), withheld=(0, 4), unreadable=True)), + (replace(_ROWS, org=None), Decision("ReadBudgets", indices=(1, 2, 4))), + (replace(_ROWS, org=""), Decision("ReadBudgets", indices=(4,))), + ), +) +def test_owner_team_selection_keeps_order_scope_and_withheld_rows( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected + + +def test_later_incomplete_budget_is_not_hidden_by_an_allowed_team() -> None: + with pytest.raises(NativeAuthFailure): + owner_teams(replace(_ROWS, budgets=(OwnerTeamBudget(1, ALLOW), OwnerTeamBudget(4, INITIAL)))) diff --git a/tests/unit/_auth/test_personal_budget.py b/tests/unit/_auth/test_personal_budget.py new file mode 100644 index 00000000000..d22c35e44e5 --- /dev/null +++ b/tests/unit/_auth/test_personal_budget.py @@ -0,0 +1,149 @@ +import math +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.common_budgets import BudgetFacts, MemberLimitFacts +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.personal_budget import PersonalTeamBudgetDecision as Decision +from litellm._auth.personal_budget import PersonalTeamBudgetFacts as Facts +from litellm._auth.personal_budget import personal_team_budget + + +def denied(budget: BudgetFacts) -> bool: + if budget.limit is None or budget.spend is None or not math.isfinite(budget.limit): + return False + return budget.spend > budget.limit if budget.scope == "Team" else budget.spend >= budget.limit + + +def reference(facts: Facts) -> Decision: + org: Final = facts.token_org or facts.team_org + if not facts.source_ready: + return Decision("Source", org_id=org) + if facts.io_failed: + return Decision("Complete") + if facts.team is None: + return Decision("ReadTeam") + if denied(facts.team): + return Decision("Complete") if facts.team_alerted else Decision("AlertTeam") + if facts.window_count is None: + return Decision("ReadWindows") + for index in range(facts.window_count): + if index >= len(facts.windows): + return Decision("ReadWindow", index=index) + if denied(facts.windows[index]): + return Decision("Complete") + if facts.check_db_only and not facts.user_present: + return Decision("Complete") + if facts.user_present and (facts.team_id_present or facts.check_db_only) and not facts.membership_loaded: + return Decision("ReadMembership") + if facts.user_present and facts.team_id_present: + limit: Final = facts.member_limit + if limit is None: + return Decision("ReadMemberLimit") + if limit.member_limit is None and facts.default_id is not None and not facts.default_loaded: + return Decision("ReadDefault", id=facts.default_id) + base: Final = ( + limit.member_limit + if limit.member_limit is not None + else limit.default_limit + if limit.default_limit is not None and not limit.default_limit <= 0 + else None + ) + if ( + base is not None + and limit.temporary_increase is not None + and limit.expires_at is not None + and limit.now_naive + ): + return Decision("Complete", log_failure=True) + increase: Final = ( + limit.temporary_increase + if limit.temporary_increase is not None and limit.expires_at is not None and limit.expires_at > limit.now + else 0.0 + ) + if facts.member is None: + return Decision("ReadMember", limit=base + increase if base is not None else None) + if facts.member.limit is not None and math.isfinite(facts.member.limit) and not facts.member_alerted: + return Decision("AlertMember") + if denied(facts.member): + return Decision("Complete") + if org is not None: + if not facts.org_loaded: + return Decision("ReadOrg", id=org) + if not facts.org_failed and facts.org is not None and denied(facts.org): + return Decision("Complete") if facts.org_alerted else Decision("AlertOrg", id=org) + return Decision("Complete", allowed=True) + + +INITIAL: Final = Facts(user_present=True, token_org="org") +TEAM: Final = replace(INITIAL, source_ready=True, team=BudgetFacts("Team", None, token_present=True)) +WINDOWS: Final = replace(TEAM, window_count=0) +MEMBERSHIP: Final = replace(WINDOWS, membership_loaded=True) +LIMIT: Final = replace(MEMBERSHIP, member_limit=MemberLimitFacts(None, None, None, None, 100)) +MEMBER: Final = replace(LIMIT, member=BudgetFacts("TeamMember", None)) +ALLOW: Final = replace(MEMBER, org_loaded=True) + + +@pytest.mark.parametrize("policy", (reference, personal_team_budget)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (INITIAL, Decision("Source", org_id="org")), + (replace(INITIAL, token_org="", team_org="team-org"), Decision("Source", org_id="team-org")), + (replace(INITIAL, token_org=None, team_org=""), Decision("Source", org_id="")), + (replace(INITIAL, source_ready=True), Decision("ReadTeam")), + (replace(TEAM, io_failed=True), Decision("Complete")), + (replace(TEAM, team=BudgetFacts("Team", 1, 2, token_present=True)), Decision("AlertTeam")), + (replace(TEAM, team=BudgetFacts("Team", 1, 2, token_present=True), team_alerted=True), Decision("Complete")), + (replace(TEAM, team=BudgetFacts("Team", 1, 1, token_present=True)), Decision("ReadWindows")), + (TEAM, Decision("ReadWindows")), + (replace(TEAM, window_count=2), Decision("ReadWindow")), + (replace(TEAM, window_count=2, windows=(BudgetFacts("TeamWindow", 1, 0),)), Decision("ReadWindow", index=1)), + (replace(TEAM, window_count=2, windows=(BudgetFacts("TeamWindow", 1, 1),)), Decision("Complete")), + (replace(WINDOWS, user_present=False, check_db_only=True), Decision("Complete")), + (replace(WINDOWS, user_present=False), Decision("ReadOrg", id="org")), + (WINDOWS, Decision("ReadMembership")), + (replace(WINDOWS, team_id_present=False), Decision("ReadOrg", id="org")), + (replace(WINDOWS, team_id_present=False, check_db_only=True), Decision("ReadMembership")), + (replace(MEMBERSHIP, team_id_present=False, check_db_only=True), Decision("ReadOrg", id="org")), + (MEMBERSHIP, Decision("ReadMemberLimit")), + (replace(LIMIT, default_id=""), Decision("ReadDefault")), + ( + replace(LIMIT, default_id="default", member_limit=MemberLimitFacts(0, None, None, None, 100)), + Decision("ReadMember", limit=0), + ), + ( + replace( + LIMIT, default_id="default", default_loaded=True, member_limit=MemberLimitFacts(None, 10, 2, 101, 100) + ), + Decision("ReadMember", limit=12), + ), + (replace(LIMIT, member_limit=MemberLimitFacts(10, None, 2, 100, 100)), Decision("ReadMember", limit=10)), + ( + replace(LIMIT, member_limit=MemberLimitFacts(10, None, 2, 101, 100, True)), + Decision("Complete", log_failure=True), + ), + (LIMIT, Decision("ReadMember")), + (replace(MEMBER, member=BudgetFacts("TeamMember", 1, 1)), Decision("AlertMember")), + (replace(MEMBER, member=BudgetFacts("TeamMember", 1, 1), member_alerted=True), Decision("Complete")), + (replace(MEMBER, member=BudgetFacts("TeamMember", float("inf"), 100)), Decision("ReadOrg", id="org")), + (MEMBER, Decision("ReadOrg", id="org")), + (replace(MEMBER, token_org=None), Decision("Complete", allowed=True)), + (ALLOW, Decision("Complete", allowed=True)), + (replace(ALLOW, org_failed=True), Decision("Complete", allowed=True)), + (replace(ALLOW, org=BudgetFacts("Organization", 1, 1)), Decision("AlertOrg", id="org")), + (replace(ALLOW, org=BudgetFacts("Organization", 1, 1), org_alerted=True), Decision("Complete")), + ), +) +def test_personal_budget_read_and_denial_order( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected + + +def test_incomplete_budget_read_is_terminal_protocol_failure() -> None: + with pytest.raises(NativeAuthFailure): + personal_team_budget(replace(TEAM, team=BudgetFacts("Team", 1))) diff --git a/tests/unit/_auth/test_personal_model_listing.py b/tests/unit/_auth/test_personal_model_listing.py new file mode 100644 index 00000000000..0d40963c03f --- /dev/null +++ b/tests/unit/_auth/test_personal_model_listing.py @@ -0,0 +1,66 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_model_listing import PersonalModelListingDecision as Decision +from litellm._auth.personal_model_listing import PersonalModelListingFacts as Facts +from litellm._auth.personal_model_listing import personal_model_listing +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm._auth.personal_owner import PersonalOwnerFacts +from tests.unit._auth.test_personal_models import ALLOW, INITIAL +from tests.unit._auth.test_personal_models import reference as scope_reference +from tests.unit._auth.test_personal_owner import reference as owner_reference + + +def reference(facts: Facts) -> Decision: + owner: Final = owner_reference(facts.owner) + if owner.kind != "Complete": + return Decision("ReadOwner") + if facts.owner.teams is None or not facts.owner.teams.loads: + return Decision("Complete") + if facts.names is None: + return Decision("ReadNames") + if len(facts.models) < len(facts.names): + return Decision("ReadModel", index=len(facts.models), teams=owner.selected) + return Decision( + "Complete", + selected=tuple( + index + for index, model in enumerate(facts.models) + if any(tuple(scope_reference(scope).allowed for scope in model.scopes)) + ), + ) + + +_BASE: Final = Facts(PersonalOwnerFacts(True, True, True, True)) +_OWNER: Final = replace( + _BASE.owner, + owner_loaded=True, + owner_present=True, + teams=OwnerTeamFacts("Models", ("team",), "ui", loaded=True, loads=(OwnerTeamLoad(True),)), +) +_LOADED: Final = replace(_BASE, owner=_OWNER, names=("first", "second", "first")) +_ALLOWED: Final = LoadedTeamModelFacts((True,), (ALLOW,)) +_DENIED: Final = LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)) + + +@pytest.mark.parametrize("policy", (reference, personal_model_listing)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadOwner")), + (replace(_BASE, owner=replace(_BASE.owner, database_present=False)), Decision("Complete")), + (replace(_BASE, owner=replace(_BASE.owner, owner_loaded=True, owner_failed=True)), Decision("Complete")), + (replace(_BASE, owner=_OWNER), Decision("ReadNames")), + (_LOADED, Decision("ReadModel", teams=(0,))), + (replace(_LOADED, models=(_DENIED,)), Decision("ReadModel", index=1, teams=(0,))), + (replace(_LOADED, models=(_ALLOWED, _DENIED, _ALLOWED)), Decision("Complete", selected=(0, 2))), + ), +) +def test_listing_preserves_order_duplicates_and_checks_after_denial( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_personal_model_request.py b/tests/unit/_auth/test_personal_model_request.py new file mode 100644 index 00000000000..40b2edf752f --- /dev/null +++ b/tests/unit/_auth/test_personal_model_request.py @@ -0,0 +1,57 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_model_request import PersonalModelRequestDecision as Decision +from litellm._auth.personal_model_request import PersonalModelRequestFacts as Facts +from litellm._auth.personal_model_request import personal_model_request +from litellm._auth.personal_models import LoadedTeamModelFacts +from tests.unit._auth.test_owner_teams import reference as teams_reference +from tests.unit._auth.test_personal_models import ALLOW, INITIAL, reference as scope_reference + + +def reference(facts: Facts) -> Decision: + if facts.owner_present and facts.teams is None: + return Decision("ReadTeams") + teams: Final = teams_reference(facts.teams).selected if facts.teams is not None else () + for index, model in enumerate(facts.requested): + if not facts.owner_present: + return Decision("Deny", model=model) + if index >= len(facts.models): + return Decision("ReadModel", index=index, teams=teams) + if not any(scope_reference(scope).allowed for scope in facts.models[index].scopes): + return Decision("Deny", model=model) + return Decision("Allow") + + +_BASE: Final = Facts(("first", "second"), True) +_TEAMS: Final = replace( + _BASE, teams=OwnerTeamFacts("Models", ("team",), "ui", loaded=True, loads=(OwnerTeamLoad(True),)) +) +_ONE: Final = replace(_TEAMS, models=(LoadedTeamModelFacts((True,), (ALLOW,)),)) + + +@pytest.mark.parametrize("policy", (reference, personal_model_request)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadTeams")), + (replace(_BASE, requested=()), Decision("ReadTeams")), + (replace(_BASE, owner_present=False), Decision("Deny", model="first")), + (replace(_BASE, owner_present=False, requested=()), Decision("Allow")), + (_TEAMS, Decision("ReadModel", teams=(0,))), + (_ONE, Decision("ReadModel", index=1, teams=(0,))), + (replace(_ONE, models=(*_ONE.models, LoadedTeamModelFacts((True,), (ALLOW,)))), Decision("Allow")), + ( + replace(_TEAMS, models=(LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)),)), + Decision("Deny", model="first"), + ), + ), +) +def test_personal_requests_read_all_teams_but_stop_at_first_model_denial( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_personal_models.py b/tests/unit/_auth/test_personal_models.py new file mode 100644 index 00000000000..5e0dffaa0bb --- /dev/null +++ b/tests/unit/_auth/test_personal_models.py @@ -0,0 +1,110 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts, personal_model_access_applies +from litellm._auth.personal_models import LoadedTeamModelFacts, loaded_team_models, personal_model_scope +from litellm._auth.personal_models import PersonalModelDecision as Decision +from litellm._auth.personal_models import PersonalModelScope as Scope +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Scope) -> Decision: + if facts.io_failed: + return Decision("Complete") + if facts.team is None: + return Decision("ReadTeam") + team: Final = grouped(facts.team, merge=True) + if team.kind == "ReadGroups": + return Decision("ReadGroups") + if team.kind != "Allow": + return Decision("Complete") + if not facts.user_present or not facts.team_id_present: + return Decision("Complete", allowed=True) + if facts.member_grants is None: + return Decision("ReadMembership") + if not facts.member_grants: + return Decision("Complete", allowed=True) + if facts.member is None: + return Decision("ReadMember") + return Decision("Complete", allowed=model_access(facts.member).kind == "Allow") + + +INITIAL: Final = Scope(True, True) +MODEL: Final = ModelAccessFacts(("model",), ("model",), object_type="team") +TEAM: Final = replace(INITIAL, team=GroupedModelScope(MODEL)) +ALLOW: Final = replace(TEAM, member_grants=()) + + +@pytest.mark.parametrize("policy", (reference, personal_model_scope)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (INITIAL, Decision("ReadTeam")), + (replace(INITIAL, io_failed=True), Decision("Complete")), + (TEAM, Decision("ReadMembership")), + (replace(TEAM, team=GroupedModelScope(replace(MODEL, grants=("other",)))), Decision("Complete")), + (replace(TEAM, team=GroupedModelScope(replace(MODEL, grants=("other",)), ("group",))), Decision("ReadGroups")), + ( + replace(TEAM, team=GroupedModelScope(replace(MODEL, grants=("other",)), ("group",), ("model",))), + Decision("ReadMembership"), + ), + (replace(TEAM, team_id_present=False), Decision("Complete", allowed=True)), + (replace(TEAM, user_present=False), Decision("Complete", allowed=True)), + (ALLOW, Decision("Complete", allowed=True)), + (replace(TEAM, member_grants=("member",)), Decision("ReadMember")), + (replace(TEAM, member_grants=("other",), member=replace(MODEL, grants=("other",))), Decision("Complete")), + (replace(TEAM, member_grants=("model",), member=MODEL), Decision("Complete", allowed=True)), + ), +) +def test_team_model_scope_orders_grants_before_membership( + policy: Callable[[Scope], Decision], facts: Scope, expected: Decision +) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (LoadedTeamModelFacts(()), Decision("Complete")), + (LoadedTeamModelFacts((False, True, False, True)), Decision("ReadScopes", indices=(1, 3))), + ( + LoadedTeamModelFacts((True, True), (ALLOW, replace(INITIAL, io_failed=True))), + Decision("Complete", allowed=True), + ), + (LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)), Decision("Complete")), + ), +) +def test_loaded_teams_are_all_read_before_union(facts: LoadedTeamModelFacts, expected: Decision) -> None: + assert loaded_team_models(facts) == expected + + +def test_later_native_failure_is_not_hidden_by_first_team_allow() -> None: + with pytest.raises(NativeAuthFailure): + loaded_team_models(LoadedTeamModelFacts((True, True), (ALLOW, INITIAL))) + + +@pytest.mark.parametrize( + ("api_key_present", "inheritance", "expected"), ((False, False, True), (False, True, False), (True, True, True)) +) +def test_inheritance_requires_virtual_key_not_just_owner( + api_key_present: bool, inheritance: bool, expected: bool +) -> None: + assert ( + personal_model_access_applies( + enabled=True, + token_present=True, + team_id=None, + user_id="", + master=False, + role=None, + api_key_present=api_key_present, + inheritance=inheritance, + ) + is expected + ) diff --git a/tests/unit/_auth/test_personal_owner.py b/tests/unit/_auth/test_personal_owner.py new file mode 100644 index 00000000000..02e75764fcd --- /dev/null +++ b/tests/unit/_auth/test_personal_owner.py @@ -0,0 +1,61 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_owner import PersonalOwnerDecision as Decision +from litellm._auth.personal_owner import PersonalOwnerFacts as Facts +from litellm._auth.personal_owner import personal_owner +from tests.unit._auth.test_owner_teams import reference as select_teams + + +def reference(facts: Facts) -> Decision: + if not all((facts.user_present, facts.database_present, facts.cache_present, facts.logging_present)): + return Decision("Complete") + if not facts.owner_loaded: + return Decision("ReadOwner") + if facts.owner_failed or not facts.owner_present: + return Decision("Complete") + if facts.teams is None: + return Decision("ReadTeams") + teams: Final = select_teams(facts.teams) + return Decision("Complete", teams.selected, teams.withheld, teams.unreadable) + + +_BASE: Final = Facts(True, True, True, True) +_OWNER: Final = replace(_BASE, owner_loaded=True, owner_present=True) + + +@pytest.mark.parametrize("policy", (reference, personal_owner)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadOwner")), + (replace(_BASE, user_present=False), Decision("Complete")), + (replace(_BASE, database_present=False), Decision("Complete")), + (replace(_BASE, cache_present=False), Decision("Complete")), + (replace(_BASE, logging_present=False), Decision("Complete")), + (replace(_BASE, owner_loaded=True), Decision("Complete")), + (replace(_OWNER, owner_failed=True), Decision("Complete")), + (_OWNER, Decision("ReadTeams")), + ( + replace( + _OWNER, + teams=OwnerTeamFacts( + "Grants", + (), + "ui", + loaded=True, + loads=(OwnerTeamLoad(True, blocked=True), OwnerTeamLoad(False, failed=True)), + ), + ), + Decision("Complete", withheld=(0,), unreadable=True), + ), + ), +) +def test_owner_lookup_and_team_status_order( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_python_core.py b/tests/unit/_auth/test_python_core.py new file mode 100644 index 00000000000..3d37c7b557d --- /dev/null +++ b/tests/unit/_auth/test_python_core.py @@ -0,0 +1,637 @@ +import math +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone +from typing import Final +from zoneinfo import ZoneInfo + +import pytest + +from litellm._auth.budget_pricing import ExemptionFacts, RouteBudgetFacts +from litellm._auth.decision import ( + AdminKey, + Allow, + AuthCore, + CachedAdminKey, + Deny, + ExpiredKey, + KeyBlocked, + KeyBudgetExceeded, + KeyFacts, + KeyWithoutToken, + MalformedKey, + MasterKey, + MasterKeyOnlyRoute, + MissingKey, + ModelBudgetExceeded, + ModelBudgetRead, + ModelBudgetScope, + NoDatabase, + NoMasterKey, + NonAdminKey, + NotAVirtualKey, + Outcome, + OwnerDeactivated, + OwnerFacts, + PresentedKey, + RequestFacts, + Settings, + Step, + TeamMemberBudget, + TeamMemberBudgetExceeded, + UserAuthDisabled, + UserAuthRoute, + VirtualKey, +) +from litellm._auth.driver import drive +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.key import hash_token +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts +from tests.unit._auth.raw_answers import budget_read, fallback_read + +MASTER: Final = "sk-master" +NOW: Final = datetime(2026, 11, 6, 12, tzinfo=timezone.utc) +SETTINGS: Final = Settings( + master_key=MASTER, + allow_user_auth=False, + has_db=True, + master_key_only_routes=frozenset({"/global/spend/reset"}), + proxy_max_budget=0.0, +) +PLAIN_KEY: Final = KeyFacts( + has_token=True, + blocked=False, + expires=None, + spend=0.0, + soft_budget=None, + max_budget=None, + team_member_spend=None, + user_id=None, + team_id=None, + budget_throttle_pct=None, +) +CHECKED: Final = ( + "prepare_key", + "inspect_key", + "enforce_model_access", + "check_budget_exemption", + "send_key_budget_alerts", + "read_model_budget_scope", + "load_key_context", +) + +NO_MODEL_BUDGETS: Final = ModelBudgetScope( + models=(), + rewritable_model=None, + key_has_model_budgets=False, + key_budget_fallbacks={}, + user_has_model_budgets=False, +) + +Call = tuple[str, object] + + +@dataclass +class RecordingHandler: + presented: PresentedKey | None = None + key: KeyFacts = PLAIN_KEY + owner: OwnerFacts | None = None + exempt: bool = False + team_member_budget: TeamMemberBudget | None = None + key_spend: float = 0.0 + proxy_spend: float | None = None + model_budget_scope: ModelBudgetScope = NO_MODEL_BUDGETS + user_model_budget_exceeded: ModelBudgetExceeded | None = None + calls: tuple[Call, ...] = () + + def _record(self, name: str, payload: object = None) -> None: + self.calls = (*self.calls, (name, payload)) + + async def resolve_end_user(self, prefetch_key: bool) -> None: + self._record("resolve_end_user", prefetch_key) + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: + self._record("resolve_presented_key", hashed_token) + return self.presented + + async def evict_presented_key(self, hashed_token: str) -> None: + self._record("evict_presented_key", hashed_token) + + async def load_key(self, hashed_token: str) -> None: + self._record("load_key", hashed_token) + + async def prepare_key(self, loaded_from_db: bool) -> None: + self._record("prepare_key", loaded_from_db) + + async def inspect_key(self) -> KeyFacts: + self._record("inspect_key") + return self.key + + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: + self._record("enforce_model_access") + return KeyAndFallbackModelAccessFacts(True, ()) + + async def prefetch_auth_objects(self) -> None: + return None + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: + self._record("load_key_owner", user_id) + return self.owner + + async def check_budget_exemption(self) -> ExemptionFacts: + self._record("check_budget_exemption") + return ExemptionFacts(RouteBudgetFacts(discovery_route=self.exempt, llm_route=True), False, False, 0) + + async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: + self._record("read_team_member_budget", (user_id, team_id)) + return self.team_member_budget + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: + self._record("alert_team_member_budget", budget) + + async def send_key_budget_alerts(self) -> None: + self._record("send_key_budget_alerts") + + async def read_key_spend(self, max_budget: float) -> float: + self._record("read_key_spend", max_budget) + return self.key_spend + + async def alert_key_budget(self, spend: float) -> None: + self._record("alert_key_budget", spend) + + async def alert_key_soft_budget(self) -> None: + self._record("alert_key_soft_budget") + + async def read_model_budget_scope(self) -> ModelBudgetScope: + self._record("read_model_budget_scope") + return self.model_budget_scope + + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: + self._record("check_key_model_budget", model) + return budget_read() + + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: + self._record("authorize_budget_fallback", model) + return fallback_read(model) + + async def rewrite_request_model(self, model: str) -> None: + self._record("rewrite_request_model", model) + + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: + self._record("check_user_model_budget", model) + return budget_read(self.user_model_budget_exceeded) + + async def load_key_context(self) -> None: + self._record("load_key_context") + + async def read_proxy_spend(self) -> float | None: + self._record("read_proxy_spend") + return self.proxy_spend + + async def alert_proxy_budget(self, spend: float) -> None: + self._record("alert_proxy_budget", spend) + + +def _facts(token: str | None, route: str = "/chat/completions", is_llm_api_route: bool = True) -> RequestFacts: + return RequestFacts(route=route, token=token, is_llm_api_route=is_llm_api_route) + + +async def _run(step: Step, handler: RecordingHandler) -> tuple[Outcome, tuple[str, ...]]: + outcome: Final = await drive(step, handler) + return outcome, tuple(name for name, _ in handler.calls) + + +async def _decide( + core: AuthCore, + token: str | None, + *, + route: str = "/chat/completions", + settings: Settings = SETTINGS, + handler: RecordingHandler | None = None, +) -> tuple[Outcome, tuple[str, ...]]: + return await _run(core.start(_facts(token, route), settings, NOW), handler or RecordingHandler()) + + +async def _check( + core: AuthCore, + key: KeyFacts, + *, + is_llm_api_route: bool = True, + settings: Settings = SETTINGS, + owner: OwnerFacts | None = None, + exempt: bool = False, + team_member_budget: TeamMemberBudget | None = None, + key_spend: float = 0.0, +) -> tuple[Outcome, tuple[str, ...]]: + handler: Final = RecordingHandler( + key=key, owner=owner, exempt=exempt, team_member_budget=team_member_budget, key_spend=key_spend + ) + facts: Final = _facts("hash", is_llm_api_route=is_llm_api_route) + return await _run(core.start_resolved_key(facts, settings, NOW), handler) + + +def _allowed(budget_throttle_pct: float | None = None) -> Outcome: + return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "token", "settings", "expected"), + [ + ("/chat/completions", None, Settings(None, False, True, frozenset(), 0.0), Allow(NoMasterKey(token=None))), + ("/chat/completions", "sk-x", Settings(None, False, False, frozenset(), 0.0), Allow(NoMasterKey(token="sk-x"))), + ("/chat/completions", None, SETTINGS, Deny(MissingKey())), + ("/chat/completions", "", SETTINGS, Deny(MalformedKey())), + ("/user/auth", "", Settings(MASTER, True, True, frozenset(), 0.0), Deny(MalformedKey())), + ("/user/auth", "sk-x", Settings(MASTER, True, True, frozenset(), 0.0), Allow(UserAuthRoute())), + ("/user/auth", "sk-x", SETTINGS, Deny(UserAuthDisabled())), + ("/user/auth", MASTER, SETTINGS, Deny(UserAuthDisabled())), + ], +) +async def test_credential_decisions_do_no_io( + core: AuthCore, route: str, token: str | None, settings: Settings, expected: Outcome +): + assert await _decide(core, token, route=route, settings=settings) == (expected, ()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/chat/completions", "/global/spend/reset"]) +async def test_master_key_resolves_end_user_without_key_prefetch(core: AuthCore, route: str): + handler: Final = RecordingHandler() + assert await _decide(core, MASTER, route=route, handler=handler) == (Allow(MasterKey()), ("resolve_end_user",)) + assert handler.calls == (("resolve_end_user", False),) + + +@pytest.mark.asyncio +async def test_unknown_virtual_key_is_loaded_from_the_db_then_checked( + core: AuthCore, +): + handler: Final = RecordingHandler() + outcome, calls = await _decide(core, "sk-x", handler=handler) + assert outcome == _allowed() + assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", *CHECKED) + assert handler.calls[2] == ("load_key", hash_token("sk-x")) + assert handler.calls[3] == ("prepare_key", True) + + +@pytest.mark.asyncio +async def test_cached_virtual_key_is_checked_without_a_db_load( + core: AuthCore, +): + handler: Final = RecordingHandler(presented=NonAdminKey()) + outcome, calls = await _decide(core, "sk-x", handler=handler) + assert outcome == _allowed() + assert calls == ("resolve_end_user", "resolve_presented_key", *CHECKED) + assert handler.calls[2] == ("prepare_key", False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW]) +async def test_unexpired_cached_admin_key_skips_route_db_and_key_checks(core: AuthCore, expires: datetime | None): + no_db: Final = Settings(MASTER, False, False, frozenset({"/global/spend/reset"}), 0.0) + outcome, calls = await _decide( + core, + "sk-admin", + route="/global/spend/reset", + settings=no_db, + handler=RecordingHandler(presented=AdminKey(expires=expires)), + ) + assert (outcome, calls) == (Allow(CachedAdminKey()), ("resolve_end_user", "resolve_presented_key")) + + +@pytest.mark.asyncio +async def test_expired_cached_admin_key_is_evicted_then_denied_with_the_presented_key( + core: AuthCore, +): + expires: Final = NOW - timedelta(seconds=1) + handler: Final = RecordingHandler(presented=AdminKey(expires=expires)) + outcome, calls = await _decide(core, "sk-admin", handler=handler) + assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label="sk-admin")) + assert calls == ("resolve_end_user", "resolve_presented_key", "evict_presented_key") + assert handler.calls[-1] == ("evict_presented_key", hash_token("sk-admin")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("presented", [None, NonAdminKey()]) +async def test_non_admin_key_cannot_use_master_key_only_routes(core: AuthCore, presented: PresentedKey | None): + outcome, calls = await _decide( + core, "sk-x", route="/global/spend/reset", handler=RecordingHandler(presented=presented) + ) + assert (outcome, calls) == (Deny(MasterKeyOnlyRoute()), ("resolve_end_user", "resolve_presented_key")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("presented", [None, NonAdminKey()]) +async def test_virtual_key_needs_a_database(core: AuthCore, presented: PresentedKey | None): + no_db: Final = Settings(MASTER, False, False, frozenset(), 0.0) + outcome, _ = await _decide(core, "sk-x", settings=no_db, handler=RecordingHandler(presented=presented)) + assert outcome == Deny(NoDatabase()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token", ["eyJhbGciOiJIUzI1NiJ9.e30.sig", "a" * 64, "sk"]) +async def test_uncached_token_without_virtual_key_prefix_is_rejected_before_the_db(core: AuthCore, token: str): + assert await _decide(core, token) == (Deny(NotAVirtualKey()), ("resolve_end_user", "resolve_presented_key")) + + +@pytest.mark.asyncio +async def test_expired_virtual_key_is_labelled_by_the_hash_it_was_loaded_with( + core: AuthCore, +): + expires: Final = NOW - timedelta(seconds=1) + outcome, _ = await _decide(core, "sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires))) + assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x"))) + + +@pytest.mark.asyncio +async def test_blocked_key_is_denied_before_model_access( + core: AuthCore, +): + assert await _check(core, replace(PLAIN_KEY, blocked=True)) == (Deny(KeyBlocked()), ("inspect_key",)) + + +@pytest.mark.asyncio +async def test_deactivated_owner_is_denied_after_model_access( + core: AuthCore, +): + outcome, calls = await _check(core, replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True)) + assert (outcome, calls) == ( + Deny(OwnerDeactivated(user_id="u1")), + ("inspect_key", "enforce_model_access", "load_key_owner"), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner", [None, OwnerFacts(deactivated=False)]) +async def test_active_or_missing_owner_continues(core: AuthCore, owner: OwnerFacts | None): + outcome, _ = await _check(core, replace(PLAIN_KEY, user_id="u1"), owner=owner) + assert outcome == _allowed() + + +@pytest.mark.asyncio +async def test_exempt_request_skips_every_budget_check_but_not_expiry( + core: AuthCore, +): + key: Final = replace( + PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=99.0, max_budget=1.0, spend=5.0, soft_budget=1.0 + ) + assert await _check(core, key, exempt=True) == ( + _allowed(), + ("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption", "load_key_context"), + ) + expired: Final = replace(key, expires=NOW - timedelta(seconds=1)) + outcome, _ = await _check(core, expired, exempt=True) + assert outcome == Deny(ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("budget", "denied"), + [ + (TeamMemberBudget(spend=10.0, max_budget=10.0), True), + (TeamMemberBudget(spend=10.1, max_budget=10.0), True), + (TeamMemberBudget(spend=9.9, max_budget=10.0), False), + (None, False), + ], +) +async def test_team_member_budget_alerts_then_denies_at_the_limit( + core: AuthCore, budget: TeamMemberBudget | None, denied: bool +): + key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0) + outcome, calls = await _check(core, key, team_member_budget=budget) + if denied: + assert budget is not None + assert outcome == Deny( + TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=budget.spend, max_budget=budget.max_budget) + ) + assert calls[-2:] == ("read_team_member_budget", "alert_team_member_budget") + else: + assert outcome == _allowed() + assert "alert_team_member_budget" not in calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key", + [ + replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=None), + replace(PLAIN_KEY, user_id=None, team_id="t1", team_member_spend=0.0), + replace(PLAIN_KEY, user_id="u1", team_id=None, team_member_spend=0.0), + ], +) +async def test_team_member_budget_is_read_only_for_a_member_with_tracked_spend(core: AuthCore, key: KeyFacts): + _, calls = await _check(core, key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert "read_team_member_budget" not in calls + + +@pytest.mark.asyncio +async def test_team_member_budget_needs_a_database( + core: AuthCore, +): + key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0) + no_db: Final = Settings(MASTER, False, False, frozenset(), 0.0) + _, calls = await _check(core, key, settings=no_db, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert "read_team_member_budget" not in calls + + +@pytest.mark.asyncio +async def test_team_member_budget_is_checked_before_expiry( + core: AuthCore, +): + key: Final = replace( + PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0, expires=NOW - timedelta(seconds=1) + ) + outcome, _ = await _check(core, key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert outcome == Deny(TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=1.0, max_budget=0.5)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend", "max_budget", "throttle", "expected"), + [ + (5.0, 5.0, None, Deny(KeyBudgetExceeded(spend=5.0, max_budget=5.0))), + (6.0, 5.0, None, Deny(KeyBudgetExceeded(spend=6.0, max_budget=5.0))), + (4.9, 5.0, None, _allowed()), + (6.0, 5.0, 0.5, _allowed(budget_throttle_pct=0.5)), + (4.9, 5.0, 0.5, _allowed()), + ( + 1e9, + math.inf, + None, + _allowed(), + ), + ( + 1e9, + math.nan, + None, + _allowed(), + ), + ], +) +async def test_key_budget_denies_or_throttles_at_the_limit( + core: AuthCore, spend: float, max_budget: float, throttle: float | None, expected: Outcome +): + key: Final = replace(PLAIN_KEY, max_budget=max_budget, budget_throttle_pct=throttle) + handler: Final = RecordingHandler(key=key, key_spend=spend) + outcome, calls = await _run(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + assert outcome == expected + assert calls[3:6] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget") + assert handler.calls[5] == ("alert_key_budget", spend) + assert ("read_model_budget_scope" in calls) == isinstance(expected, Allow) + + +@pytest.mark.asyncio +async def test_key_budget_is_enforced_only_on_llm_routes( + core: AuthCore, +): + key: Final = replace(PLAIN_KEY, max_budget=1.0) + outcome, calls = await _check(core, key, is_llm_api_route=False, key_spend=2.0) + assert outcome == _allowed() + assert "read_key_spend" not in calls + assert "read_model_budget_scope" in calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend", "soft_budget", "alerted"), + [ + (5.0, 5.0, True), + (6.0, 5.0, True), + (4.9, 5.0, False), + (10.0, None, False), + (10.0, 0.0, False), + (10.0, math.nan, False), + ], +) +async def test_soft_budget_alerts_at_the_threshold_before_model_budgets( + core: AuthCore, spend: float, soft_budget: float | None, alerted: bool +): + outcome, calls = await _check(core, replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget)) + assert outcome == _allowed() + soft_alert: Final = ("alert_key_soft_budget",) if alerted else () + assert calls[3:] == ("send_key_budget_alerts", *soft_alert, "read_model_budget_scope", "load_key_context") + + +@pytest.mark.asyncio +async def test_throttled_key_still_gets_soft_and_model_budget_checks( + core: AuthCore, +): + key: Final = replace(PLAIN_KEY, max_budget=5.0, budget_throttle_pct=0.5, spend=6.0, soft_budget=1.0) + outcome, calls = await _check(core, key, key_spend=6.0) + assert outcome == _allowed(budget_throttle_pct=0.5) + assert calls[-4:] == ("alert_key_budget", "alert_key_soft_budget", "read_model_budget_scope", "load_key_context") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exempt", [False, True]) +@pytest.mark.parametrize( + ("proxy_max_budget", "has_db", "proxy_spend", "tail"), + [ + (0.0, True, 50.0, (("load_key_context", None),)), + (-1.0, True, 50.0, (("load_key_context", None),)), + (math.nan, True, 50.0, (("load_key_context", None),)), + (10.0, False, 50.0, (("load_key_context", None),)), + (10.0, True, None, (("load_key_context", None), ("read_proxy_spend", None))), + ( + 10.0, + True, + 50.0, + (("load_key_context", None), ("read_proxy_spend", None), ("alert_proxy_budget", 50.0)), + ), + ], +) +async def test_proxy_budget_only_alerts_and_runs_for_exempt_requests_too( + core: AuthCore, + exempt: bool, + proxy_max_budget: float, + has_db: bool, + proxy_spend: float | None, + tail: tuple[Call, ...], +): + settings: Final = replace(SETTINGS, proxy_max_budget=proxy_max_budget, has_db=has_db) + handler: Final = RecordingHandler(exempt=exempt, proxy_spend=proxy_spend) + outcome, _ = await _run(core.start_resolved_key(_facts("hash"), settings, NOW), handler) + assert outcome == _allowed() + assert handler.calls[-len(tail) :] == tail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exempt", [False, True]) +async def test_key_without_token_is_denied_after_its_context_loads(core: AuthCore, exempt: bool): + settings: Final = replace(SETTINGS, proxy_max_budget=10.0) + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, has_token=False), exempt=exempt, proxy_spend=1.0) + outcome, calls = await _run(core.start_resolved_key(_facts("hash"), settings, NOW), handler) + assert outcome == Deny(KeyWithoutToken()) + assert calls[-3:] == ("load_key_context", "read_proxy_spend", "alert_proxy_budget") + + +@pytest.mark.asyncio +async def test_a_model_budget_denial_stops_the_key_chain_before_its_context_loads( + core: AuthCore, +): + exceeded: Final = ModelBudgetExceeded( + message="over", current_cost=2.0, max_budget=1.0, entity_type="user", entity_id="u-1" + ) + scope: Final = replace(NO_MODEL_BUDGETS, models=("gpt-4o",), user_has_model_budgets=True) + handler: Final = RecordingHandler(model_budget_scope=scope, user_model_budget_exceeded=exceeded) + outcome, calls = await _run(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + assert outcome == Deny(exceeded) + assert calls[-2:] == ("read_model_budget_scope", "check_user_model_budget") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("now", "expires", "expired"), + [ + (NOW, NOW, False), + (NOW, NOW - timedelta(microseconds=1), True), + (NOW, NOW.astimezone(timezone(timedelta(hours=9))), False), + (NOW.replace(tzinfo=None), NOW.replace(tzinfo=None) - timedelta(microseconds=1), True), + ( + datetime(2026, 11, 1, 1, 20, fold=1, tzinfo=ZoneInfo("America/New_York")), + datetime(2026, 11, 1, 1, 30, fold=0, tzinfo=ZoneInfo("America/New_York")), + False, + ), + ( + datetime(2026, 11, 1, 1, 20, fold=1, tzinfo=ZoneInfo("America/New_York")), + datetime(2026, 11, 1, 1, 30, fold=0, tzinfo=ZoneInfo.no_cache("America/New_York")), + True, + ), + ], +) +async def test_datetime_comparisons_preserve_timezone_identity_and_public_values( + core: AuthCore, now: datetime, expires: datetime, expired: bool +): + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, expires=expires)) + outcome: Final = await drive(core.start_resolved_key(_facts("hash"), SETTINGS, now), handler) + if expired: + assert outcome == Deny(ExpiredKey(expires, now, "hash")) + assert isinstance(outcome, Deny) and isinstance(outcome.denial, ExpiredKey) + assert outcome.denial.expires is expires + assert outcome.denial.now is now + else: + assert outcome == _allowed() + + +@pytest.mark.asyncio +async def test_mixed_naive_and_aware_expiry_is_terminal_at_the_original_check(core: AuthCore): + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, expires=NOW.replace(tzinfo=None))) + with pytest.raises(TypeError, match="offset-naive and offset-aware"): + await drive(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + assert handler.calls == (("inspect_key", None), ("enforce_model_access", None), ("check_budget_exemption", None)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget", [math.nan, math.inf, -math.inf, 0.0]) +async def test_team_member_nonfinite_budgets_keep_the_reference_comparison(core: AuthCore, budget: float): + handler: Final = RecordingHandler( + key=replace(PLAIN_KEY, user_id="u", team_id="t", team_member_spend=0), + team_member_budget=TeamMemberBudget(1.0, budget), + ) + outcome: Final = await drive(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + if budget == math.inf: + assert outcome == _allowed() + else: + assert isinstance(outcome, Deny) and isinstance(outcome.denial, TeamMemberBudgetExceeded) + assert outcome.denial.spend == 1.0 + assert outcome.denial.max_budget == budget or (math.isnan(outcome.denial.max_budget) and math.isnan(budget)) + assert handler.calls[-1][0] == "alert_team_member_budget" diff --git a/tests/unit/_auth/test_reservation.py b/tests/unit/_auth/test_reservation.py new file mode 100644 index 00000000000..23398b47c0b --- /dev/null +++ b/tests/unit/_auth/test_reservation.py @@ -0,0 +1,123 @@ +from dataclasses import replace +from itertools import product +from math import inf, nan +from typing import Final + +import pytest + +from litellm._auth import reservation + +_SCOPES: Final[tuple[reservation.Scope, ...]] = ( + "Key", + "Team", + "User", + "EndUser", + "Tag", + "AccessGroup", + "Member", + "Organization", + "Project", +) +_SOURCES: Final = tuple( + reservation.Sources(key, team, team_id, user, user_id, apply) + for key, team, team_id, user, user_id, apply in product( + (False, True), (False, True), (None, "", "team"), (False, True), (None, "", "user"), (False, True) + ) +) + + +@pytest.mark.parametrize("scope", _SCOPES) +@pytest.mark.parametrize("limit", (None, -inf, -1.0, 0.0, 1.0, inf, nan)) +def test_reservation_counter_reference(scope: reservation.Scope, limit: float | None) -> None: + expected: Final = limit is not None and ( + limit > 0 + if scope in ("Key", "Team", "User") + else limit > 0 and limit < inf + if scope == "Project" + else not limit <= 0 + ) + assert reservation.counter(scope, limit) == expected + + +@pytest.mark.parametrize("facts", _SOURCES) +def test_source_selection_preserves_team_key_user_opt_in(facts: reservation.Sources) -> None: + team: Final = facts.team_present and facts.team_id is not None + assert reservation.sources(facts) == reservation.SourcePlan( + facts.key_present, + team, + (not team or facts.apply_user_budget_to_team_keys) and facts.user_present and facts.user_id is not None, + ) + + +@pytest.mark.parametrize("cost", (None, -inf, -1.0, 0.0, 1.0, inf, nan)) +def test_estimate_and_reservation_staging(cost: float | None) -> None: + initial: Final = reservation.Gate(True) + route: Final = replace(initial, llm_route=True) + billable: Final = replace(route, unbilled=False) + model: Final = replace(billable, model_present=True) + counters: Final = replace(model, counter_count=1) + assert tuple(reservation.action(facts).kind for facts in (initial, route, billable, model, counters)) == ( + "ReadRoute", + "ReadUnbilled", + "ReadModel", + "ReadCounters", + "ReadEstimate", + ) + assert reservation.action(replace(counters, cost_read=True, cost=cost)).kind == ( + "Skip" if cost is None or cost <= 0 else "Reserve" + ) + + +@pytest.mark.parametrize("duration", (False, True)) +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("parsed", (None, False, True)) +def test_window_parse_and_warning_timing(duration: bool, limit: float | None, parsed: bool | None) -> None: + expected: Final = ( + "Skip" + if not duration or limit is None or limit <= 0 + else "ReadStart" + if parsed is None + else "AppendWindow" + if parsed + else "WarnWindow" + ) + assert reservation.action(reservation.Window(duration, limit, parsed)).kind == expected + + +@pytest.mark.parametrize("scope", ("EndUser", "Organization", "Project")) +@pytest.mark.parametrize("primary", (None, "", "primary")) +@pytest.mark.parametrize("fallback", (None, "", "fallback")) +def test_reservation_identity_keeps_empty_scope_semantics( + scope: str, primary: str | None, fallback: str | None +) -> None: + from pydantic import TypeAdapter + + facts: Final = TypeAdapter(reservation.Identity).validate_python( + {"scope": scope, "primary": primary, "fallback": fallback} + ) + expected: Final = ( + primary or fallback + if scope == "EndUser" + else primary + if scope == "Project" or primary is not None + else fallback + ) + assert reservation.identity(facts) == expected + + +@pytest.mark.parametrize("team", (False, True)) +@pytest.mark.parametrize("user", (False, True)) +@pytest.mark.parametrize("team_id", (None, "", "team")) +@pytest.mark.parametrize("user_id", (None, "", "user")) +def test_reservation_membership_requires_both_subjects( + team: bool, user: bool, team_id: str | None, user_id: str | None +) -> None: + expected: Final = "ReadMember" if team and team_id is not None and user and user_id is not None else "Skip" + assert reservation.action(reservation.MemberGate(team, team_id, user, user_id)).kind == expected + + +@pytest.mark.parametrize("cap", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("read", (False, True)) +def test_member_default_read_only_when_own_cap_absent(cap: float | None, read: bool) -> None: + expected: Final = "ReadDefault" if cap is None and not read else "ReadCap" + assert reservation.action(reservation.MemberCap(cap, read)).kind == expected diff --git a/tests/unit/_auth/test_resolved_models.py b/tests/unit/_auth/test_resolved_models.py new file mode 100644 index 00000000000..7b301e9b3be --- /dev/null +++ b/tests/unit/_auth/test_resolved_models.py @@ -0,0 +1,156 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm._auth.resolved_models import ResolvedModelDecision as Decision +from litellm._auth.resolved_models import ResolvedModelFacts as Facts +from litellm._auth.resolved_models import ResolvedTeamRead, resolved_models +from litellm._auth.team_override import TeamOverrideFacts +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access +from tests.unit._auth.test_member_models import reference as member_reference +from tests.unit._auth.test_personal_model_request import reference as personal_reference +from tests.unit._auth.test_team_override import reference as team_reference + + +def reference(facts: Facts) -> Decision: + if not facts.config_nonempty and "all-team-models" not in facts.key_models: + if facts.key is None: + return Decision("ReadKey") + key: Final = grouped(facts.key, merge=False) + if key.kind == "ReadGroups": + return Decision("ReadKeyGroups") + if key.kind != "Allow": + return Decision("Deny", key) + if facts.team_id is not None: + if facts.team_read is None: + return Decision("ReadTeamObject") + if facts.team_read.failed and not facts.embedded_ready: + return Decision("ReadEmbeddedTeam") + if facts.team_read.found or facts.team_read.failed: + if facts.team is None: + return Decision("ReadTeamScope") + team: Final = team_reference(facts.team) + if team.kind == "Deny": + return Decision("Deny", team.denial) + if facts.team_read.found and facts.user_id is not None: + if facts.member is None: + return Decision("ReadMember") + member: Final = member_reference(facts.member) + if member.kind == "Deny": + return Decision("Deny", member.denial, member.member_models) + if ( + facts.personal_enabled + and facts.team_id is None + and facts.user_id is not None + and not facts.master + and facts.role != "proxy_admin" + ): + if not facts.owner_loaded: + return Decision("ReadOwner") + if facts.owner_failed or not facts.owner_present: + return Decision("Deny", ModelAccessDecision("PersonalTeam", model=facts.model)) + if facts.personal is None: + return Decision("ReadPersonal") + personal: Final = personal_reference(facts.personal) + if personal.kind == "Deny": + return Decision("Deny", ModelAccessDecision("PersonalTeam", model=personal.model)) + if facts.project_present: + if not facts.project_loaded: + return Decision("ReadProjectObject") + if facts.project_grants: + if facts.project is None: + return Decision("ReadProjectScope") + project: Final = model_access(facts.project) + if project.kind != "Allow": + return Decision("Deny", project) + if facts.customer_present and facts.database_present: + if not facts.customer_loaded: + return Decision("ReadCustomerObject") + if facts.customer_grants: + if facts.customer is None: + return Decision("ReadCustomerScope") + customer: Final = model_access(facts.customer) + if customer.kind != "Allow": + return Decision("Deny", customer) + return Decision("Allow") + + +_MODEL: Final = ModelAccessFacts(("model",), ("model",)) +_PASS: Final = Facts(model="model", config_nonempty=True) +_TEAM: Final = replace( + _PASS, + team_id="team", + team_read=ResolvedTeamRead(found=True), + team=TeamOverrideFacts(GroupedModelScope(replace(_MODEL, object_type="team"))), +) +_PERSONAL: Final = replace(_PASS, user_id="owner", personal_enabled=True) +_PROJECT: Final = replace(_PASS, project_present=True, project_loaded=True, project_grants=("model",)) +_CUSTOMER: Final = replace( + _PASS, customer_present=True, database_present=True, customer_loaded=True, customer_grants=("model",) +) + + +@pytest.mark.parametrize("policy", (reference, resolved_models)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (Facts(), Decision("ReadKey")), + (_PASS, Decision("Allow")), + (Facts(key_models=("all-team-models",)), Decision("Allow")), + (Facts(key=GroupedModelScope(_MODEL)), Decision("Allow")), + ( + Facts(key=GroupedModelScope(replace(_MODEL, grants=("other",)), ("group",)), project_present=True), + Decision("ReadKeyGroups"), + ), + (replace(_TEAM, team_read=None), Decision("ReadTeamObject")), + (replace(_TEAM, team_read=ResolvedTeamRead(failed=True)), Decision("ReadEmbeddedTeam")), + (replace(_TEAM, team=None), Decision("ReadTeamScope")), + (_TEAM, Decision("Allow")), + (replace(_TEAM, user_id="member", project_present=True), Decision("ReadMember")), + ( + replace(_TEAM, user_id="member", member=MemberModelFacts(True, False), project_present=True), + Decision("ReadProjectObject"), + ), + ( + replace( + _TEAM, + user_id="member", + team_read=ResolvedTeamRead(failed=True), + embedded_ready=True, + project_present=True, + ), + Decision("ReadProjectObject"), + ), + (_PERSONAL, Decision("ReadOwner")), + (replace(_PERSONAL, master=True), Decision("Allow")), + (replace(_PERSONAL, role="proxy_admin"), Decision("Allow")), + (replace(_PERSONAL, owner_loaded=True), Decision("Deny", ModelAccessDecision("PersonalTeam", model="model"))), + ( + replace(_PERSONAL, owner_loaded=True, owner_present=True, owner_failed=True), + Decision("Deny", ModelAccessDecision("PersonalTeam", model="model")), + ), + (replace(_PERSONAL, owner_loaded=True, owner_present=True), Decision("ReadPersonal")), + (replace(_PROJECT, project_loaded=False), Decision("ReadProjectObject")), + (_PROJECT, Decision("ReadProjectScope")), + (replace(_PROJECT, project_grants=()), Decision("Allow")), + ( + replace( + _PROJECT, project=replace(_MODEL, object_type="project"), customer_present=True, database_present=True + ), + Decision("ReadCustomerObject"), + ), + (replace(_CUSTOMER, database_present=False), Decision("Allow")), + (replace(_CUSTOMER, customer_loaded=False), Decision("ReadCustomerObject")), + (_CUSTOMER, Decision("ReadCustomerScope")), + (replace(_CUSTOMER, customer_grants=()), Decision("Allow")), + (replace(_CUSTOMER, customer=replace(_MODEL, object_type="customer")), Decision("Allow")), + ), +) +def test_resolved_scope_and_lookup_order(policy: Callable[[Facts], Decision], facts: Facts, expected: Decision) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_route_enable.py b/tests/unit/_auth/test_route_enable.py new file mode 100644 index 00000000000..bdc2ab66335 --- /dev/null +++ b/tests/unit/_auth/test_route_enable.py @@ -0,0 +1,75 @@ +from collections.abc import Callable +from typing import Final, Literal + +import pytest + +from litellm._auth.route_enable import ( + LlmRouteDecision, + RouteEnableDecision, + RouteEnableFacts, + RouteSetting, + llm_route, + route_enable, +) +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_enable_reference import route_enable as reference + + +@pytest.mark.parametrize("policy", (reference, route_enable)) +@pytest.mark.parametrize( + "route", ("/unknown", "/models", "/v1/models", "/chat/completions", "/key/generate", "/user/info") +) +@pytest.mark.parametrize("imported", (False, True)) +@pytest.mark.parametrize("failed", (None, "Other", "Http")) +@pytest.mark.parametrize("exemptions", (None, (), ("/models", "/v1/models"))) +@pytest.mark.parametrize( + "setting", + ( + None, + RouteSetting(True, False), + RouteSetting(False, False), + RouteSetting(True, True), + RouteSetting(True, True, True), + RouteSetting(True, True, False), + ), +) +def test_optional_route_enable_decision_table( + policy: Callable[[RouteEnableFacts], RouteEnableDecision], + route: str, + imported: bool, + failed: Literal["Other", "Http"] | None, + exemptions: tuple[str, ...] | None, + setting: RouteSetting | None, +) -> None: + facts: Final = RouteEnableFacts(route_context(route), imported, failed, setting, setting, exemptions) + assert policy(facts) == reference(facts) + + +@pytest.mark.parametrize( + ("route", "configured", "expected"), + ( + ("/chat/completions", None, "Allow"), + ("/unknown", None, "Deny"), + ("/eu_jev/v1/systemone", None, "ReadEuJev"), + ("/eu_jev/v1/systemone", False, "Allow"), + ("/eu_jev/v1/systemone", True, "Deny"), + ("/eu_jev/v1/systemone/", True, "Deny"), + ), +) +def test_llm_route_configuration_read_is_selected_only_at_builtin_fallback( + route: str, configured: bool | None, expected: LlmRouteDecision +) -> None: + assert llm_route(route_context(route), configured) == expected + + +@pytest.mark.parametrize("configured", (None, False, True)) +@pytest.mark.parametrize("route", ("/eu_jev/v1/systemone", "/eu_jev/v1/systemone/")) +def test_optional_route_builtin_configuration_matches_reference(configured: bool | None, route: str) -> None: + facts: Final = RouteEnableFacts( + route_context(route), + imported=True, + llm=RouteSetting(True, True, True), + exemptions=(), + eu_jev_configured=configured, + ) + assert route_enable(facts) == reference(facts) diff --git a/tests/unit/_auth/test_route_scope.py b/tests/unit/_auth/test_route_scope.py new file mode 100644 index 00000000000..10fda677d03 --- /dev/null +++ b/tests/unit/_auth/test_route_scope.py @@ -0,0 +1,61 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.proxy import proxy_server +from litellm.proxy.auth import route_scope_adapter as native + + +@dataclass(frozen=True, slots=True) +class Result: + allowed: bool + status: int | None = None + detail: object = None + error: str | None = None + + +def outcome(call: Callable[[], object]) -> Result: + try: + return Result(call() is not False) + except HTTPException as error: + return Result(False, error.status_code, error.detail) + except TypeError as error: + return Result(False, error=str(error)) + + +def reference(route: str, settings: Mapping[str, object], premium: object, admin: bool) -> object: + """RouteChecks wrapper behavior frozen from main 2f430873aead.""" + if admin: + if "admin_only_routes" not in settings or premium is not True: + return None + if route in settings["admin_only_routes"]: # pyright: ignore[reportOperatorIssue] # freeze malformed configuration TypeError as well as valid membership + raise HTTPException( + status_code=403, detail=f"user not allowed to access this route. Route={route} is an admin only route" + ) + return None + if "allowed_routes" not in settings: + return True + return route in settings["allowed_routes"] # pyright: ignore[reportOperatorIssue] # frozen Python membership preserves string, list, mapping and invalid-scalar semantics + + +@pytest.mark.parametrize("admin", (False, True)) +@pytest.mark.parametrize("premium", (False, True, 1, None)) +@pytest.mark.parametrize("route", ("/v1/responses", "/missing", "responses", "")) +@pytest.mark.parametrize("present", (False, True)) +@pytest.mark.parametrize( + "value", + (None, True, 3, (), ("/v1/responses",), "/v1/responses", {"/v1/responses": False}, (None, "/v1/responses", 3)), +) +def test_route_scope_matches_frozen_main( + monkeypatch: pytest.MonkeyPatch, admin: bool, premium: object, route: str, present: bool, value: object +) -> None: + settings: Final[Mapping[str, object]] = ( + {("admin_only_routes" if admin else "allowed_routes"): value} if present else {} + ) + monkeypatch.setattr(proxy_server, "general_settings", dict(settings)) + monkeypatch.setattr(proxy_server, "premium_user", premium) + policy: Final = native.custom_admin_only_route_check if admin else native.is_allowed_on_proxy + assert outcome(lambda: policy(route)) == outcome(lambda: reference(route, settings, premium, admin)) diff --git a/tests/unit/_auth/test_routes.py b/tests/unit/_auth/test_routes.py new file mode 100644 index 00000000000..31f7d5b6b37 --- /dev/null +++ b/tests/unit/_auth/test_routes.py @@ -0,0 +1,191 @@ +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException, Request + +from litellm._auth import routes +from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_reference import RouteChecks as ReferenceRoutes +from tests.unit._auth.support.route_reference import managed_agent_route_allowed + + +@dataclass(frozen=True, slots=True) +class RouteResult: + allowed: bool + status: int | None = None + detail: object = None + + +def outcome(call: Callable[[], object]) -> RouteResult: + try: + result: Final = call() + return RouteResult(result is not False) + except HTTPException as error: + return RouteResult(False, error.status_code, error.detail) + + +_ROUTES: Final = tuple( + dict.fromkeys( + ( + "/unknown", + "/chat/completions", + "/v1/mcp/server", + "/v1/mcp/server/id", + "/v1/mcp/server/id/more", + "/auto_router/session", + "/models", + "/v1/responses/id", + "/agents", + "/v1/agents/id", + "/key/id/regenerate", + "/user/update", + "/user/password/change", + "/session/logout", + "/user/info", + "/organization/new", + "/team/new", + "/openai/deployments/model/versions/chat/completions", + "/engines/model/chat/completions", + "/foo/count_tokens", + "/v1beta/models/model:countTokens", + "/unknown\n", + "/key/generate", + "/key/list", + "/session/list", + ) + ) +) + + +@pytest.mark.parametrize("route", _ROUTES) +@pytest.mark.parametrize("method", (None, "GET", "POST", "DELETE", "HEAD", "OPTIONS")) +def test_route_classifications_match_reference(route: str, method: str | None) -> None: + request: Final = Request({"type": "http", "method": method or "GET", "path": route, "headers": []}) + assert RouteChecks.is_llm_api_route(route) == ReferenceRoutes.is_llm_api_route(route) + assert RouteChecks.is_unbilled_route(route) == ReferenceRoutes.is_unbilled_route(route) + assert RouteChecks.is_self_scoped_credential_route( + route, request + ) == ReferenceRoutes.is_self_scoped_credential_route(route, request) + assert RouteChecks._is_get_mcp_server_discovery_route( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route, request + ) == ReferenceRoutes._is_get_mcp_server_discovery_route(route, request) # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + + +@pytest.mark.parametrize("route", _ROUTES) +@pytest.mark.parametrize( + "allowed", + ( + None, + (), + ("llm_api_routes",), + ("llm_api_routes_without_mcp",), + ("management_routes",), + ("/v1/*",), + ("/user",), + ("*",), + ("/no",), + ), +) +@pytest.mark.parametrize("method", ("GET", "POST")) +def test_virtual_route_allow_and_denial_envelope(route: str, allowed: tuple[str, ...] | None, method: str) -> None: + request: Final = Request({"type": "http", "method": method, "path": route, "headers": []}) + token: Final = UserAPIKeyAuth(allowed_routes=list(allowed) if allowed is not None else None) + assert outcome(lambda: RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request)) == outcome( + lambda: ReferenceRoutes.is_virtual_key_allowed_to_call_route(route, token, request) + ) + + +@pytest.mark.parametrize("route", _ROUTES) +@pytest.mark.parametrize("method", ("GET", "POST", "PATCH", "HEAD")) +@pytest.mark.parametrize("keys", ((), ("user_email",), ("user_id", "user_email"))) +def test_viewer_role_denial_order_and_envelope(route: str, method: str, keys: tuple[str, ...]) -> None: + request: Final = Request({"type": "http", "method": method, "path": route, "headers": []}) + data: Final = {key: "value" for key in keys} + assert outcome( + lambda: RouteChecks._check_proxy_admin_viewer_access( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # compare the preserved private route behavior against native projection + route=route, _user_role="proxy_admin_viewer", request_data=data, request=request + ) + ) == outcome( + lambda: ReferenceRoutes._check_proxy_admin_viewer_access( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route=route, _user_role="proxy_admin_viewer", request_data=data, request=request + ) + ) + + +@pytest.mark.parametrize("group", tuple(LiteLLMRoutes)) +def test_full_catalog_classification(group: LiteLLMRoutes) -> None: + assert tuple(RouteChecks.is_llm_api_route(route) for route in group.value) == tuple( + ReferenceRoutes.is_llm_api_route(route) for route in group.value + ) + + +@pytest.mark.parametrize( + ("route", "pattern"), + ( + ("/a/id", "/a/{id}"), + ("/a/x/y", "/a/{id:path}"), + ("/a/x/y:action", "/a/{id:path}:action"), + ("/a/x\ny", "/a/{id:path}"), + ("/a/id\n", "/a/{id}"), + ("/a/b", "/a/*"), + ("/ab", "/a"), + ("/a", ""), + ("/a", "/"), + ("/a/id:action", "/a/{id}:action"), + ("/a/id", "/a/{id:path}"), + ), +) +def test_route_pattern_modes(route: str, pattern: str) -> None: + assert RouteChecks._route_matches_pattern(route, pattern) == ReferenceRoutes._route_matches_pattern(route, pattern) # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + assert RouteChecks._route_matches_allowed_route(route, pattern) == ReferenceRoutes._route_matches_allowed_route( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route, pattern + ) + assert RouteChecks.route_matches_wildcard_pattern(route, pattern) == ReferenceRoutes.route_matches_wildcard_pattern( + route, pattern + ) + assert RouteChecks._route_matches_denied_route(route, pattern) == ReferenceRoutes._route_matches_denied_route( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route, pattern + ) + + +@pytest.mark.parametrize( + "route", + ( + "/agents", + "/v1/agents", + "/chat/completions", + "/v1/chat/completions", + "/v1/responses/id", + "/realtime", + "/v1/realtime", + "/v1/mcp/server/id", + "/mcp/foo", + "/user/info", + "/v1beta/models/x:generateContent", + ), +) +@pytest.mark.parametrize("method", (None, "GET", "HEAD", "POST", "DELETE", "get")) +def test_managed_actor_route_methods(route: str, method: str | None) -> None: + assert routes.classify_route(route_context(route, method), "Managed") == managed_agent_route_allowed(route, method) + + +@pytest.mark.parametrize( + ("key", "team", "jwt", "denied", "expected"), + ( + (("/p",), (), (), (), True), + ((), ("/p",), (), (), True), + (("/q",), ("/p",), (), (), False), + ((), (), ("/p/*",), (), True), + ((), (), ("*",), (), False), + (("/p",), (), (), ("/p",), False), + ), +) +def test_passthrough_precedence( + key: tuple[str, ...], team: tuple[str, ...], jwt: tuple[str, ...], denied: tuple[str, ...], expected: bool +) -> None: + facts: Final = routes.PassThroughFacts(key, team, denied, ("/p/resource",), jwt) + assert routes.passthrough_route(routes.RouteContext("/p/resource"), facts, required=True).allowed == expected diff --git a/tests/unit/_auth/test_runtime.py b/tests/unit/_auth/test_runtime.py new file mode 100644 index 00000000000..be773a265c1 --- /dev/null +++ b/tests/unit/_auth/test_runtime.py @@ -0,0 +1,22 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.decision import Deny, InspectKey, KeyBlocked, RequestFacts +from litellm._auth.runtime import get_auth_core +from tests.unit._auth.test_python_core import NOW, PLAIN_KEY, SETTINGS + + +@pytest.mark.parametrize("backend", (None, "rust", "python", "shadow")) +def test_default_and_native_modes_consume_their_decision_once(backend: str | None, monkeypatch: pytest.MonkeyPatch): + if backend is None: + monkeypatch.delenv("LITELLM_AUTH_CORE", raising=False) + else: + monkeypatch.setenv("LITELLM_AUTH_CORE", backend) + step: Final = get_auth_core().start_resolved_key(RequestFacts("/chat/completions", "hash", True), SETTINGS, NOW) + assert isinstance(step, InspectKey) + blocked: Final = replace(PLAIN_KEY, blocked=True) + assert step.resume(blocked) == Deny(KeyBlocked()) + with pytest.raises(RuntimeError, match="already consumed"): + step.resume(blocked) diff --git a/tests/unit/_auth/test_rust_core.py b/tests/unit/_auth/test_rust_core.py new file mode 100644 index 00000000000..8db6fda188c --- /dev/null +++ b/tests/unit/_auth/test_rust_core.py @@ -0,0 +1,127 @@ +from dataclasses import replace +from typing import Final, cast + +import pytest + +from litellm._auth.decision import ( + Allow, + AuthCore, + Deny, + KeyBlocked, + NoMasterKey, + RequestFacts, + ResolveEndUser, + Step, + WithinModelBudgets, +) +from litellm._auth.driver import drive +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts +from litellm._auth.rust_core import RustCore +from litellm.rust_bridge import _native +from tests.unit._auth.shadow import ShadowCore, ShadowMismatch +from tests.unit._auth.test_python_core import NOW, PLAIN_KEY, SETTINGS, RecordingHandler + + +class ScriptedCore: + def __init__(self, step: Step) -> None: + self.step: Final = step + self.calls = 0 + + def check_model_budgets(self) -> Step: + self.calls += 1 + return self.step + + +FACTS: Final = RequestFacts("/chat/completions", "hash", True) + + +def test_native_continuations_are_one_shot_and_own_their_pending_decision(): + payload, continuation = _native.request_auth_budgets("Key") + assert payload == {"kind": "ReadModelBudgetScope"} + assert isinstance(continuation, _native.NativeAuthContinuation) + outcome, done = continuation.resume( + { + "kind": "ModelScope", + "value": { + "models": [], + "rewritable_model": None, + "key_has_model_budgets": False, + "key_budget_fallbacks": {}, + "user_has_model_budgets": False, + }, + } + ) + assert outcome == {"kind": "Allow", "grant": {"kind": "WithinModelBudgets"}} + assert done is None + with pytest.raises(RuntimeError, match="already consumed"): + continuation.resume({"kind": "Unit"}) + + +def test_invalid_native_answer_consumes_the_continuation_without_fallback(): + _, continuation = _native.request_auth_budgets("Key") + assert continuation is not None + with pytest.raises(RuntimeError, match="invalid auth effect answer"): + continuation.resume({"kind": "Bool", "value": False}) + with pytest.raises(RuntimeError, match="already consumed"): + continuation.resume({"kind": "Unit"}) + + +@pytest.mark.asyncio +async def test_shadow_replays_only_after_the_primary_finishes_and_uses_one_io_trace(): + primary: Final = ScriptedCore(ResolveEndUser(False, lambda: Allow(WithinModelBudgets()))) + reference: Final = ScriptedCore(ResolveEndUser(False, lambda: Allow(WithinModelBudgets()))) + step: Final = ShadowCore(cast(AuthCore, primary), cast(AuthCore, reference)).check_model_budgets() + assert reference.calls == 0 + handler: Final = RecordingHandler() + assert await drive(step, handler) == Allow(WithinModelBudgets()) + assert handler.calls == (("resolve_end_user", False),) + assert (primary.calls, reference.calls) == (1, 1) + + +@pytest.mark.asyncio +async def test_shadow_catches_effect_payload_mismatch_without_a_second_io(): + primary: Final = ScriptedCore(ResolveEndUser(True, lambda: Allow(WithinModelBudgets()))) + reference: Final = ScriptedCore(ResolveEndUser(False, lambda: Allow(WithinModelBudgets()))) + handler: Final = RecordingHandler() + with pytest.raises(ShadowMismatch, match="effect mismatch"): + await drive(ShadowCore(cast(AuthCore, primary), cast(AuthCore, reference)).check_model_budgets(), handler) + assert handler.calls == (("resolve_end_user", True),) + + +@pytest.mark.asyncio +async def test_shadow_catches_outcome_mismatch_without_exposing_credentials(): + with pytest.raises(ShadowMismatch) as error: + ShadowCore( + cast(AuthCore, ScriptedCore(Allow(NoMasterKey("private-credential")))), + cast(AuthCore, ScriptedCore(Deny(KeyBlocked()))), + ).check_model_budgets() + assert str(error.value) == "auth shadow mismatch at Allow" + + +@pytest.mark.asyncio +async def test_native_requests_keep_independent_continuations(): + core: Final = RustCore() + first: Final = core.start_resolved_key(replace(FACTS, token="one"), SETTINGS, NOW) + second: Final = core.start_resolved_key(replace(FACTS, token="two"), SETTINGS, NOW) + blocked: Final = RecordingHandler(key=replace(PLAIN_KEY, blocked=True)) + plain: Final = RecordingHandler() + assert await drive(second, blocked) == Deny(KeyBlocked()) + assert isinstance(await drive(first, plain), Allow) + assert blocked.calls == (("inspect_key", None),) + assert plain.calls[-1] == ("load_key_context", None) + + +@pytest.mark.asyncio +async def test_effect_failure_after_native_entry_is_never_replayed(core: AuthCore): + failure: Final = RuntimeError("effect failed") + + class FailingHandler(RecordingHandler): + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: + self._record("enforce_model_access") + raise failure + + handler: Final = FailingHandler() + with pytest.raises(RuntimeError) as error: + await drive(core.start_resolved_key(FACTS, SETTINGS, NOW), handler) + assert error.value is failure + assert handler.calls == (("inspect_key", None), ("enforce_model_access", None)) diff --git a/tests/unit/_auth/test_spend_counter_policy.py b/tests/unit/_auth/test_spend_counter_policy.py new file mode 100644 index 00000000000..e62dfc009da --- /dev/null +++ b/tests/unit/_auth/test_spend_counter_policy.py @@ -0,0 +1,81 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth import spend_counter_policy as core + +_SPEND: Final = core.Spend(2.0, False, 0.0, 10.0, False, False) + + +@pytest.mark.parametrize( + "facts, expected", + ( + (_SPEND, core.Action("ReadStrictBefore")), + (replace(_SPEND, limit=None), core.SpendValue("Value", 2.0)), + (replace(_SPEND, current=10.0), core.SpendValue("Value", 10.0)), + (replace(_SPEND, current=11.0), core.SpendValue("Value", 11.0)), + (replace(_SPEND, fallback=3.0), core.Action("ReadFloor")), + (replace(_SPEND, window=True), core.Action("ReadFloor")), + (replace(_SPEND, strict_before=True), core.Action("ReadFloor")), + (replace(_SPEND, strict_before=False), core.Action("ReadStrictAfter")), + (replace(_SPEND, strict_before=False, verified=True), core.SpendValue("Value", 2.0)), + (replace(_SPEND, strict_before=False, fallback_authoritative=True), core.SpendValue("Value", 2.0)), + (replace(_SPEND, strict_before=True, authoritative_read=True), core.Action("ReadStrictAfter")), + ( + replace(_SPEND, strict_before=True, authoritative_read=True, authoritative=3.0), + core.SpendValue("Repair", 3.0), + ), + ( + replace(_SPEND, strict_before=True, authoritative_read=True, authoritative=1.0), + core.SpendValue("Value", 2.0), + ), + (replace(_SPEND, fallback=3.0, authoritative_read=True), core.SpendValue("Value", 3.0)), + (replace(_SPEND, strict_before=False, strict_after=True), core.Action("Unverifiable")), + ( + replace(_SPEND, strict_before=True, authoritative_read=True, strict_after=False), + core.SpendValue("Value", 2.0), + ), + ), +) +def test_spend_selection_preserves_lazy_flag_reads(facts: core.Spend, expected: core.Decision) -> None: + assert core.decide(facts) == expected + + +@pytest.mark.parametrize("current", (10.0 + 1e-13, 10.5, 11.0, 12.0, 13.0)) +@pytest.mark.parametrize("strict", (False, True)) +@pytest.mark.parametrize("throttled", (False, True)) +def test_overflow_matches_reference(current: float, strict: bool, throttled: bool) -> None: + remaining: Final = 10.0 - (current - 2.0) + expected: Final = ( + core.Action("Release") + if throttled + else core.Rejected("Reject", current, None) + if remaining <= 1e-12 + else core.Rejected("Reject", current - 2.0, 2.0) + if strict and current - 10.0 > 1e-12 + else core.Resize("Resize", remaining) + ) + assert core.decide(core.OverBudget(throttled, current, 10.0, 2.0, strict)) == expected + + +@pytest.mark.parametrize("token_present", (False, True)) +@pytest.mark.parametrize("token", (None, "", "hash")) +@pytest.mark.parametrize("counter", ("spend:key:hash", "spend:key:None", "spend:key:", "spend:team:hash")) +def test_throttle_is_limited_to_the_key_counter(token_present: bool, token: str | None, counter: str) -> None: + expected: Final = "ReadThrottle" if token_present and counter == f"spend:key:{token}" else "Skip" + assert core.decide(core.ThrottleKey(counter, token_present, token)) == core.Action(expected) + + +@pytest.mark.parametrize("currents, limits", (((), ()), ((0.0, 1.0), (2.0, 2.0)), ((0.0, 2.0), (2.0, 2.0)))) +def test_batching_preserves_order_and_overflow_boundaries( + currents: tuple[float, ...], limits: tuple[float, ...] +) -> None: + expected: Final = ( + () + if not currents + else (tuple(range(len(currents))),) + if all(current + 1.0 <= limit for current, limit in zip(currents, limits, strict=True)) + else tuple((index,) for index in range(len(currents))) + ) + assert core.decide(core.ReservationGroups(currents, limits, 1.0)) == core.Groups("ReservationGroups", expected) diff --git a/tests/unit/_auth/test_team_aliases.py b/tests/unit/_auth/test_team_aliases.py new file mode 100644 index 00000000000..b61a34204dc --- /dev/null +++ b/tests/unit/_auth/test_team_aliases.py @@ -0,0 +1,49 @@ +from collections.abc import Callable +from typing import Final + +import pytest + +from litellm._auth.team_aliases import TeamAliasDecision as Decision +from litellm._auth.team_aliases import TeamAliasFacts as Facts +from litellm._auth.team_aliases import team_aliases + + +def reference(facts: Facts) -> Decision: + targets: Final = tuple(facts.aliases.get(name) for name in facts.requested) + count: Final = sum( + bool(target is not None and facts.router_present and target.startswith(f"model_name_{facts.team_id}_")) + for target in targets + ) + if count and facts.names is None: + return Decision("ReadNames", count=count) + reads: Final = iter(facts.names or ()) + return Decision( + "Complete", + models=tuple( + name + if target is None + or ( + facts.router_present and target.startswith(f"model_name_{facts.team_id}_") and target not in next(reads) + ) + else target + for name, target in zip(facts.requested, targets, strict=True) + ), + ) + + +@pytest.mark.parametrize("policy", (reference, team_aliases)) +@pytest.mark.parametrize( + "facts", + ( + Facts(("alias",), {}, "team", True), + Facts(("alias",), {"alias": "target"}, "team", True), + Facts(("alias",), {"alias": "model_name_team_target"}, "team", False), + Facts(("alias",), {"alias": "model_name_team_target"}, "team", True), + Facts(("alias", "alias"), {"alias": "model_name_team_target"}, "team", True, (("model_name_team_target",), ())), + Facts(("alias",), {"alias": "model_name_None_target"}, None, True, ((),)), + ), +) +def test_aliases_preserve_deleted_targets_and_repeated_registry_reads( + policy: Callable[[Facts], Decision], facts: Facts +) -> None: + assert policy(facts) == reference(facts) diff --git a/tests/unit/_auth/test_team_override.py b/tests/unit/_auth/test_team_override.py new file mode 100644 index 00000000000..8e956b54e3e --- /dev/null +++ b/tests/unit/_auth/test_team_override.py @@ -0,0 +1,74 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.team_override import TeamOverrideDecision as Decision +from litellm._auth.team_override import TeamOverrideFacts as Facts +from litellm._auth.team_override import team_override +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Facts) -> Decision: + if facts.team is None: + return Decision("ReadTeam") + denied: Final = grouped(facts.team, merge=True) + if denied.kind == "Allow": + return Decision("Allow") + if denied.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if denied.kind != "Deny" or denied.object_type != "team": + return Decision("Deny", denied) + if facts.assigned_models is None: + return Decision("ReadAssignedGroups") + if not facts.assigned_models: + return Decision("Deny", denied) + if facts.assigned_scope is None: + return Decision("ReadAssignedScope") + assigned: Final = model_access(facts.assigned_scope) + if assigned.kind == "Allow": + return Decision("Allow") + return Decision("Deny", denied if assigned.kind == "Deny" else assigned) + + +_MODEL: Final = ModelAccessFacts(("model",), ("other",), object_type="team") +_TEAM: Final = Facts(GroupedModelScope(_MODEL)) +_ASSIGNED: Final = replace(_TEAM, assigned_models=("model",)) + + +@pytest.mark.parametrize("policy", (reference, team_override)) +@pytest.mark.parametrize( + "facts", + ( + Facts(), + _TEAM, + Facts(GroupedModelScope(replace(_MODEL, grants=()))), + Facts(GroupedModelScope(_MODEL, ("team-group",))), + Facts(GroupedModelScope(_MODEL, ("team-group",), ("model",))), + Facts(GroupedModelScope(_MODEL, ("team-group",), ("no-access",))), + Facts(GroupedModelScope(replace(_MODEL, object_type="key"))), + Facts(GroupedModelScope(replace(_MODEL, fallback_depth=10))), + replace(_TEAM, assigned_models=()), + _ASSIGNED, + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key", grants=("model",))), + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key")), + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key", fallback_depth=10)), + ), +) +def test_team_override_reads_groups_only_after_team_denial(policy: Callable[[Facts], Decision], facts: Facts) -> None: + assert policy(facts) == reference(facts) + + +def test_assigned_denial_keeps_team_error_envelope() -> None: + facts: Final = replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key")) + assert team_override(facts) == Decision("Deny", model_access(_MODEL)) + + +def test_bad_assigned_scope_is_terminal_not_team_denial() -> None: + with pytest.raises(NativeAuthFailure): + team_override(replace(_ASSIGNED, assigned_scope=replace(_MODEL, fallback_depth=-1))) diff --git a/tests/unit/_auth/test_unpriced_models.py b/tests/unit/_auth/test_unpriced_models.py new file mode 100644 index 00000000000..fac21104611 --- /dev/null +++ b/tests/unit/_auth/test_unpriced_models.py @@ -0,0 +1,106 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.unpriced_models import DeploymentPricing, PriceEntry, UnpricedModelFacts, unpriced_model +from litellm.proxy.auth.unpriced_model_adapter import price_entry +from tests.unit._auth.support.model_registry_reference import entry_has_priced_metric + +_BASE: Final = UnpricedModelFacts("model", True) +_ROW: Final = DeploymentPricing((), True) + + +@pytest.mark.parametrize( + ("entry", "expected"), + ( + ({}, False), + ({"input_cost_per_token": 0}, False), + ({"input_cost_per_token": -1}, False), + ({"cost_per_image": 0.1}, True), + ({"cost_per_image": True}, False), + ({"cost_per_image": "0.1"}, False), + ({"cost_per_image": {"small": 0, "large": 0.2}}, True), + ({"cost_per_image": {"large": True}}, False), + ({"cost_per_image": {"large": {"nested": 1}}}, False), + ({"tiered_pricing": None}, False), + ({"tiered_pricing": ()}, True), + ({"tiered_pricing": False}, True), + ({"maximum_tokens": 1}, False), + ({"input_cost_per_token": float("nan")}, False), + ({"input_cost_per_token": float("inf")}, True), + ), +) +def test_native_price_metrics_match_reference(entry: Mapping[str, object], expected: bool) -> None: + facts: Final = replace( + _BASE, + group_found=True, + deployments=(replace(_ROW, info_loaded=True, info=price_entry(entry)),), + cost_deployments=(), + ) + assert entry_has_priced_metric(entry) == expected + assert unpriced_model(facts).unpriced == (not expected) + + +@pytest.mark.parametrize( + ("facts", "kind", "index", "unpriced"), + ( + (replace(_BASE, model=None), "Complete", 0, False), + (replace(_BASE, model=""), "Complete", 0, False), + (replace(_BASE, router_present=False), "Complete", 0, False), + (_BASE, "ReadGroup", 0, False), + (replace(_BASE, group_found=False), "Complete", 0, False), + (replace(_BASE, group_found=True), "ReadDeployments", 0, False), + (replace(_BASE, group_found=True, deployments=(_ROW,)), "ReadInfo", 0, False), + ( + replace(_BASE, group_found=True, deployments=(replace(_ROW, parameter_keys=("input_cost_per_token",)),)), + "Complete", + 0, + False, + ), + ( + replace(_BASE, group_found=True, deployments=(replace(_ROW, parameter_keys=("tiered_pricing",)),)), + "Complete", + 0, + False, + ), + (replace(_BASE, group_found=True, deployments=()), "ReadCostDeployments", 0, False), + (replace(_BASE, group_found=True, deployments=(), cost_deployments=(_ROW,)), "ReadCost", 0, False), + ( + replace( + _BASE, + group_found=True, + deployments=(), + cost_deployments=(replace(_ROW, cost_keys=("input_cost_per_token",)),), + ), + "Complete", + 0, + False, + ), + ( + replace( + _BASE, + group_found=True, + deployments=(), + cost_deployments=(replace(_ROW, cost_keys=("cost_per_second",)),), + ), + "Complete", + 0, + True, + ), + ( + replace( + _BASE, group_found=True, deployments=(replace(_ROW, info_loaded=True, info=PriceEntry(False, ())), _ROW) + ), + "ReadInfo", + 1, + False, + ), + ), +) +def test_unpriced_reads_preserve_legacy_short_circuit_order( + facts: UnpricedModelFacts, kind: str, index: int, unpriced: bool +) -> None: + decision: Final = unpriced_model(facts) + assert (decision.kind, decision.index, decision.unpriced) == (kind, index, unpriced) diff --git a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 4692088e536..bb967550c96 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -57,6 +57,32 @@ async def _single_context(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyA return [user_api_key_auth] +@pytest.mark.asyncio +@pytest.mark.parametrize("writer", [False, True]) +async def test_native_legacy_agent_group_lookup_reads_selected_pool_without_prisma(writer): + from types import SimpleNamespace + + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{"agent_id": "agent-a"}, {"agent_id": "agent-b"}])) + ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups( + SimpleNamespace(native_auth_database=native), ["legacy-group"], check_db_only=writer + ) + assert ids == {"agent-a", "agent-b"} + native.auth_read.assert_awaited_once_with("agent_access_group", ["legacy-group"], writer=writer) + + +@pytest.mark.asyncio +async def test_native_legacy_agent_group_error_is_not_empty_grants(): + from types import SimpleNamespace + + from litellm._auth.failure import NativeAuthFailure + + native = SimpleNamespace(auth_read=AsyncMock(side_effect=OSError("offline"))) + with pytest.raises(NativeAuthFailure): + await AgentRequestHandler._get_db_agent_ids_for_access_groups( + SimpleNamespace(native_auth_database=native), ["legacy-group"] + ) + + @pytest.mark.asyncio class TestAgentRequestHandler: """ @@ -1369,7 +1395,7 @@ def personal_key_owner_teams(monkeypatch: pytest.MonkeyPatch, request: pytest.Fi monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache) monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) - monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) monkeypatch.setattr( auth_checks, "get_user_object", @@ -1428,7 +1454,9 @@ async def test_empty_owner_team_id_does_not_restrict_agent_access( assert owner_teams.unreadable is False assert tuple(team.team_id for team in owner_teams.teams) == ("team-open",) - assert await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + assert ( + await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + ) assert all(call.kwargs["team_id"] == "team-open" for call in team_loader.await_args_list) @@ -1671,7 +1699,6 @@ async def test_personal_key_agent_billing_keeps_team_controls_when_org_lookup_fa org_loader: Final = AsyncMock(side_effect=RuntimeError("org lookup failed")) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) - monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader) @@ -1717,7 +1744,9 @@ async def test_personal_key_agent_billing_keeps_team_controls_when_membership_lo AsyncMock(side_effect=RuntimeError("membership lookup failed")), ) - billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(personal_key_owner_teams, "agent-b") + billed: Final = await AgentRequestHandler._personal_key_agent_billing_auth_for_team( # pyright: ignore[reportPrivateUsage] # exercise billing after owner-team selection has completed + personal_key_owner_teams, "agent-b", await _get_team_object("team-b"), strict=False + ) assert billed.team_id == "team-b" assert billed.team_rpm_limit == 5 @@ -1761,9 +1790,7 @@ async def test_personal_key_agent_billing_uses_fresh_membership_for_strict_polic ), ) - async def _get_team_membership( - user_id: str, team_id: str, **kwargs: object - ) -> LiteLLM_TeamMembership: + async def _get_team_membership(user_id: str, team_id: str, **kwargs: object) -> LiteLLM_TeamMembership: database_value: Final = bool(kwargs.get("check_db_only")) return LiteLLM_TeamMembership( user_id=user_id, @@ -1774,7 +1801,6 @@ async def test_personal_key_agent_billing_uses_fresh_membership_for_strict_polic team_loader: Final = AsyncMock(side_effect=_get_team_object) monkeypatch.setattr(auth_checks, "get_team_object", team_loader) monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(side_effect=_get_team_membership)) - monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) auth: Final = personal_key_owner_teams.model_copy(update={"requires_fresh_policy": strict}) @@ -2135,9 +2161,7 @@ async def test_strict_personal_key_agent_access_reads_owner_teams_from_the_datab return LiteLLM_TeamTable( team_id=team_id, object_permission_id=f"perm-{team_id}", - object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id=f"perm-{team_id}", agents=["agent-a"] - ), + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id=f"perm-{team_id}", agents=["agent-a"]), ) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) diff --git a/tests/unit/proxy/agent_endpoints/test_identity_store.py b/tests/unit/proxy/agent_endpoints/test_identity_store.py index 005f0b4c074..4ea0126e0ce 100644 --- a/tests/unit/proxy/agent_endpoints/test_identity_store.py +++ b/tests/unit/proxy/agent_endpoints/test_identity_store.py @@ -74,13 +74,47 @@ def setup_store( ) ) return ( - AgentIdentityStore(AgentsRepository(db), AgentIdentityRepository(db), VerifiedSubjectRepository(db), cache=cache), + AgentIdentityStore( + AgentsRepository(db), AgentIdentityRepository(db), VerifiedSubjectRepository(db), cache=cache + ), agents, identities, humans, ) +@pytest.mark.asyncio +async def test_native_agent_policy_reads_writer_and_preserves_binding_without_prisma() -> None: + expected: Final = stored_agent(enabled=False, object_permission={"models": ["approved"]}) + native: Final = SimpleNamespace(auth_read=AsyncMock(return_value=[expected.model_dump()])) + store: Final = AgentIdentityStore.from_client(SimpleNamespace(native_auth_database=native)) + actual: Final = await store.agent("agent-one") + assert actual == expected + native.auth_read.assert_awaited_once_with("agent", ["agent-one"], writer=True) + + +@pytest.mark.asyncio +async def test_native_retired_agent_lookup_preserves_missing_and_found_without_prisma() -> None: + native: Final = SimpleNamespace(auth_read=AsyncMock(side_effect=[[], [{"original_agent_id": "agent-one"}]])) + store: Final = AgentIdentityStore.from_client(SimpleNamespace(native_auth_database=native)) + assert await store.retired_agent("missing") is False + assert await store.retired_agent("agent-one") is True + native.auth_read.assert_awaited_with("retired_agent", ["agent-one"], writer=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [OSError("offline"), RuntimeError("invalid native state")]) +async def test_native_agent_failure_cannot_enter_legacy_fallback(error: Exception) -> None: + from litellm._auth.failure import NativeAuthFailure + + native: Final = SimpleNamespace(auth_read=AsyncMock(side_effect=error)) + store: Final = AgentIdentityStore.from_client(SimpleNamespace(native_auth_database=native)) + with pytest.raises(NativeAuthFailure): + await store.agent("agent-one") + with pytest.raises(NativeAuthFailure): + await store.retired_agent("agent-one") + + @pytest.mark.asyncio async def test_application_authentication_has_no_fabricated_human() -> None: store, _, _, humans = setup_store() diff --git a/tests/unit/proxy/auth/budget_alert_reference.py b/tests/unit/proxy/auth/budget_alert_reference.py new file mode 100644 index 00000000000..a47fd364db6 --- /dev/null +++ b/tests/unit/proxy/auth/budget_alert_reference.py @@ -0,0 +1,42 @@ +from collections.abc import Mapping + +from pydantic import TypeAdapter, ValidationError + +from litellm._auth.budget_alerts import Facts + + +def reference_recipients(raw: object) -> list[str]: + if isinstance(raw, list): + return [email.strip() for email in raw if isinstance(email, str) and email.strip()] + if isinstance(raw, str): + return [email.strip() for email in raw.split(",") if email.strip()] + return [] + + +def reference_alert(facts: Facts, global_config: object, local_config: object) -> tuple[bool, object]: + if not facts.member and not (facts.limit is not None and facts.spend is not None and facts.spend > 0): + return False, None + if facts.member: + try: + config = TypeAdapter(Mapping[str, object]).validate_python(local_config) + except ValidationError: + config = {} + local_config = { + key: value for key, value in config.items() if key.isdigit() and len(key) <= 3 and 1 <= int(key) <= 100 + } + global_entries = {key: reference_recipients(value) for key, value in (global_config or {}).items()} + local_entries = {key: reference_recipients(value) for key, value in (local_config or {}).items()} + merged = { + key: list(dict.fromkeys(global_entries.get(key, []) + local_entries.get(key, []))) + for key in set(global_entries) | set(local_entries) + } + if merged: + if facts.member and facts.spend <= 0: + return False, None + minimum = min((int(key) for key in merged if key.isdigit()), default=None) + if minimum is None or facts.spend < facts.limit * (minimum / 100.0): + return False, None + return True, tuple((key, tuple(value)) for key, value in merged.items()) + return ( + not facts.member and facts.spend >= facts.limit * facts.default_fraction and facts.spend < facts.limit + ), None diff --git a/tests/unit/proxy/auth/budget_reference.py b/tests/unit/proxy/auth/budget_reference.py new file mode 100644 index 00000000000..a04a997ab46 --- /dev/null +++ b/tests/unit/proxy/auth/budget_reference.py @@ -0,0 +1,44 @@ +from collections.abc import Mapping +from typing import Final + +import litellm +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + entry_names_model, + resolve_model_budget, +) + + +async def model_allowed( + limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, + entity_type: Litellm_EntityType, + entity_id: str | None, + team_id: str | None, + budgets: Mapping[str, object] | None, + model: str, + message: str, +) -> bool: + if not budgets: + return True + resolved: Final = resolve_model_budget(model, budgets, limiter._model_access_groups, team_id) + if resolved is None: + return True + limit: Final = resolved.budget_config.max_budget + if limit is None or limit < 0: + return True + local: Final = await limiter._get_spend_for_model_budget(entity_type, entity_id, model, resolved) + current: Final = ( + local if local >= limit else await limiter._with_shared_spend(entity_type, entity_id, resolved, local) + ) + if current >= limit: + raise litellm.BudgetExceededError( + message=message + if entry_names_model(resolved.budget_model, model) + else f"{message}, budget_model={resolved.budget_model}", + current_cost=current, + max_budget=limit, + entity_type=entity_type.value, + entity_id=entity_id, + ) + return True diff --git a/tests/unit/proxy/auth/phase1_budget_auth_reference.py b/tests/unit/proxy/auth/phase1_budget_auth_reference.py new file mode 100644 index 00000000000..42c4673d874 --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_auth_reference.py @@ -0,0 +1,915 @@ +"""Policy frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9, with independently supplied raw IO.""" + +from __future__ import annotations + +import asyncio +import json +import math +import re +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final, Optional, Protocol + +from opentelemetry.trace import Span +from pydantic import ConfigDict, TypeAdapter + +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import ( + MODEL_LIST_ROUTES, + CallInfo, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + Litellm_EntityType, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _dictable_table, # pyright: ignore[reportPrivateUsage] # data-only Prisma row projection +) +from litellm.proxy.common_utils.user_api_key_cache import ( + AUTH_OBJECTS_TARGET, + UserApiKeyCache, + get_management_object_ttl, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.budget_repository import BudgetRepository +from litellm.router import Router +from litellm.router_api import RegistryReader +from litellm.types.router import DeploymentTypedDict + +from .phase1_budget_route_reference import RouteChecks +from .phase1_budget_throttle_reference import budget_throttle_percentage, should_throttle_budget_exceeded + + +class _RawCacheRead(Protocol): + async def async_get_cache(self, *, key: str, local_only: bool = ...) -> object: ... + + +def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead: + return cache + + +def _get_router_zero_cost_cache(llm_router: RegistryReader) -> dict[str, bool] | None: # mutable-ok: Phase 1 public collection and identity contract + + cache: Final = llm_router.registry().zero_cost() + return cache if isinstance(cache, dict) else None + + +def is_model_cost_zero(model: str | list[str] | None, llm_router: RegistryReader | None) -> bool: # mutable-ok: Phase 1 public collection and identity contract + + if model is None or llm_router is None: + return False + + model_list: Final = [model] if isinstance(model, str) else model + + zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router) + + for model_name in model_list: + if zero_cost_cache is not None: + cached = zero_cost_cache.get(model_name) + if cached is not None: + if cached is False: + return False + continue + try: + model_group_info = llm_router.registry().group(model_name, include_hidden=True) + + if model_group_info is None: + verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + input_cost = model_group_info.input_cost_per_token + output_cost = model_group_info.output_cost_per_token + + if input_cost is None or output_cost is None: + verbose_proxy_logger.debug( + "Model %s has undefined cost (input: %s, output: %s), assuming it has cost", + model_name, + input_cost, + output_cost, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if input_cost > 0 or output_cost > 0: + verbose_proxy_logger.debug( + "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + safe_name = str(model_name).replace("\n", "").replace("\r", "") + served_deployments = _deployments_served_for(model_name, llm_router) + if not _is_cost_explicitly_configured(served_deployments): + verbose_proxy_logger.debug( + "Model %s has zero cost but no explicit cost " + "configuration in model_cost entry — treating as unknown " + "cost (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if _has_explicit_positive_cost(served_deployments): + verbose_proxy_logger.debug( + "Model %s routes to a deployment with an explicit positive per-token price (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if _has_ptu_flat_cost(served_deployments): + verbose_proxy_logger.debug( + "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " + "rate is not a free model (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + verbose_proxy_logger.debug( + "Model %s has zero cost explicitly configured (input: %s, output: %s)", + safe_name, + input_cost, + output_cost, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = True + + except Exception as e: + verbose_proxy_logger.debug("Error checking cost for model %s: %s, assuming it has cost", model_name, e) + return False + + return True + + +_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: + + registry: Final = llm_router.registry() + routed_group: Final = registry.resolve(model) + if registry.resolve(routed_group) == routed_group: + return registry.deployments(routed_group) or () + named: Final = registry.deployments(model, source="alias") + if named: + return named + return tuple( + DeploymentTypedDict(**deployment) # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + for deployment in registry.patterns().get_deployments_by_pattern(model=routed_group) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + ) + + +def _has_ptu_flat_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_info = deployment.get("model_info") or _NO_MODEL_INFO # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return True + return False + + +def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_id is None: + continue + raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: + return True + return False + + +def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_id is None: + continue + raw_entry = as_str_mapping(litellm.model_cost.get(model_id)) or _EMPTY_COST_ENTRY # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + if ( + _is_positive_cost(raw_entry.get("input_cost_per_token")) + or _is_positive_cost(raw_entry.get("output_cost_per_token")) + or _is_positive_cost(_deployment_query_price(deployment["litellm_params"], raw_entry)) + ): + return True + return False + + +def _deployment_query_price(litellm_params: Mapping[str, object], cost_entry: Mapping[str, object]) -> float | None: + provider: Final = litellm_params.get("custom_llm_provider") or str(litellm_params.get("model")).partition("/")[0] + return decisions_request_price({**cost_entry, "litellm_provider": provider}, decisions_pricing_fields(cost_entry)) + + +_EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 + + +MODEL_DISCOVERY_ROUTES: Final = frozenset( + { + *MODEL_LIST_ROUTES, + "/model/info", + "/v1/model/info", + "/v2/model/info", + "/model_group/info", + "/utils/model_info", + } +) + + +BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset( + { + "/health", + "/health/services", + "/health/test_connection", + "/auto_router/test_routing", + } +) + + +MCP_DISCOVERY_ROUTES: Final = frozenset({"/mcp-rest/tools/list", "/v1/mcp/tools"}) + + +MCP_ZERO_SPEND_JSONRPC_METHODS: Final = frozenset({"initialize", "notifications/initialized", "ping", "tools/list"}) + + +MCP_TOOL_CALL_ROUTES: Final = frozenset({"/mcp/tools/call", "/mcp-rest/tools/call"}) + + +def is_mcp_discovery_request(route: str, request_body: Mapping[str, object]) -> bool: + if route in MCP_DISCOVERY_ROUTES: + return True + if route in MCP_TOOL_CALL_ROUTES or not (route == "/mcp" or route.startswith("/mcp/")): + return False + method: Final = request_body.get("method") + return isinstance(method, str) and method in MCP_ZERO_SPEND_JSONRPC_METHODS + + +def route_skips_budget_checks(route: str) -> bool: + if route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: + return False + if route in MODEL_DISCOVERY_ROUTES: + return True + if RouteChecks.is_llm_api_route(route=route): + return False + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route) + + +def request_skips_budget_checks( + route: str, + model: str | list[str] | None, # mutable-ok: Phase 1 public collection and identity contract + llm_router: Router | None, + *, + request_body: Mapping[str, object] | None = None, +) -> bool: + if route_skips_budget_checks(route=route): + return True + if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body): + return True + return is_model_cost_zero(model=model, llm_router=llm_router) + + +KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" + + +def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: + + if key_metadata is None: + return None + budget_id: Final = key_metadata.get(KEY_END_USER_BUDGET_ID_METADATA_FIELD) + return budget_id if isinstance(budget_id, str) and budget_id != "" else None + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def get_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + budget_id: str | None = None, +) -> LiteLLM_BudgetTable | None: + + default_budget_id: Final = budget_id if budget_id is not None else litellm.max_end_user_budget_id + if prisma_client is None or default_budget_id is None: + return None + + cache_key: Final = f"default_end_user_budget:{default_budget_id}" + + cached_budget: Final = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_BudgetTable, + ) + if cached_budget is not None: + return cached_budget + + try: + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( + where={"budget_id": default_budget_id} + ) + + if budget_record is None: + verbose_proxy_logger.warning( + "Default end user budget not found in database: %s", + default_budget_id.replace("\r", "").replace("\n", ""), + ) + return None + + _budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value=_budget_obj, + model_type=LiteLLM_BudgetTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + + return _budget_obj + + except Exception as e: + verbose_proxy_logger.error("Error fetching default end user budget: %s", e) + return None + + +async def get_team_member_default_budget( + budget_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + raise_on_lookup_error: bool = False, +) -> LiteLLM_BudgetTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_team_member_default_budget") + + +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + if member_budget is not None and member_budget.max_budget is not None: + return _effective_member_max(member_budget, now) + if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: + return None + return team_default_budget.max_budget + ( + _active_member_increase(member_budget, now) if member_budget is not None else 0.0 + ) + + +async def resolve_default_end_user_budget( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + key_end_user_budget_id: str | None, + parent_otel_span: Span | None = None, +) -> LiteLLM_BudgetTable | None: + + if key_end_user_budget_id is not None: + key_budget: Final = await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + budget_id=key_end_user_budget_id, + ) + if key_budget is not None: + return key_budget + + if litellm.max_end_user_budget_id is None: + return None + + return await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) + + +async def _apply_default_budget_to_end_user( + end_user_obj: LiteLLM_EndUserTable, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable: + + if end_user_obj.budget_id is not None and end_user_obj.litellm_budget_table is not None: + return end_user_obj + + if key_end_user_budget_id is None and litellm.max_end_user_budget_id is None: + return end_user_obj + + default_budget: Final = await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + + if default_budget is None: + return end_user_obj + + verbose_proxy_logger.debug( + "Applied default budget %s to end user %s", default_budget.budget_id, end_user_obj.user_id + ) + return end_user_obj.model_copy(update=MappingProxyType({"litellm_budget_table": default_budget})) + + +async def check_end_user_budget( + end_user_obj: LiteLLM_EndUserTable, + route: str, +) -> None: + + if RouteChecks.is_info_route(route): + return + + if end_user_obj.litellm_budget_table is None: + return + + end_user_budget: Final = end_user_obj.litellm_budget_table.max_budget + if end_user_budget is None: + return + + from litellm.proxy.proxy_server import get_current_spend + + end_user_spend: Final = await get_current_spend( + counter_key=f"spend:end_user:{end_user_obj.user_id}", + fallback_spend=end_user_obj.spend or 0.0, + max_budget=end_user_budget, + fallback_authoritative=True, + ) + if end_user_spend > end_user_budget: + raise litellm.BudgetExceededError( + current_cost=end_user_spend, + max_budget=end_user_budget, + message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", + entity_type=Litellm_EntityType.END_USER.value, + entity_id=end_user_obj.user_id, + ) + + +async def _load_end_user_restricted_registry( + prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache +) -> frozenset[str] | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: _load_end_user_restricted_registry") + + +async def _end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, + key_end_user_budget_id: str | None = None, +) -> bool: + + if ( + litellm.max_end_user_budget_id is not None + or key_end_user_budget_id is not None + or litellm.validate_end_user_id_in_db + or token_end_user_max_budget is not None + ): + return False + + registry: Final = await _load_end_user_restricted_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return registry is not None and end_user_id not in registry + + +async def get_end_user_object( + end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + route: str | None = "", + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + token_end_user_max_budget: float | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_end_user_object") + + +_END_USER_VALIDATION_NEGATIVE_TTL: Final = 60 + + +_END_USER_VALIDATION_POSITIVE_TTL: Final = 300 + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def resolve_and_validate_end_user_id( + raw_end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", + key_end_user_budget_id: str | None = None, +) -> str | None: + + if raw_end_user_id is None: + return None + if not litellm.validate_end_user_id_in_db: + return raw_end_user_id + if prisma_client is None: + return raw_end_user_id + + has_default_budget: Final = bool(litellm.max_end_user_budget_id) or key_end_user_budget_id is not None + cache_key: Final = f"end_user_validation:{raw_end_user_id}" + cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) + if cached == "valid": + return raw_end_user_id + if cached == "invalid": + return raw_end_user_id if has_default_budget else None + + is_valid: Final = await _end_user_id_exists_in_db( + end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value="valid" if is_valid else "invalid", + ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL), + ) + + return raw_end_user_id if is_valid or has_default_budget else None + + +async def _end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", +) -> bool: + + try: + end_user_obj: Final = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if end_user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", e) + + try: + user_obj: Final = await get_user_object( + user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=False, + user_email=end_user_id if "@" in end_user_id else None, + ) + if user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", e) + + return False + + +async def get_team_membership( + user_id: str, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + *, + check_db_only: bool = False, +) -> Optional["LiteLLM_TeamMembership"]: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_team_membership") + + +async def get_user_object( + user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + user_id_upsert: bool, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + sso_user_id: str | None = None, + user_email: str | None = None, + check_db_only: bool | None = None, +) -> LiteLLM_UserTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_user_object") + + +_MASKED_KEY_NAME_RE: Final = re.compile(r"sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?") + + +def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: + + pct: Final = budget_throttle_percentage() + if pct is None or not should_throttle_budget_exceeded(valid_token): + return False + valid_token.budget_throttle_pct = pct # rebind-ok: preserve Phase 1 mutation order and identity + return True + + +def send_token_budget_alert( + valid_token: UserAPIKeyAuth, + spend: float, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + call_info: Final = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj is not None else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="token_budget", + user_info=call_info, + ) + ) + + +def key_budget_exceeded_error( + valid_token: UserAPIKeyAuth, spend: float, max_budget: float +) -> litellm.BudgetExceededError: + + key_label: Final = valid_token.key_alias or "key" + key_descriptor: Final = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) + else key_label + ) + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {max_budget}", + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) + + +async def virtual_key_max_budget_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + + if valid_token.max_budget is not None: + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + max_budget=valid_token.max_budget, + ) + send_token_budget_alert(valid_token, spend, proxy_logging_obj, user_obj) + + if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: + if _apply_budget_exceeded_throttle(valid_token): + return + raise key_budget_exceeded_error(valid_token, spend, valid_token.max_budget) + + +async def _virtual_key_multi_budget_check( + valid_token: UserAPIKeyAuth, +): + + if not valid_token.budget_limits: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in valid_token.budget_limits: # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + w: dict = window if isinstance(window, dict) else window.model_dump() # pyright: ignore[reportMissingTypeArgument, reportUnknownVariableType, reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_entity_type="Key", + window_entity_id=valid_token.token, + window_duration=str(w["budget_duration"]), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_start=get_budget_window_start(w), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + ) + if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + message=( + f"ExceededBudget: Key over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) + + +async def team_max_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + + if team_object is not None and team_object.max_budget is not None: + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:team:{team_object.team_id}", + fallback_spend=team_object.spend or 0.0, + max_budget=team_object.max_budget, + ) + + if math.isfinite(team_object.max_budget) and spend > team_object.max_budget: + if valid_token: + call_info: Final = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=team_object.max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.TEAM, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=spend, + max_budget=team_object.max_budget, + message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, + ) + + +async def _team_multi_budget_check( + team_object: LiteLLM_TeamTable | None, +): + + if team_object is None or not team_object.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in team_object.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() # pyright: ignore[reportMissingTypeArgument] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_entity_type="Team", + window_entity_id=team_object.team_id, + window_duration=str(w["budget_duration"]), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_start=get_budget_window_start(w), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + ) + if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + message=( + f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, + ) + + +class ModelMaxBudgetHolder(Protocol): + @property + def model_max_budget(self) -> Mapping[str, object] | None: ... + + +def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] | None: + return holder.model_max_budget + + +def _active_member_increase(self: LiteLLM_BudgetTable, now: datetime) -> float: + if self.temp_budget_increase is None or self.temp_budget_expiry is None: + return 0.0 + expiry: Final = ( + self.temp_budget_expiry.replace(tzinfo=timezone.utc) + if self.temp_budget_expiry.tzinfo is None + else self.temp_budget_expiry + ) + return 0.0 if expiry <= now else self.temp_budget_increase + + +def _effective_member_max(self: LiteLLM_BudgetTable, now: datetime) -> float | None: + if self.max_budget is None: + return None + return self.max_budget + _active_member_increase(self, now) + + +_WINDOW_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + +def _coerce_window(window: object) -> Mapping[str, object]: + if isinstance(window, Mapping): + return window # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if isinstance(window, str): + try: + return _WINDOW_OBJECT.validate_python(json.loads(window)) + except Exception: + return {} + model_dump: Final = getattr(window, "model_dump", None) + if not callable(model_dump): + return {} + dumped: Final[object] = model_dump() + return dumped if isinstance(dumped, Mapping) else {} # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + + +def get_budget_window_start(window: object) -> datetime | None: + window_dict: Final = _coerce_window(window) + budget_duration: Final = window_dict.get("budget_duration") + if budget_duration is None: + return None + try: + duration_seconds: Final = duration_in_seconds(str(budget_duration)) + except Exception: + return None + + reset_at = _coerce_datetime(window_dict.get("reset_at")) # rebind-ok: preserve Phase 1 mutation order and identity + if reset_at is None: + return datetime.now(timezone.utc) - timedelta(seconds=duration_seconds) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) # rebind-ok: preserve Phase 1 mutation order and identity + return reset_at - timedelta(seconds=duration_seconds) + + +def _coerce_datetime(value: object) -> datetime | None: + if value is None: + return None + if isinstance(value, datetime): + return value + if isinstance(value, str): + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return None + + +__all__ = ( + "_raw_cache", + "_get_router_zero_cost_cache", + "_deployments_served_for", + "_has_ptu_flat_cost", + "_is_cost_explicitly_configured", + "_has_explicit_positive_cost", + "_deployment_query_price", + "_is_positive_cost", + "_apply_default_budget_to_end_user", + "_load_end_user_restricted_registry", + "_end_user_is_known_unrestricted", + "_end_user_id_exists_in_db", + "_apply_budget_exceeded_throttle", + "_virtual_key_multi_budget_check", + "_team_multi_budget_check", + "_active_member_increase", + "_effective_member_max", + "_coerce_window", + "_coerce_datetime", +) diff --git a/tests/unit/proxy/auth/phase1_budget_context_reference.py b/tests/unit/proxy/auth/phase1_budget_context_reference.py new file mode 100644 index 00000000000..46900a23d09 --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_context_reference.py @@ -0,0 +1,352 @@ +"""Policy frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9, with independently supplied raw IO.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from typing import Final, Protocol, cast + +from opentelemetry.trace import Span + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +from .phase1_budget_auth_reference import ( + get_end_user_object, + get_key_end_user_budget_id, + get_team_member_default_budget, + get_team_membership, + model_max_budget_of, + resolve_default_end_user_budget, +) + + +class _TeamModelBudgetLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + +class _EndUserModelBudgetLimiter(Protocol): + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: ... + + +async def _check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _TeamModelBudgetLimiter, + models: list[str], # mutable-ok: Phase 1 public collection and identity contract +) -> None: + + team_model_max_budget: Final = valid_token.team_model_max_budget + if valid_token.team_id is None or not team_model_max_budget: + return + key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + for model_name in models: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + model=model_name, + ) + + +async def _end_user_model_max_budget( + end_user_id: str | None, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, +) -> Mapping[str, object] | None: + budget_table: Final = ( + end_user_object.litellm_budget_table + if end_user_object is not None + else ( + await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + if end_user_id and prisma_client is not None + else None + ) + ) + return model_max_budget_of(budget_table) if budget_table is not None else None + + +async def _check_end_user_model_budget( + end_user_id: str | None, + end_user_model_max_budget: Mapping[str, object] | None, + model_max_budget_limiter: _EndUserModelBudgetLimiter, + models: Sequence[str], + team_id: str | None, +) -> None: + if not end_user_id or not end_user_model_max_budget: + return + for model_name in models: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=end_user_model_max_budget, + model=model_name, + team_id=team_id, + ) + + +def _apply_budget_limits_to_end_user_params( + end_user_params: dict, # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + budget_info: LiteLLM_BudgetTable, + end_user_id: str | None, +) -> None: + + if budget_info.tpm_limit is not None: + end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.rpm_limit is not None: + end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.tpd_limit is not None: + end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.max_budget is not None: + end_user_params["end_user_max_budget"] = budget_info.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.model_max_budget is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + end_user_params["end_user_model_max_budget"] = budget_info.model_max_budget # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) + + +def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + valid_token.end_user_id = end_user_params.get("end_user_id") # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + if end_user_params.get("end_user_tpm_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_rpm_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_tpd_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("allowed_model_region") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.allowed_model_region = end_user_params["allowed_model_region"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_model_max_budget") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"] # rebind-ok: preserve Phase 1 mutation order and identity + return valid_token + + +async def _load_team_member_default_model_budget( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + skip_budget_checks: bool, + proxy_logging_obj: ProxyLogging | None, +) -> None: + team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None + team_metadata: Final[Mapping[str, object] | None] = ( + cast( # cast-ok: team metadata is a string-keyed JSON object + Mapping[str, object], team_metadata_value + ) + if isinstance(team_metadata_value, Mapping) + else None + ) + team_member_budget_id: Final[object | None] = ( + team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None + ) + if ( + not isinstance(user_api_key_auth_obj.user_id, str) + or not isinstance(user_api_key_auth_obj.team_id, str) + or prisma_client is None + ): + return + + try: + membership: Final = await get_team_membership( + user_id=user_api_key_auth_obj.user_id, + team_id=user_api_key_auth_obj.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + member_budget_row: Final = membership.litellm_budget_table if membership is not None else None + member_budget_value: Final[object] = ( + getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None + ) + if ( + membership is not None + and membership.budget_id is not None + and membership.budget_id != team_member_budget_id + and isinstance(member_budget_value, Mapping) + and member_budget_value + ): + user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream + member_budget_value + ) + return + if not isinstance(team_member_budget_id, str): + return + team_member_default_budget: Final = await get_team_member_default_budget( + budget_id=team_member_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + raise_on_lookup_error=True, + ) + except Exception: + if not skip_budget_checks: + raise + verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True) + return + + default_budget_value: Final[object] = ( + getattr(team_member_default_budget, "model_max_budget", None) + if team_member_default_budget is not None + else None + ) + user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream + default_budget_value if isinstance(default_budget_value, Mapping) else None + ) + + +async def _apply_key_end_user_default_budget_to_token( + valid_token: UserAPIKeyAuth, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + keep_token_limits: bool, +) -> None: + + default_budget: Final = ( + end_user_object.litellm_budget_table + if end_user_object is not None + else await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + ) + if default_budget is None: + return + + if not keep_token_limits or valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_tpm_limit is None: + valid_token.end_user_tpm_limit = default_budget.tpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_rpm_limit is None: + valid_token.end_user_rpm_limit = default_budget.rpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_tpd_limit is None: + valid_token.end_user_tpd_limit = default_budget.tpd_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_model_max_budget is None: + valid_token.end_user_model_max_budget = default_budget.model_max_budget # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + +def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): # pyright: ignore[reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values + valid_token_metadata: Final = valid_token.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata: + expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"]) # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + if expiry.tzinfo is None: + expiry = expiry.replace(tzinfo=timezone.utc) # rebind-ok: preserve Phase 1 mutation order and identity + if expiry > datetime.now(timezone.utc): + return valid_token_metadata["temp_budget_increase"] # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + return None + + +def _update_key_budget_with_temp_budget_increase( + valid_token: UserAPIKeyAuth, +) -> UserAPIKeyAuth: + if valid_token.max_budget is None: + return valid_token + temp_budget_increase: Final = _get_temp_budget_increase(valid_token) # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if not temp_budget_increase: + return valid_token + return valid_token.model_copy(update={"max_budget": valid_token.max_budget + temp_budget_increase}) + + +async def _lookup_end_user_and_apply_budget( + valid_token: UserAPIKeyAuth, + route: str, + parent_otel_span: Span | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +): + + end_user_object = None # rebind-ok: preserve Phase 1 mutation order and identity + key_end_user_budget_id: Final = get_key_end_user_budget_id(valid_token.metadata) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + try: + end_user_object = await get_end_user_object( # rebind-ok: preserve Phase 1 mutation order and identity + end_user_id=valid_token.end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + token_end_user_max_budget=valid_token.end_user_max_budget, + key_end_user_budget_id=key_end_user_budget_id, + ) + if end_user_object is not None: + end_user_params = { # rebind-ok: preserve Phase 1 mutation order and identity + "end_user_id": valid_token.end_user_id, + "allowed_model_region": end_user_object.allowed_model_region, + } + if end_user_object.litellm_budget_table is not None: + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=end_user_object.litellm_budget_table, + end_user_id=valid_token.end_user_id or "", + ) + valid_token = update_valid_token_with_end_user_params( # rebind-ok: preserve Phase 1 mutation order and identity + valid_token=valid_token, end_user_params=end_user_params + ) + elif key_end_user_budget_id is not None or litellm.max_end_user_budget_id is not None: + default_budget: Final = await resolve_default_end_user_budget( + prisma_client=prisma_client, # pyright: ignore[reportArgumentType] # frozen Phase 1 forwards the optional raw client + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + if default_budget is not None: + end_user_params = {"end_user_id": valid_token.end_user_id} # rebind-ok: preserve Phase 1 mutation order and identity + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=default_budget, + end_user_id=valid_token.end_user_id or "", + ) + valid_token = update_valid_token_with_end_user_params( # rebind-ok: preserve Phase 1 mutation order and identity + valid_token=valid_token, end_user_params=end_user_params + ) + if valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + except Exception as e: + if isinstance(e, litellm.BudgetExceededError): + raise e + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) + return valid_token, end_user_object + + +__all__ = ( + "_check_team_model_budget", + "_end_user_model_max_budget", + "_check_end_user_model_budget", + "_apply_budget_limits_to_end_user_params", + "_load_team_member_default_model_budget", + "_apply_key_end_user_default_budget_to_token", + "_get_temp_budget_increase", + "_update_key_budget_with_temp_budget_increase", + "_lookup_end_user_and_apply_budget", +) diff --git a/tests/unit/proxy/auth/phase1_budget_lookup_reference.py b/tests/unit/proxy/auth/phase1_budget_lookup_reference.py new file mode 100644 index 00000000000..6d11b7f398c --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_lookup_reference.py @@ -0,0 +1,341 @@ +"""Unchanged Phase 1 budget lookup policy with raw IO supplied by tests.""" + +from __future__ import annotations + +import asyncio +import math +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + CallInfo, + Litellm_EntityType, + LiteLLM_OrganizationTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.utils import get_utc_datetime + +from .phase1_budget_auth_reference import ( + get_team_member_default_budget, + get_team_membership, + resolve_effective_team_member_budget, +) + +TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY: Final = "team_member_max_budget_alert_emails" +_TEAM_MEMBER_ALERT_CONFIG_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +async def get_org_object( + org_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + include_budget_table: bool, +) -> LiteLLM_OrganizationTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_org_object") + + +def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """Parse emails from a list or comma-separated string.""" + if isinstance(raw, list): + return [e.strip() for e in raw if isinstance(e, str) and e.strip()] + elif isinstance(raw, str): + return [e.strip() for e in raw.split(",") if e.strip()] + return [] + + +def _merge_budget_alert_email_configs( + global_cfg: Mapping[str, str | Sequence[object] | None] | None, + per_key_cfg: Mapping[str, str | Sequence[object] | None] | None, +) -> dict[str, list[str]] | None: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """ + Per-threshold additive merge: each threshold's recipient list is the union + of global + per-key entries (deduped, global-first ordering). Missing + thresholds on one side are inherited from the other. + """ + global_cfg_normalized: Final = _normalize_alert_emails(global_cfg) + per_key_cfg_normalized: Final = _normalize_alert_emails(per_key_cfg) + if not global_cfg_normalized and (not per_key_cfg_normalized): + return None + thresholds: Final = set(global_cfg_normalized) | set(per_key_cfg_normalized) + return { + t: list(dict.fromkeys(global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []))) for t in thresholds + } + + +def _is_valid_alert_threshold_pct(pct: str) -> bool: + return pct.isdigit() and len(pct) <= 3 and (1 <= int(pct) <= 100) + + +def _alert_recipients(raw: object) -> Sequence[str] | None: + if isinstance(raw, (str, Sequence)): + return _parse_email_list(raw) # pyright: ignore[reportUnknownArgumentType] # preserve unvalidated Phase 1 oracle values and call shape + return None + + +def _valid_alert_threshold_config(raw_config: object) -> Mapping[str, str | Sequence[object] | None] | None: + try: + config: Final = _TEAM_MEMBER_ALERT_CONFIG_ADAPTER.validate_python(raw_config) + except ValidationError: + return None + return MappingProxyType( + {pct: _alert_recipients(emails) for pct, emails in config.items() if _is_valid_alert_threshold_pct(pct)} + ) + + +def team_member_max_budget_alert_check( + team_id: str, + team_alias: str | None, + team_metadata: Mapping[str, object] | None, + organization_id: str | None, + user_id: str, + user_email: str | None, + proxy_logging_obj: ProxyLogging, + spend: float, + max_budget: float, +) -> None: + raw_config: Final = team_metadata.get(TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY) if team_metadata else None + alert_email_config: Final = _merge_budget_alert_email_configs( + global_cfg=None, per_key_cfg=_valid_alert_threshold_config(raw_config) + ) + if not alert_email_config or spend <= 0: + return + min_pct: Final = min(int(pct) for pct in alert_email_config) + if spend < max_budget * (min_pct / 100.0): + return + call_info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + user_id=user_id, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_email=user_email, + event_group=Litellm_EntityType.TEAM_MEMBER, + max_budget_alert_emails=alert_email_config, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=call_info)) + + +async def _check_team_member_budget( # pyright: ignore[reportUnusedFunction] # preserve unvalidated Phase 1 oracle values and call shape + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +): + """Check if team member is over their max budget within the team.""" + if ( + team_object is not None + and team_object.team_id is not None # pyright: ignore[reportUnnecessaryComparison] # preserve unvalidated Phase 1 oracle values and call shape + and (valid_token is not None) + and (valid_token.user_id is not None) + ): + if not team_membership_loaded: + team_membership = await get_team_membership( # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None + now: Final = get_utc_datetime() + default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + team_default_budget: Final = ( + await get_team_member_default_budget( + budget_id=default_budget_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache + ) + if (member_budget_row is None or member_budget_row.max_budget is None) + and isinstance(default_budget_id, str) + else None + ) + team_member_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget_row, team_default_budget=team_default_budget, now=now + ) + if team_member_budget is not None: + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + from litellm.proxy.proxy_server import get_current_spend + + team_member_spend = await get_current_spend( # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", + fallback_spend=team_member_spend, + max_budget=team_member_budget, + ) + if not math.isfinite(team_member_budget): + return + team_member_max_budget_alert_check( + team_id=team_object.team_id, + team_alias=team_object.team_alias, + team_metadata=team_object.metadata, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # preserve unvalidated Phase 1 oracle values and call shape + organization_id=team_object.organization_id, + user_id=valid_token.user_id, + user_email=user_object.user_email if user_object is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=team_member_spend, + max_budget=team_member_budget, + ) + if team_member_spend >= team_member_budget: + raise litellm.BudgetExceededError( + current_cost=team_member_spend, + max_budget=team_member_budget, + message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{valid_token.user_id}:{team_object.team_id}", + ) + + +async def _team_soft_budget_check( # pyright: ignore[reportUnusedFunction] # preserve unvalidated Phase 1 oracle values and call shape + team_object: LiteLLM_TeamTable | None, valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging +): + """ + Triggers a budget alert if the team is over it's soft budget. + """ + if ( + team_object is not None + and team_object.soft_budget is not None + and (team_object.spend is not None) + and (team_object.spend >= team_object.soft_budget) + ): + verbose_proxy_logger.debug( + "Crossed Soft Budget for team %s, spend %s, soft_budget %s", + team_object.team_id, + team_object.spend, + team_object.soft_budget, + ) + if valid_token: + alert_emails: list[str] | None = None # mutable-ok: unchanged Phase 1 collection contract in independent oracle # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if team_object.metadata is not None and isinstance(team_object.metadata, dict): # pyright: ignore[reportUnknownMemberType, reportUnnecessaryIsInstance] # preserve unvalidated Phase 1 oracle values and call shape + soft_budget_alert_emails: Final = team_object.metadata.get("soft_budget_alerting_emails") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + if soft_budget_alert_emails is not None: + if isinstance(soft_budget_alert_emails, list): + alert_emails = [ # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip() # pyright: ignore[reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + ] + elif isinstance(soft_budget_alert_emails, str): + alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()] # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if alert_emails: + alert_emails = [email for email in alert_emails if email] # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + else: + alert_emails = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if alert_emails is None or len(alert_emails) == 0: + verbose_proxy_logger.debug( + "Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails", + team_object.team_id, + ) + return + call_info: Final = CallInfo( + token=valid_token.token, + spend=team_object.spend, + max_budget=team_object.max_budget, + soft_budget=team_object.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.TEAM, + alert_emails=alert_emails, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="soft_budget", user_info=call_info)) + + +async def organization_max_budget_check( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Check if the organization is over its max budget. + + This function checks the organization budget using: + 1. First, tries to use valid_token.org_id (if key has organization_id set) + 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) + + This ensures organization budget checks work even when keys don't have organization_id + set directly, as long as their team belongs to an organization. + + Raises: + BudgetExceededError if the organization is over its max budget. + Triggers a budget alert if the organization is over its max budget. + """ + if valid_token is None or prisma_client is None: + return + org_id: str | None = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if valid_token.org_id is not None: + org_id = valid_token.org_id # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + elif team_object is not None and team_object.organization_id is not None: + org_id = team_object.organization_id # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_id is None: + return + try: + org_table: Final = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + include_budget_table=True, + ) + except Exception: + return + if org_table is None: + return + carry_organization_budget_state(valid_token=valid_token, org_table=org_table) + org_max_budget: float | None = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_table.litellm_budget_table is not None: + org_max_budget = org_table.litellm_budget_table.max_budget # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_max_budget is None: + return + from litellm.proxy.proxy_server import get_current_spend + + org_spend: Final = await get_current_spend( + counter_key=f"spend:org:{org_id}", fallback_spend=org_table.spend or 0.0, max_budget=org_max_budget + ) + if math.isfinite(org_max_budget) and org_spend >= org_max_budget: + call_info: Final = CallInfo( + token=valid_token.token, + spend=org_spend, + max_budget=org_max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=org_id, + event_group=Litellm_EntityType.ORGANIZATION, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="organization_budget", user_info=call_info)) + raise litellm.BudgetExceededError( + current_cost=org_spend, + max_budget=org_max_budget, + message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", + entity_type=Litellm_EntityType.ORGANIZATION.value, + entity_id=org_id, + ) + + +def _normalize_alert_emails(cfg: Mapping[str, str | Sequence[object] | None] | None) -> dict[str, list[str]]: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. + + Values may legitimately arrive as list, comma-separated string, or None + from YAML/metadata; _parse_email_list tolerates all three. + """ + if not cfg: + return {} + return {k: _parse_email_list(v) for k, v in cfg.items()} diff --git a/tests/unit/proxy/auth/phase1_budget_route_reference.py b/tests/unit/proxy/auth/phase1_budget_route_reference.py new file mode 100644 index 00000000000..f409731dc7e --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_route_reference.py @@ -0,0 +1,151 @@ +"""Route classification frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9 for independent budget references.""" + +from __future__ import annotations + +import re +from collections.abc import Collection +from typing import Final + +from litellm.proxy._types import LiteLLMRoutes + +_EU_JEV_DECISION_ROUTE: Final = "/eu_jev/v1/systemone" + + +class RouteChecks: + @staticmethod + def is_llm_api_route(route: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + if route in LiteLLMRoutes.openai_routes.value: + return True + + if route in LiteLLMRoutes.anthropic_routes.value: + return True + + if route in LiteLLMRoutes.google_routes.value: + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.litellm_native_routes.value): + return True + + for openai_route in LiteLLMRoutes.openai_routes.value: + if "{" in openai_route: + if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): + return True + + if RouteChecks.is_wildcard_pattern(pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): + return True + + for google_route in LiteLLMRoutes.google_routes.value: + if "{" in google_route: + if RouteChecks._route_matches_pattern(route=route, pattern=google_route): + return True + + for anthropic_route in LiteLLMRoutes.anthropic_routes.value: + if "{" in anthropic_route: + if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): + return True + + if RouteChecks._is_azure_openai_route(route=route): + return True + + for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: + if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): + return True + return RouteChecks.is_builtin_eu_jev_decision_route(route=route) + + @staticmethod + def is_builtin_eu_jev_decision_route(route: str) -> bool: + if route.rstrip("/") != _EU_JEV_DECISION_ROUTE: + return False + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass_through_endpoints imports the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash( + route=_EU_JEV_DECISION_ROUTE + ) + + @staticmethod + def is_info_route(route: str) -> bool: + + return route in LiteLLMRoutes.info_routes.value + + @staticmethod + def _is_azure_openai_route(route: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + deployment_pattern: Final = r"^/openai/deployments/[^/]+/[^/]+/chat/completions$" + engine_pattern: Final = r"^/engines/[^/]+/chat/completions$" + + if re.match(deployment_pattern, route) or re.match(engine_pattern, route): + return True + return False + + @staticmethod + def _route_matches_pattern(route: str, pattern: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + def _placeholder_to_regex(match: re.Match[str]) -> str: + placeholder: Final = match.group(0).strip("{}") + if not placeholder.endswith(":path"): + return r"[^/]+" + + return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" + + pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) # rebind-ok: preserve Phase 1 mutation order and identity + + pattern = f"^{pattern}$" # rebind-ok: preserve Phase 1 mutation order and identity + if re.match(pattern, route): + return True + return False + + @staticmethod + def is_wildcard_pattern(pattern: str) -> bool: + + return pattern.endswith("*") + + @staticmethod + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + + if pattern.endswith("*"): + prefix: Final = pattern[:-1] + return route.startswith(prefix) + else: + return route == pattern + + @staticmethod + def check_route_access(route: str, allowed_routes: Collection[str]) -> bool: + + if route in allowed_routes: + return True + + if any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in allowed_routes + if RouteChecks.is_wildcard_pattern(pattern=allowed_route) + ): + return True + + if any( + RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes + ): + return True + + return False + + +__all__ = () diff --git a/tests/unit/proxy/auth/phase1_budget_throttle_reference.py b/tests/unit/proxy/auth/phase1_budget_throttle_reference.py new file mode 100644 index 00000000000..5b3d8469d5a --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_throttle_reference.py @@ -0,0 +1,45 @@ +""" +Throttle a key after it exceeds its own ``max_budget`` instead of blocking it. + +When a key opts in via ``throttle_on_budget_exceeded`` and a global +``budget_exceeded_throttle_percentage`` is configured, an over-budget key keeps +serving requests but at a reduced TPM/RPM (the configured percentage of its +configured limits). The decision (over budget + opted in) is made once during +auth; the scaling is recomputed from the key's original limits on every request +so it never compounds across requests. +""" + +import math +from typing import Final + +import litellm +from litellm.proxy._types import UserAPIKeyAuth + + +def budget_throttle_percentage() -> float | None: + + pct: Final = litellm.budget_exceeded_throttle_percentage + if not isinstance(pct, (int, float)) or isinstance(pct, bool): + return None + if not 0 < pct <= 1: + return None + return float(pct) + + +def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: + + if (valid_token.metadata or {}).get("throttle_on_budget_exceeded") is not True: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return False + if valid_token.tpm_limit is None and valid_token.rpm_limit is None: + return False + return budget_throttle_percentage() is not None + + +def throttled_limit(limit: int | None, pct: float | None) -> int | None: + + if limit is None or pct is None: + return limit + return max(1, math.floor(limit * pct)) + + +__all__ = () diff --git a/tests/unit/proxy/auth/route_role_reference.py b/tests/unit/proxy/auth/route_role_reference.py new file mode 100644 index 00000000000..8e0ade7c067 --- /dev/null +++ b/tests/unit/proxy/auth/route_role_reference.py @@ -0,0 +1,81 @@ +from collections.abc import Iterable +from itertools import chain +from typing import Final + +from fastapi import Request +from starlette.routing import compile_path + +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.route_checks import RouteChecks + + +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + try: + return LitellmUserRoles(user_role) + except ValueError: + return LitellmUserRoles.INTERNAL_USER + + +def get_user_role(user_obj: LiteLLM_UserTable | None) -> LitellmUserRoles | None: + return None if user_obj is None else effective_user_role(user_obj.user_role) + + +def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: + return user_obj is not None and user_obj.user_role is not None and user_obj.user_role == "proxy_admin" + + +def is_api_route_allowed( + route: str, request: Request, request_data: dict[str, object], valid_token: UserAPIKeyAuth | None, # mutable-ok: preserve the original request identity at the shared route effect + user_obj: LiteLLM_UserTable | None = None, +) -> bool: + role: Final = get_user_role(user_obj) + if valid_token is None: + raise Exception("Invalid proxy server token passed. valid_token=None.") + if not is_user_proxy_admin(user_obj): + RouteChecks.non_proxy_admin_allowed_routes_check( # pyright: ignore[reportUnknownMemberType] # Phase 1 public effect signature + user_obj=user_obj, _user_role=role, route=route, request=request, + request_data=request_data, valid_token=valid_token, + ) + return True + + +def allowed_routes_list_check(user_route: str, allowed_routes: Iterable[object]) -> bool: + for allowed in allowed_routes: + if allowed in LiteLLMRoutes.__members__: + for template in LiteLLMRoutes[allowed].value: # pyright: ignore[reportArgumentType, reportUnknownVariableType, reportUnknownMemberType] # reference intentionally preserves the untyped Phase 1 input contract + regex, _, _ = compile_path(template) # pyright: ignore[reportUnknownArgumentType] # enum values are legacy route collections + if regex.match(user_route): + return True + elif allowed.endswith("*"): # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # invalid legacy entries must raise their original AttributeError + if user_route.startswith(allowed[:-1]): # pyright: ignore[reportIndexIssue, reportUnknownArgumentType] # preserve the reference's raw entry operations + return True + elif user_route == allowed: + return True + return False + + +def allowed_routes_check(user_role: LitellmUserRoles, user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth) -> bool: + if user_role == LitellmUserRoles.PROXY_ADMIN: + return allowed_routes_list_check(user_route, litellm_proxy_roles.admin_allowed_routes) + if user_role == LitellmUserRoles.TEAM: + configured: Final = litellm_proxy_roles.team_allowed_routes + return allowed_routes_list_check( + user_route, ("openai_routes", "info_routes") if configured is None else configured # pyright: ignore[reportUnnecessaryComparison] # retain Phase 1 constructed None config behavior + ) + return False + + +def allowed_route_check_inside_route(auth: UserAPIKeyAuth, requested_user_id: str | None) -> bool: + return auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) or ( + requested_user_id is not None and auth.user_id is not None and auth.user_id == requested_user_id + ) + + +def get_actual_routes(allowed_routes: Iterable[object]) -> list[object]: # mutable-ok: public Phase 1 projection + def expanded(raw: object) -> tuple[object, ...]: + try: + return tuple(LiteLLMRoutes[raw].value) # pyright: ignore[reportArgumentType, reportUnknownArgumentType, reportUnknownMemberType] # preserve enum indexing and iterable values from Phase 1 + except KeyError: + return (raw,) + + return list(chain.from_iterable(expanded(raw) for raw in allowed_routes)) diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index b191f7835a8..cb3a9d93d45 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -2338,7 +2338,10 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) client = _WedgedPublishRedisClient() - cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000)) + cache: Final = UserApiKeyCache( + in_memory_cache=InMemoryCache(max_size_in_memory=2_000), + key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000), + ) hashed_tokens: Final = tuple( hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 78ce9728f3b..e2696f89c2a 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -20,6 +20,7 @@ import pytest from fastapi import Request, status import litellm +from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( CallInfo, Litellm_EntityType, @@ -632,6 +633,37 @@ async def test_get_key_object_should_reconnect_once_on_db_connection_error(): ) +@pytest.mark.asyncio +async def test_native_key_lookup_populates_cache_without_prisma_io(): + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{"token": "native-hash", "team_models": None, "team_blocked": None}])) + prisma = MagicMock(native_auth_database=native) + prisma.get_data = AsyncMock(side_effect=AssertionError("legacy read")) + prisma.attempt_db_reconnect = AsyncMock(side_effect=AssertionError("legacy reconnect")) + cache = UserApiKeyCache() + first = await get_key_object("native-hash", prisma, cache) + second = await get_key_object("native-hash", prisma, cache) + assert first.token == second.token == "native-hash" + assert first.team_models == second.team_models == [] + assert native.auth_read.await_count == 1 + prisma.get_data.assert_not_awaited() + prisma.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_native_key_transport_error_never_enters_prisma_reconnect(): + from litellm._auth.failure import NativeAuthFailure + + native = SimpleNamespace(auth_read=AsyncMock(side_effect=httpx.ConnectError("db offline"))) + prisma = MagicMock(native_auth_database=native) + prisma.get_data = AsyncMock() + prisma.attempt_db_reconnect = AsyncMock() + with pytest.raises(NativeAuthFailure): + await fetch_key_object_from_db_with_reconnect("hash", prisma, None, None, check_db_only=True) + native.auth_read.assert_awaited_once_with("key", ["hash"], writer=True) + prisma.get_data.assert_not_awaited() + prisma.attempt_db_reconnect.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_error(): mock_prisma_client = MagicMock() @@ -850,7 +882,7 @@ class TestAuthCacheRedisWritePolicy: """ @pytest.mark.asyncio - async def test_get_key_object_db_load_publishes_to_redis(self): + async def test_get_key_object_db_load_stays_local_and_is_reused_without_redis_publication(self): mock_prisma_client = MagicMock() mock_prisma_client.get_data = AsyncMock(return_value=UserAPIKeyAuth(token="hashed-token-db")) @@ -865,10 +897,14 @@ class TestAuthCacheRedisWritePolicy: ) assert key_obj.token == "hashed-token-db" - fake_redis.async_set_cache.assert_awaited_once() - assert ( - fake_redis.async_set_cache.await_args.kwargs.get("key") or fake_redis.async_set_cache.await_args.args[0] - ) == "hashed-token-db" + cached: Final = await get_key_object( + hashed_token="hashed-token-db", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + assert cached.token == key_obj.token + mock_prisma_client.get_data.assert_awaited_once() + fake_redis.async_set_cache.assert_not_awaited() def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values): @@ -2320,6 +2356,18 @@ async def test_every_requested_vector_store_needs_a_grant( ) +@pytest.mark.asyncio +@pytest.mark.parametrize("deny_by_default", (False, True), ids=("legacy", "strict")) +async def test_missing_database_skips_malformed_vector_request(deny_by_default: bool) -> None: + with patch("litellm.proxy.proxy_server.prisma_client", None): + assert await vector_store_access_check( + request_body={"vector_store_ids": {"malformed": "ids"}}, + valid_token=_virtual_key(), + team_object=None, + deny_by_default=deny_by_default, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "valid_token", @@ -2432,6 +2480,67 @@ async def test_file_search_inside_additional_tools_input_item_needs_a_vector_sto ) +@pytest.mark.parametrize("strict", [False, True], ids=["legacy", "strict"]) +@pytest.mark.parametrize( + ("input_items", "expected"), + [ + (None, ("top", "shared", "tool", "shared")), + ("not-a-list", ("top", "shared", "tool", "shared")), + ([{"type": "additional_tools", "tools": 1}], ("top", "shared", "tool", "shared")), + ( + [ + {"type": "message", "tools": [{"vector_store_ids": ["ignored"]}]}, + {"type": "additional_tools", "tools": [{"vector_store_ids": ["input", "shared"]}]}, + {"type": "additional_tools", "tools": [{"vector_store_ids": ["last"]}]}, + ], + ("top", "shared", "tool", "shared", "input", "shared", "last"), + ), + ], + ids=["no-input", "string-input", "malformed-tools", "ordered-duplicates"], +) +def test_vector_store_raw_reads_keep_top_level_tools_then_additional_tools( + strict: bool, input_items: object, expected: tuple[str, ...] +) -> None: + from litellm.proxy.auth.auth_checks import _legacy_requested_vector_store_ids, _strict_requested_vector_store_ids + + body: Final = { + "vector_store_ids": ["top", "shared"], + "tools": [{"vector_store_ids": ["tool", "shared"]}], + "input": input_items, + } + result: Final = ( + _strict_requested_vector_store_ids(body) if strict else _legacy_requested_vector_store_ids(body, lambda: ()) + ) + assert result == expected + + +@pytest.mark.parametrize("strict", [False, True], ids=["legacy", "strict"]) +def test_malformed_top_level_vector_ids_do_not_read_additional_tools(strict: bool) -> None: + from litellm.proxy.auth.auth_checks import ( + _legacy_requested_vector_store_ids, + _strict_requested_vector_store_ids, + _UnreadableVectorStoreIds, + ) + + with patch( + "litellm.proxy.auth.auth_checks.additional_tools_in_input", + side_effect=AssertionError("unreadable top-level IDs read input tools"), + ) as input_reader: + if strict: + with pytest.raises(ProxyException) as caught: + _strict_requested_vector_store_ids({"vector_store_ids": 42}) + assert (caught.value.code, caught.value.type, caught.value.param) == ( + "400", + "invalid_request_error", + "vector_store_ids", + ) + else: + assert isinstance( + _legacy_requested_vector_store_ids({"vector_store_ids": 42}, lambda: ()), _UnreadableVectorStoreIds + ) + input_reader.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.parametrize( "request_body", @@ -5926,11 +6035,12 @@ class TestGuardrailModificationCheck: `_get_admin_metadata` silently ignores at read time. """ - def _call(self, request_body): + def _call(self, request_body, *, allowed=False): from litellm.proxy.auth.auth_checks import guardrail_modification_check - team_object = MagicMock() - team_object.metadata = {} # no permission + team_object = LiteLLM_TeamTable( + team_id="guardrail-team", metadata={"guardrails": {"modify_guardrails": allowed}} + ) return guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): @@ -5940,47 +6050,31 @@ class TestGuardrailModificationCheck: def test_rejects_guardrails_list(self): from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"guardrails": ["custom"]}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"guardrails": ["custom"]}}) + assert exc.value.status_code == 403 def test_rejects_disable_global_guardrails_plural(self): from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"disable_global_guardrails": True}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"disable_global_guardrails": True}}) + assert exc.value.status_code == 403 def test_rejects_disable_global_guardrail_singular(self): """VERIA-28's originally-reported singular-key typo variant.""" from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"disable_global_guardrail": True}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"disable_global_guardrail": True}}) + assert exc.value.status_code == 403 def test_rejects_opted_out_global_guardrails(self): from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}}) + assert exc.value.status_code == 403 @pytest.mark.parametrize( "key", @@ -6001,46 +6095,30 @@ class TestGuardrailModificationCheck: """ from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {key: empty_value}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {key: empty_value}}) + assert exc.value.status_code == 403 def test_rejects_injection_via_litellm_metadata_key(self): """Caller can populate the OTHER metadata key; that must also 403.""" from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"litellm_metadata": {"disable_global_guardrails": True}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"litellm_metadata": {"disable_global_guardrails": True}}) + assert exc.value.status_code == 403 def test_rejects_root_level_injection(self): """Top-level injection (`request_body["disable_global_guardrails"]`) was VERIA-28's easiest variant to hit — keep it rejected.""" from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"disable_global_guardrails": True}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"disable_global_guardrails": True}) + assert exc.value.status_code == 403 def test_allows_when_team_has_permission(self): - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=True, - ): - # no-op, should not raise - self._call({"metadata": {"disable_global_guardrails": True}}) + # no-op, should not raise + self._call({"metadata": {"disable_global_guardrails": True}}, allowed=True) def test_rejects_string_encoded_metadata_bypass(self): """Regression: attacker sends metadata as JSON string to bypass the @@ -6051,13 +6129,9 @@ class TestGuardrailModificationCheck: from fastapi import HTTPException attacker_payload = {"disable_global_guardrails": True} - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": _json.dumps(attacker_payload)}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": _json.dumps(attacker_payload)}) + assert exc.value.status_code == 403 def test_rejects_string_encoded_litellm_metadata_bypass(self): """Same bypass via the litellm_metadata key.""" @@ -6066,13 +6140,9 @@ class TestGuardrailModificationCheck: from fastapi import HTTPException attacker_payload = {"guardrails": ["evaded"]} - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"litellm_metadata": _json.dumps(attacker_payload)}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"litellm_metadata": _json.dumps(attacker_payload)}) + assert exc.value.status_code == 403 def test_noop_when_string_is_not_json_object(self): """Unparseable strings should not trigger a 403 — they have no keys.""" @@ -7344,7 +7414,7 @@ class TestManagementObjectTTLHonored: @pytest.mark.asyncio async def test_uses_configured_user_api_key_cache_ttl(self): mem = _TTLCapturingInMemoryCache() - cache = UserApiKeyCache(in_memory_cache=mem, default_in_memory_ttl=300) + cache = DualCache(in_memory_cache=mem, default_in_memory_ttl=300) await _cache_management_object( key="team_id:lit-3338", @@ -7359,7 +7429,7 @@ class TestManagementObjectTTLHonored: @pytest.mark.asyncio async def test_falls_back_to_management_default_when_unconfigured(self): mem = _TTLCapturingInMemoryCache() - cache = UserApiKeyCache(in_memory_cache=mem) + cache = DualCache(in_memory_cache=mem) assert cache.default_in_memory_ttl is None await _cache_management_object( @@ -9661,6 +9731,40 @@ async def test_get_managed_vector_store_rows_by_uuids_db_fetch_validates_rows(): assert result[0].custom_llm_provider == "openai" +@pytest.mark.asyncio +async def test_native_managed_vector_store_read_preserves_owner_and_warms_cache(): + from litellm.proxy.auth.auth_checks import get_managed_vector_store_rows_by_uuids + + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{ + "vector_store_id": "vs-native", "custom_llm_provider": "openai", + "team_id": "team", "user_id": "owner", "litellm_params": {"vector_store_id": "upstream"}, + }])) + prisma = SimpleNamespace(native_auth_database=native) + cache = UserApiKeyCache() + rows = await get_managed_vector_store_rows_by_uuids(["vs-native"], prisma, cache) + cached = await get_managed_vector_store_rows_by_uuids(["vs-native"], prisma, cache) + assert len(rows) == len(cached) == 1 + assert rows[0].team_id == cached[0].team_id == "team" + assert rows[0].user_id == cached[0].user_id == "owner" + assert rows[0].litellm_params == {"vector_store_id": "upstream"} + native.auth_read.assert_awaited_once_with("managed_vector_store", ["vs-native"], writer=False) + + +@pytest.mark.asyncio +async def test_native_email_backfill_reads_the_writer_after_the_atomic_update(): + from litellm.proxy.auth.auth_checks import _backfill_null_user_email + + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{"user_id": "u", "user_email": "winner@example.test"}])) + prisma = MagicMock(native_auth_database=native) + prisma.db.litellm_usertable.update_many = AsyncMock(return_value=0) + prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=AssertionError("legacy read")) + row = await _backfill_null_user_email(prisma, UserApiKeyCache(), LiteLLM_UserTable(user_id="u"), "loser@example.test") + assert row.user_email == "winner@example.test" + native.auth_read.assert_awaited_once_with("user", ["u"], writer=True) + prisma.db.litellm_usertable.update_many.assert_awaited_once() + prisma.db.litellm_usertable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_project_object_db_fetch_returns_cached_obj(): from litellm.proxy._types import LiteLLM_ProjectTableCachedObj @@ -11850,7 +11954,9 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_ unmapped_request = _client_tag_request(client_tags, baggage=False) unmapped = _client_tag_body(client_tags, client_slot) LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(unmapped, route) - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key")) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key") + ) mapped_request = _client_tag_request(client_tags, baggage=baggage) mapped = _client_tag_body(client_tags, client_slot) diff --git a/tests/unit/proxy/auth/test_auth_object_prefetch.py b/tests/unit/proxy/auth/test_auth_object_prefetch.py index a80c637d3df..0bf448f6bc4 100644 --- a/tests/unit/proxy/auth/test_auth_object_prefetch.py +++ b/tests/unit/proxy/auth/test_auth_object_prefetch.py @@ -1,11 +1,15 @@ """Counts the Redis round trips and DB queries auth object reads cost per cache regime, and checks that the per-object getters still enforce on their own when the prefetch cannot help.""" +import json +from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from litellm._auth.failure import NativeAuthFailure from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy._types import ( LiteLLM_OrganizationTable, @@ -26,6 +30,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, end_user_restricted_registry_cache_key, + team_membership_reservation_cache_key, ) from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache @@ -75,6 +80,80 @@ ALL_ROWS = { } +@pytest.mark.asyncio +async def test_native_prefetch_fills_all_getters_with_one_read_and_preserves_optional_ids(): + native = SimpleNamespace(auth_read=AsyncMock(return_value=[ALL_ROWS])) + prisma = _prisma() + prisma.native_auth_database = native + cache = _cache(None) + + await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=cache) + user, team, member, org = await _read_all_through_getters(cache, prisma) + + assert user is not None and user.user_id == USER_ID + assert team.team_id == TEAM_ID + assert member is not None and member.spend == 3.0 + assert org is not None and org.litellm_budget_table.max_budget == 1000.0 + assert native.auth_read.await_count == 1 + kind, ids = native.auth_read.call_args.args + assert kind == "prefetch" + assert len(ids) == 1 and json.loads(ids[0]) == [USER_ID, TEAM_ID, USER_ID, ORG_ID, None] + assert prisma.db.mock_calls == [] + + await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=cache) + assert native.auth_read.await_count == 1 + + +@pytest.mark.asyncio +async def test_native_owner_getters_do_not_cache_rows_invalidated_during_their_read(): + cache: Final = _cache(None) + + async def read(kind: str, ids: list[str], writer: bool = False) -> list[object]: + cache.delete_local( + { + "user": USER_ID, + "team": f"team_id:{TEAM_ID}", + "member": team_membership_reservation_cache_key(USER_ID, TEAM_ID), + "org": f"org_id:{ORG_ID}:with_budget", + }[kind] + ) + return [{"user": USER_ROW, "team": TEAM_ROW, "member": MEMBERSHIP_ROW, "org": ORG_ROW}[kind]] + + native: Final = SimpleNamespace(auth_read=AsyncMock(side_effect=read)) + prisma: Final = _prisma() + prisma.native_auth_database = native + user, team, member, org = await _read_all_through_getters(cache, prisma) + + assert user is not None and team is not None and member is not None and org is not None + assert cache.get_local(USER_ID) is None + assert cache.get_local(f"team_id:{TEAM_ID}") is None + assert cache.get_local(team_membership_reservation_cache_key(USER_ID, TEAM_ID)) is None + assert cache.get_local(f"org_id:{ORG_ID}:with_budget") is None + assert native.auth_read.await_count == 4 + assert prisma.db.mock_calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rows", [[{"user_row": {"spend": "bad"}}], {"user_row": USER_ROW}]) +async def test_native_prefetch_invalid_projection_is_terminal_without_legacy_query(rows): + prisma = _prisma() + prisma.native_auth_database = SimpleNamespace(auth_read=AsyncMock(return_value=rows)) + with pytest.raises(NativeAuthFailure): + await prefetch_auth_objects( + refs=AuthObjectRefs(user_id=USER_ID), prisma_client=prisma, user_api_key_cache=_cache(None) + ) + assert prisma.db.mock_calls == [] + + +@pytest.mark.asyncio +async def test_native_prefetch_transport_failure_is_not_best_effort(): + prisma = _prisma() + prisma.native_auth_database = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("offline"))) + with pytest.raises(NativeAuthFailure): + await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=_cache(None)) + assert prisma.db.mock_calls == [] + + def _prisma(rows: dict[str, object] | None = ALL_ROWS) -> MagicMock: prisma = MagicMock(name="prisma_client") prisma.db.query_first = AsyncMock(return_value=rows) diff --git a/tests/unit/proxy/auth/test_budget_alert_adapter.py b/tests/unit/proxy/auth/test_budget_alert_adapter.py new file mode 100644 index 00000000000..8582074fc6a --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_alert_adapter.py @@ -0,0 +1,125 @@ +import asyncio +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth import budget_alerts +from litellm._auth.budget_alerts import Facts +from litellm._auth.failure import NativeAuthFailure, NativeAuthValueError +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.proxy._types import CallInfo, Litellm_EntityType, UserAPIKeyAuth +from litellm.proxy.auth import budget_alert_adapter as native +from litellm.rust_bridge import _native +from tests.unit.proxy.auth.budget_alert_reference import reference_alert + + +def observed_reference(facts: Facts, global_config: object, local_config: object) -> object: + try: + return reference_alert(facts, global_config, local_config) + except (ValueError, OverflowError, AttributeError) as error: + return type(error).__name__, str(error) + + +def observed_native(facts: Facts, global_config: object, local_config: object) -> object: + try: + decision: Final = budget_alerts.plan(native.read_alert_facts(facts, global_config, local_config)) + native.raise_alert_failure(decision) + config: Final = native.alert_emails(decision) if decision.emit else None + return decision.emit, None if config is None else tuple((key, tuple(value)) for key, value in config.items()) + except (ValueError, OverflowError, AttributeError) as error: + return type(error).__name__, str(error) + + +@pytest.mark.parametrize("member", (False, True)) +@pytest.mark.parametrize("spend", (-1.0, 0.0, 7.9, 8.0, 9.0, 10.0, float("nan"), float("inf"))) +@pytest.mark.parametrize("limit", (-1.0, 0.0, 10.0, float("nan"), float("inf"))) +@pytest.mark.parametrize( + "configs", + ( + (None, None), + ({"80": " a@x.test,b@x.test "}, {"80": [" b@x.test ", "c@x.test", "", None]}), + (None, {"bad": "a@x.test"}), + (None, {"0": [], "101": None, "080": [" a@x.test "]}), + (None, {"80": ("a@x.test",)}), + (None, {"٨٠": "a@x.test"}), + (None, {"²": []}), + (None, {"9" * 400: [], "80": "a@x.test"}), + (None, {"9" * 400: []}), + (None, "invalid"), + ), +) +def test_native_threshold_merge_and_trigger_match_reference( + member: bool, spend: float, limit: float, configs: tuple[object, object] +) -> None: + facts: Final = Facts(member, spend, limit, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) + global_config, local_config = configs + global_value: Final = None if member else global_config + assert observed_native(facts, global_value, local_config) == observed_reference(facts, global_value, local_config) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("config", (None, {"80": " a@x.test, b@x.test "})) +async def test_key_alert_schedules_once_without_awaiting_and_preserves_projection( + monkeypatch: pytest.MonkeyPatch, config: object +) -> None: + monkeypatch.setattr(litellm, "default_key_max_budget_alert_emails", None) + key: Final = UserAPIKeyAuth( + token="token", + spend=8, + max_budget=10, + metadata={"max_budget_alert_emails": config}, + user_id="owner", + team_id="team", + ) + logger: Final = Mock(budget_alerts=AsyncMock()) + await native.virtual_key_max_budget_alert_check(key, logger) + assert logger.budget_alerts.call_count == 1 + assert logger.budget_alerts.await_count == 0 + call_info: Final = logger.budget_alerts.call_args.kwargs["user_info"] + expected: Final = CallInfo( + token="token", + spend=8, + max_budget=10, + soft_budget=None, + user_id="owner", + team_id="team", + team_alias=None, + organization_id=None, + user_email=None, + key_alias=None, + event_group=Litellm_EntityType.KEY, + ) + projected: Final = ( + expected + if config is None + else expected.model_copy(update={"max_budget_alert_emails": {"80": ["a@x.test", "b@x.test"]}}) + ) + assert call_info.model_dump() == projected.model_dump() + assert call_info.model_fields_set == projected.model_fields_set + await asyncio.sleep(0) + assert logger.budget_alerts.await_count == 1 + + +def test_native_failure_is_terminal_before_alert_config_reads(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = ValueError("native alert decoder failed") + backend: Final = Mock(side_effect=failure) + monkeypatch.setattr(_native, "request_auth_budget_alert_plan", backend) + with pytest.raises(NativeAuthFailure, match="native alert decoder failed") as error: + native.read_alert_facts(Facts(False, 8.0, 10.0, 0.8), None, {"80": []}) + assert isinstance(error.value, ValueError) + assert backend.call_count == 1 + + +def test_pricing_native_failure_never_becomes_lookup_failure() -> None: + from litellm.proxy.auth.budget_pricing_adapter import read_model_price_facts + + failure: Final = NativeAuthValueError("native registry failed") + registry: Final = Mock(group=Mock(side_effect=failure)) + router: Final = Mock(registry=Mock(return_value=registry)) + with pytest.raises(NativeAuthFailure) as error: + read_model_price_facts("model", router) + assert error.value is failure + registry.group.assert_called_once() + registry.deployments.assert_not_called() diff --git a/tests/unit/proxy/auth/test_budget_context_adapter.py b/tests/unit/proxy/auth/test_budget_context_adapter.py new file mode 100644 index 00000000000..662bd369f0d --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_context_adapter.py @@ -0,0 +1,335 @@ +from collections.abc import Mapping +from typing import Final, Literal +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth import budget_context as core +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks +from litellm.proxy.auth import budget_context_adapter as adapter +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from tests.unit.proxy.auth import phase1_budget_auth_reference as auth_reference +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy + + +@pytest.fixture(params=(False, True), ids=("reference", "native")) +def native(request: pytest.FixtureRequest) -> bool: + return request.param is True # pyright: ignore[reportAny] # fixture parameters are the two literal booleans above + + +@pytest.mark.parametrize("scope", ("Team", "EndUser", "CustomEndUser")) +@pytest.mark.parametrize("identity", (None, "", "entity")) +@pytest.mark.parametrize("budget_count", (0, 1, 2)) +@pytest.mark.parametrize("model_count", (0, 1, 3)) +def test_native_model_read_selection( + scope: Literal["Team", "EndUser", "CustomEndUser"], identity: str | None, budget_count: int, model_count: int +) -> None: + facts: Final = core.ModelCheckFacts(scope, identity, budget_count, model_count) + expected: Final = ( + tuple(range(model_count)) if budget_count and identity is not None and (scope != "EndUser" or identity) else () + ) + assert core.model_checks(facts) == expected + + +@pytest.mark.parametrize("keep", (False, True)) +@pytest.mark.parametrize("present", (False, True)) +@pytest.mark.parametrize("limit", (None, 0, -1, 1)) +@pytest.mark.parametrize("models", (None, 0, 2)) +def test_native_token_assignment_table(keep: bool, present: bool, limit: int | None, models: int | None) -> None: + facts: Final = core.TokenDefaultsFacts(core.BudgetFields(limit, limit, limit, limit, models), keep, present) + expected: Final = tuple( + name + for name, value in (("MaxBudget", limit), ("Tpm", limit), ("Rpm", limit), ("Tpd", limit), ("Models", models)) + if present and (not keep or value is None) + ) + assert core.token_assignments(facts) == expected + + +@pytest.mark.parametrize("mode", ("Models", "Token", "Default")) +@pytest.mark.parametrize("object_present", (False, True)) +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("db", (False, True)) +@pytest.mark.parametrize("key_id", (None, "", "key")) +def test_end_user_initial_table( + mode: Literal["Models", "Token", "Default"], + object_present: bool, + identity: str | None, + db: bool, + key_id: str | None, +) -> None: + facts: Final = core.EndUserFacts(mode, object_present, identity, db, key_id) + expected: Final = ( + "Own" + if mode != "Default" and object_present + else "None" + if mode == "Models" and (not identity or not db) + else "ReadKey" + if key_id is not None + else "ReadGlobalConfig" + ) + assert core.end_user_source(facts) == expected + + +@pytest.mark.parametrize("key_row", (False, True)) +@pytest.mark.parametrize("global_id", (None, "", "proxy")) +@pytest.mark.parametrize("global_read", (False, True)) +@pytest.mark.parametrize("global_row", (False, True)) +def test_end_user_fallback_stages(key_row: bool, global_id: str | None, global_read: bool, global_row: bool) -> None: + facts: Final = core.EndUserFacts( + "Default", + False, + None, + True, + "key", + key_read=True, + key_budget=core.BudgetFields() if key_row else None, + global_config_read=True, + global_budget_id=global_id, + global_read=global_read, + global_budget=core.BudgetFields() if global_row else None, + ) + expected: Final = ( + "Key" + if key_row + else "None" + if global_id is None + else "ReadGlobal" + if not global_read + else "Global" + if global_row + else "None" + ) + assert core.end_user_source(facts) == expected + + +@pytest.mark.parametrize("default_id", (None, "", "default")) +@pytest.mark.parametrize("member_id", (None, "", "member", "default")) +@pytest.mark.parametrize("member_models", (None, 0, 2)) +@pytest.mark.parametrize("default_models", (None, 0, 2)) +@pytest.mark.parametrize("skip", (False, True)) +def test_member_snapshot_table( + default_id: str | None, member_id: str | None, member_models: int | None, default_models: int | None, skip: bool +) -> None: + facts: Final = core.MemberFacts( + "", + "", + True, + default_id, + skip, + membership_read=True, + membership=core.MembershipSnapshot(member_id, member_models), + default_read=True, + default_budget=core.BudgetFields(model_budget_count=default_models), + ) + expected: Final = ( + "Member" + if member_id is not None and member_id != default_id and member_models + else "Done" + if default_id is None + else "Default" + if default_models is not None + else "Clear" + ) + assert core.member_source(facts) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", (None, "", "entity")) +@pytest.mark.parametrize("scope", ("Team", "EndUser")) +@pytest.mark.parametrize("fail", (False, True)) +async def test_selected_model_reads_keep_order_and_stop_on_failure( + native: bool, identity: str | None, scope: Literal["Team", "EndUser"], fail: bool +) -> None: + failure: Final = RuntimeError("counter failed") + read: Final = AsyncMock(side_effect=(True, failure, True) if fail else None, return_value=True) + limiter: Final = Mock() + limiter.is_team_within_model_budget = read + limiter.is_end_user_within_model_budget = read + caps: Final[Mapping[str, object]] = {"a": {"max_budget": 1}} + token: Final = UserAPIKeyAuth(team_id=identity, team_model_max_budget=dict(caps)) + models: Final = ["a", "a", "b"] + coroutine: Final = ( + (adapter.check_team_model_budget if native else legacy._check_team_model_budget)(token, limiter, models) # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + if scope == "Team" + else (adapter.check_end_user_model_budget if native else legacy._check_end_user_model_budget)( # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + identity, caps, limiter, models, "team" + ) + ) + active: Final = identity is not None and (scope == "Team" or bool(identity)) + if fail and active: + with pytest.raises(RuntimeError, match="counter failed") as raised: + await coroutine + assert raised.value is failure + else: + await coroutine + assert read.await_count == (2 if fail else 3) if active else read.await_count == 0 + assert tuple(item.kwargs["model"] for item in read.await_args_list) == tuple(models[: read.await_count]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own", (False, True)) +@pytest.mark.parametrize("own_row", (False, True)) +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("key_found", (False, True)) +async def test_end_user_source_read_order_and_identity( + native: bool, monkeypatch: pytest.MonkeyPatch, own: bool, own_row: bool, identity: str | None, key_found: bool +) -> None: + budget: Final = LiteLLM_BudgetTable(model_max_budget={"m": {"max_budget": 1}}) + end_user: Final = ( + LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget if own_row else None) + if own + else None + ) + read: Final = AsyncMock(side_effect=(budget,) if key_found else (None, budget)) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", read) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", read) + monkeypatch.setattr(litellm, "max_end_user_budget_id", "proxy") + fn: Final = adapter.end_user_model_max_budget if native else legacy._end_user_model_max_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + value: Final = await fn(identity, end_user, "key", Mock(), UserApiKeyCache(), None) + expected: Final = getattr(budget, "model_max_budget", None) if (own and own_row) or (not own and identity) else None + assert value is expected + assert read.await_count == (0 if own or not identity else 1 if key_found else 2) + if read.await_count: + assert read.await_args_list[0].kwargs["budget_id"] == "key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("keep", (False, True)) +@pytest.mark.parametrize("present", (False, True)) +async def test_token_defaults_preserve_caps_and_projection_identity(native: bool, keep: bool, present: bool) -> None: + budget: Final = LiteLLM_BudgetTable( + max_budget=9, tpm_limit=8, rpm_limit=7, tpd_limit=6, model_max_budget={"m": {"max_budget": 5}} + ) + token: Final = UserAPIKeyAuth( + end_user_max_budget=0 if present else None, + end_user_tpm_limit=0 if present else None, + end_user_model_max_budget={} if present else None, + ) + original_models: Final = token.end_user_model_max_budget + own: Final = LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget) + fn: Final = ( + adapter.apply_key_end_user_default_budget_to_token + if native + else legacy._apply_key_end_user_default_budget_to_token # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + await fn(token, own, "key", Mock(), UserApiKeyCache(), None, keep) + assert token.end_user_max_budget == (0 if keep and present else 9) + assert token.end_user_tpm_limit == (0 if keep and present else 8) + assert (token.end_user_rpm_limit, token.end_user_tpd_limit) == (7, 6) + assert token.end_user_model_max_budget is ( + original_models if keep and present else getattr(budget, "model_max_budget", None) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("default_id", (None, "", "default")) +@pytest.mark.parametrize("member_id", (None, "default", "override")) +@pytest.mark.parametrize("member_caps", (None, {}, {"member": {"max_budget": 1}})) +@pytest.mark.parametrize("default_caps", (None, {}, {"default": {"max_budget": 2}})) +async def test_member_default_read_and_mutation_order( + native: bool, + monkeypatch: pytest.MonkeyPatch, + default_id: str | None, + member_id: str | None, + member_caps: Mapping[str, object] | None, + default_caps: Mapping[str, object] | None, +) -> None: + member_budget: Final = LiteLLM_BudgetTable(model_max_budget=dict(member_caps) if member_caps is not None else None) + default_budget: Final = LiteLLM_BudgetTable( + model_max_budget=dict(default_caps) if default_caps is not None else None + ) + membership: Final = LiteLLM_TeamMembership( + user_id="", team_id="", budget_id=member_id, litellm_budget_table=member_budget + ) + member_read: Final = AsyncMock(return_value=membership) + default_read: Final = AsyncMock(return_value=default_budget) + monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(auth_reference, "get_team_membership", member_read) + monkeypatch.setattr(legacy, "get_team_membership", member_read) + monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(auth_reference, "get_team_member_default_budget", default_read) + monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) + token: Final = UserAPIKeyAuth(user_id="", team_id="", team_member_model_max_budget={"old": {}}) + prior: Final = token.team_member_model_max_budget + team: Final = LiteLLM_TeamTable(team_id="", metadata={"team_member_budget_id": default_id}) + fn: Final = ( + adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + await fn(token, team, Mock(), UserApiKeyCache(), False, None) + override: Final = member_id is not None and member_id != default_id and bool(member_caps) + expected: Final = ( + getattr(member_budget, "model_max_budget", None) + if override + else getattr(default_budget, "model_max_budget", None) + if default_id is not None + else prior + ) + assert token.team_member_model_max_budget is expected + member_read.assert_awaited_once() + assert default_read.await_count == (0 if override or default_id is None else 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("skip", (False, True)) +@pytest.mark.parametrize("during_default", (False, True)) +async def test_member_lookup_error_policy( + native: bool, monkeypatch: pytest.MonkeyPatch, skip: bool, during_default: bool +) -> None: + error: Final = RuntimeError("budget read failed") + member_read: Final = AsyncMock(return_value=None, side_effect=None if during_default else error) + default_read: Final = AsyncMock(side_effect=error) + monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(auth_reference, "get_team_membership", member_read) + monkeypatch.setattr(legacy, "get_team_membership", member_read) + monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(auth_reference, "get_team_member_default_budget", default_read) + monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) + token: Final = UserAPIKeyAuth(user_id="u", team_id="t", team_member_model_max_budget={"old": {}}) + prior: Final = token.team_member_model_max_budget + team: Final = LiteLLM_TeamTable(team_id="t", metadata={"team_member_budget_id": "default"}) + fn: Final = ( + adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + pending: Final = fn(token, team, Mock(), UserApiKeyCache(), skip, None) + if skip: + await pending + else: + with pytest.raises(RuntimeError, match="budget read failed") as raised: + await pending + assert raised.value is error + assert token.team_member_model_max_budget is prior + + +@pytest.mark.asyncio +async def test_member_native_failure_never_becomes_skippable_lookup_error(monkeypatch: pytest.MonkeyPatch) -> None: + error: Final = NativeAuthRuntimeError("native planner failed") + read: Final = AsyncMock(side_effect=error) + monkeypatch.setattr(auth_checks, "get_team_membership", read) + monkeypatch.setattr(auth_reference, "get_team_membership", read) + with pytest.raises(NativeAuthRuntimeError, match="native planner failed") as raised: + await adapter.load_team_member_default_model_budget( + UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None + ) + assert raised.value is error + read.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_member_planner_failure_stops_before_io(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge import _native + + native: Final = Mock(side_effect=RuntimeError("decode failed")) + read: Final = AsyncMock() + monkeypatch.setattr(_native, "request_auth_context_member", native) + monkeypatch.setattr(auth_checks, "get_team_membership", read) + monkeypatch.setattr(auth_reference, "get_team_membership", read) + with pytest.raises(NativeAuthRuntimeError, match="decode failed"): + await adapter.load_team_member_default_model_budget( + UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None + ) + read.assert_not_awaited() + native.assert_called_once() diff --git a/tests/unit/proxy/auth/test_budget_lookup_adapter.py b/tests/unit/proxy/auth/test_budget_lookup_adapter.py new file mode 100644 index 00000000000..38707a68983 --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_lookup_adapter.py @@ -0,0 +1,276 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from math import inf, isnan, nan +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_OrganizationTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks, common_budget_adapter, common_budget_lookups +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +from . import phase1_budget_lookup_reference as reference +from .test_common_budget_adapter import observe + +_NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) +_CALL: Final = TypeAdapter(tuple[str, tuple[object, ...], Mapping[str, object]]) + + +def normalized_calls(events: Mock) -> tuple[tuple[str, tuple[object, ...], Mapping[str, object]], ...]: + return tuple( + ( + name, + args, + MappingProxyType( + {key: "NaN" if isinstance(value, float) and isnan(value) else value for key, value in kwargs.items()} + ), + ) + for name, args, kwargs in (_CALL.validate_python(value) for value in events.mock_calls) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("token_id", "team_id"), ((None, None), (None, "team-org"), ("", "team-org"), ("key-org", "team-org")) +) +@pytest.mark.parametrize(("token_present", "client_present"), ((True, True), (False, True), (True, False))) +@pytest.mark.parametrize("row_state", ("present", "missing", "failed")) +@pytest.mark.parametrize("limit", (None, 1.0, inf)) +async def test_organization_source_and_raw_reads_match_phase1( + monkeypatch: pytest.MonkeyPatch, + token_id: str | None, + team_id: str | None, + token_present: bool, + client_present: bool, + row_state: str, + limit: float | None, +) -> None: + row: Final = LiteLLM_OrganizationTable.model_construct( + organization_id="row", spend=2.0, budget_id="budget", litellm_budget_table=LiteLLM_BudgetTable(max_budget=limit) + ) + reader: Final = AsyncMock( + return_value=row if row_state == "present" else None, + side_effect=RuntimeError("unavailable") if row_state == "failed" else None, + ) + counter: Final = AsyncMock(return_value=2.0) + monkeypatch.setattr(reference, "get_org_object", reader) + monkeypatch.setattr(auth_checks, "get_org_object", reader) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + old: Final = UserAPIKeyAuth(org_id=token_id, token="token", key_alias="alias") + new: Final = old.model_copy(deep=True) + team: Final = LiteLLM_TeamTable(team_id="team", organization_id=team_id) + prisma: Final = Mock(spec=PrismaClient) if client_present else None + cache: Final = UserApiKeyCache() + alerts: Final = AsyncMock() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + new_log: Final = old_log + before: Final = await observe( + lambda: reference.organization_max_budget_check(old if token_present else None, team, prisma, cache, old_log), + old_log, + old, + ) + old_reads: Final = (tuple(reader.await_args_list), tuple(counter.await_args_list)) + reader.reset_mock() + counter.reset_mock() + alerts.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.organization_max_budget_check( + new if token_present else None, team, prisma, cache, new_log + ), + new_log, + new, + ) + assert after == before + assert (tuple(reader.await_args_list), tuple(counter.await_args_list)) == old_reads + assert new.model_dump() == old.model_dump() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loaded", (False, True)) +@pytest.mark.parametrize("member_limit", (None, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("default_id", (None, "", "default", 1)) +@pytest.mark.parametrize("missing_member", (False, True)) +async def test_member_default_and_counter_order_match_phase1( + monkeypatch: pytest.MonkeyPatch, + loaded: bool, + member_limit: float | None, + default_id: object, + missing_member: bool, +) -> None: + row: Final = LiteLLM_TeamMembership( + team_id="team", user_id="user", spend=2.0, litellm_budget_table=LiteLLM_BudgetTable(max_budget=member_limit) + ) + membership: Final = None if missing_member else row + member_read: Final = AsyncMock(return_value=membership) + default_read: Final = AsyncMock(return_value=LiteLLM_BudgetTable(max_budget=1)) + counter: Final = AsyncMock(return_value=2.0) + clock: Final = Mock(return_value=_NOW) + events: Final = Mock() + events.attach_mock(member_read, "membership") + events.attach_mock(default_read, "default") + events.attach_mock(counter, "counter") + events.attach_mock(clock, "clock") + for module in (reference, auth_checks): + monkeypatch.setattr(module, "get_team_membership", member_read) + monkeypatch.setattr(module, "get_team_member_default_budget", default_read) + monkeypatch.setattr(module, "get_utc_datetime", clock) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + team: Final = LiteLLM_TeamTable( + team_id="team", + metadata={ + "team_member_budget_id": default_id, + "team_member_max_budget_alert_emails": {"50": ["owner"]}, + }, + ) + key: Final = UserAPIKeyAuth(user_id="user") + prisma: Final = Mock(spec=PrismaClient) + cache: Final = UserApiKeyCache() + alerts: Final = AsyncMock() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + new_log: Final = old_log + before: Final = await observe( + lambda: reference._check_team_member_budget(team, None, key, prisma, cache, old_log, membership, loaded), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + old_reads: Final = normalized_calls(events) + events.reset_mock() + alerts.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.check_team_member_budget( + team, None, key, prisma, cache, new_log, membership, loaded + ), + new_log, + key, + ) + assert after == before + assert normalized_calls(events) == old_reads + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raw", (None, "", " a,b , ,a", "\x1ca\x1f,\u0085b\u00a0", [" a ", "", "\x1c", 4, "a"], ("a",), 4, {"a": 1}) +) +@pytest.mark.parametrize( + ("limit", "spend", "token"), + ((1, 2, True), (1, 0, True), (0, 0, True), (1, 2, False), (None, 2, True), (nan, 2, True)), +) +async def test_soft_email_normalization_matches_phase1( + raw: object, + limit: float | None, + spend: float, + token: bool, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="team", soft_budget=limit, spend=spend, metadata={"soft_budget_alerting_emails": raw} + ) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe( + lambda: reference._team_soft_budget_check(team, key if token else None, old_log), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + after: Final = await observe( + lambda: common_budget_adapter.team_soft_budget_check(team, key if token else None, new_log), + new_log, + key, + ) + assert after == before + + +@pytest.mark.asyncio +async def test_organization_native_failure_from_lookup_is_terminal(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = NativeAuthRuntimeError("native failed") + reader: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(auth_checks, "get_org_object", reader) + with pytest.raises(NativeAuthRuntimeError) as raised: + await common_budget_lookups.organization_max_budget_check( + UserAPIKeyAuth(org_id="org"), None, Mock(spec=PrismaClient), UserApiKeyCache(), Mock(spec=ProxyLogging) + ) + assert raised.value is failure + assert reader.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_present", "team_id", "token_present", "user_id"), + ( + (False, "team", True, "user"), + (True, None, True, "user"), + (True, "team", False, "user"), + (True, "team", True, None), + (True, "", True, ""), + ), +) +async def test_member_identity_stages_match_phase1( + monkeypatch: pytest.MonkeyPatch, + team_present: bool, + team_id: str | None, + token_present: bool, + user_id: str | None, +) -> None: + team: Final = LiteLLM_TeamTable.model_construct(team_id=team_id) if team_present else None + key: Final = UserAPIKeyAuth(user_id=user_id) + reader: Final = AsyncMock(return_value=None) + clock: Final = Mock(return_value=_NOW) + for module in (reference, auth_checks): + monkeypatch.setattr(module, "get_team_membership", reader) + monkeypatch.setattr(module, "get_utc_datetime", clock) + logging: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + cache: Final = UserApiKeyCache() + before: Final = await observe( + lambda: reference._check_team_member_budget(team, None, key if token_present else None, None, cache, logging), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + logging, + key, + ) + reads: Final = (tuple(reader.await_args_list), clock.call_count) + reader.reset_mock() + clock.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.check_team_member_budget( + team, None, key if token_present else None, None, cache, logging + ), + logging, + key, + ) + assert after == before + assert (tuple(reader.await_args_list), clock.call_count) == reads + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", (False, True)) +async def test_soft_budget_skip_does_not_read_metadata(native: bool) -> None: + team: Final = LiteLLM_TeamTable(team_id="team", soft_budget=10, spend=0) + object.__setattr__(team, "metadata", object()) + alerts: Final = AsyncMock() + logging: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + policy: Final = common_budget_adapter.team_soft_budget_check if native else reference._team_soft_budget_check # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + result: Final = await policy(team, UserAPIKeyAuth(), logging) + assert result is None + alerts.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw", (None, [], 1, "metadata", {1: "ignored", "soft_budget_alerting_emails": " a "})) +async def test_soft_metadata_shape_preserves_phase1_projection(raw: object) -> None: + team: Final = LiteLLM_TeamTable(team_id="team", soft_budget=1, spend=2) + object.__setattr__(team, "metadata", raw) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe( + lambda: reference._team_soft_budget_check(team, key, old_log), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + after: Final = await observe(lambda: common_budget_adapter.team_soft_budget_check(team, key, new_log), new_log, key) + assert after == before diff --git a/tests/unit/proxy/auth/test_budget_pricing_adapter.py b/tests/unit/proxy/auth/test_budget_pricing_adapter.py new file mode 100644 index 00000000000..ff5fdd62a7d --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_pricing_adapter.py @@ -0,0 +1,121 @@ +from types import SimpleNamespace +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm.proxy.auth import budget_pricing_adapter as native +from tests.unit.proxy.auth import phase1_budget_auth_reference as reference + + +def router(input_price: float | None, output_price: float | None, deployments: tuple[object, ...] = ()) -> Mock: + view: Final = Mock() + view.group.return_value = SimpleNamespace(input_cost_per_token=input_price, output_cost_per_token=output_price) + view.zero_cost.return_value = {} + view.deployments.return_value = deployments + return Mock(registry=Mock(return_value=view)) + + +@pytest.mark.parametrize(("input_price", "output_price"), ((None, 0), (0, None), (1, 0), (0, 1))) +def test_priced_or_undefined_group_does_not_read_deployments( + input_price: float | None, output_price: float | None +) -> None: + fake: Final = router(input_price, output_price) + facts: Final = native.read_model_price_facts("model", fake) + assert facts == ModelPriceFacts(group_found=True, input_price=input_price, output_price=output_price) + fake.registry().deployments.assert_not_called() + + +@pytest.mark.parametrize( + "route", + ( + "/v1/models", + "/model/info", + "/health/services", + "/chat/completions", + "/mcp", + "/mcp-rest/tools/list", + "/mcp/tools/call", + "/key/info", + ), +) +@pytest.mark.parametrize("body", (None, {}, {"method": "tools/list"}, {"method": "tools/call"}, {"method": 1})) +def test_route_and_mcp_decision_matches_production_reference(route: str, body: dict[str, object] | None) -> None: + assert native.request_skips_budget_checks( + route, None, None, request_body=body + ) == reference.request_skips_budget_checks(route, None, None, request_body=body) + assert native.route_skips_budget_checks(route) == reference.route_skips_budget_checks(route) + + +def test_missing_models_and_empty_model_list_preserve_exemption_behavior() -> None: + fake: Final = router(0, 0) + assert native.is_model_cost_zero(None, fake) is False + assert native.is_model_cost_zero("model", None) is False + assert native.is_model_cost_zero([], fake) is True + + +@pytest.mark.parametrize( + ("prices", "reads", "exempt"), (((0, 1, 0), 2, False), ((0, 0, 0), 3, True), ((1, 0, 0), 1, False)) +) +def test_raw_exemption_prices_only_native_selected_models_and_preserves_cache_writes( + monkeypatch: pytest.MonkeyPatch, prices: tuple[float, ...], reads: int, exempt: bool +) -> None: + from litellm._auth import budget_pricing + + old: Final = router(0, 0) + new: Final = router(0, 0) + monkeypatch.setitem( + litellm.model_cost, "budget-zero-fixture", {"input_cost_per_token": 0, "output_cost_per_token": 0} + ) + deployment: Final = { + "model_name": "model", + "litellm_params": {"model": "openai/fixture"}, + "model_info": {"id": "budget-zero-fixture"}, + } + old.registry().deployments.return_value = (deployment,) + new.registry().deployments.return_value = (deployment,) + old.registry().resolve.side_effect = lambda model: model + new.registry().resolve.side_effect = lambda model: model + old.registry().group.side_effect = tuple( + SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices + ) + new.registry().group.side_effect = tuple( + SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices + ) + old_cache: Final[dict[str, bool]] = {} + new_cache: Final[dict[str, bool]] = {} + old.registry().zero_cost.return_value = old_cache + new.registry().zero_cost.return_value = new_cache + models: Final = ["first", "second", "third"] + before: Final = reference.request_skips_budget_checks("/chat/completions", models, old) + facts: Final = native.read_budget_exemption_facts("/chat/completions", models, new) + assert budget_pricing.exemption(facts).exempt is before is exempt + assert len(facts.prices) == reads + assert old.registry().group.call_count == new.registry().group.call_count == reads + assert old_cache == new_cache + + +def test_raw_exemption_route_skip_does_not_touch_model_registry() -> None: + fake: Final = router(0, 0) + facts: Final = native.read_budget_exemption_facts("/v1/models", ["first", "second"], fake) + assert facts.prices == () + fake.registry.assert_not_called() + + +def test_raw_cached_price_does_not_read_model_group_or_write_again() -> None: + fake: Final = router(0, 0) + fake.registry().zero_cost.return_value = {"model": False} + facts: Final = native.read_model_cost_facts("model", fake) + assert facts.cached is False + fake.registry().group.assert_not_called() + + +def test_empty_model_list_still_reads_zero_cost_cache() -> None: + fake: Final = router(0, 0) + facts: Final = native.read_budget_exemption_facts("/chat/completions", [], fake) + assert facts.prices == () + assert facts.cache_read is True + fake.registry().zero_cost.assert_called_once_with() + fake.registry().group.assert_not_called() diff --git a/tests/unit/proxy/auth/test_common_budget_adapter.py b/tests/unit/proxy/auth/test_common_budget_adapter.py new file mode 100644 index 00000000000..ace55803f6b --- /dev/null +++ b/tests/unit/proxy/auth/test_common_budget_adapter.py @@ -0,0 +1,297 @@ +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timezone +from math import inf +from typing import Final, Literal, Protocol +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth.common_budgets import BudgetFacts +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + BudgetLimitEntry, + LiteLLM_EndUserTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth import budget_windows_adapter, common_budget_lookups +from litellm.proxy.auth import common_budget_adapter as native +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import ProxyLogging +from tests.unit.proxy.auth import phase1_budget_auth_reference as reference + + +@dataclass(frozen=True, slots=True) +class Observed: + error: tuple[object, ...] | None + alerts: tuple[object, ...] + throttle: float | None + + +class AlertLog(Protocol): + @property + def budget_alerts(self) -> AsyncMock: ... + + +async def observe(run: Callable[[], Awaitable[None]], logging: AlertLog, key: UserAPIKeyAuth) -> Observed: + try: + await run() + except litellm.BudgetExceededError as error: + failure: Final = (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id) + await asyncio.sleep(0) + return Observed(failure, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct) + await asyncio.sleep(0) + return Observed(None, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2))) +async def test_team_error_envelope_and_alert_projection_match_reference( + monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + team: Final = LiteLLM_TeamTable(team_id="team", max_budget=limit, spend=0) + key: Final = UserAPIKeyAuth(token="hash", team_id="team") + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.team_max_budget_check(team, key, old_log), old_log, key) + after: Final = await observe(lambda: native.team_max_budget_check(team, key, new_log), new_log, key) + assert after == before + assert reader.await_count == (0 if limit is None else 2) + assert reader.await_args_list[:1] == reader.await_args_list[1:] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2))) +@pytest.mark.parametrize("throttle", (None, 0.5, 2)) +async def test_key_error_envelope_alert_and_throttle_mutation_match_reference( + monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float, throttle: float | None +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", throttle) + old: Final = UserAPIKeyAuth( + token="hash", key_alias="alias", max_budget=limit, tpm_limit=10, metadata={"throttle_on_budget_exceeded": True} + ) + new: Final = old.model_copy(deep=True) + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.virtual_key_max_budget_check(old, old_log), old_log, old) + after: Final = await observe(lambda: native.virtual_key_max_budget_check(new, new_log), new_log, new) + assert after == before + assert reader.await_args_list[:1] == reader.await_args_list[1:] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("/chat/completions", "/key/info")) +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 1), (1, 2), (0, 0), (-1, 0))) +async def test_end_user_info_skip_and_strict_ceiling_match_reference( + monkeypatch: pytest.MonkeyPatch, route: str, limit: float | None, spend: float +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + user: Final = LiteLLM_EndUserTable( + user_id="end", blocked=False, litellm_budget_table=LiteLLM_BudgetTable(max_budget=limit) + ) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.check_end_user_budget(user, route), old_log, key) + after: Final = await observe(lambda: native.check_end_user_budget(user, route), new_log, key) + assert after == before + assert reader.await_args_list[:1] == reader.await_args_list[1:] + + +@pytest.mark.asyncio +async def test_spend_io_failure_is_terminal_and_does_not_replay() -> None: + failure: Final = RuntimeError("counter unavailable") + reader: Final = AsyncMock(side_effect=failure) + with pytest.raises(RuntimeError, match="counter unavailable") as raised: + await native.read_budget(BudgetFacts("Key", 1), native.CounterRead("spend:key:hash", 0, 1), reader) + assert raised.value is failure + assert reader.await_count == 1 + + +@pytest.mark.parametrize( + ("member", "default", "increase", "expiry"), + ( + (1, 5, 2, datetime(2025, 1, 2)), + (None, 5, 2, datetime(2025, 1, 2)), + (1, 5, 2, datetime(2025, 1, 1, tzinfo=timezone.utc)), + (None, 0, 2, datetime(2025, 1, 2)), + (None, -1, 2, None), + (0, 5, None, None), + ), +) +def test_member_budget_selection_and_temporary_expiry_match_reference( + member: float | None, default: float | None, increase: float | None, expiry: datetime | None +) -> None: + now: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) + row: Final = LiteLLM_BudgetTable(max_budget=member, temp_budget_increase=increase, temp_budget_expiry=expiry) + team: Final = LiteLLM_BudgetTable(max_budget=default) + assert native.resolve_effective_team_member_budget( + row, team, now + ) == reference.resolve_effective_team_member_budget(row, team, now) + + +@pytest.mark.parametrize( + "policy", (reference.resolve_effective_team_member_budget, native.resolve_effective_team_member_budget) +) +def test_naive_now_preserves_datetime_error(policy: Callable[..., float | None]) -> None: + row: Final = LiteLLM_BudgetTable(max_budget=1, temp_budget_increase=2, temp_budget_expiry=datetime(2025, 1, 2)) + with pytest.raises(TypeError, match="offset-naive"): + policy(row, None, datetime(2025, 1, 1)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Key", "Team")) +@pytest.mark.parametrize("as_model", (False, True)) +@pytest.mark.parametrize("spend", (1.0, 10.0)) +async def test_budget_window_objects_preserve_counter_contract_and_denial( + monkeypatch: pytest.MonkeyPatch, + entity: str, + as_model: bool, + spend: float, +) -> None: + entry: Final = BudgetLimitEntry( + budget_duration="24h", max_budget=10.0, reset_at=datetime(2025, 1, 2, tzinfo=timezone.utc) + ) + raw: Final = entry if as_model else entry.model_dump() + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + key: Final = UserAPIKeyAuth(token="key") + team: Final = LiteLLM_TeamTable(team_id="team") + object.__setattr__(key, "budget_limits", [raw]) + object.__setattr__(team, "budget_limits", [raw]) + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + old_call: Final = ( + ( + lambda: reference._virtual_key_multi_budget_check(key) # pyright: ignore[reportPrivateUsage] # Phase 1 reference + ) + if entity == "Key" + else ( + lambda: reference._team_multi_budget_check(team) # pyright: ignore[reportPrivateUsage] # Phase 1 reference + ) + ) + new_call: Final = ( + (lambda: budget_windows_adapter.virtual_key_multi_budget_check(key)) + if entity == "Key" + else (lambda: budget_windows_adapter.team_multi_budget_check(team)) + ) + before: Final = await observe(old_call, old_log, key) + expected_read: Final = reader.await_args + reader.reset_mock() + after: Final = await observe(new_call, new_log, key) + assert after == before + assert (after.error is not None) is (spend == 10.0) + assert reader.await_args_list == [expected_read] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entity", ("Key", "Team", "User", "Project", "TeamMember", "Organization", "EndUser", "Tag", "AccessGroup") +) +async def test_non_window_spend_reads_omit_unselected_options( + monkeypatch: pytest.MonkeyPatch, + entity: str, +) -> None: + reader: Final = AsyncMock(return_value=3.0) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + authoritative: Final = entity in {"EndUser", "Tag", "AccessGroup"} + assert ( + await native.read_counter(native.CounterRead(f"spend:{entity}:id", 2.0, 10.0, authoritative=authoritative)) + == 3.0 + ) + reader.assert_awaited_once_with( + counter_key=f"spend:{entity}:id", + fallback_spend=2.0, + max_budget=10.0, + **({"fallback_authoritative": True} if authoritative else {}), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Key", "Team")) +async def test_window_spend_reads_keep_explicit_missing_metadata( + monkeypatch: pytest.MonkeyPatch, + entity: str, +) -> None: + reader: Final = AsyncMock(return_value=3.0) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + assert ( + await native.read_counter( + native.CounterRead( + "window", + 0.0, + 10.0, + window_entity_type=entity, + window_duration="invalid", + ) + ) + == 3.0 + ) + reader.assert_awaited_once_with( + counter_key="window", + fallback_spend=0.0, + max_budget=10.0, + window_entity_type=entity, + window_entity_id=None, + window_duration="invalid", + window_start=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Team", "Project", "TeamMember")) +@pytest.mark.parametrize("spend", (1.0, 11.0)) +async def test_primary_scope_denial_with_baseline_spend_signature( + monkeypatch: pytest.MonkeyPatch, + entity: Literal["Team", "Project", "TeamMember"], + spend: float, +) -> None: + async def current_spend(counter_key: str, fallback_spend: float, max_budget: float | None = None) -> float: + return fallback_spend + + reader: Final = AsyncMock(side_effect=current_spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + team: Final = LiteLLM_TeamTable(team_id="team", spend=spend, max_budget=10.0) + key: Final = UserAPIKeyAuth(user_id="user", team_id="team") + row: Final = LiteLLM_BudgetTable(max_budget=10.0) + logging: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + project: Final = LiteLLM_ProjectTableCachedObj( + project_id="project", team_id="team", spend=spend, litellm_budget_table=row + ) + member: Final = LiteLLM_TeamMembership(user_id="user", team_id="team", spend=spend, litellm_budget_table=row) + invocation: Final = { + "Team": lambda: native.team_max_budget_check(team, None, logging), + "Project": lambda: native.project_max_budget_check(project, None, logging), + "TeamMember": lambda: common_budget_lookups.check_team_member_budget( + team, + None, + key, + None, + UserApiKeyCache(), + logging, + team_membership=member, + team_membership_loaded=True, + ), + }[entity] + result: Final = await observe(invocation, logging, key) + assert (result.error is not None) is (spend > 10.0) + reader.assert_awaited_once_with( + counter_key={ + "Team": "spend:team:team", + "Project": "spend:project:project", + "TeamMember": "spend:team_member:user:team", + }[entity], + fallback_spend=spend, + max_budget=10.0, + ) diff --git a/tests/unit/proxy/auth/test_common_budget_lookups.py b/tests/unit/proxy/auth/test_common_budget_lookups.py new file mode 100644 index 00000000000..b034dafcf47 --- /dev/null +++ b/tests/unit/proxy/auth/test_common_budget_lookups.py @@ -0,0 +1,44 @@ +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +from litellm._auth import common_budgets +from litellm.proxy import proxy_server +from litellm.proxy.auth.common_budget_lookups import _read_access_group_rows +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("limit", "reads", "denied"), ((None, 0, False), (0, 0, False), (-1, 0, False), (1, 1, True), (2, 2, False)) +) +async def test_raw_access_group_snapshot_stops_reads_at_native_denial( + monkeypatch: pytest.MonkeyPatch, + limit: float | None, + reads: int, + denied: bool, +) -> None: + counter: Final = AsyncMock(return_value=1.0) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + rows: Final = { + "first": ModelAccessGroupBudget(access_group_name="first", max_budget=limit), + "second": ModelAccessGroupBudget(access_group_name="second", max_budget=limit), + } + snapshots: Final = await _read_access_group_rows(("first", "second"), rows) + assert common_budgets.check(snapshots[0].facts).denied is denied + assert len(snapshots) == (1 if denied else 2) + assert snapshots[0].name == "first" + assert counter.await_count == reads + + +@pytest.mark.asyncio +async def test_access_group_counter_failure_is_terminal_before_next_group(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = RuntimeError("counter unavailable") + counter: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + row: Final = ModelAccessGroupBudget(access_group_name="first", max_budget=1) + with pytest.raises(RuntimeError, match="counter unavailable") as raised: + await _read_access_group_rows(("first", "second"), {"first": row, "second": row}) + assert raised.value is failure + assert counter.await_count == 1 diff --git a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py index 0c3487255f6..030a8570913 100644 --- a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py @@ -115,7 +115,7 @@ async def test_custom_auth_enforces_end_user_budget_when_common_checks_skipped() with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + "litellm.proxy.auth.auth_checks.get_end_user_object", new_callable=AsyncMock, return_value=over_budget_end_user, ), diff --git a/tests/unit/proxy/auth/test_end_user_context_adapter.py b/tests/unit/proxy/auth/test_end_user_context_adapter.py new file mode 100644 index 00000000000..8eaee4db5a9 --- /dev/null +++ b/tests/unit/proxy/auth/test_end_user_context_adapter.py @@ -0,0 +1,437 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks +from litellm.proxy.auth import end_user_context_adapter as adapter +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from tests.unit.proxy.auth import phase1_budget_auth_reference as auth_reference +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy + + +@pytest.fixture(params=(False, True), ids=("reference", "native")) +def native(request: pytest.FixtureRequest) -> bool: + return request.param is True # pyright: ignore[reportAny] # boolean fixture parameter + + +@pytest.mark.parametrize( + "value,expected", ((None, None), ("", None), (" ", " "), ("budget", "budget"), (0, None), (False, None), ({}, None)) +) +def test_metadata_id_keeps_only_nonempty_strings(native: bool, value: object, expected: str | None) -> None: + function: Final = adapter.get_key_end_user_budget_id if native else auth_reference.get_key_end_user_budget_id + assert function({"end_user_budget_id": value}) == expected + assert function(None) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("explicit", (None, "", "key")) +@pytest.mark.parametrize("global_id", (None, "", "global")) +@pytest.mark.parametrize("db_present", (False, True)) +async def test_default_budget_id_and_cache_selection( + native: bool, explicit: str | None, global_id: str | None, db_present: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + cache: Final = UserApiKeyCache() + selected: Final = explicit if explicit is not None else global_id + budget: Final = LiteLLM_BudgetTable(budget_id=selected, max_budget=0) + await cache.async_set_cache(key=f"default_end_user_budget:{selected}", value=budget, model_type=LiteLLM_BudgetTable) + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget + result: Final = await function(Mock() if db_present else None, cache, budget_id=explicit) + assert result == (budget if db_present and selected is not None else None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own_id", (None, "", "own")) +@pytest.mark.parametrize("own_present", (False, True)) +@pytest.mark.parametrize("key_id", (None, "", "key")) +@pytest.mark.parametrize("global_id", (None, "", "global")) +@pytest.mark.parametrize("found", (False, True)) +async def test_default_application_preserves_own_budget_and_copies_only_resolved_defaults( + native: bool, + own_id: str | None, + own_present: bool, + key_id: str | None, + global_id: str | None, + found: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + own: Final = LiteLLM_BudgetTable(budget_id="own", max_budget=2) + default: Final = LiteLLM_BudgetTable(budget_id="default", max_budget=3) if found else None + row: Final = LiteLLM_EndUserTable( + user_id="eu", blocked=False, budget_id=own_id, litellm_budget_table=own if own_present else None + ) + load: Final = AsyncMock(return_value=default) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", load) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", load) + function: Final = ( + adapter.apply_default_budget_to_end_user if native else auth_reference._apply_default_budget_to_end_user # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + ) # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + result: Final = await function(row, Mock(), UserApiKeyCache(), key_end_user_budget_id=key_id) + unchanged: Final = (own_id is not None and own_present) or (key_id is None and global_id is None) or not found + assert (result is row) is unchanged + assert result.litellm_budget_table is (row.litellm_budget_table if unchanged else default) + assert row.litellm_budget_table is (own if own_present else None) + expected_ids: Final = ( + () + if own_id is not None and own_present + else ( + ((key_id,) if key_id is not None else ()) + + ((global_id,) if global_id is not None and (key_id is None or not found) else ()) + ) + ) + assert tuple(call.kwargs.get("budget_id", global_id) for call in load.await_args_list) == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("global_id", (None, "", "global")) +@pytest.mark.parametrize("key_id", (None, "", "key")) +@pytest.mark.parametrize("validate", (False, True)) +@pytest.mark.parametrize("max_budget", (None, 0)) +@pytest.mark.parametrize("ids", (None, frozenset[str](), frozenset(("eu",)))) +async def test_unrestricted_registry_requires_absence_of_all_budget_and_validation_constraints( + native: bool, + global_id: str | None, + key_id: str | None, + validate: bool, + max_budget: int | None, + ids: frozenset[str] | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", validate) + load: Final = AsyncMock(return_value=ids) + monkeypatch.setattr(auth_checks, "_load_end_user_restricted_registry", load) + monkeypatch.setattr(auth_reference, "_load_end_user_restricted_registry", load) + function: Final = adapter.end_user_is_known_unrestricted if native else auth_reference._end_user_is_known_unrestricted # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + result: Final = await function("eu", Mock(), UserApiKeyCache(), max_budget, key_id) + blocked: Final = global_id is not None or key_id is not None or validate or max_budget is not None + assert result is (not blocked and ids is not None and "eu" not in ids) + assert load.await_count == (0 if blocked else 1) + + +@pytest.mark.parametrize("value", (None, 0, 1)) +@pytest.mark.parametrize("models", (None, {}, {"model": {"budget_limit": 0}})) +def test_budget_parameter_projection_keeps_zero_and_empty_model_limits( + native: bool, + value: int | None, + models: Mapping[str, object] | None, +) -> None: + budget: Final = LiteLLM_BudgetTable( + tpm_limit=value, + rpm_limit=value, + tpd_limit=value, + max_budget=value, + model_max_budget=dict(models) if models is not None else None, + ) + params: Final[dict[str, object]] = {"existing": True} # mutable-ok: legacy output argument + if native: + adapter.apply_budget_limits_to_end_user_params(params, budget, "eu") + else: + legacy._apply_budget_limits_to_end_user_params(params, budget, "eu") # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # Phase 1 unparameterized reference dictionary + expected: Final = { + "existing": True, + **( + { + name: value + for name in ("end_user_tpm_limit", "end_user_rpm_limit", "end_user_tpd_limit", "end_user_max_budget") + } + if value is not None + else {} + ), + **({"end_user_model_max_budget": models} if models is not None else {}), + } + assert params == expected + assert tuple(params) == tuple(expected) + if models is not None: + assert params["end_user_model_max_budget"] is getattr(budget, "model_max_budget") + + +@pytest.mark.parametrize("value", (None, 0, "opaque")) +def test_token_projection_preserves_unset_limits_and_never_copies_max_budget(native: bool, value: object) -> None: + token: Final = UserAPIKeyAuth( + end_user_id="old", + end_user_tpm_limit=9, + end_user_rpm_limit=9, + end_user_tpd_limit=9, + end_user_max_budget=9, + end_user_model_max_budget={"old": {}}, + allowed_model_region="eu", + ) + params: Final = { + name: value + for name in ( + "end_user_tpm_limit", + "end_user_rpm_limit", + "end_user_tpd_limit", + "end_user_max_budget", + "end_user_model_max_budget", + "allowed_model_region", + ) + } + result: Final = ( + adapter.update_valid_token_with_end_user_params(token, params) + if native + else legacy.update_valid_token_with_end_user_params(token, params) # pyright: ignore[reportUnknownMemberType] # Phase 1 unparameterized reference dictionary + ) + assert result is token + assert token.end_user_id is None + assert token.end_user_tpm_limit == (9 if value is None else value) + assert token.end_user_rpm_limit == (9 if value is None else value) + assert token.end_user_tpd_limit == (9 if value is None else value) + assert token.allowed_model_region == ("eu" if value is None else value) + assert token.end_user_model_max_budget == ({"old": {}} if value is None else value) + assert token.end_user_max_budget == 9 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cached", ("valid", "invalid", None, 0)) +@pytest.mark.parametrize("global_id,key_id", ((None, None), ("", None), ("global", None), (None, ""))) +@pytest.mark.parametrize("end_found,user_found", ((False, False), (True, False), (False, True))) +@pytest.mark.parametrize("identity", ("eu", "email@example.test")) +async def test_validation_uses_raw_row_presence_and_preserves_unknown_ids_with_defaults( + native: bool, + cached: object, + global_id: str | None, + key_id: str | None, + end_found: bool, + user_found: bool, + identity: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=f"end_user_validation:{identity}", value=cached) + end: Final = AsyncMock(return_value=LiteLLM_EndUserTable(user_id=identity, blocked=False) if end_found else None) + user: Final = AsyncMock(return_value=LiteLLM_UserTable(user_id=identity) if user_found else None) + monkeypatch.setattr(auth_checks, "get_end_user_object", end) + monkeypatch.setattr(auth_reference, "get_end_user_object", end) + monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) + function: Final = ( + adapter.resolve_and_validate_end_user_id if native else auth_reference.resolve_and_validate_end_user_id + ) + result: Final = await function(identity, Mock(), cache, key_end_user_budget_id=key_id) + valid: Final = cached == "valid" or (cached != "invalid" and (end_found or user_found)) + assert result == (identity if valid or bool(global_id) or key_id is not None else None) + assert end.await_count == (0 if cached in ("valid", "invalid") else 1) + assert user.await_count == (0 if cached in ("valid", "invalid") or end_found else 1) + if user.await_count: + assert user.await_args is not None + assert user.await_args.kwargs["user_email"] == (identity if "@" in identity else None) + assert await cache.async_get_cache(key=f"end_user_validation:{identity}") == ("valid" if valid else "invalid") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("enabled,db", ((False, False), (False, True), (True, False))) +async def test_validation_disabled_or_without_db_skips_all_lookups( + native: bool, + identity: str | None, + enabled: bool, + db: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", enabled) + cache: Final = Mock(spec=UserApiKeyCache) + read: Final = AsyncMock(side_effect=AssertionError("cache must not be read")) + cache.async_get_cache = read + function: Final = ( + adapter.resolve_and_validate_end_user_id if native else auth_reference.resolve_and_validate_end_user_id + ) + assert await function(identity, Mock() if db else None, cache) == identity + read.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("row_mode", ("missing", "empty", "budget")) +@pytest.mark.parametrize("initial_max", (None, 0, 7)) +@pytest.mark.parametrize( + "key_id,global_id,found", + ((None, None, False), ("key", None, True), (None, "global", True), ("key", "global", False)), +) +async def test_lookup_uses_row_before_default_and_only_fills_missing_max_budget( + native: bool, + row_mode: str, + initial_max: int | None, + key_id: str | None, + global_id: str | None, + found: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + budget: Final = LiteLLM_BudgetTable(budget_id="budget", max_budget=5, tpm_limit=0, model_max_budget={}) + row: Final = ( + LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget if row_mode == "budget" else None) + if row_mode != "missing" + else None + ) + end: Final = AsyncMock(return_value=row) + default: Final = AsyncMock(return_value=budget if found else None) + monkeypatch.setattr(auth_checks, "get_end_user_object", end) + monkeypatch.setattr(auth_reference, "get_end_user_object", end) + monkeypatch.setattr(legacy, "get_end_user_object", end) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", default) + token: Final = UserAPIKeyAuth( + end_user_id="eu", end_user_max_budget=initial_max, end_user_tpm_limit=9, metadata={"end_user_budget_id": key_id} + ) + result: Final = ( + await adapter.lookup_end_user_and_apply_budget( + token, "/chat/completions", None, Mock(), UserApiKeyCache(), None + ) + if native + else await legacy._lookup_end_user_and_apply_budget( # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # Phase 1 untyped reference boundary + token, "/chat/completions", None, Mock(), UserApiKeyCache(), None + ) # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # Phase 1 untyped reference boundary + ) + actual_token, actual_row = result + assert actual_token is token and actual_row is row + default_used: Final = row is None and (key_id is not None or global_id is not None) and found + assert token.end_user_max_budget == (5 if default_used and initial_max is None else initial_max) + assert token.end_user_tpm_limit == (0 if row_mode == "budget" or default_used else 9) + assert token.end_user_model_max_budget == ({} if row_mode == "budget" or default_used else None) + assert end.await_count == 1 + assert default.await_count == ( + 0 + if row is not None + else (int(key_id is not None) + int(global_id is not None and (key_id is None or not found))) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_failure", (False, True)) +async def test_lookup_io_failure_does_not_select_defaults_and_native_failure_is_terminal( + native_failure: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + failure: Final = NativeAuthRuntimeError("native") if native_failure else RuntimeError("io") + load: Final = AsyncMock(side_effect=failure) + default: Final = AsyncMock(side_effect=AssertionError("no defaults after lookup error")) + monkeypatch.setattr(auth_checks, "get_end_user_object", load) + monkeypatch.setattr(auth_reference, "get_end_user_object", load) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", default) + token: Final = UserAPIKeyAuth(end_user_id="eu", metadata={"end_user_budget_id": "key"}) + if native_failure: + with pytest.raises(NativeAuthRuntimeError) as raised: + await adapter.lookup_end_user_and_apply_budget(token, "", None, Mock(), UserApiKeyCache(), None) + assert raised.value is failure + else: + assert await adapter.lookup_end_user_and_apply_budget(token, "", None, Mock(), UserApiKeyCache(), None) == ( + token, + None, + ) + default.assert_not_called() + + +@pytest.mark.asyncio +async def test_native_validation_failure_stops_before_user_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + failure: Final = NativeAuthRuntimeError("native") + monkeypatch.setattr(auth_checks, "get_end_user_object", AsyncMock(side_effect=failure)) + monkeypatch.setattr(auth_reference, "get_end_user_object", AsyncMock(side_effect=failure)) + user: Final = AsyncMock() + monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) + with pytest.raises(NativeAuthRuntimeError) as raised: + await adapter.resolve_and_validate_end_user_id("eu", Mock(), UserApiKeyCache()) + assert raised.value is failure + user.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("found", (False, True)) +async def test_default_budget_cache_miss_reads_database_and_caches_only_found_rows( + native: bool, + found: bool, +) -> None: + cache: Final = UserApiKeyCache() + row: Final = LiteLLM_BudgetTable(budget_id="budget", max_budget=0) if found else None + read: Final = AsyncMock(return_value=row) + db: Final = Mock(db=Mock(litellm_budgettable=Mock(find_unique=read))) + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget + assert await function(db, cache, budget_id="budget") == row + read.assert_awaited_once_with(where={"budget_id": "budget"}) + assert await cache.async_get_cache(key="default_end_user_budget:budget", model_type=LiteLLM_BudgetTable) == row + if found: + assert await function(db, cache, budget_id="budget") == row + read.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_failure", (False, True)) +async def test_default_budget_database_failure_keeps_native_fault_terminal(native_failure: bool) -> None: + failure: Final = NativeAuthRuntimeError("native") if native_failure else RuntimeError("io") + read: Final = AsyncMock(side_effect=failure) + db: Final = Mock(db=Mock(litellm_budgettable=Mock(find_unique=read))) + if native_failure: + with pytest.raises(NativeAuthRuntimeError) as raised: + await adapter.get_default_end_user_budget(db, UserApiKeyCache(), budget_id="budget") + assert raised.value is failure + else: + assert await adapter.get_default_end_user_budget(db, UserApiKeyCache(), budget_id="budget") is None + read.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_default_budget_cache_failure_is_not_swallowed(native: bool) -> None: + failure: Final = RuntimeError("cache unavailable") + cache: Final = Mock(spec=UserApiKeyCache) + read: Final = AsyncMock(side_effect=failure) + cache.async_get_cache = read + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget + with pytest.raises(RuntimeError) as raised: + await function(Mock(), cache, budget_id="budget") + assert raised.value is failure + read.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own_failure", (False, True)) +@pytest.mark.parametrize("user_found", (False, True)) +async def test_validation_lookup_errors_select_secondary_identity_read( + native: bool, + own_failure: bool, + user_found: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + own: Final = AsyncMock(side_effect=RuntimeError("end-user read")) if own_failure else AsyncMock(return_value=None) + user: Final = ( + AsyncMock(return_value=LiteLLM_UserTable(user_id="eu")) + if user_found + else AsyncMock(side_effect=RuntimeError("user read")) + ) + monkeypatch.setattr(auth_checks, "get_end_user_object", own) + monkeypatch.setattr(auth_reference, "get_end_user_object", own) + monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) + function: Final = adapter.end_user_id_exists_in_db if native else auth_reference._end_user_id_exists_in_db # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + assert await function("eu", Mock(), UserApiKeyCache()) is user_found + own.assert_awaited_once() + user.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_lookup_budget_denial_preserves_exception_identity(native: bool, monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = litellm.BudgetExceededError(current_cost=2, max_budget=1) + read: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(auth_checks, "get_end_user_object", read) + monkeypatch.setattr(auth_reference, "get_end_user_object", read) + monkeypatch.setattr(legacy, "get_end_user_object", read) + invocation: Final = ( + adapter.lookup_end_user_and_apply_budget(UserAPIKeyAuth(), "", None, Mock(), UserApiKeyCache(), None) + if native + else legacy._lookup_end_user_and_apply_budget(UserAPIKeyAuth(), "", None, Mock(), UserApiKeyCache(), None) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # Phase 1 untyped reference boundary + ) + with pytest.raises(litellm.BudgetExceededError) as raised: + await invocation + assert raised.value is failure + read.assert_awaited_once() diff --git a/tests/unit/proxy/auth/test_key_budget_context_adapter.py b/tests/unit/proxy/auth/test_key_budget_context_adapter.py new file mode 100644 index 00000000000..843f74a30f5 --- /dev/null +++ b/tests/unit/proxy/auth/test_key_budget_context_adapter.py @@ -0,0 +1,108 @@ +import math +from datetime import datetime, timezone +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm._auth import budget_context as core +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth import key_budget_context_adapter as adapter +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy +from tests.unit.proxy.auth import phase1_budget_throttle_reference as budget_throttle + +_NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) + + +@pytest.mark.parametrize("native", (False, True), ids=("reference", "native")) +@pytest.mark.parametrize( + "increase", (None, False, True, 0, -3, 3, float("nan"), float("inf"), "", "3", [], [3], {}, {"n": 3}) +) +@pytest.mark.parametrize( + "expiry", ("2024-12-31", "2025-01-01T00:00:00+00:00", "2025-01-01T00:00:00.000001Z", "2026-01-01T09:00:00+09:00") +) +@pytest.mark.parametrize("limit", (None, 0, 10)) +def test_temporary_budget_reference_native( + native: bool, monkeypatch: pytest.MonkeyPatch, increase: object, expiry: str, limit: float | None +) -> None: + now: Final = Mock(return_value=_NOW) + clock: Final = Mock(wraps=datetime, now=now) + monkeypatch.setattr(adapter, "datetime", clock) + monkeypatch.setattr(legacy, "datetime", clock) + token: Final = UserAPIKeyAuth( + max_budget=limit, metadata={"temp_budget_increase": increase, "temp_budget_expiry": expiry} + ) + fn: Final = ( + adapter.update_key_budget_with_temp_budget_increase + if native + else legacy._update_key_budget_with_temp_budget_increase # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + active: Final = expiry.startswith("2026") or expiry == "2025-01-01T00:00:00.000001Z" + changes: Final = limit is not None and active and bool(increase) + if changes and not isinstance(increase, (int, float)): + with pytest.raises(TypeError, match="unsupported operand type"): + fn(token) + return + result: Final = fn(token) + assert (result is token) is not changes + if changes: + assert limit is not None and isinstance(increase, (int, float)) + expected: Final = limit + increase + assert result.max_budget is not None + assert result.max_budget == expected or (math.isnan(result.max_budget) and math.isnan(expected)) + assert token.max_budget == limit + assert now.call_count == (0 if limit is None else 1) + + +@pytest.mark.parametrize("native", (False, True), ids=("reference", "native")) +@pytest.mark.parametrize("expiry", (None, 5, "", "invalid")) +@pytest.mark.parametrize("increase_present", (False, True)) +@pytest.mark.parametrize("limit", (None, 1)) +def test_temporary_metadata_error_timing( + native: bool, monkeypatch: pytest.MonkeyPatch, expiry: object, increase_present: bool, limit: float | None +) -> None: + token: Final = UserAPIKeyAuth( + max_budget=limit, + metadata={"temp_budget_expiry": expiry, **({"temp_budget_increase": 0} if increase_present else {})}, + ) + now: Final = Mock(return_value=_NOW) + clock: Final = Mock(wraps=datetime, now=now) + monkeypatch.setattr(adapter, "datetime", clock) + monkeypatch.setattr(legacy, "datetime", clock) + fn: Final = ( + adapter.update_key_budget_with_temp_budget_increase + if native + else legacy._update_key_budget_with_temp_budget_increase # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + if limit is not None and increase_present: + with pytest.raises(ValueError if isinstance(expiry, str) else TypeError): + fn(token) + else: + assert fn(token) is token + now.assert_not_called() + + +@pytest.mark.parametrize("pct", (None, False, True, "", "0.5", 0, -0.5, 0.25, 1, 2, float("nan"), float("inf"))) +@pytest.mark.parametrize("opt_in", (None, False, True, 1, "true")) +@pytest.mark.parametrize("limit", (None, 0, 1)) +def test_throttle_reference_native( + monkeypatch: pytest.MonkeyPatch, pct: object, opt_in: object, limit: int | None +) -> None: + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", pct) + token: Final = UserAPIKeyAuth(metadata={"throttle_on_budget_exceeded": opt_in}, tpm_limit=limit) + assert adapter.budget_throttle_percentage() == budget_throttle.budget_throttle_percentage() + assert adapter.should_throttle_budget_exceeded(token) == budget_throttle.should_throttle_budget_exceeded(token) + expected: Final = ( + budget_throttle.budget_throttle_percentage() if budget_throttle.should_throttle_budget_exceeded(token) else None + ) + assert adapter.inspect_key_budget_throttle(token) == expected + + +@pytest.mark.parametrize("tpm", (None, 0, 1)) +@pytest.mark.parametrize("rpm", (None, 0, 1)) +@pytest.mark.parametrize("opt_in", (False, True)) +def test_throttle_native_selects_config_read(tpm: int | None, rpm: int | None, opt_in: bool) -> None: + assert core.throttle(core.ThrottleFacts(opt_in, tpm, rpm)).read_percentage is ( + opt_in and (tpm is not None or rpm is not None) + ) diff --git a/tests/unit/proxy/auth/test_managed_model_adapter.py b/tests/unit/proxy/auth/test_managed_model_adapter.py new file mode 100644 index 00000000000..d96e57ae375 --- /dev/null +++ b/tests/unit/proxy/auth/test_managed_model_adapter.py @@ -0,0 +1,84 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ModelAccessDeniedProxyException, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling +from litellm.proxy.auth import managed_model_adapter +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import ProxyLogging +from litellm.types.agents import AgentResponse + + +@pytest.mark.asyncio +@pytest.mark.parametrize("models", (None, "model", (), ("allowed",))) +async def test_managed_model_grants_preserve_public_denial(models: object) -> None: + token: Final = UserAPIKeyAuth().model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": models} + ) + } + ) + cache: Final = UserApiKeyCache() + authorize: Final = managed_model_adapter.can_agent_call_model + if models == ("allowed",): + with pytest.raises(ModelAccessDeniedProxyException) as denied: + await authorize("denied", token, None, None, cache, ProxyLogging(user_api_key_cache=cache)) + assert ( + denied.value.internal_message + == "agent not allowed to access model. This agent can only access models=['allowed']. Tried to access denied" + ) + else: + with pytest.raises(HTTPException) as denied: + await authorize("denied", token, None, None, cache, ProxyLogging(user_api_key_cache=cache)) + assert (denied.value.status_code, denied.value.detail) == (403, "This agent has no model grants") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("grants", ((), ("allowed",))) +async def test_agent_empty_ceiling_and_model_denial_envelopes(grants: tuple[str, ...]) -> None: + async def resolve(_agent: str) -> AgentAccessGroupCeiling: + return AgentAccessGroupCeiling(("group",), frozenset(grants), frozenset(), frozenset()) + + authorize: Final = managed_model_adapter.check_agent_access_group_model_access + with pytest.raises(ModelAccessDeniedProxyException) as denied: + await authorize("denied", UserAPIKeyAuth(agent_id="agent"), None, resolve) + assert denied.value.internal_message == ( + "agent agent access groups grant no models" + if not grants + else "agent not allowed to access model. This agent can only access models=['allowed']. Tried to access denied" + ) + + +@pytest.mark.asyncio +async def test_managed_alias_dispatch_before_scope_check() -> None: + token: Final = UserAPIKeyAuth(team_model_aliases={"alias": "allowed"}).model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": ("allowed",)} + ), + } + ) + assert await managed_model_adapter.can_agent_call_model("alias", token, None, None, None, None) is None + + +@pytest.mark.parametrize( + ("models", "expected"), + ( + (None, frozenset[str]()), + ("model", frozenset[str]()), + ((), frozenset[str]()), + (("a", 42, "a", "b"), frozenset(("a", "b"))), + ), +) +def test_managed_listing_filters_non_string_grants(models: object, expected: frozenset[str]) -> None: + token: Final = UserAPIKeyAuth().model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": models} + ) + } + ) + assert managed_model_adapter.managed_agent_model_grants(token) == expected diff --git a/tests/unit/proxy/auth/test_model_budget_adapter.py b/tests/unit/proxy/auth/test_model_budget_adapter.py new file mode 100644 index 00000000000..de22d4c78dc --- /dev/null +++ b/tests/unit/proxy/auth/test_model_budget_adapter.py @@ -0,0 +1,307 @@ +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass, field +from typing import Final +from unittest.mock import AsyncMock, call + +import pytest + +import litellm +from litellm._auth import model_budget_data +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.fallback_candidate import ( + CandidateBudgetRead, + CandidateBudgetScope, + FallbackCandidateFacts, + fallback_candidate, +) +from litellm._auth.model_budget_data import ModelBudgetEntity +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.auth.model_budget_adapter import ( + ModelBudgetIO, + enforce_model_budget, + read_model_budget, + resolve_model_budget, +) +from litellm.proxy.hooks import model_max_budget_limiter as reference +from litellm.proxy.hooks.model_budget_shared_windows import SharedWindow + +from .budget_reference import model_allowed + + +@dataclass(frozen=True, slots=True) +class Windows: + value: SharedWindow | None = None + read: AsyncMock = field(default_factory=AsyncMock) + + async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: + raise AssertionError("request auth must not open a window") + + async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: + raise AssertionError("request auth must not write a counter") + + +def groups(model: str, team_id: str | None) -> frozenset[str]: + return frozenset(("group",)) + + +@pytest.mark.parametrize( + ("model", "budgets"), + ( + ( + "provider/model", + {"model": {"max_budget": 1, "budget_duration": "1h"}, "*": {"max_budget": 2, "budget_duration": "1h"}}, + ), + ("model", {"model": {"budget_duration": "bad"}, "*": {"max_budget": 2, "budget_duration": "1h"}}), + ("model", {"model": {}, "group": {"max_budget": 2, "budget_duration": "1h"}}), + ( + "model", + {"m*": {"max_budget": 2, "budget_duration": "1h"}, "group": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ( + "mødél", + {"m*é*": {"max_budget": 2, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ( + "model", + {"model": {"max_budget": None, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ), +) +def test_resolution_matches_production_reference(model: str, budgets: Mapping[str, object]) -> None: + io: Final = ModelBudgetIO(DualCache(), groups, Windows(), 0) + old: Final = reference.resolve_model_budget(model, budgets, groups) + result: Final = resolve_model_budget(model, budgets, io, None) + assert result.read_groups is False + assert (result.entry.name if result.entry else None) == (old.budget_model if old else None) + assert (result.entry.max_budget if result.entry else None) == (old.budget_config.max_budget if old else None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("entity", "scope"), + ( + ("Key", Litellm_EntityType.KEY), + ("User", Litellm_EntityType.USER), + ("Team", Litellm_EntityType.TEAM), + ("TeamMember", Litellm_EntityType.TEAM_MEMBER), + ("EndUser", Litellm_EntityType.END_USER), + ), +) +@pytest.mark.parametrize( + ("local", "shared", "now", "limit"), + ((0, 2, 100, 1), (1, 0, 100, 1), (0, 2, 3700, 1), (0, 0, 100, 0), (0, 2, 100, -1), (0, 2, 100, None)), +) +async def test_counter_and_shared_window_match_production_reference( + entity: ModelBudgetEntity, scope: Litellm_EntityType, local: float, shared: float, now: float, limit: float | None +) -> None: + cache: Final = DualCache() + key: Final = reference.model_budget_spend_cache_key(scope, "id", "model", "1h") + cache.set_cache(key, local) + windows: Final = Windows(read=AsyncMock(return_value=SharedWindow(100, shared))) + budgets: Final = {"model": {"max_budget": limit, "budget_duration": "1h"}} + old: Final = reference._PROXY_VirtualKeyModelMaxBudgetLimiter( + cache, clock=lambda: now, model_access_groups=groups, shared_windows=windows + ) + io: Final = ModelBudgetIO(cache, groups, windows, 0, clock=lambda: now, monotonic=lambda: 0) + before: Final = await observed_model_result(model_allowed(old, scope, "id", None, budgets, "model", "exceeded")) + after: Final = await observed_model_result( + enforce_model_budget(io, entity, "id", None, "model", budgets, "exceeded") + ) + assert after == before + + +async def observed_model_result(call: Awaitable[bool]) -> tuple[object, ...]: + try: + return (await call,) + except litellm.BudgetExceededError as error: + return (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id) + + +@pytest.mark.asyncio +async def test_legacy_counter_is_added_to_the_current_counter_before_shared_read() -> None: + cache: Final = DualCache() + cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "model", "1h"), 0.4) + cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "provider/model", "1h"), 0.6) + windows: Final = Windows(read=AsyncMock(side_effect=AssertionError("already over local limit"))) + io: Final = ModelBudgetIO(cache, groups, windows, 0, monotonic=lambda: 1) + reading: Final = await read_model_budget( + io, "Key", "id", None, "provider/model", {"model": {"max_budget": 1, "budget_duration": "1h"}} + ) + assert reading.snapshot.local_spend == 0.4 + assert reading.snapshot.legacy_spend == 0.6 + assert model_budget_data.check(reading.snapshot).denied is True + windows.read.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shared_read_failure_is_terminal_and_not_replayed() -> None: + failure: Final = RuntimeError("database unavailable") + windows: Final = Windows(read=AsyncMock(side_effect=failure)) + io: Final = ModelBudgetIO(DualCache(), groups, windows, 0, monotonic=lambda: 0) + with pytest.raises(RuntimeError, match="database unavailable") as raised: + await read_model_budget(io, "User", "id", None, "model", {"model": {"max_budget": 1, "budget_duration": "1h"}}) + assert raised.value is failure + assert windows.read.await_count == 1 + + +@pytest.fixture +def alias_catalog(monkeypatch: pytest.MonkeyPatch) -> Mapping[str, object]: + catalog: Final = {"budget-latest": AliasedCostEntry({}, ("budget-target", "budget-latest"))} + monkeypatch.setattr(litellm, "model_cost", catalog) + return catalog + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Key", "User", "Team", "TeamMember", "EndUser")) +@pytest.mark.parametrize( + ("alias_spend", "target_spend", "expected"), + ((1.0, 2.0, "budget-latest"), (0.0, 2.0, "budget-target"), (0.0, 0.0, None)), +) +async def test_raw_alias_scope_preserves_first_denial_and_counter_order( + alias_catalog: Mapping[str, object], + entity: ModelBudgetEntity, + alias_spend: float, + target_spend: float, + expected: str | None, +) -> None: + cache: Final = DualCache() + scope: Final = { + "Key": Litellm_EntityType.KEY, + "User": Litellm_EntityType.USER, + "Team": Litellm_EntityType.TEAM, + "TeamMember": Litellm_EntityType.TEAM_MEMBER, + "EndUser": Litellm_EntityType.END_USER, + }[entity] + alias_key: Final = reference.model_budget_spend_cache_key(scope, "id", "budget-latest", "1h") + target_key: Final = reference.model_budget_spend_cache_key(scope, "id", "budget-target", "1h") + cache.in_memory_cache.set_cache(alias_key, alias_spend) # pyright: ignore[reportUnknownMemberType] # legacy cache setter has no annotations + cache.in_memory_cache.set_cache(target_key, target_spend) # pyright: ignore[reportUnknownMemberType] # legacy cache setter has no annotations + shared_read: Final = AsyncMock(return_value=None) + windows: Final = Windows(read=shared_read) + io: Final = ModelBudgetIO( + cache, groups, windows, 0, clock=lambda: 100, monotonic=lambda: 0, price_entry=alias_catalog.get + ) + reading: Final = await read_model_budget( + io, + entity, + "id", + None, + "budget-latest", + { + "budget-latest": {"max_budget": 1, "budget_duration": "1h"}, + "budget-target": {"max_budget": 2, "budget_duration": "1h"}, + }, + ) + decision: Final = model_budget_data.check(reading.snapshot) + assert decision.denied is (expected is not None) + assert reading.budget_model == (expected or "budget-target") + assert decision.current_spend == (alias_spend if expected == "budget-latest" else target_spend) + candidate: Final = fallback_candidate( + FallbackCandidateFacts( + True, + True, + True, + membership="Loaded", + price=ModelPriceFacts(cached=False), + budget_models=("budget-latest",), + budget_scopes=(CandidateBudgetScope(entity, True, True),), + budgets=(CandidateBudgetRead(0, entity, reading.snapshot),), + ) + ) + assert candidate.kind == ("Complete" if expected is not None else "ReadAuthorization") + assert candidate.usable is False + expected_shared: Final = ( + [] if entity == "EndUser" or alias_spend == 1 else [alias_key] if target_spend == 2 else [alias_key, target_key] + ) + assert shared_read.call_args_list == [call(spend_key=key, budget_duration="1h") for key in expected_shared] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("overrides", (("budget-latest",), ("budget-target",), ("budget-latest", "budget-target"))) +async def test_team_overrides_apply_across_alias_and_target( + alias_catalog: Mapping[str, object], overrides: tuple[str, ...] +) -> None: + io: Final = ModelBudgetIO(DualCache(), groups, Windows(), 0, price_entry=alias_catalog.get) + budgets: Final = {name: {"max_budget": 0, "budget_duration": "1h"} for name in ("budget-latest", "budget-target")} + key: Final = {name: {"max_budget": 1, "budget_duration": "1h"} for name in overrides} + reading: Final = await read_model_budget(io, "Team", "id", "id", "budget-latest", budgets, key) + expected: Final = next((name for name in budgets if name not in overrides), None) + assert model_budget_data.check(reading.snapshot).denied is (expected is not None) + assert reading.budget_model == expected + + +@pytest.mark.asyncio +async def test_shared_native_model_flow_enforces_target_before_user_scope(alias_catalog: Mapping[str, object]) -> None: + from starlette.requests import Request + + from litellm._auth.decision import Deny, ModelBudgetExceeded + from litellm._auth.driver import drive_model_budgets + from litellm._auth.rust_core import RustCore + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import ( + _ModelBudgetEffects, # pyright: ignore[reportPrivateUsage] # exercise the production raw effect boundary + ) + + token: Final = UserAPIKeyAuth( + token="id", + user_id="owner", + key_alias="alias-key", + model_max_budget={"budget-target": {"max_budget": 0, "budget_duration": "1h"}}, + user_model_max_budget={"budget-latest": {"max_budget": 0, "budget_duration": "1h"}}, + ) + effects: Final = _ModelBudgetEffects( + request=Request({"type": "http", "headers": [], "path": "/chat/completions"}), + request_data={"model": "budget-latest"}, + valid_token=token, + route="/chat/completions", + limiter=reference._PROXY_VirtualKeyModelMaxBudgetLimiter(DualCache(), clock=lambda: 100), # pyright: ignore[reportPrivateUsage] # concrete limiter integration + llm_router=None, + llm_model_list=None, + ) + outcome: Final = await drive_model_budgets(RustCore().check_model_budgets(), effects) + assert isinstance(outcome, Deny) and isinstance(outcome.denial, ModelBudgetExceeded) + assert (outcome.denial.entity_type, outcome.denial.entity_id, outcome.denial.max_budget) == ("key", "id", 0.0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("authorization_denied", (False, True)) +async def test_fallback_probe_retains_original_denial_without_authorization_retries( + alias_catalog: Mapping[str, object], authorization_denied: bool +) -> None: + from starlette.requests import Request + + from litellm._auth.decision import Deny, ModelBudgetExceeded + from litellm._auth.driver import drive_model_budgets + from litellm._auth.rust_core import RustCore + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import ( + _ModelBudgetEffects, # pyright: ignore[reportPrivateUsage] # exercise the production raw effect boundary + ) + + data: Final[dict[str, object]] = {"model": "initial"} + token: Final = UserAPIKeyAuth( + token="id", + models=["initial", "later"], + model_max_budget={ + "initial": {"max_budget": 0, "budget_duration": "1h"}, + "budget-target": {"max_budget": 1 if authorization_denied else 0, "budget_duration": "1h"}, + }, + budget_fallbacks={"initial": ["budget-latest", "later"] if authorization_denied else ["budget-latest"]}, + ) + effects: Final = _ModelBudgetEffects( + request=Request({"type": "http", "headers": [], "path": "/chat/completions"}), + request_data=data, + valid_token=token, + route="/chat/completions", + limiter=reference._PROXY_VirtualKeyModelMaxBudgetLimiter( # pyright: ignore[reportPrivateUsage] # concrete limiter integration + DualCache(), clock=lambda: 100, shared_windows=Windows(read=AsyncMock(return_value=None)) + ), + llm_router=None, + llm_model_list=None, + ) + outcome: Final = await drive_model_budgets(RustCore().check_model_budgets(), effects) + assert isinstance(outcome, Deny) and isinstance(outcome.denial, ModelBudgetExceeded) + assert "model=initial" in outcome.denial.message + assert data == {"model": "initial"} diff --git a/tests/unit/proxy/auth/test_model_policy_read.py b/tests/unit/proxy/auth/test_model_policy_read.py new file mode 100644 index 00000000000..0bafa72e79b --- /dev/null +++ b/tests/unit/proxy/auth/test_model_policy_read.py @@ -0,0 +1,18 @@ +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy.auth.model_policy_read import require_native_fact + + +@pytest.mark.parametrize("value", (False, 0, "", (), object())) +def test_read_precondition_preserves_falsy_and_opaque_facts(value: object) -> None: + assert require_native_fact(value) is value + + +def test_missing_read_fact_is_terminal() -> None: + with pytest.raises(NativeAuthRuntimeError) as raised: + require_native_fact(None) + error: Final = raised.value + assert isinstance(error, RuntimeError) diff --git a/tests/unit/proxy/auth/test_model_projection_adapter.py b/tests/unit/proxy/auth/test_model_projection_adapter.py new file mode 100644 index 00000000000..60817b2fc82 --- /dev/null +++ b/tests/unit/proxy/auth/test_model_projection_adapter.py @@ -0,0 +1,203 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from queue import SimpleQueue +from typing import Final, Literal, cast + +import pytest + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import grouped_model_access +from litellm.proxy._types import ModelAccessDeniedProxyException +from litellm.proxy.auth import model_access_adapter, model_projection_adapter +from litellm.router_api import RegistryReader +from tests.unit._auth.support import model_registry_reference as reference + + +@dataclass(frozen=True, slots=True) +class _Info: + team_id: str | None = "team" + team_public_model_name: str | None = "public" + + +@dataclass(frozen=True, slots=True) +class _Deployment: + model_info: _Info = field(default_factory=_Info) + + +@dataclass(frozen=True, slots=True) +class _Registry: + events: SimpleQueue[tuple[str, str | None]] + entries: tuple[str, ...] = ("registered",) + router_aliases: Mapping[str, str] = field(default_factory=lambda: {"alias": "target"}) + groups: Mapping[str, tuple[str, ...]] = field(default_factory=lambda: {"group": ("m",)}) + deployment: _Deployment | None = field(default_factory=_Deployment) + fail: str | None = None + + def observe(self, kind: str, model: str | None = None) -> None: + self.events.put((kind, model)) + if self.fail == kind: + raise RuntimeError(f"failed {kind}") + + def names(self) -> list[str]: # mutable-ok: the legacy sentinel resolver concatenates the registry's concrete list + self.observe("names") + return list(self.entries) + + def aliases(self) -> Mapping[str, str]: + self.observe("aliases") + return self.router_aliases + + def resolve(self, model: str, *, kind: Literal["alias"]) -> str | None: + self.observe("resolve", model) + return self.router_aliases.get(model) + + def access_groups(self, model_name: str, team_id: str | None) -> Mapping[str, tuple[str, ...]]: + self.observe("groups", model_name) + return self.groups + + def deployments(self, model: str, *, source: Literal["first_in_group"]) -> _Deployment | None: + self.observe("deployment", model) + return self.deployment + + +@dataclass(frozen=True, slots=True) +class _Router: + view: _Registry + + def registry(self) -> _Registry: + self.view.observe("registry") + return self.view + + +def _reader(view: _Registry) -> RegistryReader: + return cast(RegistryReader, _Router(view)) # noqa: TID251 # cast-ok: typed IO double implements only the registry methods consumed by auth + + +def _events(events: SimpleQueue[tuple[str, str | None]]) -> tuple[tuple[str, str | None], ...]: + return tuple(events.get() for _ in range(events.qsize())) + + +def _run(native: bool, model: str | tuple[str, ...], grants: tuple[str, ...], registry: _Registry) -> str: + router: Final = _reader(registry) + aliases: Final = {"z": "model_name_team_deploy", "a": "model_name_team_deploy", "team-alias": "target"} + try: + if native: + model_access_adapter.can_object_call_model(model, router, grants, aliases, "team", {"key": "target"}) + else: + reference.can_object_call_model( + model if isinstance(model, str) else list(model), + router, + list(grants), + aliases, + "team", + {"key": "target"}, + ) + return "allowed" + except ModelAccessDeniedProxyException as error: + return str((error.message, error.internal_message, error.type, error.param, error.code)) + + +@pytest.mark.parametrize( + "model", ("m", "alias", "global", "key", "team-alias", "model_name_team_deploy", ("m", "m"), ("denied", "m"), ()) +) +@pytest.mark.parametrize( + "grants", + ( + (), + ("m",), + ("target",), + ("group",), + ("all-team-models",), + ("all-team-models", "other"), + ("*",), + ("other/*", "provider/*"), + ("a",), + ("public",), + ), +) +def test_registry_observation_order_and_denials_match_legacy( + monkeypatch: pytest.MonkeyPatch, model: str | tuple[str, ...], grants: tuple[str, ...] +) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + registry: Final = _Registry(events) + + def provider(model: str) -> tuple[str, str, None, None]: + events.put(("provider", model)) + return "leaf", "provider", None, None + + monkeypatch.setattr(litellm, "get_llm_provider", provider) + monkeypatch.setattr(litellm, "model_alias_map", {"global": "target"}) + expected: Final = _run(False, model, grants, registry) + expected_events: Final = _events(events) + assert _run(True, model, grants, registry) == expected + assert _events(events) == expected_events + + +@pytest.mark.parametrize("native", (False, True)) +@pytest.mark.parametrize("failed", ("aliases", "groups", "names", "deployment")) +def test_registry_failure_stops_before_later_reads(native: bool, failed: str) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + registry: Final = _Registry(events, fail=failed) + with pytest.raises(RuntimeError, match=f"failed {failed}") as error: + _run(native, "model_name_team_deploy", ("all-team-models",), registry) + assert not isinstance(error.value, NativeAuthFailure) + assert _events(events)[-1][0] == failed + + +@pytest.mark.parametrize("native", (False, True)) +def test_depth_denial_never_reads_registry(native: bool) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + router: Final = _reader(_Registry(events, fail="registry")) + policy: Final = model_access_adapter.can_object_call_model if native else reference.can_object_call_model + with pytest.raises(Exception, match="max fallback depth exceeded"): + policy("m", router, (), fallback_depth=100) + assert _events(events) == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("group_models", (("other",), ("other/*",), ("m",))) +async def test_group_retry_gets_fresh_registry_reads_even_for_duplicate_grants(group_models: tuple[str, ...]) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + router: Final = _reader(_Registry(events)) + facts: Final = model_access_adapter.project_model_facts("m", router, ("other",)) + assert _events(events) == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) + scope: Final = await model_projection_adapter.read_grouped_scope( + GroupedModelScope(facts, ("group",), group_models), router, merge=False + ) + assert _events(events)[:4] == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) + assert grouped_model_access(scope.facts, scope.group_ids, scope.group_models).kind == ( + "Allow" if group_models == ("m",) else "Deny" + ) + + +def test_native_provider_failure_is_not_a_missing_provider(monkeypatch: pytest.MonkeyPatch) -> None: + def provider(model: str) -> tuple[str, str, None, None]: + raise NativeAuthRuntimeError("terminal provider failure") + + monkeypatch.setattr(litellm, "get_llm_provider", provider) + with pytest.raises(NativeAuthRuntimeError, match="terminal provider failure"): + model_access_adapter.can_object_call_model("model", None, ("provider/*",)) + + +@pytest.mark.parametrize("native", (False, True)) +def test_compaction_parent_is_read_only_after_child_denial(native: bool) -> None: + from litellm.router_strategy.complexity_router.context_compaction import native_compaction_call + + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + with native_compaction_call("parent", "child"): + assert _run(native, "child", ("parent",), _Registry(events)) == "allowed" + assert _events(events) == ( + ("registry", None), + ("aliases", None), + ("registry", None), + ("groups", "child"), + ("registry", None), + ("groups", "parent"), + ) + + +def test_large_model_lists_do_not_consume_python_recursion_depth() -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + assert _run(True, ("m",) * 180, ("m",), _Registry(events)) == "allowed" + assert _events(events) == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) * 180 diff --git a/tests/unit/proxy/auth/test_org_route_adapter.py b/tests/unit/proxy/auth/test_org_route_adapter.py new file mode 100644 index 00000000000..ebae5ba615e --- /dev/null +++ b/tests/unit/proxy/auth/test_org_route_adapter.py @@ -0,0 +1,101 @@ +from collections.abc import Iterator, Mapping +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.proxy._types import LiteLLM_UserTable, ProxyException +from litellm.proxy.auth import org_route_adapter as native + +from .test_route_role_adapter import observed + +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) + + +def reference_admin(body: Mapping[str, object], user: LiteLLM_UserTable | None) -> bool: + if user is None or user.organization_memberships is None: + return False + singular: Final = body.get("organization_id") + plural: Final = body.get("organizations") + candidates: Final = ( + *((singular,) if singular is not None else ()), + *(_OBJECTS.validate_python(plural) if isinstance(plural, list) else ()), + ) + admins: Final = frozenset( + item.organization_id for item in user.organization_memberships + if item.user_role == "org_admin" and item.organization_id is not None # pyright: ignore[reportUnnecessaryComparison] # preserve Phase 1 constructed membership handling + ) + return bool(candidates) and all(candidate in admins for candidate in candidates) + + +def user(kind: str) -> LiteLLM_UserTable | None: + if kind == "absent": + return None + return LiteLLM_UserTable.model_validate({ + "user_id": "u", + "user_role": "proxy_admin" if kind == "proxy" else "internal_user", + "organization_memberships": None if kind == "none" else ( + [] if kind == "empty" else [{ + "organization_id": "org", "user_role": "org_admin", "user_id": "u", + "created_at": "2025-01-01T00:00:00Z", "updated_at": "2025-01-01T00:00:00Z", + }] + ), + }) + + +@pytest.mark.parametrize("kind", ("absent", "none", "empty", "member", "proxy")) +@pytest.mark.parametrize("singular", (None, "org", "other", "", 0, False, 1.5, [], {}, ())) +@pytest.mark.parametrize("plural", (None, [], ["org"], ["other", []], ["org", []], [0], "org")) +def test_raw_singular_and_plural_keep_membership_failure_and_short_circuit( + kind: str, singular: object, plural: object +) -> None: + member: Final = user(kind) + body: Final = {"organization_id": singular, "organizations": plural} + assert observed(lambda: native.user_is_org_admin(body, member)) == observed(lambda: reference_admin(body, member)) + + +@pytest.mark.parametrize("identifier", (0, False, 1.5, (), "")) +@pytest.mark.parametrize("route", ("/organization/delete", "/team/new")) +def test_nonstring_present_identifier_is_not_changed_to_missing(identifier: object, route: str) -> None: + with pytest.raises(ProxyException) as captured: + native.organization_role_based_access_check({"organization_id": identifier}, user("member"), route) + expected: Final = ( + f"You do not have a role within the selected organization. Passed organization_id: {identifier}. Please contact the organization admin to request access." + if route == "/organization/delete" else + f"You do not have the required role to call {route}. Your role is None in Organization {identifier}" + ) + assert captured.value.message == expected + assert captured.value.code == "401" + + +@pytest.mark.parametrize("identifier", ([], {})) +@pytest.mark.parametrize("route", ("/organization/delete", "/team/new")) +def test_unhashable_identifier_raises_only_when_membership_lookup_is_selected(identifier: object, route: str) -> None: + with pytest.raises(TypeError, match="unhashable type"): + native.organization_role_based_access_check({"organization_id": identifier}, user("member"), route) + assert native.organization_role_based_access_check({"organization_id": identifier}, user("proxy"), route) is None + assert native.organization_role_based_access_check({"organization_id": identifier}, user("absent"), route) is None + with pytest.raises(ProxyException, match="not a member of any organization"): + native.organization_role_based_access_check({"organization_id": identifier}, user("none"), "/organization/delete") + + +class UnreadableBody(Mapping[str, object]): + def __getitem__(self, key: str) -> object: + raise AssertionError(f"unexpected body read: {key}") + + def __iter__(self) -> Iterator[str]: + return iter(()) + + def __len__(self) -> int: + return 0 + + +class UnreadableHash: + def __hash__(self) -> int: + raise AssertionError("unexpected candidate hash") + + +def test_no_user_never_reads_request_and_denied_candidate_never_hashes_later_candidate() -> None: + assert native.user_is_org_admin(UnreadableBody(), None) is False + assert native.organization_role_based_access_check(UnreadableBody(), None, "/organization/delete") is None + assert native.user_is_org_admin({"organizations": ["not-allowed", UnreadableHash()]}, user("member")) is False diff --git a/tests/unit/proxy/auth/test_personal_key_team_model_access.py b/tests/unit/proxy/auth/test_personal_key_team_model_access.py index 38f34adb48e..31b5c304bfc 100644 --- a/tests/unit/proxy/auth/test_personal_key_team_model_access.py +++ b/tests/unit/proxy/auth/test_personal_key_team_model_access.py @@ -320,14 +320,12 @@ async def test_model_listing_loads_each_team_once(monkeypatch: pytest.MonkeyPatc @pytest.mark.asyncio async def test_team_check_error_grants_nothing_for_that_team(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: _enable(monkeypatch) - check_team: Final = auth_checks.can_team_access_model - async def _failing_for_team_b(model: str, team_object: LiteLLM_TeamTableCachedObj, **_: object) -> bool: - if team_object.team_id == "team-b": - raise RuntimeError("access group lookup failed") - return await check_team(model=model, team_object=team_object, llm_router=_ROUTER) + async def _failing_for_team_b(*, team_id: str, **_: object) -> None: + if team_id == "team-b": + raise RuntimeError("membership lookup failed") - monkeypatch.setattr(auth_checks, "can_team_access_model", _failing_for_team_b) + monkeypatch.setattr(auth_checks, "get_team_membership", _failing_for_team_b) owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) token: Final = UserAPIKeyAuth(token="k1", user_id="u1") assert await _common_checks("gpt-a", owner, token) is True diff --git a/tests/unit/proxy/auth/test_resolvers_store.py b/tests/unit/proxy/auth/test_resolvers_store.py index 607162000ab..bd0fd9a452d 100644 --- a/tests/unit/proxy/auth/test_resolvers_store.py +++ b/tests/unit/proxy/auth/test_resolvers_store.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import Dict, Optional +from types import SimpleNamespace +from typing import Dict, Final, Optional import pytest @@ -14,6 +15,8 @@ from litellm.proxy.auth.resolvers.exceptions import ( from litellm.proxy.auth.auth_method import AuthMethod from litellm.proxy.auth.resolvers.models import Principal, PrincipalType from litellm.proxy.auth.resolvers.store import IdentityStore +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm._auth.failure import NativeAuthRuntimeError class _FakeCache: @@ -92,3 +95,39 @@ async def test_resolve_raises_a_generic_key_not_found_carrying_the_hash_when_no_ assert exc_info.value.message == INVALID_API_KEY_ERROR_MESSAGE assert missing not in str(exc_info.value) assert exc_info.value.hashed_token == missing + + +class _InvalidatedNativeRead: + def __init__(self, cache: UserApiKeyCache, *, always_invalidate: bool = False) -> None: + self.cache = cache + self.calls = 0 + self.always_invalidate = always_invalidate + + async def auth_read(self, kind: str, ids: list[str], writer: bool = False) -> list[dict[str, object]]: + self.calls += 1 + if self.calls == 1 or self.always_invalidate: + self.cache.delete_local(ids[0]) + return [{"token": ids[0], "blocked": self.calls > 1, "models": ["allowed"]}] + + +async def test_invalidated_key_fill_reloads_instead_of_returning_or_caching_the_old_identity(): + cache: Final = UserApiKeyCache() + handle: Final = _InvalidatedNativeRead(cache) + token: Final = hash_token("sk-fill-fence") + store: Final = IdentityStore(SimpleNamespace(native_auth_database=handle), cache) + principal: Final = await store.resolve(token) + assert IdentityStore.key_from_principal(principal).blocked is True + assert handle.calls == 2 + cached: Final = await cache.async_get_cache(token, model_type=UserAPIKeyAuth) + assert cached is not None and cached.blocked is True + + +async def test_continuously_invalidated_key_fills_stop_without_caching_a_stale_identity(): + cache: Final = UserApiKeyCache() + handle: Final = _InvalidatedNativeRead(cache, always_invalidate=True) + token: Final = hash_token("sk-contended-fill") + store: Final = IdentityStore(SimpleNamespace(native_auth_database=handle), cache) + with pytest.raises(NativeAuthRuntimeError, match="changed during lookup"): + await store.resolve(token) + assert handle.calls == 2 + assert cache.get_local(token) is None diff --git a/tests/unit/proxy/auth/test_route_enable_adapter.py b/tests/unit/proxy/auth/test_route_enable_adapter.py new file mode 100644 index 00000000000..6f3f48d321c --- /dev/null +++ b/tests/unit/proxy/auth/test_route_enable_adapter.py @@ -0,0 +1,140 @@ +from queue import SimpleQueue +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.route_enable import RouteFeature, RouteSetting +from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.auth.route_enable_adapter import RouteEnableReaders, check_route_enabled + + +@pytest.mark.parametrize("route", ("/models", "/chat/completions")) +def test_enterprise_class_exemptions_preserve_enabled_route_policy(route: str) -> None: + class EnterpriseRoutes: + LLM_API_EXEMPT_ROUTES: Final = ("/models",) + + readers: Final = RouteEnableReaders( + enterprise=lambda: EnterpriseRoutes, + setting=lambda _: RouteSetting(True, True), + secret=lambda _: True, + ) + if route == "/models": + assert check_route_enabled(route, readers=readers) is None + else: + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + assert (error.value.status_code, error.value.detail) == ( + 403, + "LLM API routes are disabled for this instance.", + ) + + +def test_missing_optional_class_exemptions_preserve_ignored_read_failure() -> None: + readers: Final = RouteEnableReaders(enterprise=lambda: object()) + assert check_route_enabled("/chat/completions", readers=readers) is None + + +@pytest.mark.parametrize( + ("route", "expected"), + ( + ("/unknown", ("import",)), + ("/models", ("import", "exemptions")), + ("/chat/completions", ("import", "exemptions", "Llm", "secret-Llm")), + ("/key/generate", ("import", "Admin", "secret-Admin")), + ), +) +def test_optional_route_reads_are_lazy(route: str, expected: tuple[str, ...]) -> None: + events: Final[SimpleQueue[str]] = SimpleQueue() + + def enterprise() -> object: + events.put("import") + return object() + + def setting(feature: RouteFeature) -> RouteSetting: + events.put(feature) + return RouteSetting(False, False) + + def secret(feature: RouteFeature) -> bool: + events.put(f"secret-{feature}") + return False + + def exemptions(_enterprise: object) -> tuple[str, ...]: + events.put("exemptions") + return ("/models",) + + check_route_enabled(route, readers=RouteEnableReaders(enterprise, setting, secret, exemptions)) + assert tuple(events.get() for _ in range(events.qsize())) == expected + + +@pytest.mark.parametrize(("route", "label"), (("/key/generate", "ADMIN"), ("/chat/completions", "LLM API"))) +@pytest.mark.parametrize("premium", (False, True)) +def test_optional_route_denial_envelopes(route: str, label: str, premium: bool) -> None: + readers: Final = RouteEnableReaders( + enterprise=lambda: object(), + setting=lambda _: RouteSetting(True, premium), + secret=lambda _: True, + exemptions=lambda _: (), + ) + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + if premium: + disabled: Final = "Management" if label == "ADMIN" else label + assert (error.value.status_code, error.value.detail) == ( + 403, + f"{disabled} routes are disabled for this instance.", + ) + else: + assert (error.value.status_code, error.value.detail) == ( + 500, + f"\U0001f6a8\U0001f6a8\U0001f6a8 DISABLING {label} ENDPOINTS is an Enterprise feature\n\U0001f6a8 {CommonProxyErrors.not_premium_user.value}", + ) + + +@pytest.mark.parametrize( + "error", + ( + ImportError("optional package"), + RuntimeError("settings unavailable"), + HTTPException(418, "read failure"), + NativeAuthRuntimeError("terminal native failure"), + ), +) +def test_optional_import_failure_policy(error: Exception) -> None: + def enterprise() -> object: + raise error + + readers: Final = RouteEnableReaders(enterprise=enterprise) + if isinstance(error, (HTTPException, NativeAuthRuntimeError)): + with pytest.raises(type(error)) as caught: + check_route_enabled("/key/generate", readers=readers) + assert caught.value is error + else: + assert check_route_enabled("/key/generate", readers=readers) is None + + +@pytest.mark.parametrize("configured", (False, True)) +@pytest.mark.parametrize("route", ("/eu_jev/v1/systemone", "/eu_jev/v1/systemone/")) +def test_configured_eu_jev_route_skips_llm_disable_setting(configured: bool, route: str) -> None: + events: Final[SimpleQueue[str]] = SimpleQueue() + + def eu_jev() -> bool: + events.put("configuration") + return configured + + def setting(feature: RouteFeature) -> RouteSetting: + events.put(feature) + return RouteSetting(True, True) + + readers: Final = RouteEnableReaders( + enterprise=lambda: object(), setting=setting, secret=lambda _: True, exemptions=lambda _: (), eu_jev=eu_jev + ) + if configured: + assert check_route_enabled(route, readers=readers) is None + assert tuple(events.get() for _ in range(events.qsize())) == ("configuration",) + else: + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + assert (error.value.status_code, error.value.detail) == (403, "LLM API routes are disabled for this instance.") + assert tuple(events.get() for _ in range(events.qsize())) == ("configuration", "Llm") diff --git a/tests/unit/proxy/auth/test_route_role_adapter.py b/tests/unit/proxy/auth/test_route_role_adapter.py new file mode 100644 index 00000000000..dce460a6b37 --- /dev/null +++ b/tests/unit/proxy/auth/test_route_role_adapter.py @@ -0,0 +1,126 @@ +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException, Request + +from litellm._auth import route_roles as core +from litellm._auth.failure import NativeAuthFailure +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth import route_role_adapter as native +from litellm.rust_bridge import _native + +from . import route_role_reference as reference + +_ROLES: Final = (None, "", "unknown", "PROXY_ADMIN", *(role.value for role in LitellmUserRoles)) + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: type[Exception] + message: str + status: int | None + detail: object + + +def observed(call: Callable[[], object]) -> object: + try: + return call() + except Exception as error: + return Failure( + type(error), str(error), + error.status_code if isinstance(error, HTTPException) else None, + error.detail if isinstance(error, HTTPException) else None, + ) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("present", (False, True)) +def test_role_and_admin_projection_matches_reference(role: str | None, present: bool) -> None: + user: Final = LiteLLM_UserTable(user_id="u", user_role=role) if present else None + assert native.effective_user_role(role) is reference.effective_user_role(role) + assert native.get_user_role(user) is reference.get_user_role(user) + assert native.is_user_proxy_admin(user) is reference.is_user_proxy_admin(user) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("user_id", (None, "", "u")) +@pytest.mark.parametrize("requested", (None, "", "u", "other")) +def test_inside_route_admin_viewer_or_non_none_self(role: str | None, user_id: str | None, requested: str | None) -> None: + auth: Final = UserAPIKeyAuth.model_construct(user_role=role, user_id=user_id) + assert native.allowed_route_check_inside_route(auth, requested) is reference.allowed_route_check_inside_route(auth, requested) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("token_present", (False, True)) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/team/new")) +def test_api_route_role_precedes_token_and_nonadmin_check(role: str | None, token_present: bool, route: str) -> None: + user: Final = LiteLLM_UserTable(user_id="u", user_role=role) + token: Final = UserAPIKeyAuth(user_id="u") if token_present else None + request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + assert observed(lambda: native.is_api_route_allowed(route, request, {}, token, user)) == observed( + lambda: reference.is_api_route_allowed(route, request, {}, token, user) + ) + + +@pytest.mark.parametrize("role", tuple(LitellmUserRoles)) +@pytest.mark.parametrize("configured", (None, (), ("info_routes",), ("/x*",))) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/x/y", "/not-configured")) +def test_jwt_admin_team_none_default_and_empty_deny( + role: LitellmUserRoles, configured: tuple[str, ...] | None, route: str +) -> None: + config: Final = LiteLLM_JWTAuth.model_construct( + admin_allowed_routes=None if configured is None else list(configured), + team_allowed_routes=None if configured is None else list(configured), + ) + assert observed(lambda: native.allowed_routes_check(role, route, config)) == observed( + lambda: reference.allowed_routes_check(role, route, config) + ) + + +@pytest.mark.parametrize("entries", ( + (), ("/x",), ("/x*",), ("/x/{id}",), ("info_routes",), + ("/x", None), (None, "/x"), (1,), (False,), (list[object](),), (dict[str, object](),), + (set[object](),), ((list[object](),),), (("x",),), ("missing_group", "info_routes"), ("/x", list[object]()), +)) +@pytest.mark.parametrize("route", ("/x", "/x/a", "/x/", "/user/info", "")) +def test_listed_route_invalid_entries_and_short_circuit_match_reference(entries: tuple[object, ...], route: str) -> None: + assert observed(lambda: native.allowed_routes_list_check(route, entries)) == observed( + lambda: reference.allowed_routes_list_check(route, entries) + ) + + +@pytest.mark.parametrize("group", tuple(LiteLLMRoutes.__members__)) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/key/a/reset_spend", "/credentials/a/b", "/x\n", "")) +def test_every_route_group_uses_starlette_template_semantics(group: str, route: str) -> None: + assert observed(lambda: native.allowed_routes_list_check(route, (group,))) == observed( + lambda: reference.allowed_routes_list_check(route, (group,)) + ) + + +@pytest.mark.parametrize("entries", (("info_routes", "info_routes"), ("/x", None, 1, ("x",)), ([],), ({},))) +def test_expansion_preserves_order_duplicates_nonstring_values_and_errors(entries: tuple[object, ...]) -> None: + assert observed(lambda: native.get_actual_routes(entries)) == observed(lambda: reference.get_actual_routes(entries)) + + +@pytest.mark.parametrize("compiled", ("^/thing/(?P[^/]+)$", "^/thing/(?P.*)$", "^/v1\\.0/$")) +@pytest.mark.parametrize("route", ("/thing/", "/thing/a", "/thing/a/b", "/thing/a\n", "/thing/a\nb", "/v1.0/", "/v1x0/")) +def test_compiled_group_regex_keeps_empty_path_literal_escaping_and_final_newline(compiled: str, route: str) -> None: + import re + + decision: Final = core.listed(core.ListedFacts(route, core.Entry("group", "str", templates=("unused",)), compiled=compiled)) + assert (decision.action == "Allow") is (re.match(compiled, route) is not None) + + +def test_native_fault_is_terminal_before_reading_first_allowed_entry(monkeypatch: pytest.MonkeyPatch) -> None: + def broken(_facts: object) -> object: + raise ValueError("decode failed") + + def entries() -> Iterable[object]: + raise AssertionError("entry read after native failure") + yield "unused" + + monkeypatch.setattr(_native, "request_auth_route_role_listed", broken) + with pytest.raises(NativeAuthFailure, match="decode failed"): + native.allowed_routes_list_check("/x", entries()) diff --git a/tests/unit/proxy/auth/test_router_override_fallback_auth.py b/tests/unit/proxy/auth/test_router_override_fallback_auth.py index d35a5851659..549b7740375 100644 --- a/tests/unit/proxy/auth/test_router_override_fallback_auth.py +++ b/tests/unit/proxy/auth/test_router_override_fallback_auth.py @@ -5,15 +5,23 @@ override is promoted to per-request kwargs after auth and lets a caller execute requests against models their API key cannot call. """ -from typing import List -from unittest.mock import AsyncMock, patch +from typing import Final, List import pytest -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm import Router + +from litellm.proxy._types import LitellmUserRoles, ModelAccessDeniedProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import fallback_target_model_name, iter_request_fallback_targets from litellm.proxy.auth.user_api_key_auth import enforce_key_and_fallback_model_access +_ROUTER: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/unit-model", "api_key": "sk-unit-test"}} + for name in ("gpt-3.5-turbo", "real-fallback", "string-fallback") + ] +) + def _fallback_model_names(fallbacks): """Model names the auth check validates for a top-level ``fallbacks`` value.""" @@ -38,9 +46,11 @@ def _key_with_models(models: List[str]) -> UserAPIKeyAuth: def test_fallback_model_names_router_config_shape(): """Router-config shape: ``[{primary: [fallback_list]}]``.""" - assert _fallback_model_names( - [{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}] - ) == ["gpt-4", "claude-3", "o1"] + assert _fallback_model_names([{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}]) == [ + "gpt-4", + "claude-3", + "o1", + ] def test_fallback_model_names_simple_string_shape(): @@ -58,9 +68,10 @@ def test_fallback_model_names_client_side_shape(): def test_fallback_model_names_nested_deployment_fallbacks(): """A deployment target's own nested fallback field is unrolled too.""" - assert _fallback_model_names( - [{"primary": [{"model": "gpt-4", "fallbacks": [{"gpt-4": ["deepseek-chat"]}]}]}] - ) == ["gpt-4", "deepseek-chat"] + assert _fallback_model_names([{"primary": [{"model": "gpt-4", "fallbacks": [{"gpt-4": ["deepseek-chat"]}]}]}]) == [ + "gpt-4", + "deepseek-chat", + ] def test_fallback_model_names_empty_or_none(): @@ -85,34 +96,16 @@ async def test_router_override_fallbacks_validated_against_key_allowlist(): }, } - seen_models: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen_models.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - # Both the primary model and the override-nested fallback must be - # checked against the API key's allowlist. - assert "gpt-3.5-turbo" in seen_models - assert "unauthorized-model" in seen_models + assert denied.value.code == "403" @pytest.mark.asyncio @@ -136,31 +129,16 @@ async def test_router_override_all_fallback_fields_validated(fallback_field): }, } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - assert "smuggled-model" in seen + assert denied.value.code == "403" @pytest.mark.asyncio @@ -183,31 +161,16 @@ async def test_top_level_fallback_fields_validated(fallback_field): fallback_field: [{"gpt-3.5-turbo": ["top-level-smuggled"]}], } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - assert "top-level-smuggled" in seen + assert denied.value.code == "403" @pytest.mark.asyncio @@ -229,31 +192,16 @@ async def test_nested_deployment_fallback_inner_model_validated(): ], } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - assert "deep-smuggled-model" in seen + assert denied.value.code == "403" @pytest.mark.asyncio @@ -274,32 +222,16 @@ async def test_model_less_fallback_dict_is_skipped_never_passed_as_none(): ], } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route="/v1/chat/completions", - request=None, - llm_model_list=None, - llm_router=None, - ) - - assert None not in seen - assert seen == ["gpt-3.5-turbo", "real-fallback", "string-fallback"] + valid_token.models = ["gpt-3.5-turbo", "real-fallback", "string-fallback"] + result = await enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=_ROUTER, + ) + assert result is None @pytest.mark.asyncio @@ -312,28 +244,12 @@ async def test_router_override_without_fallbacks_does_not_break_auth(): "router_settings_override": {"num_retries": 3, "timeout": 30}, } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route="/v1/chat/completions", - request=None, - llm_model_list=None, - llm_router=None, - ) - - assert seen == ["gpt-3.5-turbo"] + result = await enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=_ROUTER, + ) + assert result is None diff --git a/tests/unit/proxy/auth/test_spend_counter_adapter.py b/tests/unit/proxy/auth/test_spend_counter_adapter.py new file mode 100644 index 00000000000..e576a583cb0 --- /dev/null +++ b/tests/unit/proxy/auth/test_spend_counter_adapter.py @@ -0,0 +1,49 @@ +from typing import Final, NoReturn +from unittest.mock import AsyncMock, Mock, call + +import pytest + +from litellm._auth.spend_counter_policy import Spend +from litellm.proxy.auth.spend_counter_adapter import SpendIO, resolve_spend + + +def unavailable() -> NoReturn: + raise RuntimeError("unverified spend") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallback, window, floor, strict, expected, calls", + ( + (0.0, False, None, (False, False), 2.0, (call.read_strict(), call.read_strict())), + (0.0, False, 4.0, (True,), 4.0, (call.read_strict(), call.read_floor(), call.repair(4.0))), + (5.0, False, None, (), 5.0, (call.read_floor(),)), + (0.0, True, 1.0, (), 2.0, (call.read_floor(),)), + ), +) +async def test_native_plan_preserves_floor_and_flag_read_order( + fallback: float, + window: bool, + floor: float | None, + strict: tuple[bool, ...], + expected: float, + calls: tuple[object, ...], +) -> None: + effects: Final = Mock() + effects.attach_mock(Mock(side_effect=strict), "read_strict") + effects.attach_mock(AsyncMock(return_value=floor), "read_floor") + effects.attach_mock(AsyncMock(), "repair") + io: Final = SpendIO(effects.read_strict, effects.read_floor, effects.repair, unavailable) + result: Final = await resolve_spend(Spend(2.0, False, fallback, 10.0, window, False), io) + assert (result, tuple(effects.mock_calls)) == (expected, calls) + + +@pytest.mark.asyncio +async def test_unverified_spend_reads_the_latest_fail_closed_flag() -> None: + strict: Final = Mock(side_effect=(False, True)) + floor: Final = AsyncMock() + with pytest.raises(RuntimeError, match="unverified spend"): + await resolve_spend( + Spend(2.0, False, 0.0, 10.0, False, False), SpendIO(strict, floor, AsyncMock(), unavailable) + ) + assert (strict.call_count, floor.await_count) == (2, 0) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 19f9d0dcc73..df9a5a1bca8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -15,10 +15,12 @@ from litellm._logging import verbose_proxy_logger import json import logging import litellm +from tests.unit._auth.reference_core import PythonCore from litellm.proxy.auth.user_api_key_auth import ( + _enforce_model_budgets, + _ModelBudgetEffects, user_api_key_auth, UserAPIKeyAuth, - get_api_key_from_custom_header, ) from fastapi import WebSocket, HTTPException, status @@ -26,6 +28,23 @@ from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles from tests._master_key import MASTER_KEY +async def _check_user_model_budgets(valid_token: UserAPIKeyAuth, limiter: object, model: str) -> None: + request: Final = MagicMock() + request.scope = {} + await _enforce_model_budgets( + PythonCore().check_user_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data={"model": model}, + valid_token=valid_token, + route="/chat/completions", + limiter=limiter, + llm_router=None, + llm_model_list=None, + ), + ) + + class Request: def __init__(self, client_ip: Optional[str] = None, headers: Optional[dict] = None): self.client = MagicMock() @@ -506,50 +525,6 @@ async def test_auth_not_connected_to_db(): assert valid_token.token == "failed-to-connect-to-db" -def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api_key): - verbose_proxy_logger.setLevel(logging.DEBUG) - request = MagicMock(spec=Request) - request.headers = headers - api_key = get_api_key_from_custom_header(request=request, custom_litellm_key_header_name=custom_header_name) - assert api_key == expected_api_key - - -def test_get_api_key_from_custom_header_bearer_token(): - token = "sk-" + "1" * 8 - _assert_api_key_from_custom_header( - headers={"x-custom-api-key": f"Bearer {token}"}, - custom_header_name="x-custom-api-key", - expected_api_key=token, - ) - - -def test_get_api_key_from_custom_header_empty_value(): - _assert_api_key_from_custom_header( - headers={"x-custom-api-key": ""}, - custom_header_name="x-custom-api-key", - expected_api_key="", - ) - - -def test_get_api_key_from_custom_header_missing_header(): - _assert_api_key_from_custom_header( - headers={}, - custom_header_name="X-Custom-API-Key", - expected_api_key="", - ) - - -def test_get_api_key_from_custom_header_different_casing(): - token = "sk-" + "1" * 8 - _assert_api_key_from_custom_header( - headers={"X-CUSTOM-API-KEY": f"Bearer {token}"}, - custom_header_name="X-Custom-API-Key", - expected_api_key=token, - ) - - - - @pytest.mark.parametrize( "user_role, auth_user_id, requested_user_id, expected_result", [ @@ -1426,27 +1401,25 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_ stored on LiteLLM_UserTable, accepted by /user/new and /user/update, and read by nothing, so a user-level per-model budget never blocked anything. """ - from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget - calls = [] class _Limiter: - async def is_user_within_model_budget( - self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None - ) -> bool: - calls.append((user_id, user_model_max_budget, model)) - return True + async def read_model_budget_raw( + self, entity: str, entity_id: str | None, team_id: str | None, model: str, + budgets: Mapping[str, object] | None, key_budgets: Mapping[str, object] | None = None, + ): + from litellm._auth.model_budget_data import ModelBudgetSnapshot + from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading + + calls.append((entity_id, budgets, model)) + return ModelBudgetReading(ModelBudgetSnapshot(None, None, None, None, None, None, 0), model, model) valid_token = UserAPIKeyAuth( token="hash", user_id=user_id, user_model_max_budget=user_model_max_budget, ) - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=_Limiter(), - models=["gpt-4"], - ) + await _check_user_model_budgets(valid_token, _Limiter(), model="gpt-4") assert len(calls) == expected_calls if expected_calls: assert calls[0] == ("u-1", user_model_max_budget, "gpt-4") @@ -1588,7 +1561,6 @@ async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(ove nothing ever reads it, which looks enforced and is not. """ from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth - from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key from litellm.proxy.proxy_server import model_max_budget_limiter @@ -1616,18 +1588,11 @@ async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(ove if expect_refusal: with pytest.raises(litellm.BudgetExceededError) as exc: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=[model], - ) + await _check_user_model_budgets(valid_token, model_max_budget_limiter, model=model) assert exc.value.entity_type == Litellm_EntityType.USER.value + assert exc.value.entity_id == user_id else: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=[model], - ) + await _check_user_model_budgets(valid_token, model_max_budget_limiter, model=model) def test_jwt_path_enforces_the_user_model_budget_before_returning(): @@ -1645,7 +1610,7 @@ def test_jwt_path_enforces_the_user_model_budget_before_returning(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) def calls_before_each_return(node): @@ -1654,7 +1619,7 @@ def test_jwt_path_enforces_the_user_model_budget_before_returning(): if isinstance(child, ast.Call): fn = child.func name = getattr(fn, "id", None) or getattr(fn, "attr", None) - if name == "_check_user_model_budget": + if name == "_enforce_model_budgets": seen_check.append(child.lineno) return seen_check @@ -1690,7 +1655,7 @@ def test_every_jwt_branch_carries_the_user_model_budget(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) assignments = [ @@ -1708,7 +1673,7 @@ def test_every_jwt_branch_carries_the_user_model_budget(): assert "auto_registered" in targets, ( f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}" ) - assert "valid_token" in targets, "the virtual-key path must carry the user's model budget" + assert "key" in targets, "the virtual-key path must carry the user's model budget" @pytest.mark.asyncio @@ -1815,8 +1780,10 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): whether the caller used a JWT or a virtual key, and makes that log a lie. Structural rather than behavioural on purpose: the defect is a call site - sitting outside a guard, and driving both auth paths to a zero-cost model - would prove it for the two requests exercised rather than for every site. + sitting outside a guard, and driving the JWT path to a zero-cost model + would prove it for the request exercised rather than for every site. The + virtual-key path's exemption is decided by the auth core and covered in + tests/unit/_auth/test_python_core.py. """ import ast import inspect @@ -1824,11 +1791,7 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): from litellm.proxy.auth import user_api_key_auth as auth_module - tree = ast.parse( - textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) - + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) - ) + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool: for parent in ast.walk(node): @@ -1848,69 +1811,41 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): calls = [ node for node in ast.walk(tree) - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget" + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_enforce_model_budgets" ] - assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}" + assert len(calls) == 1, f"expected the JWT call site, found {len(calls)}" unguarded = [c for c in calls if not guarded_by_skip(tree, c)] assert not unguarded, ( - f"{len(unguarded)} _check_user_model_budget call(s) run even when " + f"{len(unguarded)} _enforce_model_budgets call(s) run even when " "skip_budget_checks is set, so a zero-cost model is enforced on one auth path and not the other" ) -def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): - """ - The custom-auth helper runs its own key, user and end-user per-model budget - checks. If it does not honour the zero-cost skip that the JWT and - virtual-key paths honour, the same free request is refused under one auth - method and served under the others. - - Asserted structurally, on the same reasoning as the sibling test: the defect - is a check sitting outside a guard, and it must hold for checks added later - rather than only for whichever request a behavioural test happened to drive. - """ - import ast - import inspect - import textwrap - +@pytest.mark.asyncio +async def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): from litellm.proxy.auth import user_api_key_auth as auth_module - src = textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks)) - tree = ast.parse(src) - - assert "skip_budget_checks" in src, "the custom-auth path never computes the zero-cost skip flag" - - budget_calls = ( - "_check_key_model_budget_with_fallback", - "_check_user_model_budget", - "is_end_user_within_model_budget", + router = litellm.Router(model_list=[{ + "model_name": "free-model", + "litellm_params": { + "model": "openai/free-model", "api_key": "unused", + "input_cost_per_token": 0, "output_cost_per_token": 0, + }, + "model_info": {"id": "free-id"}, + }]) + token = UserAPIKeyAuth( + token="test-key", + model_max_budget={"free-model": {"max_budget": 0}}, ) - - def guarding_ifs(target: ast.AST) -> list[ast.If]: - return [ - node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is target for sub in ast.walk(node)) - ] - - def mentions_skip(node: ast.If) -> bool: - return any(isinstance(sub, ast.Name) and sub.id == "skip_budget_checks" for sub in ast.walk(node.test)) - - for call_name in budget_calls: - calls = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and ( - (isinstance(node.func, ast.Name) and node.func.id == call_name) - or (isinstance(node.func, ast.Attribute) and node.func.attr == call_name) - ) - ] - assert calls, f"{call_name} is no longer called here; update this invariant" - for call in calls: - assert any(mentions_skip(node) for node in guarding_ifs(call)), ( - f"{call_name} runs even for a zero-cost model, so custom auth refuses " - "requests the JWT and virtual-key paths serve" - ) + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + result = await auth_module._run_post_custom_auth_checks( + token, None, {"model": "free-model"}, "/v1/chat/completions", None + ) + assert result is token def test_custom_auth_attaches_the_user_budget_even_when_it_does_not_enforce(): @@ -1972,7 +1907,7 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) # Half one: the shared block copies the user row's budget onto the token. @@ -2154,7 +2089,7 @@ async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods( with ( patch("litellm.proxy.proxy_server.premium_user", True), patch("litellm.proxy.auth.handle_jwt.JWTAuthManager.auth_builder", return_value=jwt_response), - patch("litellm.proxy.auth.user_api_key_auth._check_user_model_budget", user_model_budget_check), + patch("litellm.proxy.auth.user_api_key_auth._enforce_model_budgets", user_model_budget_check), ): result: Final = await user_api_key_auth_builder( request=mock_request, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 8440d3d2ab2..8b87f985da8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -74,10 +74,12 @@ from litellm.proxy.auth.auth_checks import ( cache_key_object, jwt_key_mapping_cache_key, ) +from tests.unit._auth.reference_core import PythonCore from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model from litellm.proxy.auth.user_api_key_auth import ( - _check_key_model_budget_with_fallback, + _enforce_model_budgets, + _ModelBudgetEffects, _check_team_member_model_budget_with_fallback, _ensure_litellm_received_at_on_request_state, _ensure_parent_otel_span_on_request_state, @@ -449,26 +451,15 @@ async def test_should_not_reuse_cached_key_object_for_request_state(): @pytest.mark.asyncio async def test_custom_auth_does_not_enforce_key_model_access_by_default(): valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch( - "litellm.proxy.proxy_server.general_settings", - {}, - ), - ): - await _run_post_custom_auth_checks( + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = await _run_post_custom_auth_checks( valid_token=valid_token, request=None, - request_data=request_data, + request_data={"model": "gpt-4o"}, route="/v1/chat/completions", parent_otel_span=None, ) - mock_can_key.assert_not_awaited() + assert result is valid_token @pytest.mark.asyncio @@ -497,29 +488,20 @@ async def test_custom_auth_honors_key_level_model_access_restriction_allowed_wit request_data = {"model": "gpt-4o-mini"} with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, ), ): - await _run_post_custom_auth_checks( + result = await _run_post_custom_auth_checks( valid_token=valid_token, request=None, request_data=request_data, route="/v1/chat/completions", parent_otel_span=None, ) - mock_can_key.assert_awaited_once_with( - model="gpt-4o-mini", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) + assert result is valid_token @pytest.mark.asyncio @@ -528,63 +510,37 @@ async def test_custom_auth_enforces_key_model_access_from_file_route_header_with request = _RoutingRequest(headers={"x-litellm-model": "restricted-model"}) with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, ), ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=request, - request_data={}, - route="/v1/files", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="restricted-model", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) + with pytest.raises(ProxyException) as denial: + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=request, + request_data={}, + route="/v1/files", + parent_otel_span=None, + ) + assert denial.value.type == ProxyErrorTypes.key_model_access_denied + assert "restricted-model" in str(denial.value) @pytest.mark.asyncio async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in(): valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - mock_can_key.side_effect = ProxyException( - message="Key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=401, - ) + with patch("litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}): with pytest.raises(ProxyException) as exc: await _run_post_custom_auth_checks( valid_token=valid_token, request=None, - request_data=request_data, + request_data={"model": "gpt-4o"}, route="/v1/chat/completions", parent_otel_span=None, ) - - assert exc.value.type == ProxyErrorTypes.key_model_access_denied + assert exc.value.type == ProxyErrorTypes.key_model_access_denied def _proxy_server_attrs_for_custom_auth(*, user_custom_auth): @@ -753,6 +709,61 @@ async def _authenticate_through_custom_auth() -> UserAPIKeyAuth: ) +async def _authenticate_master_key(request_data: dict[str, object]) -> UserAPIKeyAuth: + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder + + request: Final = Request(scope={"type": "http", "headers": []}) + request._url = URL(url="/chat/completions") + return await user_api_key_auth_builder( + request=request, + api_key="Bearer sk-master-key", + azure_api_key_header=None, + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data=request_data, + ) + + +@pytest.mark.asyncio +async def test_master_key_request_does_not_inherit_previous_requests_end_user_limits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as _proxy_server_mod + + proxy_attrs: Final = _proxy_server_attrs_for_custom_auth(user_custom_auth=None) + for attr, val in proxy_attrs.items(): + monkeypatch.setattr(_proxy_server_mod, attr, val, raising=False) + monkeypatch.setattr(_proxy_server_mod, "user_api_key_cache", UserApiKeyCache()) + capped: Final = LiteLLM_EndUserTable( + user_id="eu-capped", + blocked=False, + litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=5, rpm_limit=7), + ) + + async def end_user_object(end_user_id: str, **_: object) -> LiteLLM_EndUserTable | None: + return capped if end_user_id == "eu-capped" else None + + with patch("litellm.proxy.auth.user_api_key_auth.get_end_user_object", side_effect=end_user_object): + capped_request: Final = await _authenticate_master_key({"user": "eu-capped"}) + await asyncio.sleep(0.05) + next_request: Final = await _authenticate_master_key({}) + + assert (capped_request.end_user_id, capped_request.end_user_tpm_limit, capped_request.end_user_rpm_limit) == ( + "eu-capped", + 5, + 7, + ) + assert (next_request.end_user_id, next_request.end_user_tpm_limit, next_request.end_user_rpm_limit) == ( + None, + None, + None, + ) + assert next_request.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + + @pytest.mark.asyncio async def test_user_custom_auth_returning_one_shared_token_keeps_each_requests_end_user_apart( monkeypatch: pytest.MonkeyPatch, @@ -1839,7 +1850,13 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th from litellm.proxy.proxy_server import hash_token api_key = "sk-prefetch-order-test" - valid_token = UserAPIKeyAuth(api_key=api_key, token=hash_token(api_key), user_id="u1", team_id="t1") + valid_token = UserAPIKeyAuth( + api_key=api_key, + token=hash_token(api_key), + user_id="u1", + team_id="t1", + models=["gpt-4o" if model_allowed else "other-model"], + ) mock_cache = AsyncMock() mock_cache.async_get_cache = AsyncMock(return_value=None) @@ -1855,7 +1872,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th "proxy_logging_obj": mock_proxy_logging_obj, "master_key": "sk-master-key", "general_settings": {}, - "llm_model_list": [], + "llm_model_list": [{"model_name": "gpt-4o"}], "llm_router": None, "open_telemetry_logger": None, "model_max_budget_limiter": MagicMock(), @@ -1864,12 +1881,6 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th "litellm_proxy_admin_name": "admin", } _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} - denied = ProxyException( - message="Key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=401, - ) try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1882,11 +1893,6 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th new_callable=AsyncMock, return_value=valid_token, ), - patch( # test-quality-ok: the observable is whether the prefetch runs before or after this check - "litellm.proxy.auth.user_api_key_auth.enforce_key_and_fallback_model_access", - new_callable=AsyncMock, - side_effect=None if model_allowed else denied, - ), patch( # test-quality-ok: counting prefetch calls on a denied request IS the regression being pinned "litellm.proxy.auth.user_api_key_auth.prefetch_auth_objects", new_callable=AsyncMock ) as mock_prefetch, @@ -5487,8 +5493,6 @@ async def _run_centralized_checks_for_end_user( setattr(_proxy_server_mod, attribute, value) - - @pytest.mark.parametrize( ("make_token", "make_end_user", "gpt_4o_spend", "model", "route", "must_raise"), [ @@ -9540,10 +9544,53 @@ class TestJWTAuthUserEmail: assert result.user_email == "admin@example.com" -class TestCheckKeyModelBudgetWithFallback: - """`_check_key_model_budget_with_fallback` must reroute a request to the - first configured `budget_fallbacks` entry still within its own budget, - and only raise `BudgetExceededError` when no fallback is available.""" +def _key_over_budget_for(*models: str) -> AsyncMock: + from litellm._auth.model_budget_data import ModelBudgetSnapshot + from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading + + async def read_model_budget_raw( + entity: str, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> ModelBudgetReading: + return ModelBudgetReading( + ModelBudgetSnapshot(5, 10 if model in models else 0, None, None, None, None, 0), model, model + ) + + limiter: Final = AsyncMock() + limiter.read_model_budget_raw.side_effect = read_model_budget_raw + return limiter + + +async def _check_key_model_budgets( + valid_token: UserAPIKeyAuth, + limiter: object, + request_data: dict[str, object], + request: object, + llm_router: litellm.Router | None = None, +) -> None: + await _enforce_model_budgets( + PythonCore().check_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=valid_token, + route="/chat/completions", + limiter=limiter, + llm_router=llm_router, + llm_model_list=None, + ), + ) + + +_KEY_MODEL_BUDGET: Final = {"gpt-4o": {"budget_limit": 5.0, "time_period": "1d"}} + + +class TestKeyModelBudgetFallback: + """An over-budget key reroutes to the first `budget_fallbacks` entry it may call and can still afford.""" def _make_request(self): request = MagicMock() @@ -9552,162 +9599,66 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio async def test_within_budget_does_not_reroute(self): - valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) - limiter = AsyncMock() - limiter.is_key_within_model_budget.return_value = True + valid_token = UserAPIKeyAuth( + token="test-key", model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]} + ) request_data = {"model": "gpt-4o"} request = self._make_request() - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for(), request_data, request) assert request_data["model"] == "gpt-4o" - limiter.get_fallback_model_within_budget.assert_not_awaited() assert "parsed_body" not in request.scope @pytest.mark.asyncio - async def test_exceeded_budget_reroutes_to_fallback(self): + async def test_exceeded_budget_reroutes_to_fallback_visible_to_a_later_body_read(self): valid_token = UserAPIKeyAuth( token="test-key", + model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini", "claude-haiku"]}, ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for("gpt-4o"), request_data, request) assert request_data["model"] == "gpt-4o-mini" - limiter.get_fallback_model_within_budget.assert_awaited_once_with(user_api_key_dict=valid_token, model="gpt-4o") - # the rerouted model must be visible to a later, separate - # `_read_request_body` call on the same `request` (route handlers - # re-parse the body from this cache instead of reusing the dict). - cached_keys, cached_body = request.scope["parsed_body"] + _, cached_body = request.scope["parsed_body"] assert cached_body["model"] == "gpt-4o-mini" @pytest.mark.asyncio - async def test_raises_when_every_fallback_also_exceeded(self): - valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = None + @pytest.mark.parametrize( + ("token_fields", "over_budget"), + [ + ({"budget_fallbacks": {"gpt-4o": ["gpt-4o-mini"]}}, ("gpt-4o", "gpt-4o-mini")), + ({"models": ["gpt-4o"], "budget_fallbacks": {"gpt-4o": ["restricted-model"]}}, ("gpt-4o",)), + ( + { + "models": ["gpt-4o", "restricted-model"], + "team_id": "team-1", + "team_models": ["gpt-4o"], + "budget_fallbacks": {"gpt-4o": ["restricted-model"]}, + }, + ("gpt-4o",), + ), + ], + ids=["every_fallback_over_budget", "fallback_outside_key_models", "fallback_outside_team_models"], + ) + async def test_raises_the_original_budget_error_when_no_fallback_is_usable( + self, token_fields: dict[str, object], over_budget: tuple[str, ...] + ): + valid_token = UserAPIKeyAuth(token="test-key", model_max_budget=_KEY_MODEL_BUDGET, **token_fields) request_data = {"model": "gpt-4o"} - request = self._make_request() with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, + await _check_key_model_budgets( + valid_token, _key_over_budget_for(*over_budget), request_data, self._make_request() ) - assert exc_info.value is original_error - assert request_data["model"] == "gpt-4o" - - @pytest.mark.asyncio - async def test_raises_when_fallback_not_authorized(self): - """If the fallback model is within budget but the key is not allowed - to call it, the original BudgetExceededError must be raised instead - of rerouting to an unauthorized model.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o"], - budget_fallbacks={"gpt-4o": ["restricted-model"]}, + assert ( + exc_info.value.message == "LiteLLM Virtual Key: test-key, key_alias: None, exceeded budget for model=gpt-4o" ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = "restricted-model" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert exc_info.value is original_error - assert request_data["model"] == "gpt-4o" - - @pytest.mark.asyncio - async def test_reroute_succeeds_when_fallback_is_authorized(self): - """If the fallback model is within budget AND authorized, the request - must be rerouted to it.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "gpt-4o-mini"], - budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, - ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert request_data["model"] == "gpt-4o-mini" - - @pytest.mark.asyncio - async def test_raises_when_fallback_blocked_by_team_models(self): - """If the key allows the fallback but the team does not, the original - BudgetExceededError must be raised.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "restricted-model"], - team_id="team-1", - team_models=["gpt-4o"], - budget_fallbacks={"gpt-4o": ["restricted-model"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = "restricted-model" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert exc_info.value is original_error + assert (exc_info.value.current_cost, exc_info.value.max_budget) == (10, 5) assert request_data["model"] == "gpt-4o" @pytest.mark.asyncio @@ -9718,56 +9669,18 @@ class TestCheckKeyModelBudgetWithFallback: valid_token = UserAPIKeyAuth( token="test-key", models=["gpt-4o", "gpt-4o-mini"], + model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() request.scope["path_params"] = {"model": "gpt-4o"} - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for("gpt-4o"), request_data, request) assert request_data["model"] == "gpt-4o-mini" assert request.scope["path_params"]["model"] == "gpt-4o-mini" - @pytest.mark.asyncio - async def test_raises_when_model_from_nested_field(self): - """Budget fallback must not attempt rewrite when the checked model - came from a nested field (e.g. session.model) rather than the - top-level request_data['model'].""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "gpt-4o-mini"], - budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - request_data = {"session": {"model": "gpt-4o"}} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) - - assert exc_info.value is original_error - assert "model" not in request_data - class TestCheckTeamMemberModelBudgetWithFallback: @pytest.fixture(autouse=True) @@ -10786,13 +10699,7 @@ class TestCheckTeamMemberModelBudgetWithFallback: request_data: Final = {"model": "gpt-4o"} request: Final = self._request("gpt-4o") - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, limiter, request_data, request) assert request_data["model"] == "gpt-4o-mini" diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index d41c4595a0d..1d30cbd7b1e 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -145,6 +145,8 @@ class _ScriptedPubSubRedisClient(Redis): class _FakeRedisCache: + store_settings = None + def __init__(self, client: object, namespace: str | None = None) -> None: self._client = client self.namespace = namespace @@ -207,8 +209,8 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: instead of serving the stale object until the TTL expires. """ cache = UserApiKeyCache() - cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) - assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + cache.set_local("project_id:p-1", {"models": []}, None) + assert cache.get_local("project_id:p-1") is not None pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")]) subscriber = AuthCacheInvalidationSubscriber( @@ -221,7 +223,7 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: finally: await subscriber.stop() - assert cache.in_memory_cache.get_cache("project_id:p-1") is None + assert cache.get_local("project_id:p-1") is None assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] @@ -357,7 +359,7 @@ async def test_subscriber_writes_and_drops_the_native_spend_counter_copies_a_bro @pytest.mark.asyncio async def test_subscriber_ignores_malformed_messages() -> None: cache = UserApiKeyCache() - cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + cache.set_local("project_id:p-1", {"models": []}, None) subscriber = AuthCacheInvalidationSubscriber( redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), @@ -368,7 +370,7 @@ async def test_subscriber_ignores_malformed_messages() -> None: subscriber._apply_message("raw string") subscriber._apply_message(None) - assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + assert cache.get_local("project_id:p-1") is not None @pytest.mark.asyncio diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index 6c9104bb711..c5c31dcfa21 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -80,6 +80,19 @@ async def test_resync_exception_returns_false_without_negative_caching(): assert spy.calls == ["new-model", "new-model"] +@pytest.mark.asyncio +async def test_native_read_through_failure_is_terminal_and_not_negative_cached(): + from litellm._auth.failure import NativeAuthRuntimeError + + spy: Final = ResyncSpy(error=NativeAuthRuntimeError("native connection lost")) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded) + with pytest.raises(NativeAuthRuntimeError): + await read_through.attempt("agent") + with pytest.raises(NativeAuthRuntimeError): + await read_through.attempt("agent") + assert spy.calls == ["agent", "agent"] + + @pytest.mark.asyncio async def test_concurrent_attempts_for_missing_key_resync_once(): class SlowResyncSpy(ResyncSpy): @@ -257,6 +270,69 @@ async def test_get_agent_with_read_through_recovers_agent_by_name(clean_agent_re ) +@pytest.mark.asyncio +@pytest.mark.parametrize("by_name", [False, True]) +async def test_native_agent_read_through_preserves_identity_and_permission_without_prisma( + by_name, clean_agent_registry, fresh_agent_read_through, monkeypatch +): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + binding = { + "agent_id": "native-agent", + "provider": "microsoft_entra", + "tenant_id": "tenant", + "client_id": "client", + "issuer": "https://login.microsoftonline.com/tenant/v2.0", + "revision": "revision", + } + row = { + **FakeAgentRow("native-agent", "Native agent name").model_dump(), + "identity": binding, + "enabled": False, + "access_group_ids": ["group"], + "object_permission": {"models": ["approved-model"]}, + } + native = SimpleNamespace(auth_read=AsyncMock(side_effect=[[], [row]] if by_name else [[row]])) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(native_auth_database=native)) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + lookup = "Native agent name" if by_name else "native-agent" + + agent = await get_agent_with_read_through(lookup) + assert agent is not None and agent.identity is not None + assert agent.identity.model_dump(include=set(binding)) == binding + assert agent.object_permission == {"models": ["approved-model"]} + assert agent.access_group_ids == ["group"] and agent.enabled is False + assert clean_agent_registry.get_agent_by_id(agent_id="native-agent") == agent + assert [call.args[0] for call in native.auth_read.call_args_list] == ( + ["agent", "agent_name"] if by_name else ["agent"] + ) + + await get_agent_with_read_through(lookup) + assert native.auth_read.await_count == (2 if by_name else 1) + + +@pytest.mark.asyncio +async def test_native_agent_read_through_error_does_not_retry_by_name_or_return_miss( + clean_agent_registry, fresh_agent_read_through, monkeypatch +): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm._auth.failure import NativeAuthFailure + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + native = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("offline"))) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(native_auth_database=native)) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + with pytest.raises(NativeAuthFailure): + await get_agent_with_read_through("native-failed-agent") + native.auth_read.assert_awaited_once_with("agent", ["native-failed-agent"], writer=False) + assert clean_agent_registry.get_agent_by_id(agent_id="native-failed-agent") is None + + @pytest.mark.asyncio async def test_get_agent_with_read_through_returns_none_for_unknown_agent( clean_agent_registry, fresh_agent_read_through, monkeypatch @@ -1083,7 +1159,9 @@ async def test_resync_model_deployments_by_ids_reloads_credentials_only_for_one_ model_params: Final = {"model": "openai/gpt-backup"} row: Final = _stored_model_row( "backup-id", - model_params if stored_credential_name is None else model_params | {"litellm_credential_name": stored_credential_name}, + model_params + if stored_credential_name is None + else model_params | {"litellm_credential_name": stored_credential_name}, ) stored_credential: Final = CredentialItem( credential_name="shared-openai", diff --git a/tests/unit/proxy/common_utils/test_user_api_key_cache.py b/tests/unit/proxy/common_utils/test_user_api_key_cache.py index d3c9043e4b6..1f95815d1c5 100644 --- a/tests/unit/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/unit/proxy/common_utils/test_user_api_key_cache.py @@ -1,6 +1,6 @@ import datetime import hashlib -from typing import Any, Final +from typing import Final from unittest.mock import MagicMock import pytest @@ -17,7 +17,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( get_management_object_ttl, is_user_key_cache_key, ) -from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum from litellm.redis.stores.auth_cache import AuthCacheStore from litellm.rust_bridge._native import NativeAuthCache from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache @@ -25,18 +24,6 @@ from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cach HASHED_TOKEN = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest() -class CapturingInMemoryCache(InMemoryCache): - """Records ``ttl`` passed into ``set_cache`` (what DualCache injects).""" - - def __init__(self) -> None: - super().__init__() - self.last_ttl: Any = None - - def set_cache(self, key, value, **kwargs): # type: ignore[override] - self.last_ttl = kwargs.get("ttl") - super().set_cache(key, value, **kwargs) - - class _ManualClock: def __init__(self) -> None: self.seconds = 1000.0 @@ -55,38 +42,73 @@ def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: class TestUserApiKeyCache: @pytest.mark.asyncio - async def test_async_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + @pytest.mark.parametrize("invalidate", (True, False)) + async def test_native_db_fill_is_local_and_generation_fenced(self, invalidate: bool) -> None: + cache, fake = shared_cache() + generation: Final = cache.current_generation() + if invalidate: + cache.delete_local(HASHED_TOKEN) + await cache.async_set_cache( + HASHED_TOKEN, + _make_key_obj(HASHED_TOKEN), + model_type=UserAPIKeyAuth, + fill_generation=generation, + ) + assert fake.redis == {} + assert fake.commands == [] + assert cache.get_cache(HASHED_TOKEN, local_only=True, model_type=UserAPIKeyAuth) == ( + None if invalidate else _make_key_obj(HASHED_TOKEN) + ) + + def test_native_fill_fence_and_request_local_projection_without_redis(self): + cache: Final = UserApiKeyCache() + stale: Final = cache.begin_fill(HASHED_TOKEN) + cache.delete_local(HASHED_TOKEN) + assert not cache.set_if_generation(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), stale) + assert cache.set_if_generation( + HASHED_TOKEN, UserAPIKeyAuth(token=HASHED_TOKEN, models=["original"]), cache.begin_fill(HASHED_TOKEN) + ) + first: Final = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert first is not None + first.models.append("request-only") + second: Final = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert second is not None and second is not first + assert second.models == ["original"] + cache.flush_cache() + assert cache.get_local(HASHED_TOKEN) is None + + @pytest.mark.asyncio + async def test_async_set_with_zero_default_ttl_expires_without_a_python_copy( self, ): - """ - If ``general_settings.user_api_key_cache_ttl`` is absent, the proxy never - calls ``update_cache_ttl``; ``user_api_key_cache`` keeps - ``default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl``. - DualCache must forward that as the in-memory ``ttl`` kwarg on each set. - """ - mem = CapturingInMemoryCache() + mem = InMemoryCache() cache = UserApiKeyCache( in_memory_cache=mem, - default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + default_in_memory_ttl=0, ) + cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=None) await cache.async_set_cache( "k", _make_key_obj("t"), model_type=UserAPIKeyAuth, ) - expected = UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value - assert mem.last_ttl == expected + assert cache.get_local("k") is None + assert mem.get_cache("k") is None + cache.set_cache("k", {"token": "t"}, ttl=60) + assert cache.get_local("k") == {"token": "t"} - def test_sync_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + def test_sync_set_with_zero_default_ttl_expires_without_a_python_copy( self, ): - mem = CapturingInMemoryCache() + mem = InMemoryCache() cache = UserApiKeyCache( in_memory_cache=mem, - default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + default_in_memory_ttl=0, ) + cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=None) cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth) - assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value + assert cache.get_local("sk") is None + assert mem.get_cache("sk") is None @pytest.mark.asyncio async def test_the_shared_store_gets_default_in_memory_ttl_and_writes_without_a_ttl(self): @@ -173,7 +195,9 @@ class TestUserApiKeyCache: async def test_a_model_holding_a_datetime_is_cached_as_its_json_dump(self): cache = UserApiKeyCache( redis_cache=MagicMock(spec=RedisCache), - auth_cache=lambda _redis, settings: AuthCacheStore(NativeAuthCache(None, memory_entries=settings.memory_entries)), + auth_cache=lambda _redis, settings: AuthCacheStore( + NativeAuthCache(None, memory_entries=settings.memory_entries) + ), ) key_obj = UserAPIKeyAuth(token="abc", expires=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)) @@ -273,13 +297,14 @@ class TestUserKeyObjectPartition: @pytest.mark.asyncio async def test_update_cache_ttl_applies_to_key_object_partition(self): - key_partition = CapturingInMemoryCache() + key_partition = InMemoryCache() cache = UserApiKeyCache(default_in_memory_ttl=60, key_object_in_memory_cache=key_partition) - cache.update_cache_ttl(default_in_memory_ttl=7, default_redis_ttl=7) + cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=7) await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) - assert key_partition.last_ttl == 7 + assert cache.get_local(HASHED_TOKEN) is None + assert key_partition.get_cache(HASHED_TOKEN) is None @pytest.mark.asyncio async def test_attach_redis_cache_moves_both_partitions_onto_the_shared_store(self): @@ -372,8 +397,10 @@ class TestUserKeyObjectPartition: cache.set_local(HASHED_TOKEN, {"token": HASHED_TOKEN}, None) cache.set_local(end_user_cache_key("u1"), {"user_id": "u1"}, 5) - assert cache.key_object_cache.in_memory_cache.get_cache(HASHED_TOKEN) == {"token": HASHED_TOKEN} - assert cache.in_memory_cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"} + assert cache.get_local(HASHED_TOKEN) == {"token": HASHED_TOKEN} + assert cache.get_local(end_user_cache_key("u1")) == {"user_id": "u1"} + assert cache.key_object_cache.in_memory_cache.get_cache(HASHED_TOKEN) is None + assert cache.in_memory_cache.get_cache(end_user_cache_key("u1")) is None cache.delete_local(HASHED_TOKEN) assert cache.get_local(HASHED_TOKEN) is None @@ -461,11 +488,7 @@ class TestManagementObjectTTL: @pytest.mark.asyncio async def test_management_write_uses_configured_ttl_over_constant(self): - mem = CapturingInMemoryCache() - cache = UserApiKeyCache( - in_memory_cache=mem, - default_in_memory_ttl=300, - ) + cache, fake = shared_cache(default_in_memory_ttl=300) assert get_management_object_ttl(cache) != (DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL) await cache.async_set_cache( @@ -475,4 +498,4 @@ class TestManagementObjectTTL: ttl=get_management_object_ttl(cache), ) - assert mem.last_ttl == 300 + assert fake.commands == [("SET", "team_id:abc", 300)] diff --git a/tests/unit/proxy/db/test_spend_counter_reseed.py b/tests/unit/proxy/db/test_spend_counter_reseed.py index 7314bfbb7c5..cfb915db8e4 100644 --- a/tests/unit/proxy/db/test_spend_counter_reseed.py +++ b/tests/unit/proxy/db/test_spend_counter_reseed.py @@ -16,6 +16,7 @@ from unittest.mock import AsyncMock import pytest +from litellm._auth.failure import NativeAuthFailure from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY @@ -27,6 +28,101 @@ from litellm.redis.stores.spend_counters import SpendCounterStore WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("counter", "kind", "entity_id"), + [ + ("spend:key:k", "key", "k"), + ("spend:user:u", "user", "u"), + ("spend:team:t", "team", "t"), + ("spend:org:o", "org", "o"), + ("spend:project:p", "project", "p"), + ("spend:team_member:u:with:colon:t", "member", "u:with:colon:t"), + ], +) +async def test_native_reseed_reads_authoritative_zero_without_a_prisma_handle(counter, kind, entity_id): + native = SimpleNamespace(auth_spend=AsyncMock(return_value=0.0)) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.from_db(prisma, counter) == 0.0 + native.auth_spend.assert_awaited_once_with(kind, entity_id) + + +@pytest.mark.asyncio +async def test_native_reseed_distinguishes_missing_from_invariant_failure(): + native = SimpleNamespace(auth_spend=AsyncMock(side_effect=[None, RuntimeError("offline")])) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.from_db(prisma, "spend:user:u") is None + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.from_db(prisma, "spend:user:u") + assert native.auth_spend.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [OSError("connection lost"), asyncio.TimeoutError()]) +async def test_native_spend_io_failure_preserves_cached_fallback_without_prisma(error): + native = SimpleNamespace(auth_spend=AsyncMock(side_effect=error)) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.from_db(prisma, "spend:key:hash") is None + assert await SpendCounterReseed.end_user_from_db(prisma, "spend:end_user:u") is None + assert native.auth_spend.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("aggregate", [[{"spend": 4.5}], OSError("connection still lost")]) +async def test_native_window_io_failure_uses_only_native_aggregate_and_preserves_unavailable(aggregate): + native = SimpleNamespace(auth_read=AsyncMock(side_effect=[OSError("connection lost"), aggregate])) + result = await SpendCounterReseed.window_from_db( + SimpleNamespace(native_auth_database=native), "Key", "hash", "1d", WINDOW_START + ) + assert result == (None if isinstance(aggregate, OSError) else 4.5) + assert [call.args[0] for call in native.auth_read.call_args_list] == ["budget_window", "window_spend"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", ["not a number", {}, []]) +async def test_native_invalid_spend_cannot_become_cached_fallback(payload): + native = SimpleNamespace(auth_spend=AsyncMock(return_value=payload)) + prisma = SimpleNamespace(native_auth_database=native) + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.from_db(prisma, "spend:key:hash") + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.end_user_from_db(prisma, "spend:end_user:u") + + +@pytest.mark.asyncio +async def test_native_end_user_floor_preserves_zero(): + native = SimpleNamespace(auth_spend=AsyncMock(return_value=0.0)) + assert ( + await SpendCounterReseed.end_user_from_db(SimpleNamespace(native_auth_database=native), "spend:end_user:u") + == 0.0 + ) + native.auth_spend.assert_awaited_once_with("end_user", "u") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("row", [None, {"spend": 99.0, "window_start": WINDOW_START - timedelta(days=1)}]) +async def test_native_stale_or_missing_window_uses_native_spend_aggregate(row): + native = SimpleNamespace(auth_read=AsyncMock(side_effect=[[row] if row else [], [{"spend": 4.5}]])) + value = await SpendCounterReseed.window_from_db( + SimpleNamespace(native_auth_database=native), "Key", "hash", "1d", WINDOW_START + ) + assert value == 4.5 + assert native.auth_read.call_args_list[0].args == ("budget_window", ["Key", "hash", "1d"]) + assert native.auth_read.call_args_list[1].args == ("window_spend", ["Key", "hash", WINDOW_START.isoformat()]) + + +@pytest.mark.asyncio +async def test_native_current_window_does_not_range_scan_and_failure_does_not_fallback(): + native = SimpleNamespace( + auth_read=AsyncMock(side_effect=[[{"spend": 0.0, "window_start": WINDOW_START}], RuntimeError("offline")]) + ) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.window_from_db(prisma, "Key", "hash", "1d", WINDOW_START) == 0.0 + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.window_from_db(prisma, "Key", "hash", "1d", WINDOW_START) + assert native.auth_read.await_count == 2 + + class _FakeFindUniqueTable: def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: self._row = row diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 21cbe8667ff..e892fe43e65 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -131,26 +131,29 @@ def test_resolve_model_budget(): assert resolve_model_budget(model="non-existent", model_max_budget=model_max_budget) is None -# Test is_key_within_model_budget @pytest.mark.asyncio -async def test_is_key_within_model_budget(budget_limiter): - # Mock user API key dict - user_api_key = UserAPIKeyAuth( +@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0)) +async def test_is_key_within_model_budget( + budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, + spend: float, +) -> None: + user_api_key: Final = UserAPIKeyAuth( token="test-key", key_alias="test-alias", model_max_budget={"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, ) - - # Test when model is within budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): - assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True - - # Test when model exceeds budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): + counter: Final = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.KEY, + entity_id="test-key", + budget_model="gpt-4", + budget_duration="1d", + ) + await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend) + if spend >= 100.0: with pytest.raises(litellm.BudgetExceededError): await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") - - # Test model not in budget config + else: + assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True assert await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") is True @@ -224,38 +227,26 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim assert call_kwargs["response_cost"] == 0.05 -# Test is_end_user_within_model_budget @pytest.mark.asyncio -async def test_is_end_user_within_model_budget(budget_limiter): - # Test when model is within budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): - assert ( - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "gpt-4", - ) - is True - ) - - # Test when model exceeds budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): - with pytest.raises(litellm.BudgetExceededError): - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "gpt-4", - ) - - # Test model not in budget config - assert ( - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "non-existent", - ) - is True +@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0)) +async def test_is_end_user_within_model_budget( + budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, + spend: float, +) -> None: + counter: Final = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.END_USER, + entity_id="test-user", + budget_model="gpt-4", + budget_duration="1d", ) + await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend) + budgets: Final = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}} + if spend >= 100.0: + with pytest.raises(litellm.BudgetExceededError): + await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4") + else: + assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4") is True + assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "non-existent") is True # Test _get_spend_for_model_budget for the end-user scope @@ -510,70 +501,6 @@ async def test_model_budget_limiter_without_redis_tracks_spend_locally_and_queue assert limiter.redis_increment_operation_queue == [] -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( - budget_limiter, -): - user_api_key = UserAPIKeyAuth(token="test-key", budget_fallbacks={}) - assert await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") is None - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_first_within_budget( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}}, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=1.0): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result == "gpt-4o-mini" - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_skips_exhausted_fallback( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={ - "gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}, - "claude-haiku": {"budget_limit": 100.0, "time_period": "1d"}, - }, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - - async def _spend_for_model(entity_type, entity_id, model, resolved): - return 150.0 if resolved.budget_model == "gpt-4o-mini" else 1.0 - - with patch.object( - budget_limiter, - "_get_spend_for_model_budget", - side_effect=_spend_for_model, - ): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result == "claude-haiku" - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_none_when_chain_exhausted( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={ - "gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}, - "claude-haiku": {"budget_limit": 100.0, "time_period": "1d"}, - }, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result is None - - @pytest.mark.asyncio async def test_async_log_success_event_skips_redis_push_without_redis(budget_limiter): """When dual_cache has no Redis backend, do not await _push_in_memory_increments_to_redis.""" diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index af8bc254e2c..a64dab0bf9b 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -51,7 +51,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_session_id_for_spend_log, _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, - is_master_key, redact_logged_api_key, _redact_prompt_leaks_in_error_string, sanitize_error_information_for_spend_logs, @@ -3484,37 +3483,6 @@ def test_get_logging_payload_includes_request_duration_ms(): assert payload["request_duration_ms"] == 3000 -class TestIsMasterKey: - """Tests for _is_master_key handling None inputs without raising TypeError.""" - - def test_none_api_key_returns_false(self): - """Regression: _is_master_key(None, 'sk-master') should return False, not raise TypeError.""" - assert is_master_key(api_key=None, _master_key="sk-master-key") is False - - def test_none_master_key_returns_false(self): - assert is_master_key(api_key="sk-some-key", _master_key=None) is False - - def test_both_none_returns_false(self): - assert is_master_key(api_key=None, _master_key=None) is False - - def test_matching_key_returns_true(self): - assert is_master_key(api_key="sk-master", _master_key="sk-master") is True - - def test_non_matching_key_returns_false(self): - assert is_master_key(api_key="sk-other", _master_key="sk-master") is False - - def test_master_key_hash_is_rejected(self): - """ - ``_is_master_key`` must not accept ``hash_token(master_key)`` as - equivalent to the raw master key — only the raw value matches. - """ - from litellm.proxy.utils import hash_token - - master = "sk-master-key-123" - hashed = hash_token(master) - assert is_master_key(api_key=hashed, _master_key=master) is False - - def test_sanitize_request_body_strips_secret_fields(): """ secret_fields contains raw HTTP headers (including Authorization Bearer diff --git a/tests/unit/proxy/test_tools_allowlist_enforcement.py b/tests/unit/proxy/test_tools_allowlist_enforcement.py index efaf25c480c..0fafaf0ef4e 100644 --- a/tests/unit/proxy/test_tools_allowlist_enforcement.py +++ b/tests/unit/proxy/test_tools_allowlist_enforcement.py @@ -6,14 +6,14 @@ Covers: - extract_request_tool_names: OpenAI chat, responses, Anthropic, generate_content, MCP """ -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Final +from unittest.mock import MagicMock, patch import pytest from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import check_tools_allowlist from litellm.proxy.guardrails.tool_name_extraction import ( - TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names, ) @@ -58,9 +58,7 @@ class TestExtractRequestToolNames: {"type": "function", "name": "get_current_weather", "description": "x"}, ] } - assert extract_request_tool_names("/v1/responses", data) == [ - "get_current_weather" - ] + assert extract_request_tool_names("/v1/responses", data) == ["get_current_weather"] def test_openai_responses_mcp_tools(self): data = { @@ -129,9 +127,7 @@ class TestExtractRequestToolNames: }, ] } - assert extract_request_tool_names("/generate_content", data) == [ - "schedule_meeting" - ] + assert extract_request_tool_names("/generate_content", data) == ["schedule_meeting"] def test_mcp_call_tool_name(self): data = {"name": "my_tool", "arguments": {}} @@ -164,6 +160,57 @@ class TestExtractRequestToolNames: class TestCheckToolsAllowlist: """Test allowlist enforcement in auth (no DB in hot path).""" + @pytest.mark.asyncio + @pytest.mark.parametrize("key", [None, [], "forbidden", {"tool": "forbidden"}]) + @pytest.mark.parametrize("team", [None, [], "forbidden", {"tool": "forbidden"}]) + async def test_disabled_allowlists_never_extract_tool_names(self, key: object, team: object) -> None: + with patch( + "litellm.proxy.auth.auth_checks.extract_request_tool_names", + side_effect=AssertionError("disabled allowlist read request tools"), + ) as extractor: + await check_tools_allowlist( + request_body={"tools": object()}, + valid_token=_token(metadata={"allowed_tools": key}, team_metadata={"allowed_tools": team}), + team_object=None, + route="/v1/responses", + ) + extractor.assert_not_called() + + @pytest.mark.asyncio + async def test_empty_tool_names_do_not_convert_allowlist_members(self) -> None: + member: Final = MagicMock() + with patch("litellm.proxy.auth.auth_checks.extract_request_tool_names", return_value=[]) as extractor: + await check_tools_allowlist( + request_body={}, + valid_token=_token(metadata={"allowed_tools": [member]}), + team_object=None, + route="/v1/responses", + ) + assert extractor.call_count == 1 + member.__str__.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("key", [None, [], "ignored", {"tool": "ignored"}]) + async def test_team_members_are_stringified_only_after_tool_names(self, key: object) -> None: + member: Final = MagicMock() + member.__str__.return_value = "allowed" + + def names(route: str, body: object) -> list[str]: + member.__str__.assert_not_called() + return ["allowed", "123", "None", "False"] + + with patch("litellm.proxy.auth.auth_checks.extract_request_tool_names", side_effect=names): + await check_tools_allowlist( + request_body={}, + valid_token=_token( + metadata={"allowed_tools": key}, + team_metadata={"allowed_tools": [member, 123, None, False]}, + ), + team_object=None, + route="/v1/responses", + ) + member.__str__.assert_called_once_with() + @pytest.mark.asyncio async def test_no_allowlist_passes(self): token = _token(metadata={}, team_metadata={}) diff --git a/tests/unit/redis/stores/auth_cache_double.py b/tests/unit/redis/stores/auth_cache_double.py index 056770f1860..d440d1257a6 100644 --- a/tests/unit/redis/stores/auth_cache_double.py +++ b/tests/unit/redis/stores/auth_cache_double.py @@ -44,6 +44,16 @@ class FakeAuthCache: commands: list[tuple[object, ...]] = field(default_factory=list) settings: list[AuthCacheSettings] = field(default_factory=list) fail: BaseException | None = None + generation: int = 0 + + def current_generation(self) -> int: + return self.generation + + def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: + if generation != self.generation: + return False + self.set_local(key, value, ttl_seconds) + return True def _round_trip(self, *command: object) -> None: self.commands.append(command) @@ -147,9 +157,11 @@ class FakeAuthCache: self.memory[key] = (value, ttl_seconds) def delete_local(self, key: str) -> None: + self.generation += 1 self.memory.pop(key, None) def flush_local(self) -> None: + self.generation += 1 self.memory.clear() def len_local(self) -> LocalSizes: diff --git a/tests/unit/repositories/test_database_settings.py b/tests/unit/repositories/test_database_settings.py new file mode 100644 index 00000000000..8dac4a6a4cb --- /dev/null +++ b/tests/unit/repositories/test_database_settings.py @@ -0,0 +1,62 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm.repositories.database_settings import NativeDatabaseSettings, NativeDatabaseSettingsError + + +def test_projects_prisma_connection_settings_without_forwarding_url() -> None: + settings: Final = NativeDatabaseSettings.from_url( + "postgresql://auth%40example.test:p%3Aa%2Fs%25s@db.test:5433/gateway%20db" + "?schema=auth%20schema&connection_limit=7&pool_timeout=3&connect_timeout=4" + "&max_idle_connection_lifetime=60&pgbouncer=true&sslmode=require&sslaccept=strict&sslcert=/ca.pem" + ) + assert settings == NativeDatabaseSettings( + host="db.test", + port=5433, + username="auth@example.test", + password="p:a/s%s", + database="gateway db", + schema="auth schema", + ssl_mode="verify-full", + ssl_root_cert="/ca.pem", + max_connections=7, + acquire_timeout=3, + idle_timeout=60, + statement_cache_capacity=0, + connect_timeout=4, + ) + + +@pytest.mark.parametrize( + ("query", "ssl_mode", "certificate"), + ( + ("sslmode=disable&sslaccept=strict", "disable", None), + ("sslmode=verify-ca&sslrootcert=/root.pem", "verify-ca", "/root.pem"), + ("sslmode=require&sslaccept=accept_invalid_certs", "require", None), + ("sslmode=prefer", "prefer", None), + ), +) +def test_preserves_tls_policy(query: str, ssl_mode: str, certificate: str | None) -> None: + settings: Final = NativeDatabaseSettings.from_url(f"postgresql://user:password@db.test/gateway?{query}") + assert (settings.ssl_mode, settings.ssl_root_cert) == (ssl_mode, certificate) + + +def test_database_settings_repr_redacts_password_and_provider() -> None: + settings: Final = NativeDatabaseSettings.from_url("postgresql://user:opaque-password@db.test/gateway") + assert "opaque-password" not in repr(settings) + assert repr(replace(settings, password="another-password")) == repr(settings) + + +def test_preserves_unix_socket_and_disabled_idle_expiration() -> None: + settings: Final = NativeDatabaseSettings.from_url( + "postgresql://user@localhost/gateway?host=%2Fvar%2Frun%2Fpostgresql&max_idle_connection_lifetime=0" + ) + assert (settings.host, settings.password, settings.idle_timeout) == ("/var/run/postgresql", None, None) + + +def test_rejects_non_postgres_without_echoing_credentials() -> None: + with pytest.raises(NativeDatabaseSettingsError, match="PostgreSQL is required") as error: + NativeDatabaseSettings.from_url("mysql://user:opaque-password@db.test/gateway") + assert "opaque-password" not in str(error.value) diff --git a/tests/unit/repositories/test_native_auth_repository.py b/tests/unit/repositories/test_native_auth_repository.py new file mode 100644 index 00000000000..b0e50a288d5 --- /dev/null +++ b/tests/unit/repositories/test_native_auth_repository.py @@ -0,0 +1,149 @@ +import asyncio +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.repositories.native_auth_repository import NativeAuthRepository, native_auth_repository + + +@pytest.mark.asyncio +async def test_native_key_preserves_joined_budget_and_membership_projection(): + raw = { + "token": "token-hash", + "user_id": "user", + "organization_id": "key-org", + "org_id": "team-org", + "team_models": None, + "team_blocked": None, + "max_budget": None, + "litellm_budget_table_max_budget": 25.0, + "model_max_budget": {}, + "litellm_budget_table_model_max_budget": {"model": {"max_budget": 1.0}}, + "team_model_aliases": {"alias": "model"}, + "team_members_with_roles": [{"user_id": "user", "role": "user"}, {"user_id": "user", "role": "admin"}], + "expires": datetime(2027, 1, 1, tzinfo=timezone.utc), + } + handle = SimpleNamespace(auth_read=AsyncMock(return_value=[raw])) + row = await NativeAuthRepository(handle).key("token-hash", writer=True) + + assert row is not None + assert (row.org_id, row.team_models, row.team_blocked, row.max_budget) == ("key-org", [], False, 25.0) + assert row.model_max_budget == {"model": {"max_budget": 1.0}} + assert row.team_model_aliases == {"alias": "model"} + assert row.team_member is not None and row.team_member.role == "admin" + assert row.expires == "2027-01-01T00:00:00+00:00" + assert row.last_refreshed_at is not None + assert raw["max_budget"] is None and raw["team_models"] is None + handle.auth_read.assert_awaited_once_with("key", ["token-hash"], writer=True) + + +@pytest.mark.asyncio +async def test_native_rows_keep_null_limits_and_make_request_local_models(): + raw = {"user_id": "u", "team_id": "t", "spend": None, "litellm_budget_table": {"max_budget": None}} + handle = SimpleNamespace(auth_read=AsyncMock(return_value=[raw])) + repo = NativeAuthRepository(handle) + first = await repo.one("member", ["u", "t"], LiteLLM_TeamMembership) + second = await repo.one("member", ["u", "t"], LiteLLM_TeamMembership) + + assert first is not None and second is not None + assert first.spend is None and second.spend is None + assert first.litellm_budget_table is not None and second.litellm_budget_table is not None + first.litellm_budget_table.max_budget = 3.0 + assert second.litellm_budget_table.max_budget is None + assert raw["litellm_budget_table"] == {"max_budget": None} + + +@pytest.mark.asyncio +async def test_native_large_id_batch_is_chunked_deduplicated_and_keeps_missing_rows_absent(): + async def read(kind, ids, writer=False): + assert kind == "budget" and writer is True + assert len(ids) <= 1000 + return [{"budget_id": entity_id} for entity_id in ids if entity_id != "missing"] + + ids = tuple(f"budget-{i}" for i in range(1001)) + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=read)) + rows = await NativeAuthRepository(handle).many( + "budget", (*ids, ids[0], "missing"), LiteLLM_BudgetTable, writer=True + ) + assert tuple(row.budget_id for row in rows) == ids + assert [len(call.args[1]) for call in handle.auth_read.call_args_list] == [1000, 2] + + +@pytest.mark.asyncio +async def test_native_large_batch_error_never_returns_partial_rows(): + from litellm.repositories.native_auth_repository import NativeAuthDatabaseError + + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=[[{"budget_id": "first"}], OSError("offline")])) + with pytest.raises(NativeAuthDatabaseError): + await NativeAuthRepository(handle).many("budget", tuple(str(i) for i in range(1001)), LiteLLM_BudgetTable) + assert handle.auth_read.await_count == 2 + + +@pytest.mark.asyncio +async def test_native_empty_batch_preserves_native_empty_read_contract(): + handle = SimpleNamespace(auth_read=AsyncMock(return_value=[])) + rows = await NativeAuthRepository(handle).many("budget", (), LiteLLM_BudgetTable) + assert rows == () + handle.auth_read.assert_awaited_once_with("budget", [], writer=False) + + +@pytest.mark.asyncio +async def test_native_missing_row_is_none_and_zero_budget_is_not_missing(): + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=[[], [{"budget_id": "b", "max_budget": 0.0}]])) + repo = NativeAuthRepository(handle) + assert await repo.one("budget", ["missing"], LiteLLM_BudgetTable) is None + found = await repo.one("budget", ["b"], LiteLLM_BudgetTable) + assert found is not None and found.max_budget == 0.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("answer", [None, {}, ["bad"], [{"max_budget": "bad"}]]) +async def test_native_invalid_rows_are_terminal(answer): + handle = SimpleNamespace(auth_read=AsyncMock(return_value=answer)) + with pytest.raises(NativeAuthFailure): + await NativeAuthRepository(handle).one("budget", ["b"], LiteLLM_BudgetTable) + assert handle.auth_read.await_count == 1 + + +@pytest.mark.asyncio +async def test_native_failure_and_cancellation_are_not_clean_misses(): + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("connection closed"))) + with pytest.raises(NativeAuthFailure, match="read failed"): + await NativeAuthRepository(handle).key("hash") + cancelled = SimpleNamespace(auth_read=AsyncMock(side_effect=asyncio.CancelledError)) + with pytest.raises(asyncio.CancelledError): + await NativeAuthRepository(cancelled).key("hash") + + +@pytest.mark.asyncio +async def test_native_timeout_keeps_deadline_identity_and_terminal_native_semantics(): + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=asyncio.TimeoutError)) + with pytest.raises(NativeAuthFailure) as raised: + await NativeAuthRepository(handle).key("hash") + assert isinstance(raised.value, DBLookupDeadlineExceeded) + assert raised.value.lookup == "key" + + +@pytest.mark.asyncio +async def test_native_spend_preserves_missing_and_authoritative_zero(): + handle = SimpleNamespace(auth_spend=AsyncMock(side_effect=[None, 0.0, RuntimeError("offline")])) + repo = NativeAuthRepository(handle) + assert await repo.spend("user", "u") is None + assert await repo.spend("user", "u") == 0.0 + with pytest.raises(NativeAuthFailure): + await repo.spend("user", "u") + + +def test_repository_selection_requires_an_explicit_handle(): + assert native_auth_repository(SimpleNamespace(native_auth_database=None)) is None + assert native_auth_repository(MagicMock()) is None + handle = SimpleNamespace(auth_read=AsyncMock()) + repo = native_auth_repository(SimpleNamespace(native_auth_database=handle)) + assert repo is not None and repo.database is handle