fix(proxy): skip token-id-only chunks on guarded streams and put stop_reason on the final text chunk

litellm/litellm_core_utils/streaming_handler.py

@@ -1314,6 +1314,8 @@ class CustomStreamWrapper:
                 _apply_choice_fields(model_response.choices[0], chunk_choice_fields)
                 if sibling_finish_reason is not None:
                     self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason)
+                elif response_obj.get("is_finished"):
+                    self._apply_finish_choice_fields(model_response.choices[0])
                 self._optional_combine_thinking_block_in_choices(model_response=model_response)
 
                 return model_response

litellm/proxy/utils.py

@@ -172,7 +172,7 @@ from litellm.litellm_core_utils.served_output_texts import (
 )
 from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model
 from litellm.llms import load_guardrail_translation_mappings
-from litellm.llms.base_llm.guardrail_translation.utils import stream_item_items
+from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.proxy._types import (
     AlertType,
@@ -706,6 +706,37 @@ def _drop_stream_chunk_token_ids(chunk: object) -> None:
         drop_token_ids(choice)
 
 
+_PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object])
+_EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {})
+
+
+def _holds_stream_payload(value: object) -> bool:
+    if isinstance(value, dict):
+        nested: Final = _PLAIN_STREAM_FIELDS.validate_python(value).values()
+        return any(item not in _EMPTY_STREAM_VALUES for item in nested)
+    return value not in _EMPTY_STREAM_VALUES
+
+
+def _choice_carries_only_token_ids(choice: object) -> bool:
+    if stream_item_field(choice, "token_ids") is None or not isinstance(choice, BaseModel):
+        return False
+    fields: Final = _PLAIN_STREAM_FIELDS.validate_python(choice.model_dump())
+    return not any(_holds_stream_payload(value) for key, value in fields.items() if key not in ("index", "token_ids"))
+
+
+def _is_token_id_only_chunk(chunk: object) -> bool:
+    choices: Final = stream_item_items(chunk, "choices")
+    return bool(choices) and all(_choice_carries_only_token_ids(choice) for choice in choices)
+
+
+async def _without_token_id_only_chunks(response: AsyncIterable[object]) -> AsyncGenerator[object, None]:
+    async for chunk in response:
+        if _is_token_id_only_chunk(chunk):
+            _drop_stream_chunk_token_ids(chunk)
+            continue
+        yield chunk
+
+
 def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool:
     return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines))
 
@@ -4540,7 +4571,7 @@ class ProxyLogging:
         self,
         response,
         user_api_key_dict: UserAPIKeyAuth,
-        request_data: dict,
+        request_data: dict[str, object],
     ):
         """
         Allow user to modify outgoing streaming data -> Given a whole response iterator.
@@ -4574,7 +4605,6 @@ class ProxyLogging:
         # Merge model-level guardrails before checking which guardrails to run
         request_data = check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router)
 
-        current_response = response
         stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
 
         pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines)
@@ -4590,6 +4620,7 @@ class ProxyLogging:
         drops_token_ids: Final = bool(post_call_pipelines) or any(
             isinstance(callback, CustomGuardrail) for callback, _ in running_overrides
         )
+        current_response = _without_token_id_only_chunks(response) if drops_token_ids else response
         guarded_layers: Final[list[AsyncGenerator[object, None]]] = []  # mutable-ok: closed on disconnect
         for resolved_callback, kind in running_overrides:
             effective_kind = (

tests/unit/proxy/test_proxy_logging_hook_detection.py

@@ -400,6 +400,68 @@ async def test_post_call_stream_keeps_token_ids_when_no_guardrail_runs_on_the_re
     assert served_token_ids == ([1, 2, 3], [4])
 
 
+def _completion_chunk(text: str, token_ids: list[int], finish_reason: str | None = None) -> object:
+    from litellm.types.utils import TextChoices, TextCompletionResponse
+
+    return TextCompletionResponse(
+        object="text_completion",
+        choices=[TextChoices(index=0, text=text, finish_reason=finish_reason, token_ids=token_ids)],
+    )
+
+
+def _guarded_completions_stream(monkeypatch, action: str, held_back: tuple[object, ...]):
+    from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
+        ContentFilterGuardrail,
+    )
+    from litellm.types.guardrails import BlockedWord, ContentFilterAction
+
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [
+            ContentFilterGuardrail(
+                guardrail_name="output-filter",
+                blocked_words=[BlockedWord(keyword="7", action=ContentFilterAction(action))],
+                event_hook="post_call",
+                default_on=True,
+            )
+        ],
+    )
+
+    async def upstream():
+        for chunk in held_back:
+            yield chunk
+        yield _completion_chunk(", 6, 7", [21, 11, 22], finish_reason="stop")
+
+    return ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
+        response=upstream(),
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/completions"),
+        request_data={"model": "qwen3", "metadata": {}},
+    )
+
+
+@pytest.mark.asyncio
+async def test_post_call_stream_guardrail_blocks_completions_before_serving_token_id_only_chunks(monkeypatch):
+    stream = _guarded_completions_stream(
+        monkeypatch, "BLOCK", (_completion_chunk("", [11]), _completion_chunk("", [220]))
+    )
+
+    with pytest.raises(HTTPException):
+        await stream.__anext__()
+
+
+@pytest.mark.asyncio
+async def test_post_call_stream_guardrail_serves_no_chunk_left_empty_by_the_token_id_drop(monkeypatch):
+    held_back = (_completion_chunk("", [11]), _completion_chunk("", [220]))
+
+    served = [chunk async for chunk in _guarded_completions_stream(monkeypatch, "MASK", held_back)]
+
+    assert [(chunk.choices[0].text, chunk.choices[0].token_ids) for chunk in served] == [
+        (", 6, [KEYWORD_REDACTED]", None)
+    ]
+    assert [chunk.choices[0].token_ids for chunk in held_back] == [None, None]
+
+
 @pytest.mark.asyncio
 async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch):
     from litellm.caching.caching import DualCache

tests/unit/text_completions/test_streaming_iterator.py

@@ -131,6 +131,18 @@ def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text
     }
 
 
+def test_text_completion_stream_puts_stop_reason_on_the_final_text_chunk() -> None:
+    responses: Final = _text_completion_stream(
+        (_completion_chunk(", 5", None), _completion_chunk(", 6", "stop", stop_reason=", 7"))
+    )
+
+    assert [(response.choices[0].get("text"), response.choices[0].get("stop_reason")) for response in responses] == [
+        (", 5", None),
+        (", 6", ", 7"),
+        (None, ", 7"),
+    ]
+
+
 @pytest.mark.parametrize("final_text", [", 6", ""], ids=["text-then-finish", "finish-only"])
 def test_text_completion_stream_keeps_held_back_token_ids_and_the_stop_reason(final_text: str) -> None:
     responses: Final = _text_completion_stream(