diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fbcb35d66c9..2ae285d6eef 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4012,6 +4012,64 @@ async def get_org_object( return _org_obj +def _last_known_org_cache_key(org_id: str) -> str: + return f"org_id:{org_id}:with_budget:last_known" + + +async def _keep_last_known_org( + org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache +) -> None: + cache_key: Final = _last_known_org_cache_key(org_id) + held_locally: Final = await user_api_key_cache.async_get_cache( + key=cache_key, local_only=True, model_type=LiteLLM_OrganizationTable + ) + if held_locally is not None: + return + await user_api_key_cache.async_set_cache( + key=cache_key, + value=org, + model_type=LiteLLM_OrganizationTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + + +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: + org: Final = 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 not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): + verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) + return None + last_known_org: Final = await user_api_key_cache.async_get_cache( + key=_last_known_org_cache_key(org_id), + model_type=LiteLLM_OrganizationTable, + ) + if last_known_org is not None: + return last_known_org + if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable(): + return None + raise + if org is None: + return None + await _keep_last_known_org(org, org_id, user_api_key_cache) + return org + + 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 6f6db73f99e..de0131772bc 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -58,6 +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_for_request, get_project_object, get_team_membership, get_team_object, @@ -2611,6 +2612,47 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() +async def _inherit_org_identity( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTableCachedObj | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> None: + if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: + user_api_key_auth_obj.org_id = team_object.organization_id + already_populated: Final = any( + value is not None + for value in ( + user_api_key_auth_obj.organization_alias, + user_api_key_auth_obj.organization_max_budget, + user_api_key_auth_obj.organization_tpm_limit, + user_api_key_auth_obj.organization_rpm_limit, + user_api_key_auth_obj.organization_metadata, + ) + ) + if user_api_key_auth_obj.org_id is None or already_populated or prisma_client is None: + 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 + user_api_key_auth_obj.organization_metadata = org_object.metadata + budget: Final = org_object.litellm_budget_table + if budget is None: + return + user_api_key_auth_obj.organization_max_budget = budget.max_budget + user_api_key_auth_obj.organization_tpm_limit = budget.tpm_limit + user_api_key_auth_obj.organization_rpm_limit = budget.rpm_limit + + def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, object]) -> bool: return master_key is None and not any( general_settings.get(flag, False) @@ -2849,8 +2891,14 @@ async def _run_centralized_common_checks( user_object=user_object, ) - if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: - user_api_key_auth_obj.org_id = team_object.organization_id + await _inherit_org_identity( + user_api_key_auth_obj=user_api_key_auth_obj, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and 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 90ce821d62e..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 @@ -5739,9 +5739,15 @@ class TestMCPDcrBridgeDelegateAdmission: prisma and are swallowed (``_safe_fetch`` / the SCIM gate's fail-open), so their checks skip. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was the reload key.""" + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError + get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect) + get_org_object = AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")) 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.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 d7666f5e694..b0cda30dfe5 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 @@ -11974,9 +11974,13 @@ async def test_oauth_credential_write_keeps_virtual_key_permissions( from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.bridge_token_flow import authorize_oauth_credential_request from litellm.proxy._types import UserAPIKeyAuth, hash_token - from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError, jwt_key_mapping_cache_key handler, signing_key = jwt_oauth_identity + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.get_org_object", + AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")), + ) key: Final = "sk-oauth-permission-test" hashed: Final = hash_token(key) credential: Final = UserAPIKeyAuth( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0a6f6d8e69b..a0256e40b8c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6087,6 +6087,107 @@ async def test_organization_budget_check_carries_org_state_on_the_token(): assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0) +@pytest.mark.parametrize("warmed_by_auth_prefetch", [False, True]) +@pytest.mark.asyncio +async def test_get_org_object_for_request_serves_last_known_org_through_db_outage(warmed_by_auth_prefetch): + """A JWT whose team sits in an org resolves the org on every request, and the org row + is cached for only DEFAULT_IN_MEMORY_TTL seconds while the team and user rows ride the + 60s management-object TTL. Without a last-known copy, a DB outage a few seconds old + turned that traffic into 503s while the same request through a virtual key kept + succeeding on its cached team. The copy must exist whoever filled the short-lived entry: + this lookup's own DB read, or the virtual-key auth prefetch warming it for the same org.""" + from litellm.proxy._types import LiteLLM_OrganizationTable + from litellm.proxy.auth.auth_checks import get_org_object_for_request + + org_columns = { + "organization_id": "org-1", + "organization_alias": "platform-org", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + "litellm_budget_table": {"budget_id": "b1", "max_budget": 50.0, "tpm_limit": 700, "rpm_limit": 7}, + } + org_row = MagicMock() + org_row.model_dump = lambda: org_columns + db_outage = ConnectionRefusedError("db unavailable") + prisma_client = MagicMock() + prisma_client.db.litellm_organizationtable.find_unique = AsyncMock( + side_effect=[db_outage] if warmed_by_auth_prefetch else [org_row, db_outage] + ) + user_api_key_cache = UserApiKeyCache() + if warmed_by_auth_prefetch: + await user_api_key_cache.async_set_cache( + key="org_id:org-1:with_budget", + value=LiteLLM_OrganizationTable.model_validate(org_columns), + model_type=LiteLLM_OrganizationTable, + ) + + async def _lookup(): + return await get_org_object_for_request( + org_id="org-1", + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: the outage fallback reads this module global; no dependency injection seam exists + warm = await _lookup() + assert warm is not None and warm.organization_alias == "platform-org" + await user_api_key_cache.async_delete_cache("org_id:org-1:with_budget") + + during_outage = await _lookup() + + assert prisma_client.db.litellm_organizationtable.find_unique.await_count == (1 if warmed_by_auth_prefetch else 2) + assert during_outage is not None + assert during_outage.organization_alias == "platform-org" + assert during_outage.litellm_budget_table is not None + assert during_outage.litellm_budget_table.rpm_limit == 7 + assert during_outage.litellm_budget_table.max_budget == 50.0 + + +@pytest.mark.asyncio +async def test_get_org_object_for_request_writes_the_last_known_org_only_when_absent(): + """The last-known copy is written when this worker holds none, never per request: + with Redis attached, a write on every cached org hit would cost one SET per JWT request.""" + from litellm.proxy._types import LiteLLM_OrganizationTable + from litellm.proxy.auth.auth_checks import get_org_object_for_request + + class _WriteRecordingCache(UserApiKeyCache): + def __init__(self): + super().__init__() + self.written_keys = [] + + async def async_set_cache(self, key, value, local_only=False, **kwargs): + self.written_keys.append(key) + return await super().async_set_cache(key=key, value=value, local_only=local_only, **kwargs) + + user_api_key_cache = _WriteRecordingCache() + await user_api_key_cache.async_set_cache( + key="org_id:org-1:with_budget", + value=LiteLLM_OrganizationTable( + organization_id="org-1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + ), + model_type=LiteLLM_OrganizationTable, + ) + + for _ in range(3): + org = await get_org_object_for_request( + org_id="org-1", + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert org is not None and org.organization_alias == "platform-org" + + assert user_api_key_cache.written_keys.count("org_id:org-1:with_budget:last_known") == 1 + + @pytest.mark.parametrize( "max_budget, spend, expect_blocked", [ 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 e6975edd4cb..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 @@ -25,6 +25,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, @@ -35,6 +36,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.auth_checks import ( + OrganizationNotFoundError, TeamNotFoundError, UserNotFoundError, get_key_object, @@ -5986,6 +5988,152 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_id,team_org_id,existing_alias,existing_rpm,lookup_mode,allow_db_unavailable,expect_lookup_error,expected_org_id,expected_alias,expected_limits", + [ + (None, "t1", "org-from-team", None, None, "success", False, False, "org-from-team", "acme-org", (12.5, 700, 7)), + ("org-jwt", None, None, None, None, "success", False, False, "org-jwt", "acme-org", (12.5, 700, 7)), + ("org-pinned", None, None, "preset", None, "success", False, False, "org-pinned", "preset", (None, None, None)), + ("org-view", None, None, None, 3, "success", False, False, "org-view", None, (None, None, 3)), + ("org-missing", None, None, None, None, "missing", False, False, "org-missing", None, (None, None, None)), + ("org-db-failure-allowed", None, None, None, None, "db_failure", True, False, "org-db-failure-allowed", None, (None, None, None)), + ("org-db-failure-denied", None, None, None, None, "db_failure", False, True, "org-db-failure-denied", None, (None, None, None)), + ("org-bad-row", None, None, None, None, "bad_row", False, False, "org-bad-row", None, (None, None, None)), + ("org-nobudget", None, None, None, None, "no_budget", False, False, "org-nobudget", "acme-org", (None, None, None)), + ], +) +async def test_centralized_common_checks_inherits_org_identity( + key_org_id: str | None, + team_id: str | None, + team_org_id: str | None, + existing_alias: str | None, + existing_rpm: int | None, + lookup_mode: str, + allow_db_unavailable: bool, + expect_lookup_error: bool, + expected_org_id: str | None, + expected_alias: str | None, + expected_limits: tuple[float | None, int | None, int | None], +) -> None: + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id=team_id, + org_id=key_org_id, + organization_alias=existing_alias, + organization_rpm_limit=existing_rpm, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = ( + LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) if team_id is not None else None + ) + organization = LiteLLM_OrganizationTable( + organization_id=expected_org_id, + organization_alias="acme-org", + budget_id="budget-id", + metadata={"model_rpm_limit": {"gpt-4o": 2}}, + models=[], + created_by="test", + updated_by="test", + litellm_budget_table=( + None + if lookup_mode == "no_budget" + else LiteLLM_BudgetTable(budget_id="budget-id", max_budget=12.5, tpm_limit=700, rpm_limit=7) + ), + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + attrs["prisma_client"] = MagicMock() + attrs["general_settings"] = {"allow_requests_on_db_unavailable": allow_db_unavailable} + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( # test-quality-ok: centralized auth calls this module helper directly; no dependency injection seam exists + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + 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.auth_checks.get_org_object", + new_callable=AsyncMock, + return_value=organization, + ) as mock_get_org_object, + patch( # test-quality-ok: capture downstream token state without invoking unrelated common checks + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + if lookup_mode == "missing": + mock_get_org_object.side_effect = OrganizationNotFoundError("x") + elif lookup_mode == "db_failure": + mock_get_org_object.side_effect = ConnectionRefusedError("db unavailable") + elif lookup_mode == "bad_row": + mock_get_org_object.side_effect = ValueError("row failed validation") + + if expect_lookup_error: + with pytest.raises(ConnectionRefusedError, match="db unavailable"): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + else: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + assert token.org_id == expected_org_id + if expect_lookup_error: + mock_checks.assert_not_awaited() + assert token.organization_alias is None + assert token.organization_max_budget is None + assert token.organization_tpm_limit is None + assert token.organization_rpm_limit is None + return + + mock_checks.assert_awaited_once() + assert token.organization_alias == expected_alias + assert ( + token.organization_max_budget, + token.organization_tpm_limit, + token.organization_rpm_limit, + ) == expected_limits + checked_token = mock_checks.await_args.kwargs["valid_token"] + assert checked_token.org_id == expected_org_id + assert checked_token.organization_alias == expected_alias + if team_id is None: + mock_get_team_object.assert_not_awaited() + else: + mock_get_team_object.assert_awaited_once() + if existing_alias is not None or existing_rpm is not None: + mock_get_org_object.assert_not_awaited() + assert token.organization_metadata is None + else: + mock_get_org_object.assert_awaited_once() + assert mock_get_org_object.await_args.kwargs["org_id"] == expected_org_id + assert mock_get_org_object.await_args.kwargs["include_budget_table"] is True + if lookup_mode not in {"missing", "db_failure", "bad_row"}: + assert token.organization_metadata == {"model_rpm_limit": {"gpt-4o": 2}} + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_cli_session_token_org_backfilled_from_team(monkeypatch): """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted