diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index edee50bdfc..c4c9aea6f6 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -21,6 +21,7 @@ import httpx from httpx._types import CookieTypes, QueryParamTypes, RequestFiles import litellm +from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler @@ -390,19 +391,28 @@ def _sync_streaming( ): from litellm.utils import executor + raw_bytes: List[bytes] = [] + flush_scheduled = False try: - raw_bytes: List[bytes] = [] for chunk in response.iter_bytes(): # type: ignore raw_bytes.append(chunk) yield chunk - - executor.submit( - litellm_logging_obj.flush_passthrough_collected_chunks, - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - except Exception as e: - raise e + finally: + if not flush_scheduled and raw_bytes: + flush_scheduled = True + try: + executor.submit( + litellm_logging_obj.flush_passthrough_collected_chunks, + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + except Exception as e: + verbose_logger.exception( + "Failed to schedule passthrough spend-tracking flush " + "in _sync_streaming; %d buffered chunks dropped: %s", + len(raw_bytes), + e, + ) async def _async_streaming( @@ -411,23 +421,45 @@ async def _async_streaming( provider_config: "BasePassthroughConfig", ): iter_response = await response + try: iter_response.raise_for_status() - raw_bytes: List[bytes] = [] - - async for chunk in iter_response.aiter_bytes(): # type: ignore - raw_bytes.append(chunk) - yield chunk - - asyncio.create_task( - litellm_logging_obj.async_flush_passthrough_collected_chunks( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - ) except Exception: try: await iter_response.aclose() except Exception: pass raise + + raw_bytes: List[bytes] = [] + flush_scheduled = False + try: + async for chunk in iter_response.aiter_bytes(): # type: ignore + raw_bytes.append(chunk) + yield chunk + except Exception: + try: + await iter_response.aclose() + except Exception: + pass + raise + finally: + # GeneratorExit (raised on client disconnect) is not caught by + # `except Exception`; the finally block ensures partial usage + # still gets flushed for spend tracking. See LIT-2642. + if not flush_scheduled and raw_bytes: + flush_scheduled = True + try: + asyncio.create_task( + litellm_logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + ) + except Exception as e: + verbose_logger.exception( + "Failed to schedule passthrough spend-tracking flush " + "in _async_streaming; %d buffered chunks dropped: %s", + len(raw_bytes), + e, + ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 302d7e76ed..cbfcd34c43 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -36,21 +36,16 @@ class PassThroughStreamingHandler: passthrough_success_handler_obj: PassThroughEndpointLogging, url_route: str, ): - """ - - Yields chunks from the response - - Collect non-empty chunks for post-processing (logging) - - Inject cost into chunks if include_cost_in_streaming_usage is enabled - """ - try: - raw_bytes: List[bytes] = [] - # Extract model name for cost injection - model_name = PassThroughStreamingHandler._extract_model_for_cost_injection( - request_body=request_body, - url_route=url_route, - endpoint_type=endpoint_type, - litellm_logging_obj=litellm_logging_obj, - ) + raw_bytes: List[bytes] = [] + logging_scheduled = False + model_name = PassThroughStreamingHandler._extract_model_for_cost_injection( + request_body=request_body, + url_route=url_route, + endpoint_type=endpoint_type, + litellm_logging_obj=litellm_logging_obj, + ) + try: async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) if ( @@ -58,7 +53,6 @@ class PassThroughStreamingHandler: and model_name ): if endpoint_type == EndpointType.VERTEX_AI: - # Only handle streamRawPredict (uses Anthropic format) if "streamRawPredict" in url_route or "rawPredict" in url_route: modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( chunk, model_name @@ -73,25 +67,32 @@ class PassThroughStreamingHandler: chunk = modified_chunk yield chunk - - # After all chunks are processed, handle post-processing - end_time = datetime.now() - - asyncio.create_task( - PassThroughStreamingHandler._route_streaming_logging_to_handler( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body or {}, - endpoint_type=endpoint_type, - start_time=start_time, - raw_bytes=raw_bytes, - end_time=end_time, - ) - ) except Exception as e: verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}") raise + finally: + # GeneratorExit (raised on client disconnect) is not caught by + # `except Exception`; the finally block ensures partial usage + # still gets logged for spend tracking. See LIT-2642. + if not logging_scheduled and raw_bytes: + logging_scheduled = True + try: + asyncio.create_task( + PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body or {}, + endpoint_type=endpoint_type, + start_time=start_time, + raw_bytes=raw_bytes, + end_time=datetime.now(), + ) + ) + except Exception as e: + verbose_proxy_logger.error( + f"Error scheduling chunk_processor logging: {str(e)}" + ) @staticmethod async def _route_streaming_logging_to_handler( diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py index 2d772c4a63..75061dda94 100644 --- a/tests/otel_tests/test_prometheus.py +++ b/tests/otel_tests/test_prometheus.py @@ -115,6 +115,13 @@ async def test_proxy_failure_metrics(): "litellm_llm_api_failed_requests_metric_total{", # Deprecated but may still be used ] + # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for + # hash_token(master_key) so the master key (or its hash) never + # propagates into metrics. See PR #26484. + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS + # Check if either pattern is in metrics and contains required fields found_metric = False for pattern in expected_patterns: @@ -125,8 +132,7 @@ async def test_proxy_failure_metrics(): 'api_key_alias="None"' in line and 'exception_class="Openai.RateLimitError"' in line and 'exception_status="429"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and 'route="/chat/completions"' in line ): @@ -135,8 +141,7 @@ async def test_proxy_failure_metrics(): # For deprecated llm_api metric, check llm-specific fields elif "litellm_llm_api_failed_requests_metric_total{" in line: if ( - 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + f'hashed_api_key="{expected_hashed_api_key}"' in line and 'model="429"' in line ): # The deprecated metric uses the actual model from the request found_metric = True @@ -156,8 +161,7 @@ async def test_proxy_failure_metrics(): for line in metrics.split("\n"): if ( total_requests_pattern in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and 'status_code="429"' in line ): @@ -195,6 +199,12 @@ async def test_proxy_success_metrics(): assert END_USER_ID not in metrics + # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for + # hash_token(master_key) (PR #26484). + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS + # Check if the success metric is present and correct - use flexible matching # Check for request_total_latency_metric with required fields # Note: The model can be "gpt-3.5-turbo-0301" or similar depending on what's returned @@ -203,8 +213,7 @@ async def test_proxy_success_metrics(): if ( "litellm_request_total_latency_metric_bucket{" in line and 'api_key_alias="None"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-openai-endpoint"' in line and 'le="0.005"' in line ): @@ -221,8 +230,7 @@ async def test_proxy_success_metrics(): if ( "litellm_llm_api_latency_metric_bucket{" in line and 'api_key_alias="None"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-openai-endpoint"' in line and 'le="0.005"' in line ): @@ -298,6 +306,12 @@ async def test_proxy_fallback_metrics(): print("/metrics", metrics) + # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for + # hash_token(master_key) (PR #26484). + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS + # Check if successful fallback metric is incremented - use flexible matching found_successful_fallback = False for line in metrics.split("\n"): @@ -307,8 +321,7 @@ async def test_proxy_fallback_metrics(): and 'exception_class="Openai.RateLimitError"' in line and 'exception_status="429"' in line and 'fallback_model="fake-openai-endpoint"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and "1.0" in line ): @@ -328,8 +341,7 @@ async def test_proxy_fallback_metrics(): and 'exception_class="Openai.RateLimitError"' in line and 'exception_status="429"' in line and 'fallback_model="unknown-model"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and "1.0" in line ): diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py new file mode 100644 index 0000000000..f3fe3ae5c3 --- /dev/null +++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py @@ -0,0 +1,244 @@ +"""Regression tests for LIT-2642 — interrupted streams must still flush usage.""" + +import asyncio +from typing import List +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + + +def _make_streaming_response(chunks: List[bytes]): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + mock.headers = httpx.Headers({"content-type": "application/vnd.amazon.eventstream"}) + mock.raise_for_status = MagicMock(return_value=None) + + async def _aiter_bytes(): + for chunk in chunks: + yield chunk + + mock.aiter_bytes = _aiter_bytes + mock.aclose = AsyncMock() + return mock + + +def _make_logging_obj(): + mock = MagicMock() + mock.async_flush_passthrough_collected_chunks = AsyncMock() + return mock + + +class _ImmediateExecutor: + def submit(self, fn, *args, **kwargs): + fn(*args, **kwargs) + + +@pytest.mark.asyncio +async def test_async_streaming_flushes_on_normal_completion(): + from litellm.passthrough.main import _async_streaming + + chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] + mock_response = _make_streaming_response(chunks) + + async def response_coro(): + return mock_response + + mock_logging_obj = _make_logging_obj() + provider_config = MagicMock() + + received = [] + async for chunk in _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ): + received.append(chunk) + + assert received == chunks + + await asyncio.sleep(0) + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = ( + mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs + ) + assert call_kwargs["raw_bytes"] == chunks + assert call_kwargs["provider_config"] is provider_config + + +@pytest.mark.asyncio +async def test_async_streaming_flushes_on_client_disconnect(): + from litellm.passthrough.main import _async_streaming + + chunks = [ + b'{"chunk": 1, "outputTokens": 10}', + b'{"chunk": 2, "outputTokens": 12}', + b'{"chunk": 3, "outputTokens": 8}', + ] + mock_response = _make_streaming_response(chunks) + + async def response_coro(): + return mock_response + + mock_logging_obj = _make_logging_obj() + provider_config = MagicMock() + + gen = _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ) + + received = [await gen.__anext__()] + await gen.aclose() + + assert received == [chunks[0]] + + await asyncio.sleep(0) + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = ( + mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs + ) + assert call_kwargs["raw_bytes"] == [chunks[0]] + + +@pytest.mark.asyncio +async def test_async_streaming_does_not_flush_on_4xx(): + from litellm.passthrough.main import _async_streaming + + err_response = MagicMock(spec=httpx.Response) + err_response.status_code = 429 + + def _raise(): + raise httpx.HTTPStatusError( + "429", + request=httpx.Request("POST", "https://example.com"), + response=httpx.Response( + 429, request=httpx.Request("POST", "https://example.com") + ), + ) + + err_response.raise_for_status = _raise + err_response.aclose = AsyncMock() + + async def response_coro(): + return err_response + + mock_logging_obj = _make_logging_obj() + + with pytest.raises(httpx.HTTPStatusError): + async for _ in _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=MagicMock(), + ): + pass + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_streaming_flushes_on_upstream_exception_with_partial_data(): + from litellm.passthrough.main import _async_streaming + + partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"] + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock(return_value=None) + mock_response.aclose = AsyncMock() + + async def _aiter_bytes_then_raise(): + for c in partial_chunks: + yield c + raise httpx.ReadError("upstream disconnected") + + mock_response.aiter_bytes = _aiter_bytes_then_raise + + async def response_coro(): + return mock_response + + mock_logging_obj = _make_logging_obj() + provider_config = MagicMock() + + received = [] + with pytest.raises(httpx.ReadError): + async for chunk in _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ): + received.append(chunk) + + assert received == partial_chunks + + await asyncio.sleep(0) + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = ( + mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs + ) + assert call_kwargs["raw_bytes"] == partial_chunks + + +def test_sync_streaming_flushes_on_normal_completion(): + from litellm.passthrough.main import _sync_streaming + + chunks = [b"a", b"b", b"c"] + + mock_response = MagicMock(spec=httpx.Response) + + def _iter_bytes(): + yield from chunks + + mock_response.iter_bytes = _iter_bytes + + mock_logging_obj = MagicMock() + mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() + provider_config = MagicMock() + + with patch("litellm.utils.executor", _ImmediateExecutor()): + received = list( + _sync_streaming( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ) + ) + + assert received == chunks + mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once() + + +def test_sync_streaming_flushes_on_early_close(): + from litellm.passthrough.main import _sync_streaming + + chunks = [b"first", b"second", b"third"] + + mock_response = MagicMock(spec=httpx.Response) + + def _iter_bytes(): + yield from chunks + + mock_response.iter_bytes = _iter_bytes + + mock_logging_obj = MagicMock() + mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() + provider_config = MagicMock() + + with patch("litellm.utils.executor", _ImmediateExecutor()): + gen = _sync_streaming( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ) + + first = next(gen) + gen.close() + + assert first == chunks[0] + mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = mock_logging_obj.flush_passthrough_collected_chunks.call_args.kwargs + assert call_kwargs["raw_bytes"] == [chunks[0]] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py new file mode 100644 index 0000000000..f73aee77cc --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -0,0 +1,120 @@ +"""Regression tests for LIT-2642 — interrupted pass-through streams must still log usage.""" + +import asyncio +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + +def _make_streaming_response(chunks): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes(): + for c in chunks: + yield c + + mock.aiter_bytes = _aiter_bytes + return mock + + +@pytest.mark.asyncio +async def test_chunk_processor_logs_on_normal_completion(): + chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] + response = _make_streaming_response(chunks) + + mock_logging_obj = MagicMock() + mock_passthrough_handler = MagicMock() + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=mock_logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=mock_passthrough_handler, + url_route="/bedrock/model/claude/invoke-with-response-stream", + ): + received.append(chunk) + + await asyncio.sleep(0) + + assert received == chunks + mock_route.assert_called_once() + call_kwargs = mock_route.call_args.kwargs + assert call_kwargs["raw_bytes"] == chunks + + +@pytest.mark.asyncio +async def test_chunk_processor_logs_on_client_disconnect(): + chunks = [b"event-1", b"event-2", b"event-3"] + response = _make_streaming_response(chunks) + + mock_logging_obj = MagicMock() + mock_passthrough_handler = MagicMock() + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=mock_logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=mock_passthrough_handler, + url_route="/bedrock/model/claude/invoke-with-response-stream", + ) + + first = await gen.__anext__() + await gen.aclose() + + await asyncio.sleep(0) + + assert first == chunks[0] + mock_route.assert_called_once() + call_kwargs = mock_route.call_args.kwargs + assert call_kwargs["raw_bytes"] == [chunks[0]] + + +@pytest.mark.asyncio +async def test_chunk_processor_does_not_schedule_logging_when_no_chunks(): + response = _make_streaming_response([]) + + mock_logging_obj = MagicMock() + mock_passthrough_handler = MagicMock() + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=mock_logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=mock_passthrough_handler, + url_route="/bedrock/model/claude/invoke-with-response-stream", + ): + received.append(chunk) + + assert received == [] + mock_route.assert_not_called()