fix(policy_engine): scope pipeline-managed guardrail skips to the pipeline's mode

litellm/proxy/utils.py

@@ -11,7 +11,7 @@ import sys
 import threading
 import time
 import traceback
-from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Mapping, Sequence
+from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
 from dataclasses import dataclass, field
 from datetime import date, datetime, timedelta, timezone
 from email.mime.multipart import MIMEMultipart
@@ -446,12 +446,14 @@ def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardrail
     )
 
 
-def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[str]:
-    managed: Final = _policy_state_metadata(data).get("_pipeline_managed_guardrails")
-    return (
-        frozenset(cast("Collection[str]", managed))  # cast-ok: the policy engine wrote these guardrail names
-        if managed
-        else frozenset()
+def _pipeline_managed_guardrail_names(
+    data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
+) -> frozenset[str]:
+    return frozenset(
+        step.guardrail
+        for _policy_name, pipeline in _policy_pipelines(data)
+        if pipeline.mode == mode
+        for step in pipeline.steps
     )
 
 
@@ -1837,7 +1839,7 @@ class ProxyLogging:
             )
 
             # Get pipeline-managed guardrails to skip in normal loop
-            pipeline_managed: Final = _pipeline_managed_guardrail_names(data)
+            pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call")
 
             caps: Final = ProxyLogging._callback_capabilities()
             # Skip the per-request callback walk entirely when nothing in
@@ -2825,7 +2827,7 @@ class ProxyLogging:
         if pipeline_response is not None:
             response = pipeline_response  # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
 
-        pipeline_managed: Final = _pipeline_managed_guardrail_names(data)
+        pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
         guardrail_callbacks: Final[list[CustomGuardrail]] = []
         other_callbacks: Final[list[CustomLogger]] = []
         try:

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

@@ -977,6 +977,79 @@ async def test_post_call_pipeline_managed_default_on_guardrail_runs_exactly_once
     assert seen["count"] == 1
 
 
+@pytest.mark.asyncio
+async def test_post_call_hook_still_runs_guardrail_managed_only_by_pre_call_pipeline(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Dict[str, Any] = {"count": 0}
+
+    class DualStageGuardrail(CustomGuardrail):
+        async def async_post_call_success_hook(self, data, user_api_key_dict, response):
+            seen["count"] += 1
+            return None
+
+    pre_call_pipeline = GuardrailPipeline(
+        mode="pre_call",
+        steps=[PipelineStep(guardrail="gr-dual", on_pass="allow", on_fail="block")],
+    )
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [DualStageGuardrail(guardrail_name="gr-dual", event_hook=["pre_call", "post_call"], default_on=True)],
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = {
+        "model": "m",
+        "messages": [{"role": "user", "content": "hi"}],
+        "metadata": {
+            "_guardrail_pipelines": [("request-governance", pre_call_pipeline)],
+            "_pipeline_managed_guardrails": {"gr-dual"},
+        },
+    }
+
+    await proxy_logging.post_call_success_hook(
+        data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
+    )
+
+    assert seen["count"] == 1
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipeline(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Dict[str, Any] = {"count": 0}
+
+    class DualStageGuardrail(CustomGuardrail):
+        async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+            seen["count"] += 1
+            return data
+
+    post_call_pipeline = GuardrailPipeline(
+        mode="post_call",
+        steps=[PipelineStep(guardrail="gr-dual", on_pass="allow", on_fail="block")],
+    )
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [DualStageGuardrail(guardrail_name="gr-dual", event_hook=["pre_call", "post_call"], default_on=True)],
+    )
+    data = {
+        "model": "m",
+        "messages": [{"role": "user", "content": "hi"}],
+        "metadata": {
+            "_guardrail_pipelines": [("response-governance", post_call_pipeline)],
+            "_pipeline_managed_guardrails": {"gr-dual"},
+        },
+    }
+
+    await proxy_logging.pre_call_hook(
+        user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion"
+    )
+
+    assert seen["count"] == 1
+
+
 @pytest.mark.asyncio
 async def test_post_call_pipeline_replacement_response_reaches_caller(
     proxy_logging, make_user_api_key_auth, monkeypatch