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

@@ -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"],

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

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()),

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"

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,