test: cover private diagnostics across stream and rule paths

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -3501,11 +3501,17 @@ def test_dispatch_petals_empty_stream_after_finish_raises(
         _run_dispatch(initialized_custom_stream_wrapper, chunk=None)
 
 
+@pytest.mark.parametrize("redact", (False, True))
 def test_dispatch_cached_response_extracts_delta(
     initialized_custom_stream_wrapper: CustomStreamWrapper,
+    redact: bool,
+    capsys: pytest.CaptureFixture[str],
+    monkeypatch: pytest.MonkeyPatch,
 ):
     """cached_response replays a stored ModelResponseStream; the dispatch lifts
     its delta content, finish_reason and id back onto the live response."""
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
+    monkeypatch.setattr(litellm, "set_verbose", True)
     initialized_custom_stream_wrapper.custom_llm_provider = "cached_response"
     chunk = ModelResponseStream(
         id="chatcmpl-cache-1",
@@ -3527,6 +3533,7 @@ def test_dispatch_cached_response_extracts_delta(
     assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
     assert model_response.id == "chatcmpl-cache-1"
     assert initialized_custom_stream_wrapper.response_id == "chatcmpl-cache-1"
+    assert ("cached text" in capsys.readouterr().out) is not redact
 
 
 def test_dispatch_cached_response_without_choices_is_an_empty_chunk(
@@ -3698,11 +3705,17 @@ def test_dispatch_custom_provider_finish_only_returns_none_early(
     assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
 
 
+@pytest.mark.parametrize("redact", (False, True))
 def test_dispatch_text_completion_codestral_parses_chunk(
     initialized_custom_stream_wrapper: CustomStreamWrapper,
+    redact: bool,
+    capsys: pytest.CaptureFixture[str],
+    monkeypatch: pytest.MonkeyPatch,
 ):
     """text-completion-codestral streams raw SSE JSON strings that the dispatch
     routes through CodestralTextCompletionConfig to extract content/finish."""
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
+    monkeypatch.setattr(litellm, "set_verbose", True)
     initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-codestral"
     chunk = json.dumps(
         {"choices": [{"delta": {"content": "codestral text"}, "finish_reason": "stop"}]}
@@ -3713,6 +3726,7 @@ def test_dispatch_text_completion_codestral_parses_chunk(
     assert isinstance(result, _ProviderChunkParsed)
     assert completion_obj["content"] == "codestral text"
     assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
+    assert ("codestral text" in capsys.readouterr().out) is not redact
 
 
 def test_dispatch_text_completion_codestral_requires_string(
@@ -3726,10 +3740,16 @@ def test_dispatch_text_completion_codestral_requires_string(
         _run_dispatch(initialized_custom_stream_wrapper, {"not": "a string"})
 
 
+@pytest.mark.parametrize("redact", (False, True))
 def test_dispatch_triton_stream(
     initialized_custom_stream_wrapper: CustomStreamWrapper,
+    redact: bool,
+    capsys: pytest.CaptureFixture[str],
+    monkeypatch: pytest.MonkeyPatch,
 ):
     """triton stream chunks arrive as dicts keyed by text_output/stop_reason."""
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
+    monkeypatch.setattr(litellm, "set_verbose", True)
     initialized_custom_stream_wrapper.custom_llm_provider = "triton"
     chunk = {"text_output": "triton text", "is_finished": True, "stop_reason": "stop"}
 
@@ -3738,13 +3758,20 @@ def test_dispatch_triton_stream(
     assert isinstance(result, _ProviderChunkParsed)
     assert completion_obj["content"] == "triton text"
     assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
+    assert ("triton text" in capsys.readouterr().out) is not redact
 
 
+@pytest.mark.parametrize("redact", (False, True))
 def test_dispatch_text_completion_openai_with_usage(
     initialized_custom_stream_wrapper: CustomStreamWrapper,
+    redact: bool,
+    capsys: pytest.CaptureFixture[str],
+    monkeypatch: pytest.MonkeyPatch,
 ):
     """text-completion-openai chunks expose choices[].text and an optional usage
     block that the dispatch lifts onto the model response."""
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
+    monkeypatch.setattr(litellm, "set_verbose", True)
     initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-openai"
 
     class _Choice:
@@ -3769,6 +3796,7 @@ def test_dispatch_text_completion_openai_with_usage(
     assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
     assert model_response.usage.prompt_tokens == 5
     assert model_response.usage.total_tokens == 8
+    assert ("oai text" in capsys.readouterr().out) is not redact
 
 
 @pytest.mark.parametrize("custom_llm_provider", ["text-completion-openai", "azure_text"])
@@ -7953,3 +7981,38 @@ def test_legacy_stream_verbose_output_obeys_diagnostic_privacy(monkeypatch, caps
     monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", enabled)
     print_verbose("chunk: %s", "PRIVATE_STREAM_CONTENT")
     assert capsys.readouterr().out.strip() == ("redacted-by-litellm" if enabled else "chunk: PRIVATE_STREAM_CONTENT")
+
+
+@pytest.mark.parametrize("redact", (False, True))
+def test_predibase_stream_privacy_preserves_content_and_finish(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+    redact: bool,
+    capsys: pytest.CaptureFixture[str],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
+    monkeypatch.setattr(litellm, "set_verbose", True)
+    initialized_custom_stream_wrapper.custom_llm_provider = "predibase"
+    result, _, completion_obj = _run_dispatch(
+        initialized_custom_stream_wrapper,
+        b'data: {"token": {"text": "PRIVATE_PREDIBASE_REPLY"}, "details": {"finish_reason": "stop"}}',
+    )
+    assert isinstance(result, _ProviderChunkParsed)
+    assert completion_obj["content"] == "PRIVATE_PREDIBASE_REPLY"
+    assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
+    assert ("PRIVATE_PREDIBASE_REPLY" in capsys.readouterr().out) is not redact
+
+
+@pytest.mark.parametrize("redact", (False, True))
+def test_triton_invalid_stream_chunk_keeps_error_out_of_private_diagnostics(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+    redact: bool,
+    capsys: pytest.CaptureFixture[str],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
+    monkeypatch.setattr(litellm, "set_verbose", True)
+    initialized_custom_stream_wrapper.custom_llm_provider = "triton"
+    with pytest.raises(ValueError, match=r"Unable to parse response.*PRIVATE_TRITON_CHUNK"):
+        _run_dispatch(initialized_custom_stream_wrapper, ["PRIVATE_TRITON_CHUNK"])
+    assert ("PRIVATE_TRITON_CHUNK" in capsys.readouterr().out) is not redact

tests/unit/proxy/test_post_call_rules.py

@@ -0,0 +1,26 @@
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.proxy.post_call_rules import post_response_rule
+
+
+@pytest.mark.parametrize("redact", (False, True))
+@pytest.mark.parametrize("long_response", (False, True))
+def test_diagnostic_privacy_preserves_response_rule_decision(
+    redact: bool,
+    long_response: bool,
+    monkeypatch: pytest.MonkeyPatch,
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
+    marker: Final = "PRIVATE_RULE_INPUT"
+    response: Final = marker * (20 if long_response else 1)
+    expected: Final = (
+        {"decision": True}
+        if long_response
+        else {"decision": False, "message": "This violates LiteLLM Proxy Rules. Response too short"}
+    )
+    assert post_response_rule(response) == expected
+    assert (marker in capsys.readouterr().out) is not redact

tests/unit/responses/mcp/test_mcp_streaming_iterator.py

@@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock
 import pytest
 from mcp.types import CallToolResult, TextContent
 
-import litellm  # noqa: F401 - ensures litellm.responses.main is registered in sys.modules
+import litellm
 from litellm.responses.mcp.mcp_streaming_iterator import (
     MCPEnhancedStreamingIterator,
 )
@@ -293,14 +293,18 @@ async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch)
     assert completed[-1].response.id != "resp-interim"
 
 
+@pytest.mark.parametrize("redact", (False, True))
 @pytest.mark.asyncio
-async def test_follow_up_call_failure_emits_terminal_error_event(monkeypatch):
+async def test_follow_up_call_failure_emits_terminal_error_event(
+    monkeypatch: pytest.MonkeyPatch, redact: bool, capsys: pytest.CaptureFixture[str]
+) -> None:
     """
     Regression test for the silent-swallow path: when the follow-up LLM call
     raises, the stream must emit a terminal `error` event instead of ending
     silently after the tool events (which surfaced to clients as a successful
     but empty completion).
     """
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
     _mock_mcp_environment(monkeypatch)
 
     aresponses_mock = AsyncMock(side_effect=RuntimeError("boom from provider"))
@@ -319,16 +323,21 @@ async def test_follow_up_call_failure_emits_terminal_error_event(monkeypatch):
     assert len(error_events) == 1
     assert error_events[0].error.type == "mcp_gateway_error"
     assert "boom from provider" in error_events[0].error.message
+    assert ("boom from provider" in capsys.readouterr().err) is not redact
 
 
+@pytest.mark.parametrize("redact", (False, True))
 @pytest.mark.asyncio
-async def test_initial_call_failure_is_stashed_for_eager_reraise(monkeypatch):
+async def test_initial_call_failure_is_stashed_for_eager_reraise(
+    monkeypatch: pytest.MonkeyPatch, redact: bool, capsys: pytest.CaptureFixture[str]
+) -> None:
     """
     Regression test: a failing initial LLM call must be stashed as
     `_initial_creation_error` so `aresponses_api_with_mcp` can re-raise it as a
     real 4xx before any SSE bytes are written, instead of returning HTTP 200
     with an empty stream.
     """
+    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
     _mock_mcp_environment(monkeypatch)
 
     aresponses_mock = AsyncMock(side_effect=RuntimeError("initial boom"))
@@ -339,6 +348,7 @@ async def test_initial_call_failure_is_stashed_for_eager_reraise(monkeypatch):
 
     assert iterator._initial_creation_error is not None
     assert "initial boom" in str(iterator._initial_creation_error)
+    assert ("initial boom" in capsys.readouterr().err) is not redact
 
 
 def _reasoning_item(encrypted_content: str):