From 571ada0b0f09c0356055626160f78c583a9b2998 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:36:25 -0700 Subject: [PATCH] feat(rust_bridge): mark native streams with the x-litellm-rust header (#42758) Non-streaming responses served by the Rust core already carry x-litellm-rust: true through _hidden_params.additional_headers, which the SDK exposes and the gateway renders as a response header. Native streams did not, because the lifecycle Stream and SyncStream objects had nowhere to hold hidden params and the marker writer skips objects without them. Give both stream classes the same _hidden_params bag every other litellm response has, so the existing marker attaches without wrapping the stream or changing its identity. Co-authored-by: Yujong Lee --- litellm/rust_bridge/lifecycle.py | 2 + .../test_litellm/rust_bridge/test_runtime.py | 56 +++++++++++++++++++ .../messages/test_callbacks.py | 3 + 3 files changed, 61 insertions(+) diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index b3b7a1888c3..2f243e8c212 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -81,6 +81,7 @@ class Stream(AsyncIterator[object]): def __init__(self, execution: Execution) -> None: self._execution: Final = execution self._done = False + self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place def __aiter__(self) -> Stream: return self @@ -117,6 +118,7 @@ class SyncStream(Iterator[object]): def __init__(self, execution: Execution) -> None: self._execution: Final = execution self._done = False + self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place def __iter__(self) -> SyncStream: return self diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index bff7ded3114..bbe25e0de13 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -12,6 +12,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di from litellm.rust_bridge import bindings, configuration, runtime from litellm.rust_bridge.catalog import Delivery, Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.lifecycle import Complete, Open, Stream, SyncStream, Yield class RustBridgeDeclined(Exception): @@ -264,6 +265,61 @@ async def test_native_response_marker_reaches_caller_with_existing_metadata(shap } +class ScriptedStreamExecution: + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._steps: Final = iter((*(Yield(chunk) for chunk in chunks), Complete(None))) + self.closed = False + + def start(self) -> Open: + return Open(None) + + def resume_value(self, value: object) -> Yield | Complete: + return next(self._steps) + + def resume_error(self, error: BaseException) -> Complete: + return Complete(None) + + def close(self) -> None: + self.closed = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_native_stream_marker_reaches_caller_without_wrapping_or_consuming_the_stream( + asynchronous: bool, +) -> None: + chunks: Final = (b"event: message_start\n\n", b"event: message_stop\n\n") + execution: Final = ScriptedStreamExecution(chunks) + stream: Final[Stream | SyncStream] = Stream(execution) if asynchronous else SyncStream(execution) + bound: Final[bindings.NativeBinding[Callable[[], object]]] = bindings.NativeBinding( + "messages", validate=lambda _: None + ) + bound.override(lambda: stream) + + def python() -> object: + pytest.fail("native success must not fall back") + + async def anative(fn: Callable[[], object]) -> object: + return fn() + + async def apython() -> object: + return python() + + result: Final = ( + await runtime.arun(CONTEXT, binding=bound, native=anative, python=apython, rules=rules(Rollout.RUST_REQUIRED)) + if asynchronous + else runtime.run( + CONTEXT, binding=bound, native=lambda fn: fn(), python=python, rules=rules(Rollout.RUST_REQUIRED) + ) + ) + assert result is stream + assert get_hidden_params_dict(result) == {"additional_headers": {"x-litellm-rust": "true"}} + assert not execution.closed + delivered: Final = tuple([chunk async for chunk in result]) if isinstance(result, Stream) else tuple(result) + assert delivered == chunks + assert execution.closed + + def test_upstream_error_maps_to_api_error_without_fallback() -> None: calls: Final = recorder(RustUpstreamError(429, "rate limited")) diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index d5dc5437398..19043780eb6 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -5,6 +5,7 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import catalog from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -126,6 +127,7 @@ async def test_native_messages_stream_relays_provider_events_and_logs_success_on **arguments(messages_server, stream=True, callbacks=[recorder]) ) assert isinstance(stream, AsyncIterator) + assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} first: Final = await anext(stream) await drain_logging() assert "async_log_success_event" not in recorder.names @@ -169,6 +171,7 @@ def test_native_sync_messages_stream_relays_provider_events_and_logs_success_onc stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) assert isinstance(stream, Iterator) + assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} assert b"".join(stream) == sse_payload() assert_served_natively(messages_server)