diff --git a/litellm/proxy/auth/v2/INTEGRATION.md b/litellm/proxy/auth/v2/INTEGRATION.md index e58ad0571f5..83bfd703be5 100644 --- a/litellm/proxy/auth/v2/INTEGRATION.md +++ b/litellm/proxy/auth/v2/INTEGRATION.md @@ -45,13 +45,13 @@ independent of v1/v2. JWT, and OAuth authenticators return a thin identity (no budget/limit fields), so the hooks read `None` and enforce nothing for those logins. -**Built.** `enrichment.py` (`enrich_identity`) copies the user/team limit fields +**Built.** `stages/enrichment.py` (`enrich_identity`) copies the user/team limit fields 1:1 from the rows into the identity's distinct `user_*` / `team_*` slots, filling only unset fields so it never overrides an already-resolved value. It is wired into the inference path in `entry.py` for non-virtual-key logins (`_enrich_for_limits`), with the loaders (`get_user_object` auth_checks.py:1650, `get_team_object` auth_checks.py:1982) injected. Unit-tested in -`test_enrichment.py`. +`stages/test_enrichment.py`. **What remains (live).** Only the verification below. The mapping is additive (these logins enforce nothing today, so it cannot regress existing behavior), but @@ -74,12 +74,12 @@ limit silently over- or under-enforces a customer's spend. **Done.** Team, organization, and global proxy-spend caps live in v1's `common_checks`, which v2 does not run. Rather than move them (which would risk -v1's path), `budgets.py` (`enforce_hierarchy_budgets`) calls the *same* functions +v1's path), `stages/budgets.py` (`enforce_hierarchy_budgets`) calls the *same* functions v1 uses — `_team_max_budget_check`, `_organization_max_budget_check`, and `get_global_proxy_spend` + `_global_proxy_budget_check` — from v2's inference path. One implementation, two callers: single authority, correct counter keys, no edit to v1. A breach surfaces as the same 429 `ProxyException` v1 raises. -Unit-tested in `test_budgets.py` (team over/under budget, no-team no-op). +Unit-tested in `stages/test_budgets.py` (team over/under budget, no-team no-op). **What remains (live).** Cross-pod spend-counter accuracy. Set a team `max_budget` (and `litellm.max_budget` for the global cap) below current spend and confirm diff --git a/litellm/proxy/auth/v2/__init__.py b/litellm/proxy/auth/v2/__init__.py index 4ed1e7ad7e5..ca09fd64c55 100644 --- a/litellm/proxy/auth/v2/__init__.py +++ b/litellm/proxy/auth/v2/__init__.py @@ -6,9 +6,9 @@ from .context import ( set_auth_context, try_get_auth_context, ) -from .end_user import resolve_end_user -from .enrichment import enrich_identity from .entry import user_api_key_auth_v2 +from .stages.end_user import resolve_end_user +from .stages.enrichment import enrich_identity __all__ = [ "user_api_key_auth_v2", diff --git a/litellm/proxy/auth/v2/authn/__init__.py b/litellm/proxy/auth/v2/authn/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/auth/v2/authenticators.py b/litellm/proxy/auth/v2/authn/authenticators.py similarity index 99% rename from litellm/proxy/auth/v2/authenticators.py rename to litellm/proxy/auth/v2/authn/authenticators.py index a86a026d7b0..ca4e082a7ec 100644 --- a/litellm/proxy/auth/v2/authenticators.py +++ b/litellm/proxy/auth/v2/authn/authenticators.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Protocol, runtime_checkab from fastapi import HTTPException, status -from .context import AuthMethod +from ..context import AuthMethod if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth diff --git a/litellm/proxy/auth/v2/jwt_claims.py b/litellm/proxy/auth/v2/authn/jwt_claims.py similarity index 100% rename from litellm/proxy/auth/v2/jwt_claims.py rename to litellm/proxy/auth/v2/authn/jwt_claims.py diff --git a/litellm/proxy/auth/v2/jwt_verifier.py b/litellm/proxy/auth/v2/authn/jwt_verifier.py similarity index 100% rename from litellm/proxy/auth/v2/jwt_verifier.py rename to litellm/proxy/auth/v2/authn/jwt_verifier.py diff --git a/litellm/proxy/auth/v2/oauth2_introspection.py b/litellm/proxy/auth/v2/authn/oauth2_introspection.py similarity index 100% rename from litellm/proxy/auth/v2/oauth2_introspection.py rename to litellm/proxy/auth/v2/authn/oauth2_introspection.py diff --git a/litellm/proxy/auth/v2/authz/__init__.py b/litellm/proxy/auth/v2/authz/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/auth/v2/authorizer.py b/litellm/proxy/auth/v2/authz/authorizer.py similarity index 98% rename from litellm/proxy/auth/v2/authorizer.py rename to litellm/proxy/auth/v2/authz/authorizer.py index cf8e03cde4b..c16327f31d8 100644 --- a/litellm/proxy/auth/v2/authorizer.py +++ b/litellm/proxy/auth/v2/authz/authorizer.py @@ -1,7 +1,7 @@ import logging from typing import Any, Dict, Optional -from .principal import Principal +from ..principal import Principal from .route_map import GovernedRoute, match_route logger = logging.getLogger("litellm.proxy.auth.v2") diff --git a/litellm/proxy/auth/v2/enforcer.py b/litellm/proxy/auth/v2/authz/enforcer.py similarity index 100% rename from litellm/proxy/auth/v2/enforcer.py rename to litellm/proxy/auth/v2/authz/enforcer.py diff --git a/litellm/proxy/auth/v2/model.conf b/litellm/proxy/auth/v2/authz/model.conf similarity index 100% rename from litellm/proxy/auth/v2/model.conf rename to litellm/proxy/auth/v2/authz/model.conf diff --git a/litellm/proxy/auth/v2/policy_admin.py b/litellm/proxy/auth/v2/authz/policy_admin.py similarity index 100% rename from litellm/proxy/auth/v2/policy_admin.py rename to litellm/proxy/auth/v2/authz/policy_admin.py diff --git a/litellm/proxy/auth/v2/policy_store.py b/litellm/proxy/auth/v2/authz/policy_store.py similarity index 100% rename from litellm/proxy/auth/v2/policy_store.py rename to litellm/proxy/auth/v2/authz/policy_store.py diff --git a/litellm/proxy/auth/v2/route_map.py b/litellm/proxy/auth/v2/authz/route_map.py similarity index 100% rename from litellm/proxy/auth/v2/route_map.py rename to litellm/proxy/auth/v2/authz/route_map.py diff --git a/litellm/proxy/auth/v2/entry.py b/litellm/proxy/auth/v2/entry.py index 1bef310ad80..9450d28144e 100644 --- a/litellm/proxy/auth/v2/entry.py +++ b/litellm/proxy/auth/v2/entry.py @@ -4,16 +4,16 @@ from fastapi import HTTPException, Request, status from litellm.integrations.otel.runtime import seed_request_identity -from .authenticators import AuthContext, AuthResult, authenticate -from .authorizer import AuthorizationDenied, authorize -from .budgets import enforce_hierarchy_budgets +from .authn.authenticators import AuthContext, AuthResult, authenticate +from .authz.authorizer import AuthorizationDenied, authorize +from .authz.enforcer import CasbinEnforcer +from .authz.policy_store import load_policy_snapshot +from .authz.route_map import is_inference_route, match_route from .context import AuthMethod, RequestAuthContext, set_auth_context -from .end_user import resolve_end_user -from .enforcer import CasbinEnforcer -from .enrichment import enrich_identity -from .policy_store import load_policy_snapshot from .principal import Principal, build_principal -from .route_map import is_inference_route, match_route +from .stages.budgets import enforce_hierarchy_budgets +from .stages.end_user import resolve_end_user +from .stages.enrichment import enrich_identity async def _anonymous_identity(api_key: Optional[str]) -> Any: diff --git a/litellm/proxy/auth/v2/management_endpoints.py b/litellm/proxy/auth/v2/management_endpoints.py index 0dd55192832..fce0c80765d 100644 --- a/litellm/proxy/auth/v2/management_endpoints.py +++ b/litellm/proxy/auth/v2/management_endpoints.py @@ -6,12 +6,12 @@ from pydantic import BaseModel from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from .policy_admin import ( +from .authz.policy_admin import ( PolicyValidationError, make_assignment_rule, make_permission_rule, ) -from .policy_store import reset_cache +from .authz.policy_store import reset_cache router = APIRouter(tags=["auth_v2"]) diff --git a/litellm/proxy/auth/v2/stages/__init__.py b/litellm/proxy/auth/v2/stages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/auth/v2/budgets.py b/litellm/proxy/auth/v2/stages/budgets.py similarity index 100% rename from litellm/proxy/auth/v2/budgets.py rename to litellm/proxy/auth/v2/stages/budgets.py diff --git a/litellm/proxy/auth/v2/end_user.py b/litellm/proxy/auth/v2/stages/end_user.py similarity index 97% rename from litellm/proxy/auth/v2/end_user.py rename to litellm/proxy/auth/v2/stages/end_user.py index d2eb74940e1..f7397311244 100644 --- a/litellm/proxy/auth/v2/end_user.py +++ b/litellm/proxy/auth/v2/stages/end_user.py @@ -2,7 +2,7 @@ from typing import Awaitable, Callable, Optional from starlette.requests import Request -from .context import attach_end_user +from ..context import attach_end_user # Extraction reuses the existing request-body/header logic; validation (the # customer-table lookup) is injected so this stage is testable without a DB and diff --git a/litellm/proxy/auth/v2/enrichment.py b/litellm/proxy/auth/v2/stages/enrichment.py similarity index 100% rename from litellm/proxy/auth/v2/enrichment.py rename to litellm/proxy/auth/v2/stages/enrichment.py diff --git a/tests/test_litellm/proxy/auth/v2/test_authenticators.py b/tests/test_litellm/proxy/auth/v2/authn/test_authenticators.py similarity index 98% rename from tests/test_litellm/proxy/auth/v2/test_authenticators.py rename to tests/test_litellm/proxy/auth/v2/authn/test_authenticators.py index c1b31f44981..9afd829560a 100644 --- a/tests/test_litellm/proxy/auth/v2/test_authenticators.py +++ b/tests/test_litellm/proxy/auth/v2/authn/test_authenticators.py @@ -1,6 +1,6 @@ import pytest -from litellm.proxy.auth.v2.authenticators import ( +from litellm.proxy.auth.v2.authn.authenticators import ( OAuth2IntrospectionAuthenticator, JWTAuthenticator, MasterKeyAuthenticator, diff --git a/tests/test_litellm/proxy/auth/v2/test_jwt_claims.py b/tests/test_litellm/proxy/auth/v2/authn/test_jwt_claims.py similarity index 96% rename from tests/test_litellm/proxy/auth/v2/test_jwt_claims.py rename to tests/test_litellm/proxy/auth/v2/authn/test_jwt_claims.py index 8c5a8085b85..9a579638135 100644 --- a/tests/test_litellm/proxy/auth/v2/test_jwt_claims.py +++ b/tests/test_litellm/proxy/auth/v2/authn/test_jwt_claims.py @@ -1,6 +1,6 @@ import pytest -from litellm.proxy.auth.v2.jwt_claims import ( +from litellm.proxy.auth.v2.authn.jwt_claims import ( JWTClaimError, JWTSettings, extract_identity, diff --git a/tests/test_litellm/proxy/auth/v2/test_jwt_verifier.py b/tests/test_litellm/proxy/auth/v2/authn/test_jwt_verifier.py similarity index 97% rename from tests/test_litellm/proxy/auth/v2/test_jwt_verifier.py rename to tests/test_litellm/proxy/auth/v2/authn/test_jwt_verifier.py index 3a02ff8bb72..68bc9a04902 100644 --- a/tests/test_litellm/proxy/auth/v2/test_jwt_verifier.py +++ b/tests/test_litellm/proxy/auth/v2/authn/test_jwt_verifier.py @@ -3,7 +3,7 @@ import time import pytest from authlib.jose import JsonWebKey, jwt -from litellm.proxy.auth.v2.jwt_verifier import JWTVerificationError, verify +from litellm.proxy.auth.v2.authn.jwt_verifier import JWTVerificationError, verify ISSUER = "https://idp.example" AUDIENCE = "litellm" diff --git a/tests/test_litellm/proxy/auth/v2/test_oauth2_introspection.py b/tests/test_litellm/proxy/auth/v2/authn/test_oauth2_introspection.py similarity index 96% rename from tests/test_litellm/proxy/auth/v2/test_oauth2_introspection.py rename to tests/test_litellm/proxy/auth/v2/authn/test_oauth2_introspection.py index 315adc37213..b1c12d51fff 100644 --- a/tests/test_litellm/proxy/auth/v2/test_oauth2_introspection.py +++ b/tests/test_litellm/proxy/auth/v2/authn/test_oauth2_introspection.py @@ -1,6 +1,6 @@ import pytest -from litellm.proxy.auth.v2.oauth2_introspection import ( +from litellm.proxy.auth.v2.authn.oauth2_introspection import ( IntrospectionSettings, OAuth2IntrospectionError, parse_introspection_response, diff --git a/tests/test_litellm/proxy/auth/v2/test_authorizer.py b/tests/test_litellm/proxy/auth/v2/authz/test_authorizer.py similarity index 96% rename from tests/test_litellm/proxy/auth/v2/test_authorizer.py rename to tests/test_litellm/proxy/auth/v2/authz/test_authorizer.py index 5ee6cb498db..71c17399a47 100644 --- a/tests/test_litellm/proxy/auth/v2/test_authorizer.py +++ b/tests/test_litellm/proxy/auth/v2/authz/test_authorizer.py @@ -2,7 +2,7 @@ import logging import pytest -from litellm.proxy.auth.v2.authorizer import AuthorizationDenied, authorize +from litellm.proxy.auth.v2.authz.authorizer import AuthorizationDenied, authorize from litellm.proxy.auth.v2.principal import Principal PRINCIPAL = Principal(subject="user:u1", domain="*", groupings=[]) diff --git a/tests/test_litellm/proxy/auth/v2/test_enforcer.py b/tests/test_litellm/proxy/auth/v2/authz/test_enforcer.py similarity index 98% rename from tests/test_litellm/proxy/auth/v2/test_enforcer.py rename to tests/test_litellm/proxy/auth/v2/authz/test_enforcer.py index 822b9376a8e..11629afbd55 100644 --- a/tests/test_litellm/proxy/auth/v2/test_enforcer.py +++ b/tests/test_litellm/proxy/auth/v2/authz/test_enforcer.py @@ -1,4 +1,4 @@ -from litellm.proxy.auth.v2.enforcer import CasbinEnforcer +from litellm.proxy.auth.v2.authz.enforcer import CasbinEnforcer READER_POLICY = ["role:model_reader", "*", "model:*", "read", "allow"] ADMIN_POLICY = ["role:proxy_admin", "*", "*", "*", "allow"] diff --git a/tests/test_litellm/proxy/auth/v2/test_policy_admin.py b/tests/test_litellm/proxy/auth/v2/authz/test_policy_admin.py similarity index 97% rename from tests/test_litellm/proxy/auth/v2/test_policy_admin.py rename to tests/test_litellm/proxy/auth/v2/authz/test_policy_admin.py index b54d0cdf721..2478d4a94bb 100644 --- a/tests/test_litellm/proxy/auth/v2/test_policy_admin.py +++ b/tests/test_litellm/proxy/auth/v2/authz/test_policy_admin.py @@ -1,6 +1,6 @@ import pytest -from litellm.proxy.auth.v2.policy_admin import ( +from litellm.proxy.auth.v2.authz.policy_admin import ( PolicyValidationError, make_assignment_rule, make_permission_rule, diff --git a/tests/test_litellm/proxy/auth/v2/test_policy_store.py b/tests/test_litellm/proxy/auth/v2/authz/test_policy_store.py similarity index 96% rename from tests/test_litellm/proxy/auth/v2/test_policy_store.py rename to tests/test_litellm/proxy/auth/v2/authz/test_policy_store.py index cc6079db1d1..f0a34e60d13 100644 --- a/tests/test_litellm/proxy/auth/v2/test_policy_store.py +++ b/tests/test_litellm/proxy/auth/v2/authz/test_policy_store.py @@ -1,7 +1,7 @@ import pytest -from litellm.proxy.auth.v2 import policy_store -from litellm.proxy.auth.v2.policy_store import ( +from litellm.proxy.auth.v2.authz import policy_store +from litellm.proxy.auth.v2.authz.policy_store import ( DEFAULT_POLICIES, load_policy_snapshot, reset_cache, diff --git a/tests/test_litellm/proxy/auth/v2/test_route_map.py b/tests/test_litellm/proxy/auth/v2/authz/test_route_map.py similarity index 98% rename from tests/test_litellm/proxy/auth/v2/test_route_map.py rename to tests/test_litellm/proxy/auth/v2/authz/test_route_map.py index f9ed99c7626..ad2c11582d6 100644 --- a/tests/test_litellm/proxy/auth/v2/test_route_map.py +++ b/tests/test_litellm/proxy/auth/v2/authz/test_route_map.py @@ -1,4 +1,4 @@ -from litellm.proxy.auth.v2.route_map import is_inference_route, match_route +from litellm.proxy.auth.v2.authz.route_map import is_inference_route, match_route def test_model_routes_map_to_resource_and_action(): diff --git a/tests/test_litellm/proxy/auth/v2/test_budgets.py b/tests/test_litellm/proxy/auth/v2/stages/test_budgets.py similarity index 96% rename from tests/test_litellm/proxy/auth/v2/test_budgets.py rename to tests/test_litellm/proxy/auth/v2/stages/test_budgets.py index f01237cc758..a040cb0d23d 100644 --- a/tests/test_litellm/proxy/auth/v2/test_budgets.py +++ b/tests/test_litellm/proxy/auth/v2/stages/test_budgets.py @@ -12,7 +12,7 @@ import pytest import litellm from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.v2.budgets import enforce_hierarchy_budgets +from litellm.proxy.auth.v2.stages.budgets import enforce_hierarchy_budgets class _Logging: diff --git a/tests/test_litellm/proxy/auth/v2/test_end_user.py b/tests/test_litellm/proxy/auth/v2/stages/test_end_user.py similarity index 96% rename from tests/test_litellm/proxy/auth/v2/test_end_user.py rename to tests/test_litellm/proxy/auth/v2/stages/test_end_user.py index e1211c396d9..fbe61207d98 100644 --- a/tests/test_litellm/proxy/auth/v2/test_end_user.py +++ b/tests/test_litellm/proxy/auth/v2/stages/test_end_user.py @@ -8,7 +8,7 @@ from litellm.proxy.auth.v2.context import ( get_auth_context, set_auth_context, ) -from litellm.proxy.auth.v2.end_user import resolve_end_user +from litellm.proxy.auth.v2.stages.end_user import resolve_end_user from litellm.proxy.auth.v2.principal import Principal diff --git a/tests/test_litellm/proxy/auth/v2/test_enrichment.py b/tests/test_litellm/proxy/auth/v2/stages/test_enrichment.py similarity index 97% rename from tests/test_litellm/proxy/auth/v2/test_enrichment.py rename to tests/test_litellm/proxy/auth/v2/stages/test_enrichment.py index beb2ba2afc3..6a93eb8ce8a 100644 --- a/tests/test_litellm/proxy/auth/v2/test_enrichment.py +++ b/tests/test_litellm/proxy/auth/v2/stages/test_enrichment.py @@ -2,7 +2,7 @@ from types import SimpleNamespace import pytest -from litellm.proxy.auth.v2.enrichment import enrich_identity +from litellm.proxy.auth.v2.stages.enrichment import enrich_identity def _identity(**overrides): diff --git a/tests/test_litellm/proxy/auth/v2/test_integration_flow.py b/tests/test_litellm/proxy/auth/v2/test_integration_flow.py index 70a64e99a09..1d0b6b133b2 100644 --- a/tests/test_litellm/proxy/auth/v2/test_integration_flow.py +++ b/tests/test_litellm/proxy/auth/v2/test_integration_flow.py @@ -13,7 +13,7 @@ from fastapi.testclient import TestClient import litellm.proxy.proxy_server as ps from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.auth.v2 import policy_store +from litellm.proxy.auth.v2.authz import policy_store from litellm.proxy.auth.v2.management_endpoints import router as auth_v2_router MASTER_KEY = "sk-master-1234"