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 <yujong@berri.ai>

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

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"))
 

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)