diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index f639ad49d5e..fc62a7ad59c 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -17,7 +17,11 @@ from typing import ( ) from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, + redact_nested_match_and_regex_keys, +) from litellm.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.secret_managers.main import str_to_bool @@ -954,17 +958,8 @@ class CustomGuardrail(CustomLogger): # should not happen container[key] = [existing, slg] - if "metadata" in request_data: - if request_data["metadata"] is None: - request_data["metadata"] = {} - _append_guardrail_info(request_data["metadata"]) - elif "litellm_metadata" in request_data: - _append_guardrail_info(request_data["litellm_metadata"]) - else: - # Ensure guardrail info is always logged (e.g. proxy may not have set - # metadata yet). Attach to "metadata" so spend log / standard logging see it. - request_data["metadata"] = {} - _append_guardrail_info(request_data["metadata"]) + _, metadata_bucket = get_or_create_metadata_bucket(request_data) + _append_guardrail_info(metadata_bucket) _guardrail_self_recorded.set(True) @@ -1223,7 +1218,7 @@ def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object) """ if logging_obj is None: return - meta_src = request_data.get("metadata") or request_data.get("litellm_metadata") or {} + meta_src = request_data.get(get_metadata_variable_name_from_kwargs(request_data)) or {} slg_info = meta_src.get("standard_logging_guardrail_information") if not slg_info: return diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index fea55cd1db4..12465377b51 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -883,8 +883,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): request_data: dict, parent_span: Optional[Any], ) -> None: - """Emit ``guardrail`` spans from ``request_data["metadata"] - ["standard_logging_guardrail_information"]``. + """Emit ``guardrail`` spans from the request's proxy-internal metadata bucket + (``standard_logging_guardrail_information``). Routed through ``_create_guardrail_span`` so the dedupe state in ``_otel_internal`` is honoured — if ``_handle_failure`` already @@ -892,7 +892,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): """ from opentelemetry import trace as _trace - metadata = (request_data or {}).get("metadata") or {} + from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + ) + + request_data = request_data or {} + metadata = request_data.get(get_metadata_variable_name_from_kwargs(request_data)) or {} guardrail_information = metadata.get("standard_logging_guardrail_information") if not guardrail_information: return diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 88dddb59cc7..cecc35ee1c1 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -195,6 +195,25 @@ def get_metadata_variable_name_from_kwargs( return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" +def get_or_create_metadata_bucket( + request_data: dict, +) -> tuple[Literal["metadata", "litellm_metadata"], dict]: + """ + Return the proxy-internal metadata bucket for this request, creating it if absent. + + Batch/file routes store proxy state in ``litellm_metadata`` so the OpenAI + ``metadata`` field can remain provider-safe (string values only). Every writer and + reader of proxy-internal metadata resolves the bucket through here, so a caller that + supplies its own ``metadata`` field cannot split them across two dicts. + """ + metadata_key = get_metadata_variable_name_from_kwargs(request_data) + metadata_bucket = request_data.get(metadata_key) + if not isinstance(metadata_bucket, dict): + metadata_bucket = {} + request_data[metadata_key] = metadata_bucket + return metadata_key, metadata_bucket + + def get_litellm_metadata_from_kwargs(kwargs: dict): """ Helper to get litellm metadata from all litellm request kwargs diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index a9c2a12aff7..33bca782e0b 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,12 +1,16 @@ import copy import os -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, +) from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -406,23 +410,6 @@ def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]: return headers -def get_metadata_variable_name_from_kwargs( - kwargs: dict, -) -> Literal["metadata", "litellm_metadata"]: - """ - Helper to return what the "metadata" field should be called in the request data - - - New endpoints return `litellm_metadata` - - Old endpoints return `metadata` - - Context: - - LiteLLM used `metadata` as an internal field for storing metadata - - OpenAI then started using this field for their metadata - - LiteLLM is now moving to using `litellm_metadata` for our metadata - """ - return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - - LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset( { "applied_policies", @@ -450,23 +437,6 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset( ) -def _get_or_create_proxy_metadata_bucket( - request_data: Dict, -) -> tuple[Literal["metadata", "litellm_metadata"], dict]: - """ - Return the proxy-internal metadata bucket for this request. - - Batch/file routes store proxy state in ``litellm_metadata`` so the OpenAI - ``metadata`` field can remain provider-safe (string values only). - """ - metadata_key = get_metadata_variable_name_from_kwargs(request_data) - metadata_bucket = request_data.get(metadata_key) - if not isinstance(metadata_bucket, dict): - metadata_bucket = {} - request_data[metadata_key] = metadata_bucket - return metadata_key, metadata_bucket - - def sanitize_openai_provider_metadata( metadata: Optional[Dict[str, Any]], ) -> Optional[Dict[str, str]]: @@ -496,7 +466,7 @@ def sanitize_openai_provider_metadata( def add_guardrail_to_applied_guardrails_header(request_data: Dict, guardrail_name: Optional[str]): if guardrail_name is None: return - _, _metadata = _get_or_create_proxy_metadata_bucket(request_data) + _, _metadata = get_or_create_metadata_bucket(request_data) if "applied_guardrails" in _metadata: if guardrail_name not in _metadata["applied_guardrails"]: _metadata["applied_guardrails"].append(guardrail_name) @@ -513,7 +483,7 @@ def add_policy_to_applied_policies_header(request_data: Dict, policy_name: Optio """ if policy_name is None: return - _, _metadata = _get_or_create_proxy_metadata_bucket(request_data) + _, _metadata = get_or_create_metadata_bucket(request_data) if "applied_policies" in _metadata: if policy_name not in _metadata["applied_policies"]: _metadata["applied_policies"].append(policy_name) @@ -531,7 +501,7 @@ def add_policy_sources_to_metadata(request_data: Dict, policy_sources: Dict[str, """ if not policy_sources: return - _, _metadata = _get_or_create_proxy_metadata_bucket(request_data) + _, _metadata = get_or_create_metadata_bucket(request_data) existing = _metadata.get("policy_sources", {}) if not isinstance(existing, dict): existing = {} diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 28b9dec100f..f2c5a95202b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -30,6 +30,10 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, +) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( @@ -432,7 +436,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Override to store only the Model Armor API response, not the entire data dict. This prevents circular references in logging. """ - metadata = (request_data.get("metadata") or {}) if isinstance(request_data, dict) else {} + metadata = ( + request_data.get(get_metadata_variable_name_from_kwargs(request_data)) or {} + if isinstance(request_data, dict) + else {} + ) guardrail_response = metadata.get("_model_armor_response", {}) # Determine status – default to "success" but prefer the explicit value if present. @@ -471,7 +479,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): blocking, while fail_on_error still governs real Model Armor API errors. """ from litellm.proxy.common_utils.callback_utils import ( - _get_or_create_proxy_metadata_bucket, add_guardrail_to_applied_guardrails_header, ) @@ -491,7 +498,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) # Use the same metadata bucket the header helper writes to, so the logged Model Armor # payload and status land where _process_response reads them on every route. - _, metadata = _get_or_create_proxy_metadata_bucket(data) + _, metadata = get_or_create_metadata_bucket(data) fail_on_error = bool(self.optional_params.get("fail_on_error", True)) if unscannable_references > 0: @@ -607,7 +614,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # overwritten by another coroutine. blocked = self._should_block_content(armor_response, allow_sanitization=self.mask_request_content) if isinstance(data, dict): - metadata = data.setdefault("metadata", {}) # ensures metadata exists and is unique per request + _, metadata = get_or_create_metadata_bucket(data) # ensures metadata exists and is unique per request # Accumulate so a prior file scan on the same request is not overwritten by this text scan. metadata["_model_armor_response"] = self._append_armor_response( metadata.get("_model_armor_response"), @@ -702,7 +709,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): blocked = self._should_block_content(armor_response, allow_sanitization=self.mask_request_content) # Store the armor response for logging if isinstance(data, dict): - metadata = data.setdefault("metadata", {}) + _, metadata = get_or_create_metadata_bucket(data) # Accumulate so a prior file scan on the same request is not overwritten by this text scan. metadata["_model_armor_response"] = self._append_armor_response( metadata.get("_model_armor_response"), @@ -868,7 +875,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Attach Model Armor response & status to this request's metadata to avoid race conditions if isinstance(request_data, dict): - metadata = request_data.setdefault("metadata", {}) + _, metadata = get_or_create_metadata_bucket(request_data) metadata["_model_armor_response"] = self._build_logging_response(armor_response) metadata["_model_armor_status"] = ( "blocked" if self._should_block_content(armor_response) else "success" diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index acb2e50c79b..9364d7eae3a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -38,6 +38,10 @@ from litellm._uuid import uuid from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -668,16 +672,13 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: Optional[d """ if guardrail_data is None: return - source_metadata = guardrail_data.get("metadata") - if not isinstance(source_metadata, dict): - return + source_key = get_metadata_variable_name_from_kwargs(guardrail_data) + source_metadata = guardrail_data.get(source_key) or {} entries = source_metadata.get("standard_logging_guardrail_information") if not entries: return - metadata = request_data.get("metadata") - if not isinstance(metadata, dict): - metadata = request_data["metadata"] = {} + _, metadata = get_or_create_metadata_bucket(request_data) metadata.setdefault("standard_logging_guardrail_information", list(entries)) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 4ea79f9e2a4..0ef26c1c5f4 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -654,6 +654,53 @@ class TestGuardrailLoggingAggregation: assert len(info) == 2 assert info[1]["guardrail_name"] == "test_guardrail" + def test_caller_metadata_does_not_divert_the_entry_from_the_reader(self): + """A caller-supplied `metadata` field must not send the entry to a bucket the + spend log never reads. Routes in LITELLM_METADATA_ROUTES (/v1/messages, + /v1/responses, batches, files) seed `litellm_metadata`, and Claude Code sends + `metadata.user_id`, so both keys are present on the same request.""" + request_data = { + "metadata": {"user_id": "device-account-session"}, + "litellm_metadata": {"user_api_key_hash": "abc"}, + } + + self._invoke_add_log(request_data) + + assert ( + "standard_logging_guardrail_information" not in request_data["metadata"] + ), "entry landed in the caller's metadata, where the spend log does not read it" + info = request_data["litellm_metadata"][ + "standard_logging_guardrail_information" + ] + assert len(info) == 1 + assert info[0]["guardrail_name"] == "test_guardrail" + + def test_entry_and_applied_guardrails_header_share_one_bucket(self): + """The x-litellm-applied-guardrails writer and the guardrail-info writer must + resolve the same bucket, otherwise the response header and the spend log + disagree about whether the guardrail ran.""" + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + request_data = { + "metadata": {"user_id": "device-account-session"}, + "litellm_metadata": {}, + } + + self._invoke_add_log(request_data) + add_guardrail_to_applied_guardrails_header( + request_data=request_data, guardrail_name="test_guardrail" + ) + + buckets = { + key + for key in ("metadata", "litellm_metadata") + for field in ("standard_logging_guardrail_information", "applied_guardrails") + if field in request_data[key] + } + assert buckets == {"litellm_metadata"} + class TestGuardrailOtelSpanEmission: """Recording a guardrail emits its otel span inline, so every guardrail diff --git a/tests/test_litellm/integrations/test_guardrail_logging_sync.py b/tests/test_litellm/integrations/test_guardrail_logging_sync.py index 5dcd1114b3d..f9e1a3efbd0 100644 --- a/tests/test_litellm/integrations/test_guardrail_logging_sync.py +++ b/tests/test_litellm/integrations/test_guardrail_logging_sync.py @@ -59,8 +59,11 @@ def test_syncs_from_metadata_key(): assert result == [entry] -def test_metadata_wins_over_litellm_metadata(): - """metadata key takes precedence over litellm_metadata when both are present.""" +def test_litellm_metadata_wins_over_caller_metadata(): + """When both keys are present the helper must read the bucket the writer used, + which get_or_create_metadata_bucket resolves to litellm_metadata. Reading the + caller's metadata instead is how a guardrail entry went missing from spend logs + on the routes that seed litellm_metadata.""" entry_meta = _make_slg_entry("from-metadata") entry_lm = _make_slg_entry("from-litellm_metadata") request_data = { @@ -74,7 +77,25 @@ def test_metadata_wins_over_litellm_metadata(): result = logging_obj.litellm_params["metadata"].get( "standard_logging_guardrail_information" ) - assert result == [entry_meta] + assert result == [entry_lm] + + +def test_syncs_when_caller_sends_its_own_metadata(): + """The Claude Code shape: caller metadata present, guardrail entry in the seeded + litellm_metadata bucket. The entry must still reach the spend-log payload.""" + entry = _make_slg_entry() + request_data = { + "metadata": {"user_id": "device-account-session"}, + "litellm_metadata": {"standard_logging_guardrail_information": [entry]}, + } + logging_obj = _FakeLogging() + + _sync_guardrail_info_to_logging_obj(request_data, logging_obj) + + result = logging_obj.litellm_params["metadata"].get( + "standard_logging_guardrail_information" + ) + assert result == [entry] def test_noop_when_no_guardrail_info(): diff --git a/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py b/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py index ace9399cf53..c3e9d67ddad 100644 --- a/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py +++ b/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py @@ -279,6 +279,48 @@ class TestGuardrailSpanOnViolation(unittest.TestCase): parent_span.context.span_id, ) + def test_post_call_failure_hook_emits_span_when_caller_sends_metadata(self): + """On routes that seed ``litellm_metadata`` the guardrail entry lives there, + not in the caller's own ``metadata`` field. Reading a hard-coded ``metadata`` + key drops the span for exactly the requests that carry both.""" + otel, provider, exporter = _make_otel() + parent_span = provider.get_tracer(__name__).start_span(PROXY_SPAN_NAME) + + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", + parent_otel_span=parent_span, + request_route="/v1/messages", + ) + + request_data = { + "model": "claude-haiku", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"user_id": "device-account-session"}, + "litellm_metadata": { + "standard_logging_guardrail_information": [ + _slg_entry("guardrail_intervened", _bedrock_block_response()) + ], + }, + } + + _run( + otel.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("guardrail blocked"), + user_api_key_dict=user_api_key_dict, + ) + ) + + guardrail_spans = [ + s for s in exporter.get_finished_spans() if s.name == GUARDRAIL_SPAN_NAME + ] + self.assertEqual( + len(guardrail_spans), + 1, + "the guardrail span must be emitted from the resolved metadata bucket, " + "not from a hard-coded 'metadata' key", + ) + def test_handle_failure_and_post_call_failure_hook_dedupe(self): """When _handle_failure and async_post_call_failure_hook BOTH fire for the same request (the production flow on a guardrail block), diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index b67ea91bb0b..b4f539da286 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -4,12 +4,53 @@ import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + get_or_create_metadata_bucket, map_finish_reason, reconstruct_model_name, redact_nested_match_and_regex_keys, ) +class TestGetOrCreateMetadataBucket: + """The single owner every guardrail writer and reader shares, so the response + header and the spend log can never disagree about which dict a record lives in.""" + + def test_prefers_litellm_metadata_when_both_present(self): + request_data = {"metadata": {"user_id": "caller"}, "litellm_metadata": {}} + + key, bucket = get_or_create_metadata_bucket(request_data) + + assert key == "litellm_metadata" + assert bucket is request_data["litellm_metadata"] + + def test_uses_metadata_when_litellm_metadata_absent(self): + request_data = {"metadata": {"user_id": "caller"}} + + key, bucket = get_or_create_metadata_bucket(request_data) + + assert key == "metadata" + assert bucket is request_data["metadata"] + + def test_creates_the_bucket_in_place_when_missing(self): + request_data: dict = {} + + key, bucket = get_or_create_metadata_bucket(request_data) + + assert key == "metadata" + assert request_data["metadata"] is bucket + bucket["k"] = "v" + assert request_data["metadata"]["k"] == "v" + + def test_replaces_a_non_dict_bucket(self): + request_data = {"litellm_metadata": None} + + key, bucket = get_or_create_metadata_bucket(request_data) + + assert key == "litellm_metadata" + assert isinstance(request_data["litellm_metadata"], dict) + assert bucket is request_data["litellm_metadata"] + + def test_reconstruct_model_name_prefers_deployment_value(): """Ensure deployment metadata wins when reconstructing the model name.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index fe6cb98d1f5..9002d1f81a3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -729,10 +729,15 @@ async def test_openai_moderation_post_call_request_data_passthrough(): mock_make_request.assert_called_once() - # Guardrail info in the REAL request_data (not a throwaway) - guardrail_info_list = request_data["metadata"].get( - "standard_logging_guardrail_information" + # Guardrail info in the REAL request_data (not a throwaway). The unified hook + # seeds litellm_metadata, so read the bucket the resolver names rather than + # assuming "metadata"; the spend log reads it the same way. + from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, ) + + bucket = request_data[get_metadata_variable_name_from_kwargs(request_data)] + guardrail_info_list = bucket.get("standard_logging_guardrail_information") assert guardrail_info_list is not None assert isinstance(guardrail_info_list[0]["guardrail_response"], dict) assert "results" in guardrail_info_list[0]["guardrail_response"] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 0358ca998aa..914af0e2368 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -259,10 +259,15 @@ async def test_openai_moderation_streaming_end_of_stream_request_data_passthroug ): pass - # Verify guardrail info reached the REAL request_data (not a throwaway) - guardrail_info_list = request_data["metadata"].get( - "standard_logging_guardrail_information" + # Verify guardrail info reached the REAL request_data (not a throwaway). The + # unified hook seeds litellm_metadata, so read the bucket the resolver names + # rather than assuming "metadata"; the spend log reads it the same way. + from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, ) + + bucket = request_data[get_metadata_variable_name_from_kwargs(request_data)] + guardrail_info_list = bucket.get("standard_logging_guardrail_information") assert ( guardrail_info_list is not None ), "Guardrail info should be in request_data after streaming" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 18b5bd92411..89b6af27719 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3502,6 +3502,44 @@ async def test_single_scan_response_stays_a_dict(): assert isinstance(request_data["metadata"]["_model_armor_response"], dict) +@pytest.mark.asyncio +async def test_scan_result_reaches_the_logger_on_a_seeded_route(): + """On routes that seed `litellm_metadata` the scan result must land in that bucket + and be found by `_process_response`. Writing the file-scan result through the shared + resolver while the text-scan writers and the reader used a hard-coded `metadata` key + split the record in two, so the logged guardrail payload came back empty.""" + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + request_data = { + "model": "claude-haiku", + "messages": [_file_message(pdf_b64)], + "metadata": {"user_id": "device-account-session"}, + "litellm_metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert "_model_armor_response" not in request_data["metadata"] + assert "_model_armor_response" in request_data["litellm_metadata"] + + before = len(request_data["litellm_metadata"].get("standard_logging_guardrail_information", [])) + guardrail._process_response(response=None, request_data=request_data) + + logged = request_data["litellm_metadata"]["standard_logging_guardrail_information"] + assert len(logged) == before + 1 + assert logged[-1]["guardrail_response"], "the logger recorded an empty Model Armor payload" + + @pytest.mark.asyncio async def test_pre_call_blocks_supported_document_with_undecodable_base64(): """A supported document whose inline base64 will not decode cannot be scanned, so it fails closed."""