mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-10 23:06:52 +00:00
fix(guardrails): team-level guardrails and global policy guardrails can run together (#26466)
* fix(guardrails): apply team-level guardrails alongside global policy guardrails Two bugs prevented team-direct guardrails from being automatically applied when using a team-scoped API key: 1. Auth caching: `valid_token.team_metadata` was never refreshed from the freshly-fetched team object at the "Check 6" step in `_user_api_key_auth_builder`. Guardrails added to a team after the key was first cached were therefore invisible to `move_guardrails_to_metadata`. Fix: propagate `_team_obj.metadata` → `valid_token.team_metadata` after every "Check 6" team fetch (user_api_key_auth.py). 2. Guardrail execution: `get_guardrail_from_metadata` checked `data["litellm_metadata"]` before `data["metadata"]`. When a request carried a non-empty `litellm_metadata` without a "guardrails" key, the merged guardrail list written to `data["metadata"]` by `move_guardrails_to_metadata` was shadowed and the guardrail received an empty requested-guardrails list (custom_guardrail.py). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix merge conflict --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Shivam Rawat
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user