diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 8eaccf1ec64..5bcd959d0fa 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/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 diff --git a/tests/unit/proxy/test_post_call_rules.py b/tests/unit/proxy/test_post_call_rules.py new file mode 100644 index 00000000000..55ee5142bcf --- /dev/null +++ b/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 diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index f745b5bfadf..a9d4e7067a1 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/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):