test: wait for requested guardrail propagation

tests/e2e/guardrails/guardrails_client.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import time
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Literal
+from typing import Final, Literal
 
 from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker
 from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap
@@ -405,6 +405,26 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient:
     return GuardrailsClient(proxy=proxy)
 
 
+def poll_until_guardrail_applied(
+    call: Callable[[], StreamingResponse],
+    guardrail_name: str,
+    *,
+    timeout: float = POLL_TIMEOUT,
+    interval: float = POLL_INTERVAL,
+    now: Callable[[], float] = time.monotonic,
+    sleep: Callable[[float], None] = time.sleep,
+) -> StreamingResponse:
+    deadline: Final = now() + timeout
+    while (
+        (result := call()).ok
+        and guardrail_name
+        not in (name.strip() for name in result.headers.get("x-litellm-applied-guardrails", "").split(","))
+        and (remaining := deadline - now()) > 0
+    ):
+        sleep(min(interval, remaining))
+    return result
+
+
 def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]:
     """Retry a call that a guardrail should reject until it is, returning the last result.
 

tests/e2e/guardrails/test_guardrails_client.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass
+from itertools import chain, repeat
+from typing import Final
+
+import pytest
+
+from e2e_http import StreamingResponse
+from guardrails_client import poll_until_guardrail_applied
+
+
+@dataclass
+class Clock:
+    elapsed: float = 0.0
+
+    def now(self) -> float:
+        return self.elapsed
+
+    def sleep(self, seconds: float) -> None:
+        self.elapsed += seconds
+
+
+def _response(applied: str, status: int = 200) -> StreamingResponse:
+    return StreamingResponse(status_code=status, body="{}", headers={"x-litellm-applied-guardrails": applied})
+
+
+def test_waits_for_requested_guardrail_after_an_unrelated_global_guardrail() -> None:
+    clock: Final = Clock()
+    expected: Final = _response("global-filter, tool-permission")
+    responses: Final = iter((_response("global-filter"), expected))
+
+    result: Final = poll_until_guardrail_applied(
+        lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep
+    )
+
+    assert result is expected
+    assert clock.elapsed == 2
+
+
+@pytest.mark.parametrize("applied", ("", "global-filter", "tool-permission-sibling"))
+def test_missing_exact_guardrail_returns_failure_evidence_at_deadline(applied: str) -> None:
+    clock: Final = Clock()
+    missing: Final = _response(applied)
+
+    result: Final = poll_until_guardrail_applied(
+        lambda: missing, "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep
+    )
+
+    assert result is missing
+    assert clock.elapsed == 5
+
+
+@pytest.mark.parametrize("status", (400, 401, 429, 500))
+def test_http_failure_is_not_hidden_by_a_later_success(status: int) -> None:
+    clock: Final = Clock()
+    failed: Final = _response("", status)
+    responses: Final = iter(chain((failed,), repeat(_response("tool-permission"))))
+
+    result: Final = poll_until_guardrail_applied(
+        lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep
+    )
+
+    assert result is failed
+    assert clock.elapsed == 0

tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py

@@ -30,6 +30,7 @@ from guardrails_client import (
     ToolPermissionParamsBody,
     ToolPermissionRuleBody,
     poll_until_blocked,
+    poll_until_guardrail_applied,
 )
 from lifecycle import ResourceManager
 from models import ChatResponse, ChatTool, ChatToolFunction
@@ -84,8 +85,8 @@ def _register_tool_permission(client: GuardrailsClient, resources: ResourceManag
     resources.defer(lambda: client.delete_guardrail(guardrail_id))
 
 
-def _applied_guardrails(outcome: StreamingResponse) -> str:
-    return outcome.headers.get("x-litellm-applied-guardrails", "")
+def _applied_guardrails(outcome: StreamingResponse) -> tuple[str, ...]:
+    return tuple(name.strip() for name in outcome.headers.get("x-litellm-applied-guardrails", "").split(","))
 
 
 def _tool_call_names(response: ChatResponse) -> tuple[str, ...]:
@@ -144,14 +145,17 @@ class TestToolPermissionPreCall:
         name = f"e2e-toolperm-allow-{unique_marker()}"
         _register_tool_permission(client, resources, name=name)
 
-        outcome = client.chat_raw(
-            scoped_key,
-            MODEL,
-            TOOL_PROMPT,
-            guardrails=[name],
-            max_tokens=128,
-            tools=[ALLOWED_TOOL],
-            tool_choice="required",
+        outcome = poll_until_guardrail_applied(
+            lambda: client.chat_raw(
+                scoped_key,
+                MODEL,
+                TOOL_PROMPT,
+                guardrails=[name],
+                max_tokens=128,
+                tools=[ALLOWED_TOOL],
+                tool_choice="required",
+            ),
+            name,
         )
 
         assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}"