fix: bind post-call markers to checked output

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/custom_guardrail.py

@@ -1,6 +1,7 @@
 import contextvars
 import copy
 import hashlib
+import json
 import os
 import secrets
 from collections.abc import Iterable, Mapping, Sequence
@@ -10,7 +11,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
 
 import httpx
-from pydantic import TypeAdapter
+from pydantic import BaseModel, TypeAdapter
 
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_logger
@@ -940,14 +941,26 @@ class CustomGuardrail(CustomLogger):
         return _request_carries_marker(data, self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN))
 
     def _post_call_deployment_marker(self, response: object) -> str | None:
-        response_id: Final[object] = (
-            cast("Mapping[object, object]", response).get("id")  # cast-ok: only reading a mapping value
-            if isinstance(response, Mapping)
-            else getattr(response, "id", None)
-        )
+        try:
+            response_data: Final[Mapping[str, object]] = (
+                response.model_dump(mode="json")
+                if isinstance(response, BaseModel)
+                else _REQUEST_FIELDS.validate_python(response)
+            )
+            output: Final = tuple(
+                (key, response_data[key]) for key in ("choices", "output", "content", "data") if key in response_data
+            )
+            if not output:
+                return None
+            digest: Final = hashlib.sha256(
+                json.dumps(output, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
+            ).hexdigest()
+        except (TypeError, ValueError):
+            return None
+        response_id: Final = response_data.get("id")
         if not isinstance(response_id, str) or not response_id.strip():
             return None
-        return self._marker(f"{_POST_CALL_DEPLOYMENT_HOOK_RAN_TOKEN}:{response_id}")
+        return self._marker(f"{_POST_CALL_DEPLOYMENT_HOOK_RAN_TOKEN}:{response_id}:{digest}")
 
     def post_call_deployment_hook_ran(self, data: Mapping[str, object], response: object) -> bool:
         return _request_carries_marker(data, self._post_call_deployment_marker(response))

tests/unit/integrations/test_custom_guardrail.py

@@ -3133,9 +3133,32 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
     def test_post_call_marker_requires_usable_id(self, response_id: object) -> None:
         guardrail: Final = CustomGuardrail(guardrail_name="post-mod")
         data: Final[dict[str, object]] = {"metadata": {}}
-        guardrail._record_post_call_deployment_hook_run(data, {"id": response_id})
+        guardrail._record_post_call_deployment_hook_run(data, {"id": response_id, "choices": []})
         assert data == {"metadata": {}}
-        assert guardrail.post_call_deployment_hook_ran(data, {"id": response_id}) is False
+        assert guardrail.post_call_deployment_hook_ran(data, {"id": response_id, "choices": []}) is False
+
+    @pytest.mark.parametrize("output", [None, object(), float("nan")])
+    def test_post_call_marker_requires_serializable_output(self, output: object) -> None:
+        guardrail: Final = CustomGuardrail(guardrail_name="post-mod")
+        data: Final[dict[str, object]] = {"metadata": {}}
+        response: Final = {"id": "chatcmpl-same", **({"choices": output} if output is not None else {})}
+        guardrail._record_post_call_deployment_hook_run(data, response)
+        assert data == {"metadata": {}}
+        assert guardrail.post_call_deployment_hook_ran(data, response) is False
+
+    def test_post_call_marker_tracks_tool_calls_but_not_routing_fields(self) -> None:
+        guardrail: Final = CustomGuardrail(guardrail_name="post-mod")
+        data: Final[dict[str, object]] = {"metadata": {}}
+        response: Final = {
+            "id": "chatcmpl-same",
+            "model": "deployment",
+            "choices": [{"message": {"tool_calls": [{"function": {"arguments": "checked"}}]}}],
+        }
+        guardrail._record_post_call_deployment_hook_run(data, response)
+        assert guardrail.post_call_deployment_hook_ran(data, {**response, "model": "alias"}) is True
+        assert guardrail.post_call_deployment_hook_ran(
+            data, {**response, "choices": [{"message": {"tool_calls": [{"function": {"arguments": "rewritten"}}]}}]}
+        ) is False
 
     @pytest.mark.asyncio
     async def test_apply_guardrail_retains_request_identity(self) -> None:

tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py

@@ -14,7 +14,8 @@ from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import CallTypes, UserAPIKeyAuth
 from litellm.proxy.utils import ProxyLogging
 from litellm.types.guardrails import GuardrailEventHooks
-from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse
+from litellm.types.utils import GenericGuardrailAPIInputs, LLMResponseTypes, ModelResponse
+from litellm.utils import async_post_call_success_deployment_hook
 
 
 @pytest.fixture(autouse=True)
@@ -86,6 +87,52 @@ async def test_post_call_moderates_returned_response_once(
     assert guardrail.evaluations == (2 if scenario in ("discarded", "no-id") else 1)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("parallel", [False, True])
+@pytest.mark.parametrize("rewrite", [False, True])
+async def test_post_call_rechecks_output_after_later_deployment_callback(
+    proxy_logging: ProxyLogging, monkeypatch: pytest.MonkeyPatch, parallel: bool, rewrite: bool
+) -> None:
+    class Guardrail(CustomGuardrail):
+        def __init__(self) -> None:
+            super().__init__(guardrail_name="post-mod", event_hook=GuardrailEventHooks.post_call)
+            self.evaluations = 0
+            self.run_in_parallel = parallel
+
+        async def async_post_call_success_hook(
+            self, data: dict[str, object], response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth
+        ) -> None:
+            self.evaluations += 1
+            if isinstance(response, ModelResponse) and "restored" in response.model_dump_json():
+                raise HTTPException(status_code=400, detail="blocked restored output")
+
+    class Rewriter(CustomLogger):
+        async def async_post_call_success_deployment_hook(
+            self, request_data: dict[str, object], response: LLMResponseTypes, call_type: CallTypes | None
+        ) -> ModelResponse | None:
+            if not rewrite or not isinstance(response, ModelResponse):
+                return None
+            return ModelResponse(
+                id=response.id, choices=[{"message": {"role": "assistant", "content": "restored"}}]
+            )
+
+    guardrail: Final = Guardrail()
+    monkeypatch.setattr(litellm, "callbacks", [guardrail, Rewriter()])
+    data: Final = {"model": "m", "guardrails": ["post-mod"], "metadata": {}}
+    response: Final = ModelResponse(
+        id="chatcmpl-same", choices=[{"message": {"role": "assistant", "content": "placeholder"}}]
+    )
+    returned: Final = await async_post_call_success_deployment_hook(data, response, CallTypes.acompletion)
+    if rewrite:
+        with pytest.raises(HTTPException, match="blocked restored output"):
+            await proxy_logging.post_call_success_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), response=returned)
+    else:
+        assert await proxy_logging.post_call_success_hook(
+            data=data, user_api_key_dict=UserAPIKeyAuth(), response=returned
+        ) is returned
+    assert guardrail.evaluations == (2 if rewrite else 1)
+
+
 @pytest.mark.asyncio
 async def test_post_call_success_hook_returns_response_when_no_callbacks(proxy_logging, make_user_api_key_auth, mock_callbacks_disabled):
     response = {"original": True, "model": "m", "choices": []}