From 55569729b05d601c139e43b8faba447983e89f74 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 28 Aug 2026 18:16:35 -0700
Subject: [PATCH] fix(policy_engine): scope pipeline-managed guardrail skips to
 the pipeline's mode

---
 litellm/proxy/utils.py                        | 20 ++---
 .../proxy_logging/test_guardrail_pipeline.py  | 73 +++++++++++++++++++
 2 files changed, 84 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 75cc5c259a7..d3f2e1d7301 100644
--- a/litellm/proxy/utils.py
+++ b/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:
diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
index c92f5fa6c55..4e1ccf71c5c 100644
--- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
+++ b/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