diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index b1bf3483a9..b7dae9e9b4 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -331,8 +331,17 @@ class CustomGuardrail(CustomLogger): if "guardrails" in data: return data["guardrails"] - metadata = data.get("litellm_metadata") or data.get("metadata", {}) - return metadata.get("guardrails") or [] + # Check both metadata locations. For regular endpoints move_guardrails_to_metadata + # writes to "metadata"; for thread/assistant endpoints it writes to + # "litellm_metadata". We check the one that actually contains the "guardrails" + # key so that a non-empty litellm_metadata without guardrails does not shadow + # the merged list stored in metadata (which would cause team guardrails to be + # silently skipped while default_on=True policy guardrails still fire). + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) or {} + if isinstance(meta, dict) and "guardrails" in meta: + return meta.get("guardrails") or [] + return [] def _guardrail_is_in_requested_guardrails( self, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2ad643e297..b8db3cd2a7 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1445,6 +1445,10 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 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 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 ca8c7344ab..9c43ebcbe7 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 @@ -1737,7 +1737,112 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_team_metadata_refreshed_from_team_object_during_auth(): + """ + Regression test: when a cached API key has stale team_metadata (e.g. a + guardrail was added to the team after the key was cached), the auth flow + must update valid_token.team_metadata from the freshly fetched team object + so that move_guardrails_to_metadata picks up the new guardrail. + + Before the fix: valid_token.team_metadata was never updated from _team_obj + at the "Check 6" team-auth step in _user_api_key_auth_builder, so stale + team_metadata persisted for the lifetime of the key cache entry. + """ + from starlette.datastructures import URL + from starlette.requests import Request + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + api_key = "sk-test-team-metadata-refresh" + + # Simulate a cached key whose team_metadata was captured BEFORE the + # guardrail was added — so it has no "guardrails" key. + stale_team_metadata: dict = {"some_old_key": "some_old_value"} + valid_token = UserAPIKeyAuth( + api_key=api_key, + token=api_key, + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-guardrail-test", + team_metadata=stale_team_metadata, + ) + + # The fresh team object returned by get_team_object has the new guardrail. + fresh_team_obj = LiteLLM_TeamTableCachedObj( + team_id="team-guardrail-test", + metadata={"guardrails": ["test-guardrail-333"]}, + ) + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=valid_token) + mock_cache.async_set_cache = AsyncMock(return_value=None) + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _originals = {k: getattr(_proxy_server_mod, k, None) for k in _attrs} + + try: + for k, v in _attrs.items(): + setattr(_proxy_server_mod, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fresh_team_obj, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert result.team_metadata == {"guardrails": ["test-guardrail-333"]}, ( + f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}" + ) + + finally: + for k, v in _originals.items(): + setattr(_proxy_server_mod, k, v) + # --------------------------------------------------------------------------- + # _run_centralized_common_checks — centralized authz gate # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 1f232e410a..1e5c688e36 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -3281,3 +3281,144 @@ def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth # Even with BYOK enabled, x-api-key must be stripped when it was used # as the LiteLLM auth header (anti-replay guard). assert "x-api-key" not in result + + +# --------------------------------------------------------------------------- +# Team guardrail + global policy regression tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_team_guardrail_merges_with_global_policy(): + """ + Regression: team's direct guardrail must be present alongside guardrails + resolved from a global policy (scope='*') configured by the admin. + + The bug: get_guardrail_from_metadata checked litellm_metadata before + metadata. When the request contained a non-empty litellm_metadata field + (without a 'guardrails' key), the merged list in data["metadata"] was + shadowed and non-default guardrails silently received an empty + requested_guardrails list. + """ + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.litellm_pre_call_utils import move_guardrails_to_metadata + from litellm.types.proxy.policy_engine import ( + Policy, + PolicyAttachment, + PolicyGuardrails, + ) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + # Simulate a request that carries litellm_metadata (without guardrails) + # which previously shadowed data["metadata"]["guardrails"]. + "litellm_metadata": {"some_user_field": "some_value"}, + "metadata": {}, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={"guardrails": ["team-direct-guardrail"]}, + ) + + policy_registry = get_policy_registry() + policy_registry._policies = { + "global-policy": Policy( + guardrails=PolicyGuardrails(add=["policy-guardrail-1", "policy-guardrail-2"]), + ), + } + policy_registry._initialized = True + + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [ + PolicyAttachment(policy="global-policy", scope="*"), + ] + attachment_registry._initialized = True + + try: + with patch("litellm.proxy.utils._premium_user_check"): + await move_guardrails_to_metadata( + data=data, + _metadata_variable_name="metadata", + user_api_key_dict=user_api_key_dict, + ) + + guardrails = data["metadata"].get("guardrails", []) + + assert "team-direct-guardrail" in guardrails, \ + f"Team guardrail missing from merged list: {guardrails}" + assert "policy-guardrail-1" in guardrails, \ + f"policy-guardrail-1 missing: {guardrails}" + assert "policy-guardrail-2" in guardrails, \ + f"policy-guardrail-2 missing: {guardrails}" + assert len(guardrails) == len(set(guardrails)), \ + f"Duplicates in guardrails list: {guardrails}" + + # Verify get_guardrail_from_metadata returns the merged list even + # when litellm_metadata is present (the bug: it returned [] before fix) + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _DummyGuardrail(CustomGuardrail): + pass + + dummy = _DummyGuardrail(guardrail_name="team-direct-guardrail") + returned = dummy.get_guardrail_from_metadata(data) + assert "team-direct-guardrail" in returned, ( + f"get_guardrail_from_metadata shadowed by litellm_metadata; got: {returned}" + ) + + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + +@pytest.mark.asyncio +async def test_get_guardrail_from_metadata_prefers_metadata_over_litellm_metadata(): + """ + Unit test: get_guardrail_from_metadata must read from data["metadata"] first. + A non-empty data["litellm_metadata"] without a 'guardrails' key must not + shadow data["metadata"]["guardrails"]. + """ + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _DummyGuardrail(CustomGuardrail): + pass + + dummy = _DummyGuardrail(guardrail_name="my-guardrail") + + data = { + "metadata": {"guardrails": ["my-guardrail", "other-guardrail"]}, + "litellm_metadata": {"some_field": "some_value"}, # no 'guardrails' key + } + + result = dummy.get_guardrail_from_metadata(data) + assert result == ["my-guardrail", "other-guardrail"], ( + f"Expected guardrails from metadata, got: {result}" + ) + + +def test_get_guardrail_from_metadata_reads_litellm_metadata_when_no_metadata(): + """ + get_guardrail_from_metadata must still read from litellm_metadata when + data["metadata"] has no 'guardrails' key (thread/assistant endpoint path). + """ + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _DummyGuardrail(CustomGuardrail): + pass + + dummy = _DummyGuardrail(guardrail_name="my-guardrail") + + data = { + "metadata": {"requester_metadata": {"user": "alice"}}, # no guardrails key + "litellm_metadata": {"guardrails": ["my-guardrail"]}, + } + + result = dummy.get_guardrail_from_metadata(data) + assert result == ["my-guardrail"], ( + f"Expected guardrails from litellm_metadata fallback, got: {result}" + )