test(ocr): restore guardrail contracts

.github/workflows/test-rust.yml

@@ -9,6 +9,7 @@ on:
       - "litellm/integrations/custom_logger.py"
       - "litellm/litellm_core_utils/litellm_logging.py"
       - "litellm/litellm_core_utils/logging_worker.py"
+      - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
       - "litellm/llms/base_llm/ocr/**"
@@ -38,6 +39,7 @@ on:
       - "litellm/integrations/custom_logger.py"
       - "litellm/litellm_core_utils/litellm_logging.py"
       - "litellm/litellm_core_utils/logging_worker.py"
+      - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
       - "litellm/llms/base_llm/ocr/**"

tests/test_litellm_rust/README.md

@@ -4,7 +4,7 @@ This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR be
 
 A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions
 
-`ocr/test_requests.py` covers provider payloads, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_dispatch.py` covers public sync and async native dispatch, explicit Python dispatch, fallback, and the native compression header. `test_ocr.py` is the strict wire-level smoke test
+`ocr/test_requests.py` covers provider payloads, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. `ocr/test_dispatch.py` covers public sync and async native dispatch, explicit Python dispatch, fallback, and the native compression header. `test_ocr.py` is the strict wire-level smoke test
 
 Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports
 

tests/test_litellm_rust/conftest.py

@@ -32,6 +32,7 @@ CALLBACK_ATTRIBUTES: Final = (
 EXPECTED_FAILURE_REASONS: Final = {
     "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070",
     "ocr/test_dispatch.py": "requires the OCR native dispatch implementation from #40070",
+    "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070",
     "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070",
 }
 

tests/test_litellm_rust/ocr/test_guardrails.py

@@ -0,0 +1,74 @@
+from typing import Final
+
+import pytest
+from fastapi import HTTPException
+
+import litellm
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
+from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
+from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_aocr
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def ocr_server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
+    return recording_server
+
+
+class ReplaceOCRMarkdown(CustomGuardrail):
+    def __init__(self) -> None:
+        super().__init__(
+            guardrail_name="replace-ocr-markdown", event_hook=GuardrailEventHooks.post_call, default_on=True
+        )
+        self.call_types: list[CallTypes] = []
+
+    async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
+        self.call_types.append(call_type)
+        reviewed_page: Final = response.pages[0].model_copy(update={"markdown": "Reviewed OCR"})
+        return response.model_copy(update={"pages": [reviewed_page]})
+
+
+@pytest.mark.asyncio
+async def test_public_aocr_post_call_content_filter_blocks_matching_markdown(
+    ocr_server: RecordingServer,
+) -> None:
+    guardrail: Final = ContentFilterGuardrail(
+        guardrail_name="block-native-ocr-markdown",
+        event_hook=GuardrailEventHooks.post_call,
+        blocked_words=[BlockedWord(keyword="native OCR response", action=ContentFilterAction.BLOCK)],
+    )
+    litellm.callbacks.append(guardrail)
+
+    with pytest.raises(HTTPException, match="Content blocked") as blocked:
+        await call_aocr(ocr_server, guardrails=[guardrail.guardrail_name])
+
+    assert blocked.value.status_code == 400
+    assert len(ocr_server.requests) == 1
+
+
+@pytest.mark.asyncio
+async def test_public_aocr_post_call_replacement_reaches_caller_and_success_callback(
+    ocr_server: RecordingServer,
+) -> None:
+    guardrail: Final = ReplaceOCRMarkdown()
+    recorder: Final = RecordingLogger()
+    litellm.callbacks.append(guardrail)
+
+    response: Final = await call_aocr(
+        ocr_server,
+        callbacks=[recorder],
+        guardrails=[guardrail.guardrail_name],
+    )
+    success_events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert guardrail.call_types == [CallTypes.aocr]
+    assert response.pages[0].markdown == "Reviewed OCR"
+    assert len(success_events) == 1
+    assert success_events[0].response.pages[0].markdown == "Reviewed OCR"
+    assert "guardrails" not in ocr_server.requests[0].body