fix(mcp): keep tool attribution on guardrail-blocked REST calls (#42790)

* fix(mcp): keep tool attribution on guardrail-blocked REST calls

Co-Authored-By: bot_apk <apk@cognition.ai>

* fix(proxy): keep content enforcers in the pre-call walk when guardrails are skipped

Co-Authored-By: bot_apk <apk@cognition.ai>

* test(proxy): accept skip_guardrails kwarg in pre_call_hook test doubles

Co-Authored-By: bot_apk <apk@cognition.ai>

* test(proxy): drop section comment flagged by repo comment policy

Co-Authored-By: bot_apk <apk@cognition.ai>

* test(proxy): drop docstrings from skip_guardrails tests

Co-Authored-By: bot_apk <apk@cognition.ai>

* test(proxy): wrap pre_call_hook mocks under the line limit

Co-Authored-By: bot_apk <apk@cognition.ai>

* refactor(proxy): drop skip_guardrails docstring sentence

Co-Authored-By: bot_apk <apk@cognition.ai>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: bot_apk <apk@cognition.ai>

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -1155,6 +1155,7 @@ if MCP_AVAILABLE:
                     route_type=CallTypes.call_mcp_tool.value,
                     proxy_logging_obj=proxy_logging_obj,
                     general_settings=general_settings,
+                    skip_guardrails=True,
                 )
 
                 # Extract MCP auth headers from request and add to data dict

litellm/proxy/common_request_processing.py

@@ -1999,6 +1999,7 @@ class ProxyBaseLLMRequestProcessing:
         model: str | None = None,
         llm_router: Router | None = None,
         rate_limited_model: str | None = None,
+        skip_guardrails: bool = False,
     ) -> tuple[dict, LiteLLMLoggingObj]:
         start_time: Final = datetime.now()  # start before calling guardrail hooks
 
@@ -2187,6 +2188,7 @@ class ProxyBaseLLMRequestProcessing:
             user_api_key_dict=user_api_key_dict,
             data=self.data,
             call_type=route_type,
+            skip_guardrails=skip_guardrails,
         )
         await _enforce_guardrail_added_tag_budgets(
             data=self.data,

litellm/proxy/utils.py

@@ -2306,6 +2306,7 @@ class ProxyLogging:
         data: None,
         call_type: CallTypesLiteral,
         guardrails_only: bool = False,
+        skip_guardrails: bool = False,
     ) -> None:
         pass
 
@@ -2316,6 +2317,7 @@ class ProxyLogging:
         data: dict,
         call_type: CallTypesLiteral,
         guardrails_only: bool = False,
+        skip_guardrails: bool = False,
     ) -> dict:
         pass
 
@@ -2325,6 +2327,7 @@ class ProxyLogging:
         data: dict | None,
         call_type: CallTypesLiteral,
         guardrails_only: bool = False,
+        skip_guardrails: bool = False,
     ) -> dict | None:
         """
         Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body.
@@ -2340,6 +2343,9 @@ class ProxyLogging:
         """
         verbose_proxy_logger.debug("Inside Proxy Logging Pre-call hook!")
 
+        if guardrails_only and skip_guardrails:
+            raise ValueError("guardrails_only and skip_guardrails are mutually exclusive")
+
         if not guardrails_only:
             self._init_response_taking_too_long_task(data=data)
 
@@ -2387,16 +2393,19 @@ class ProxyLogging:
 
         try:
             # Execute guardrail pipelines before the normal callback loop
-            data, _ = await self._maybe_execute_pipelines(  # rebind-ok: pipeline edits feed the callback loop below
-                data=data,
-                user_api_key_dict=user_api_key_dict,
-                call_type=call_type,
-                event_hook="pre_call",
-                raw_request_snapshot=raw_request_snapshot,
-            )
+            if not skip_guardrails:
+                data, _ = await self._maybe_execute_pipelines(  # rebind-ok: pipeline edits feed the callback loop below
+                    data=data,
+                    user_api_key_dict=user_api_key_dict,
+                    call_type=call_type,
+                    event_hook="pre_call",
+                    raw_request_snapshot=raw_request_snapshot,
+                )
 
             # Get pipeline-managed guardrails to skip in normal loop
-            pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call")
+            pipeline_managed: Final[frozenset[str]] = (
+                frozenset() if skip_guardrails else pipeline_managed_guardrail_names(data, "pre_call")
+            )
 
             caps: Final = ProxyLogging._callback_capabilities()
             # Skip the per-request callback walk entirely when nothing in
@@ -2405,7 +2414,7 @@ class ProxyLogging:
             # ``time.time()`` x2 per registered callback for the common
             # "callbacks=[]" case on small / dev deployments.
             if (
-                not caps.has_guardrail
+                (skip_guardrails or not caps.has_guardrail)
                 and not caps.has_content_enforcer
                 and (guardrails_only or not caps.has_pre_call_override)
             ):
@@ -2413,12 +2422,16 @@ class ProxyLogging:
                     self._process_guardrail_metadata(data)
                 return data
 
-            parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
-                cb
-                for cb in caps.resolved_callbacks
-                if isinstance(cb, CustomGuardrail)
-                and getattr(cb, "run_in_parallel", False)
-                and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed)
+            parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = (
+                ()
+                if skip_guardrails
+                else tuple(
+                    cb
+                    for cb in caps.resolved_callbacks
+                    if isinstance(cb, CustomGuardrail)
+                    and getattr(cb, "run_in_parallel", False)
+                    and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed)
+                )
             )
 
             deferred_route_exc: SensitiveDataRouteException | None = None
@@ -2426,6 +2439,9 @@ class ProxyLogging:
                 start_time = time.time()
                 try:
                     if isinstance(_callback, CustomGuardrail) and data is not None:
+                        if skip_guardrails:
+                            continue
+
                         # Skip guardrails managed by a pipeline
                         if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed:
                             continue

tests/integration/mcp/test_mcp_accounting_guardrails.py

@@ -142,7 +142,7 @@ def _content_filter(gateway: Gateway, mode: str) -> Iterator[str]:
 def test_pre_mcp_call_guardrail_blocks_before_the_peer_and_still_logs_spend(
     gateway: Gateway, entry: EntryPoint
 ) -> None:
-    with _content_filter(gateway, "pre_mcp_call") as guardrail, mcp_peer() as peer, gateway.scenario() as scenario:
+    with _content_filter(gateway, "pre_mcp_call"), mcp_peer() as peer, gateway.scenario() as scenario:
         alias: Final = "guard" + uuid.uuid4().hex[:8]
         identity: Final = _priced_server(scenario, peer, alias)
         key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
@@ -158,12 +158,8 @@ def test_pre_mcp_call_guardrail_blocks_before_the_peer_and_still_logs_spend(
         assert len(rows) == 2, rows
         failures: Final = [row for row in rows if row["status"] == "failure"]
         assert len(failures) == 1, rows
-        if failures[0]["model"] == "":
-            pytest.skip(
-                f"BUG: guardrail-blocked MCP call on {entry} logs a spend row with an empty model and no tool name "
-                f"(guardrail {guardrail})"
-            )
         assert failures[0]["model"] == f"MCP: {alias}-add", failures[0]
+        assert _tool_metadata(failures[0])["mcp_server_name"] == alias, failures[0]
 
 
 def test_guardrail_blocked_call_never_reaches_peer_through_the_official_client(gateway: Gateway) -> None:

tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py

@@ -2705,7 +2705,7 @@ class TestCallToolRestAPI:
 
         pre_call_finished_at = {}
 
-        async def slow_pre_call_hook(user_api_key_dict, data, call_type):
+        async def slow_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             await asyncio.sleep(0.05)
             pre_call_finished_at["value"] = datetime.now()
             return data
@@ -2960,10 +2960,10 @@ class TestCallToolRestAPI:
                 message="Content blocked", model="mcp-tool-call", request_data={}, guardrail_name="block-all"
             )
 
-        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type):
+        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             return data
 
-        async def blocking_pre_call_hook(user_api_key_dict, data, call_type):
+        async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             raise guardrail_error
 
         async def fake_execute_mcp_tool(**kwargs):
@@ -3044,7 +3044,7 @@ class TestCallToolRestAPI:
         async def fake_add_litellm_data_to_request(**kwargs):
             return kwargs.get("data", {})
 
-        async def blocking_pre_call_hook(user_api_key_dict, data, call_type):
+        async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             raise guardrail_error
 
         failure_logging = AsyncMock(side_effect=RuntimeError("spend log db down"))
@@ -3153,7 +3153,7 @@ class TestCallToolRestAPI:
         async def fake_add_litellm_data_to_request(**kwargs):
             return kwargs.get("data", {})
 
-        async def blocking_pre_call_hook(user_api_key_dict, data, call_type):
+        async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             raise guardrail_error
 
         failure_logging = AsyncMock()

tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py

@@ -588,7 +588,9 @@ async def test_message_send_reports_an_unresolvable_entra_credential_as_internal
     user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1")
 
     mock_proxy_logging = MagicMock()
-    mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
+    mock_proxy_logging.pre_call_hook = AsyncMock(
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+    )
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
     downstream = AsyncMock()
 
@@ -956,7 +958,9 @@ async def test_subscribe_to_task_calls_pre_call_hook():
             yield chunk
 
     mock_proxy_logging = MagicMock()
-    mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
+    mock_proxy_logging.pre_call_hook = AsyncMock(
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+    )
     mock_proxy_logging.async_post_call_streaming_iterator_hook = _passthrough_iterator
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
 
@@ -1089,7 +1093,9 @@ async def test_task_method_failure_hook_uses_enriched_request_data():
     mock_handler.post = AsyncMock(side_effect=RuntimeError("upstream failed"))
 
     mock_proxy_logging = MagicMock()
-    mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
+    mock_proxy_logging.pre_call_hook = AsyncMock(
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+    )
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
 
     with ExitStack() as stack:
@@ -1154,7 +1160,9 @@ async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400()
     user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1")
 
     mock_proxy_logging = MagicMock()
-    mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
+    mock_proxy_logging.pre_call_hook = AsyncMock(
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+    )
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
 
     with ExitStack() as stack:

tests/test_litellm/proxy/test_common_request_processing.py

@@ -424,7 +424,7 @@ class TestProxyBaseLLMRequestProcessing:
         async def mock_add_litellm_data_to_request(*args, **kwargs):
             return {}
 
-        async def mock_common_processing_pre_call_logic(user_api_key_dict, data, call_type):
+        async def mock_common_processing_pre_call_logic(user_api_key_dict, data, call_type, skip_guardrails=False):
             data_copy = copy.deepcopy(data)
             return data_copy
 
@@ -520,7 +520,7 @@ class TestProxyBaseLLMRequestProcessing:
                 },
             }
 
-        async def mock_pre_call_hook(user_api_key_dict, data, call_type):
+        async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             data["messages"] = [{"role": "user", "content": "my ssn is <MASKED>"}]
             return data
 
@@ -565,7 +565,7 @@ class TestProxyBaseLLMRequestProcessing:
         async def mock_add_litellm_data_to_request(*args, **kwargs):
             return copy.deepcopy(request_body)
 
-        async def mock_pre_call_hook(user_api_key_dict, data, call_type):
+        async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             data.setdefault("metadata", {}).setdefault("tags", []).extend(guardrail_tags)
             return data
 
@@ -745,7 +745,7 @@ class TestProxyBaseLLMRequestProcessing:
         async def retry_add_litellm_data_to_request(*args, **kwargs):
             return first_pass_data
 
-        async def idempotent_pre_call_hook(user_api_key_dict, data, call_type):
+        async def idempotent_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             return data
 
         monkeypatch.setattr(
@@ -888,7 +888,7 @@ class TestProxyBaseLLMRequestProcessing:
 
         seen_metadata: dict = {}
 
-        async def mock_pre_call_hook(user_api_key_dict, data, call_type):
+        async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             seen_metadata.update(data.get("metadata") or {})
             return data
 
@@ -959,7 +959,7 @@ class TestProxyBaseLLMRequestProcessing:
         async def mock_add_litellm_data_to_request(*args, **kwargs):
             return {}
 
-        async def mock_common_processing_pre_call_logic(user_api_key_dict, data, call_type):
+        async def mock_common_processing_pre_call_logic(user_api_key_dict, data, call_type, skip_guardrails=False):
             data_copy = copy.deepcopy(data)
             return data_copy
 
@@ -1960,7 +1960,7 @@ class TestProxyBaseLLMRequestProcessing:
             data["metadata"] = data.get("metadata", {})
             return data
 
-        async def mock_pre_call_hook(user_api_key_dict, data, call_type):
+        async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             return copy.deepcopy(data)
 
         mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
@@ -6912,7 +6912,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
         limiter_models: list[str] = []
 
         async def run_limiter(
-            user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str
+            user_api_key_dict: ProxyUserAPIKeyAuth,
+            data: dict[str, object],
+            call_type: str,
+            skip_guardrails: bool = False,
         ) -> dict[str, object]:
             limiter_models.append(str(data["model"]))
             await limiter.async_pre_call_hook(
@@ -7132,7 +7135,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
         run_limiter = rig[0].pre_call_hook
 
         async def limiter_then_guardrail(
-            user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str
+            user_api_key_dict: ProxyUserAPIKeyAuth,
+            data: dict[str, object],
+            call_type: str,
+            skip_guardrails: bool = False,
         ) -> dict[str, object]:
             limited = await run_limiter(user_api_key_dict=user_api_key_dict, data=data, call_type=call_type)
             if guardrail not in (limited["metadata"].get("guardrails") or []):
@@ -7958,7 +7964,7 @@ class TestPerRequestModelGroupAlias:
         async def mock_add_litellm_data_to_request(*args, **kwargs):
             return kwargs.get("data", {})
 
-        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type):
+        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             return copy.deepcopy(data)
 
         mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
@@ -8007,7 +8013,7 @@ class TestPerRequestModelGroupAlias:
         async def mock_add_litellm_data_to_request(*args, **kwargs):
             return kwargs.get("data", {})
 
-        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type):
+        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             return copy.deepcopy(data)
 
         mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
@@ -8046,7 +8052,7 @@ class TestPerRequestModelGroupAlias:
         async def mock_add_litellm_data_to_request(*args, **kwargs):
             return kwargs.get("data", {})
 
-        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type):
+        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
             return copy.deepcopy(data)
 
         mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
@@ -9729,7 +9735,10 @@ class TestBackgroundResponseRetrievalGovernance:
             return data
 
         async def decrypting_pre_call_hook(
-            user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str
+            user_api_key_dict: ProxyUserAPIKeyAuth,
+            data: dict[str, object],
+            call_type: str,
+            skip_guardrails: bool = False,
         ) -> dict[str, object]:
             if data.get("response_id") == client_facing_response_id:
                 data["response_id"] = encoded_response_id

tests/test_litellm/proxy/test_model_level_guardrails.py

@@ -600,7 +600,7 @@ async def test_pre_call_merges_model_level_guardrails_before_pre_call_hook():
 
     captured_pre_call_guardrails: list = []
 
-    async def fake_pre_call_hook(*, user_api_key_dict, data, call_type):
+    async def fake_pre_call_hook(*, user_api_key_dict, data, call_type, skip_guardrails=False):
         # Snapshot the list rather than the dict: metadata is shared by
         # reference, so a merge that happens after this point would otherwise
         # show up here retroactively and the assertion would pass either way.

tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py

@@ -945,3 +945,53 @@ async def test_pre_call_block_keeps_request_declared_guardrail_in_applied_guardr
             call_type="completion",
         )
     assert data["metadata"]["applied_guardrails"] == ["blocker", "declared-post-call"]
+
+
+@pytest.mark.asyncio
+async def test_skip_guardrails_still_runs_non_guardrail_callbacks(proxy_logging, make_user_api_key_auth, monkeypatch):
+    accountant = _Accountant()
+    monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(), accountant])
+    proxy_logging.slack_alerting_instance = MagicMock(alerting=None)
+
+    data = _secret_request()
+    out = await proxy_logging.pre_call_hook(
+        user_api_key_dict=make_user_api_key_auth(),
+        data=data,
+        call_type="completion",
+        skip_guardrails=True,
+    )
+    assert out is data
+    assert "SECRET" in out["messages"][0]["content"]
+    assert accountant.calls == 1
+
+
+@pytest.mark.asyncio
+async def test_default_walk_still_blocks_on_the_same_setup(proxy_logging, make_user_api_key_auth, monkeypatch):
+    accountant = _Accountant()
+    monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(), accountant])
+    proxy_logging.slack_alerting_instance = MagicMock(alerting=None)
+
+    with pytest.raises(HTTPException, match="blocked"):
+        await proxy_logging.pre_call_hook(
+            user_api_key_dict=make_user_api_key_auth(),
+            data=_secret_request(),
+            call_type="completion",
+        )
+    assert accountant.calls == 0
+
+
+@pytest.mark.asyncio
+async def test_guardrails_only_and_skip_guardrails_are_mutually_exclusive(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    monkeypatch.setattr(litellm, "callbacks", [])
+    proxy_logging.slack_alerting_instance = MagicMock(alerting=None)
+
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        await proxy_logging.pre_call_hook(
+            user_api_key_dict=make_user_api_key_auth(),
+            data={"model": "m"},
+            call_type="completion",
+            guardrails_only=True,
+            skip_guardrails=True,
+        )