fix(guardrails): keep Gray Swan key metadata, scan a cut-short text stream, and pass non-chunk frames through

litellm/llms/openai/completion/guardrail_translation/handler.py

@@ -205,17 +205,18 @@ class OpenAITextCompletionHandler(BaseTranslation):
 
     async def process_output_streaming_response(
         self,
-        responses_so_far: list[TextCompletionResponse],  # mutable-ok: rewrites the caller's buffered chunks in place
+        responses_so_far: list[object],  # mutable-ok: rewrites the caller's buffered chunks in place
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
         user_api_key_dict: "UserAPIKeyAuth | None" = None,
         request_data: dict[str, object] | None = None,  # mutable-ok: the hooks share and mutate the request payload
         stream_transform_sink: StreamTransformSink | None = None,
         deliver_ended_stream_rewrites: bool = False,
-    ) -> Sequence[TextCompletionResponse]:
-        if not self._check_streaming_has_ended(responses_so_far):
+    ) -> Sequence[object]:
+        chunks: Final = _text_completion_chunks(responses_so_far)
+        if not self._check_streaming_has_ended(chunks):
             return responses_so_far
-        assembled: Final = self._assemble_ended_stream(responses_so_far)
+        assembled: Final = self._assemble_ended_stream(chunks)
         pre_guardrail_texts: Final = _choice_texts(assembled)
         await self.process_output_response(
             response=assembled,
@@ -230,16 +231,33 @@ class OpenAITextCompletionHandler(BaseTranslation):
             after for before, after in zip(pre_guardrail_texts, _choice_texts(assembled)) if after != before
         )
         for choice_index, text in rewrites:
-            self._write_choice_text(responses_so_far, choice_index, text)
+            self._write_choice_text(chunks, choice_index, text)
         return responses_so_far
 
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
-        chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, TextCompletionResponse))
+        chunks: Final = _text_completion_chunks(responses_so_far)
         if not self._check_streaming_has_ended(chunks):
             return None
         assembled: Final = self._assemble_ended_stream(chunks)
         return StreamingScanKey(texts=tuple(text for _, text in _choice_texts(assembled)), stream_ended=True)
 
+    def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]:
+        chunks: Final = _text_completion_chunks(responses_so_far)
+        streamed_choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks))
+        finished_indices: Final = frozenset(choice.index for choice in streamed_choices if choice.finish_reason)
+        unfinished_indices: Final = tuple(
+            dict.fromkeys(choice.index for choice in streamed_choices if choice.index not in finished_indices)
+        )
+        if not unfinished_indices:
+            return tuple(responses_so_far)
+        terminator: Final = TextCompletionResponse(
+            id=chunks[0].id,
+            model=chunks[0].model,
+            object="text_completion",
+            choices=[TextChoices(index=index, text="", finish_reason="length") for index in unfinished_indices],
+        )
+        return (*responses_so_far, terminator)
+
     def build_stream_error_items(
         self,
         exc: "HTTPException",
@@ -293,6 +311,10 @@ class OpenAITextCompletionHandler(BaseTranslation):
                 choice.text = ""
 
 
+def _text_completion_chunks(responses_so_far: Sequence[object]) -> tuple[TextCompletionResponse, ...]:
+    return tuple(item for item in responses_so_far if isinstance(item, TextCompletionResponse))
+
+
 def _choice_text(choice: TextChoices) -> str:
     text: Final = stream_item_field(choice, "text")
     return text if isinstance(text, str) else ""

litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py

@@ -81,6 +81,21 @@ def _as_plain_dict(item: object) -> Mapping[str, object]:
     return cast("Mapping[str, object]", item)  # cast-ok: wire rows are message/tool-call dicts
 
 
+_STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object])
+
+
+def _str_keyed(value: object) -> Mapping[str, object]:
+    return _STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else {}
+
+
+def _proxy_metadata(request_data: Mapping[str, object]) -> Mapping[str, object]:
+    litellm_metadata: Final = request_data.get("litellm_metadata")
+    if litellm_metadata is not None:
+        return _str_keyed(litellm_metadata)
+    metadata: Final = _str_keyed(request_data.get("metadata"))
+    return {key: value for key, value in metadata.items() if key.startswith("user_api_key_")}
+
+
 class GraySwanGuardrailMissingSecrets(Exception):
     """Raised when the Gray Swan API key is missing."""
 
@@ -624,10 +639,7 @@ class GraySwanGuardrail(CustomGuardrail):
         from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
         from litellm.llms import load_guardrail_translation_mappings
 
-        litellm_metadata: Final = request_data.get("litellm_metadata")
-        request_route: Final = (
-            litellm_metadata.get("user_api_key_request_route") if isinstance(litellm_metadata, Mapping) else None
-        )
+        request_route: Final = _proxy_metadata(request_data).get("user_api_key_request_route")
         route_call_types: Final = get_call_types_for_route(request_route) if isinstance(request_route, str) else None
         call_type: Final = (
             (route_call_types[0].value if route_call_types else None)
@@ -697,13 +709,12 @@ class GraySwanGuardrail(CustomGuardrail):
 
         inbound_headers: Final = self._extract_inbound_headers(request_data, logging_obj)
 
-        litellm_metadata: Final = request_data.get("litellm_metadata")
-        cleaned_litellm_metadata: Final = dict(litellm_metadata) if isinstance(litellm_metadata, dict) else {}
-        if inbound_headers:
-            existing_headers: Final = cleaned_litellm_metadata.get("headers")
-            cleaned_litellm_metadata["headers"] = (
-                {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers
-            )
+        proxy_metadata: Final = _proxy_metadata(request_data)
+        cleaned_litellm_metadata: Final = (
+            {**proxy_metadata, "headers": {**_str_keyed(proxy_metadata.get("headers")), **inbound_headers}}
+            if inbound_headers
+            else proxy_metadata
+        )
         if cleaned_litellm_metadata:
             sanitized: Final[object] = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={})
             if isinstance(sanitized, dict) and sanitized:

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

@@ -5,6 +5,7 @@ import pytest
 from fastapi import HTTPException
 
 from litellm.integrations.custom_guardrail import ModifyResponseException
+from litellm.llms.openai.completion.guardrail_translation.handler import OpenAITextCompletionHandler
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.grayswan import grayswan as grayswan_module
 from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import (
@@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import (
     GraySwanGuardrailAPIError,
 )
 from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.utils import TextChoices, TextCompletionResponse
 
 
 @pytest.fixture
@@ -884,3 +886,24 @@ async def test_pre_call_payload_unchanged() -> None:
         {"role": "user", "content": "second"},
     ]
     assert "tools" not in payload
+
+
+@pytest.mark.asyncio
+async def test_post_call_on_a_text_completion_forwards_only_the_key_and_team_metadata() -> None:
+    guardrail = _post_call_guardrail()
+    client = _CapturingClient()
+    guardrail.async_handler = client
+
+    await OpenAITextCompletionHandler().process_output_response(
+        response=TextCompletionResponse(
+            id="cmpl-1", model="gpt-4.1-mini", choices=[TextChoices(index=0, text="hello", finish_reason="stop")]
+        ),
+        guardrail_to_apply=guardrail,
+        user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key"),
+        request_data={"model": "gpt-4.1-mini", "prompt": "say hello", "metadata": {"requester_note": "private"}},
+    )
+
+    forwarded = client.calls[0]["json"]["litellm_metadata"]
+    assert forwarded["user_api_key_team_id"] == "team-456"
+    assert forwarded["user_api_key_key_alias"] == "ci-key"
+    assert "requester_note" not in forwarded

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

@@ -2084,6 +2084,73 @@ async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_a_text_co
     assert _warnings(caplog) == []
 
 
+class _FlatBlockGuardrail(CustomGuardrail):
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        if any("kumquat" in text for text in inputs.get("texts") or []):
+            raise HTTPException(status_code=400, detail="flat guardrail blocked kumquat")
+        return inputs
+
+
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_delivers_a_flat_guardrail_block_on_a_text_completion_stream_under_a_pipeline(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Final[Dict[str, Any]] = {}
+    flat: Final = _FlatBlockGuardrail(
+        guardrail_name="flat-block", event_hook=GuardrailEventHooks.post_call, default_on=True
+    )
+    monkeypatch.setattr(litellm, "callbacks", [flat, _live_scan_guardrail(seen)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"}
+    chunks: Final = [
+        _text_completion_chunk("A kum"),
+        _text_completion_chunk("quat tart"),
+        _text_completion_chunk("", "stop"),
+    ]
+    auth: Final = make_user_api_key_auth(request_route="/v1/completions")
+
+    await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True)
+    delivered: Final = [
+        item
+        async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
+        )
+    ]
+
+    assert seen["texts"] == ["A kumquat tart"]
+    assert isinstance(delivered[-1], bytes)
+    assert json.loads(delivered[-1].decode().removeprefix("data: "))["error"]["message"] == (
+        "flat guardrail blocked kumquat"
+    )
+
+
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_live_pipeline_scans_what_a_disconnected_text_completion_client_received(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Final[Dict[str, Any]] = {}
+    monkeypatch.setattr(litellm, "callbacks", [_live_scan_guardrail(seen)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data: Final = {**_post_call_pipeline_data(step=_DETECT_ONLY_STEP, stream=True), "prompt": "hi"}
+    chunks: Final = [
+        _text_completion_chunk("hello "),
+        _text_completion_chunk("world"),
+        _text_completion_chunk("", "stop"),
+    ]
+    auth: Final = make_user_api_key_auth(request_route="/v1/completions")
+
+    await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True)
+    stream: Final = proxy_logging.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
+    )
+    first: Final = await stream.__anext__()
+    await stream.aclose()
+
+    assert first.choices[0].text == "hello "
+    assert seen["texts"] == ["hello "]
+    assert _post_call_statuses(data) == ["success"]
+
+
 @pytest.mark.asyncio
 async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none(
     proxy_logging, make_user_api_key_auth, monkeypatch