From 82e3f3980d44f3822fa30ae089d0a034335a402c Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 23:59:33 +0000 Subject: [PATCH] refactor(auth): resolve org identity through an auth_checks helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 28 ++++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 29 +++++-------------- .../auth/test_user_api_key_auth_mcp.py | 2 +- .../mcp_server/test_discoverable_endpoints.py | 2 +- .../proxy/auth/test_user_api_key_auth.py | 2 +- 5 files changed, 39 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index cdada970956..ba37eed037f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4012,6 +4012,34 @@ async def get_org_object( return _org_obj +async def get_org_object_for_request( + org_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> LiteLLM_OrganizationTable | None: + try: + return await get_org_object( + org_id=org_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, + include_budget_table=True, + ) + except OrganizationNotFoundError: + return None + except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits + if ( + PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e) + and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + ): + raise + verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) + return None + + async def _get_resources_from_access_groups( access_group_ids: Sequence[str], resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b4d8648c8a9..7ef1c2775ab 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -41,7 +41,6 @@ from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, - OrganizationNotFoundError, TeamNotFoundError, _cache_key_object, _can_object_call_model, @@ -59,7 +58,7 @@ from litellm.proxy.auth.auth_checks import ( get_jwt_key_mapping_object, get_key_end_user_budget_id, get_object_permission, - get_org_object, + get_org_object_for_request, get_project_object, get_team_membership, get_team_object, @@ -2630,25 +2629,13 @@ async def _inherit_org_identity( ) if user_api_key_auth_obj.org_id is None or already_populated or prisma_client is None: return - try: - org_object: Final = await get_org_object( - org_id=user_api_key_auth_obj.org_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, - include_budget_table=True, - ) - except OrganizationNotFoundError: - return - except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits - if ( - PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e) - and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - ): - raise - verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) - return + org_object: Final = await get_org_object_for_request( + org_id=user_api_key_auth_obj.org_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 org_object is None: return user_api_key_auth_obj.organization_alias = org_object.organization_alias diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index a0fb76349b2..4380df194ed 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5746,7 +5746,7 @@ class TestMCPDcrBridgeDelegateAdmission: patchers = [ patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object), patch( # test-quality-ok: central auth now resolves org limits; this fixture models a missing org row - "litellm.proxy.auth.user_api_key_auth.get_org_object", get_org_object + "litellm.proxy.auth.auth_checks.get_org_object", get_org_object ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 200df078e00..3556722ff6e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11874,7 +11874,7 @@ async def test_oauth_credential_write_keeps_virtual_key_permissions( handler, signing_key = jwt_oauth_identity monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.get_org_object", + "litellm.proxy.auth.auth_checks.get_org_object", AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")), ) key: Final = "sk-oauth-permission-test" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 761bd454eaf..8593be751fa 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -6065,7 +6065,7 @@ async def test_centralized_common_checks_inherits_org_identity( return_value=fetched_team, ) as mock_get_team_object, patch( # test-quality-ok: centralized auth calls this module helper directly; no dependency injection seam exists - "litellm.proxy.auth.user_api_key_auth.get_org_object", + "litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock, return_value=organization, ) as mock_get_org_object,