fix(noma): keep masked choices when a later choice's scan fails open on a stream

litellm/proxy/guardrails/guardrail_hooks/noma/noma.py

@@ -826,6 +826,7 @@ class NomaGuardrail(CustomGuardrail):
         )
 
         if isinstance(assembled_model_response, ModelResponse):
+            texts_before_check: Final = self._choice_texts(assembled_model_response)
             try:
                 processed_response: Final = await self._check_llm_response(
                     request_data,
@@ -839,7 +840,11 @@ class NomaGuardrail(CustomGuardrail):
                 if self.block_failures:
                     raise
                 verbose_proxy_logger.error("Noma streaming post-call hook failed: %s", e)
-                for chunk in all_chunks:
+                if self._choice_texts(assembled_model_response) == texts_before_check:
+                    for chunk in all_chunks:
+                        yield chunk
+                    return
+                async for chunk in MockResponseIterator(model_response=assembled_model_response):
                     yield chunk
                 return
 

tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py

@@ -1,4 +1,5 @@
 import copy
+from itertools import chain
 import os
 from typing import cast
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -1951,6 +1952,103 @@ class TestNomaMultiChoiceResponses:
 
         assert [choice.message.content for choice in response.choices] == ["Mail *******", "Write *******"]
 
+    @pytest.mark.asyncio
+    async def test_fail_open_stream_keeps_redactions_made_before_a_later_choice_fails(self):
+        guardrail = NomaGuardrail(
+            api_key="test-api-key",
+            api_base="https://api.test.noma.security/",
+            anonymize_input=True,
+            monitor_mode=False,
+            block_failures=False,
+            guardrail_name="noma-post-call",
+            event_hook="post_call",
+            default_on=True,
+        )
+
+        def verdict_for_text(text):
+            if text.startswith("Write"):
+                raise httpx.ConnectError("Noma is unreachable")
+            return {
+                "aggregatedScanResult": False,
+                "scanResult": [
+                    {
+                        "role": "assistant",
+                        "type": "message",
+                        "results": {
+                            "anonymizedContent": {"anonymized": text.replace("test@example.com", "*******")},
+                            "sensitiveData": {"PII": {"result": False, "status": "SUCCESS"}},
+                        },
+                    }
+                ],
+            }
+
+        async def upstream():
+            for index, part in ((0, "Mail test@example.com"), (1, "Write back soon")):
+                yield ModelResponseStream(
+                    id="resp_1",
+                    choices=[StreamingChoices(delta=Delta(content=part), finish_reason="stop", index=index)],
+                    created=1234567890,
+                    model="gpt-4.1-mini",
+                    object="chat.completion.chunk",
+                )
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)):
+            streamed = [
+                chunk
+                async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+                    user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+                    response=upstream(),
+                    request_data={"litellm_call_id": "call-1", "metadata": {}},
+                )
+            ]
+
+        streamed_choices = tuple(chain.from_iterable(part.choices for part in streamed))
+        text_by_choice = {
+            index: "".join(choice.delta.content or "" for choice in streamed_choices if choice.index == index)
+            for index in (0, 1)
+        }
+        assert text_by_choice == {0: "Mail *******", 1: "Write back soon"}
+
+    @pytest.mark.asyncio
+    async def test_fail_open_stream_replays_the_original_chunks_when_nothing_was_redacted(self):
+        guardrail = NomaGuardrail(
+            api_key="test-api-key",
+            api_base="https://api.test.noma.security/",
+            anonymize_input=True,
+            monitor_mode=False,
+            block_failures=False,
+            guardrail_name="noma-post-call",
+            event_hook="post_call",
+            default_on=True,
+        )
+        originals = [
+            ModelResponseStream(
+                id="resp_1",
+                choices=[StreamingChoices(delta=Delta(content=part), index=0)],
+                created=1234567890,
+                model="gpt-4.1-mini",
+                object="chat.completion.chunk",
+            )
+            for part in ("Kiwis", " are", " fuzzy.")
+        ]
+
+        async def upstream():
+            for chunk in originals:
+                yield chunk
+
+        with patch.object(guardrail.async_handler, "post", side_effect=httpx.ConnectError("Noma is unreachable")):
+            streamed = [
+                chunk
+                async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+                    user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+                    response=upstream(),
+                    request_data={"litellm_call_id": "call-1", "metadata": {}},
+                )
+            ]
+
+        assert len(streamed) == len(originals)
+        assert all(sent is original for sent, original in zip(streamed, originals))
+
 
 class TestNomaAnonymizationFlow:
     """Test full anonymization flow with real Noma response objects"""