diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3a715b80a2d..5922285f643 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c40090233be..cbf4786affe 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ce2f97d6d55..e25b3bed757 100644 --- a/litellm/proxy/utils.py +++ b/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 diff --git a/tests/integration/mcp/test_mcp_accounting_guardrails.py b/tests/integration/mcp/test_mcp_accounting_guardrails.py index 4daa2c93fa1..323afad40db 100644 --- a/tests/integration/mcp/test_mcp_accounting_guardrails.py +++ b/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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index b25538d1814..e20d74ab60d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/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() diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index b9a260f5b14..8a7ab0f0001 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/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: diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 5b9cd761dda..0ef02f76a4f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/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 "}] 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 diff --git a/tests/test_litellm/proxy/test_model_level_guardrails.py b/tests/test_litellm/proxy/test_model_level_guardrails.py index a1278e399b5..9eaae49c46a 100644 --- a/tests/test_litellm/proxy/test_model_level_guardrails.py +++ b/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. diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index dbc6fba4ab1..e26eab0a759 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/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, + )