From fe9efd21bb385a355f19e86bb39319bb44fc59a8 Mon Sep 17 00:00:00 2001 From: rajeshshanmugam-smf Date: Tue, 27 Oct 2026 17:03:23 +0530 Subject: [PATCH] fix(guardrails): prevent duplicate pre_call guardrail runs when metadata.user_id is present --- litellm/integrations/custom_guardrail.py | 19 +++-- tests/e2e/guardrails/guardrails_client.py | 3 + .../guardrails/test_presidio_masking_e2e.py | 65 ++++++++++++++- tests/e2e/models.py | 6 ++ .../unit/proxy/test_model_level_guardrails.py | 83 +++++++++++++++++++ 5 files changed, 163 insertions(+), 13 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..5eee39c5f6e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -845,15 +845,16 @@ class CustomGuardrail(CustomLogger): return for meta_key in ("metadata", "litellm_metadata"): meta = data.get(meta_key) - if isinstance(meta, dict): - executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) - if isinstance(executed, list): - if marker not in executed: - executed.append(marker) - else: - meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] - return - data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} + if not isinstance(meta, dict): + continue + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list): + if marker not in executed: + executed.append(marker) + else: + meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] + if not isinstance(data.get("metadata"), dict) and not isinstance(data.get("litellm_metadata"), dict): + data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool: marker: Final = self._pre_call_marker() diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 08124a03313..0f59ee4680d 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -16,6 +16,7 @@ from lifecycle import ResourceManager from models import ( AnthropicMessagesBody, AnthropicMessagesResponse, + AnthropicRequestMetadata, ChatBody, ChatMessage, ChatMetadata, @@ -491,6 +492,7 @@ class GuardrailsClient: *, guardrails: list[str] | None = None, max_tokens: int = 16, + metadata: AnthropicRequestMetadata | None = None, ) -> Result[AnthropicMessagesResponse]: return self.proxy.messages( key, @@ -499,6 +501,7 @@ class GuardrailsClient: messages=[ChatMessage(role="user", content=text)], max_tokens=max_tokens, guardrails=guardrails, + metadata=metadata, ), ) diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py index 9fc2c7d3ea2..eea924d8b70 100644 --- a/tests/e2e/guardrails/test_presidio_masking_e2e.py +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -45,10 +45,12 @@ from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEnt from lifecycle import ResourceManager from models import ( AnthropicMessagesResponse, + AnthropicRequestMetadata, ChatResponse, GuardrailEntityMatch, GuardrailRunRecord, SpendLogRow, + SpendLogsParams, ) pytestmark = pytest.mark.e2e @@ -213,7 +215,7 @@ def _messages_text(response: AnthropicMessagesResponse) -> str: def _assert_eventually_masked[R: BaseModel]( fetch: Callable[[], Result[R]], extract: Callable[[R], str], *, email: str -) -> None: +) -> R: """Retry the call until the response comes back masked, to the propagation deadline. An unmasked early response is in-flight guardrail propagation, not a failure, and neither is a transient non-Success (a replica that has not @@ -233,7 +235,7 @@ def _assert_eventually_masked[R: BaseModel]( f"the raw phone number must be masked before the model sees it, but the " f"response echoed it: {content[:300]!r}" ) - return + return data case _: last = f"" if time.monotonic() >= deadline: @@ -289,17 +291,35 @@ class TestPresidioPreCallMasking: def test_pre_call_masks_pii_on_messages( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: + """metadata.user_id must not make the same pre_call guardrail evaluate twice.""" name = f"e2e-presidio-pre-msg-{unique_marker()}" _register_presidio(client, resources, name=name) email = _fake_email() prompt = _pii_prompt(unique_marker(), email) + user_id = f"messages-user-{unique_marker()}" - _assert_eventually_masked( - lambda: client.messages(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128), + response = _assert_eventually_masked( + lambda: client.messages( + scoped_key, + MODEL, + prompt, + guardrails=[name], + max_tokens=128, + metadata=AnthropicRequestMetadata(user_id=user_id), + ), _messages_text, email=email, ) + request_id = response.id + assert request_id, f"the served messages response must carry an id to look the spend log up by: {response}" + + pre_call_runs = _pre_call_runs_after_settle(client, request_id, name) + assert pre_call_runs, f"no pre_call record appeared for request {request_id}" + assert len(pre_call_runs) == 1, ( + "metadata.user_id must not record a second pre_call evaluation of the same guardrail, " + f"got {len(pre_call_runs)}" + ) #: Room for the model's reasoning tokens plus the three-line answer; a lower cap @@ -754,6 +774,43 @@ def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]: return tuple(metadata.guardrail_information or ()) +def _named_pre_call_runs(rows: list[SpendLogRow], guardrail_name: str) -> tuple[GuardrailRunRecord, ...]: + if not rows: + return () + return tuple( + entry + for entry in _guardrail_records(rows[0]) + if entry.guardrail_name == guardrail_name and entry.guardrail_mode == "pre_call" + ) + + +def _pre_call_runs_after_settle( + client: GuardrailsClient, request_id: str, guardrail_name: str +) -> tuple[GuardrailRunRecord, ...]: + """Re-read the spend log until the pre_call count is the same twice. + + The first row can show up before a second evaluation is appended to it. + """ + seen: Final = client.proxy.poll_logs_for_request_id( + request_id, + predicate=lambda logged: len(_named_pre_call_runs(logged, guardrail_name)) >= 1, + ) + if not seen: + return () + deadline: Final = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS + current = _named_pre_call_runs(seen, guardrail_name) + while time.monotonic() < deadline: + time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) + nxt = _named_pre_call_runs( + client.proxy.spend_logs(SpendLogsParams(request_id=request_id)), + guardrail_name, + ) + if len(nxt) == len(current): + return nxt + current = nxt + return current + + def _poll_until_guardrail_applied( client: GuardrailsClient, key: str, guardrail_name: str, prompt: str ) -> StreamingResponse: diff --git a/tests/e2e/models.py b/tests/e2e/models.py index afddc7e25d6..a81cdb8540b 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -632,6 +632,10 @@ class AnthropicToolChoice(BaseModel): name: str | None = None +class AnthropicRequestMetadata(BaseModel): + user_id: str + + class AnthropicMessagesBody(BaseModel): model: str messages: list[AnthropicMessage] @@ -641,6 +645,7 @@ class AnthropicMessagesBody(BaseModel): tool_choice: AnthropicToolChoice | None = None thinking: ThinkingParam | None = None guardrails: list[str] | None = None + metadata: AnthropicRequestMetadata | None = None cache: dict[str, bool] | None = {"no-cache": True} @@ -672,6 +677,7 @@ class AnthropicMessagesResponse(BaseModel): for triage.""" model_config = ConfigDict(extra="allow") + id: str | None = None model: str | None = None content: list[AnthropicContentBlock] | None = None choices: list[ChatChoice] | None = None diff --git a/tests/unit/proxy/test_model_level_guardrails.py b/tests/unit/proxy/test_model_level_guardrails.py index 0bc9c348127..f887c0ba781 100644 --- a/tests/unit/proxy/test_model_level_guardrails.py +++ b/tests/unit/proxy/test_model_level_guardrails.py @@ -208,6 +208,7 @@ async def test_pre_call_hook_runs_once_with_model_level_guardrails(): await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion) assert guardrail.pre_call_count == 1 + assert "litellm_metadata" not in data @pytest.mark.asyncio @@ -300,6 +301,88 @@ async def test_deployment_hook_runs_pre_call_without_proxy_loop(): assert guardrail.pre_call_count == 1 +@pytest.mark.asyncio +async def test_messages_user_id_metadata_runs_pre_call_guardrail_once(): + """ + Anthropic metadata.user_id sits beside litellm_metadata that already records + another guardrail. After this guardrail records that it ran, + translate_anthropic_to_openai keeps the masked text, the user id, and the + other guardrail's marker, and neither deployment hook scans again. + """ + from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.llms.anthropic.pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + from litellm.proxy._types import CallTypes + from litellm.types.guardrails import GuardrailEventHooks + + class MaskingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="masking-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + for message in data["messages"]: + content = message.get("content") + if isinstance(content, str): + message["content"] = content.replace("raw-pii-token", "masked-pii-token") + return data + + class PriorGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="prior-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + return data + + prior = PriorGuardrail() + prior_metadata: dict[str, object] = {} + prior.mark_pre_call_hook_ran({"litellm_metadata": prior_metadata}) + prior_executed = prior_metadata[PRE_CALL_EXECUTED_GUARDRAILS_KEY] + assert isinstance(prior_executed, list) and prior_executed + prior_marker = prior_executed[0] + + guardrail = MaskingGuardrail() + data = { + "model": "gpt-4o", + "max_tokens": 16, + "messages": [{"role": "user", "content": "repeat raw-pii-token"}], + "metadata": {"user_id": "test-user"}, + "litellm_metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [prior_marker]}, + } + + await guardrail.async_pre_call_hook(None, None, data, "anthropic_messages") + guardrail.mark_pre_call_hook_ran(data) + translated = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request=data, + )[0] + translated["guardrails"] = ["prior-guardrail", "masking-guardrail"] + await prior.async_pre_call_deployment_hook(translated, CallTypes.acompletion) + await guardrail.async_pre_call_deployment_hook(translated, CallTypes.acompletion) + + surviving = translated["metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY] + content = translated["messages"][0]["content"] + assert isinstance(surviving, list) + assert prior_marker in surviving + assert any(isinstance(item, str) and item.endswith(":masking-guardrail") for item in surviving) + assert prior.pre_call_count == 0 + assert guardrail.pre_call_count == 1 + assert content == "repeat masked-pii-token" + assert translated["user"] == "test-user" + + # --------------------------------------------------------------------------- # Integration test: post_call_success_hook with model-level guardrails # ---------------------------------------------------------------------------