fix(sagemaker): forward native streaming events as they arrive to cut TTFT

Mirror the sagemaker_chat fix on the native sagemaker/ streaming path: the
sync and async completion handlers read the invocations-response-stream body
with iter_bytes(chunk_size=1024) / aiter_bytes(chunk_size=1024), so httpx
withholds bytes until 1024 accumulate and tokens arrive in gap-then-burst
waves. Drop the fixed chunk size so each decoded event is forwarded as its
bytes arrive.

Also add a boundary-agnostic decoder test proving frames reassemble correctly
regardless of where transport reads split the stream.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/sagemaker/completion/handler.py

@@ -216,7 +216,7 @@ class SagemakerLLM(BaseAWSLLM):
 
                 decoder = AWSEventStreamDecoder(model="")
 
-                completion_stream = decoder.iter_bytes(sync_response.iter_bytes(chunk_size=1024))
+                completion_stream = decoder.iter_bytes(sync_response.iter_bytes())
                 streaming_response = CustomStreamWrapper(
                     completion_stream=completion_stream,
                     model=model,
@@ -358,7 +358,7 @@ class SagemakerLLM(BaseAWSLLM):
                 raise SagemakerError(status_code=response.status_code, message=response.text)
 
             decoder = AWSEventStreamDecoder(model="")
-            completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024))
+            completion_stream = decoder.aiter_bytes(response.aiter_bytes())
 
             return completion_stream
 

tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py

@@ -207,3 +207,29 @@ def test_signed_body_includes_stream_flag():
     )
     assert signed_body is not None
     assert json.loads(signed_body)["stream"] is True
+
+
+@pytest.mark.parametrize("split_size", [1, 3, 7, 64, 4096])
+def test_decoder_reassembles_frames_across_arbitrary_byte_boundaries(split_size):
+    """Correctness must not depend on chunk boundaries falling on frame edges.
+
+    Removing `chunk_size=1024` lets httpx yield raw transport reads, so in
+    production a single read can straddle several frames or split one frame in
+    half. This re-chunks the concatenated stream at boundaries that deliberately
+    ignore frame edges and asserts every delta still decodes, in order, exactly
+    once - the guarantee botocore's EventStreamBuffer provides.
+    """
+    from litellm.llms.sagemaker.chat.transformation import AWSEventStreamDecoder
+
+    frames = _make_frames(24)
+    blob = b"".join(frames)
+    chunks = [blob[i : i + split_size] for i in range(0, len(blob), split_size)]
+
+    decoder = AWSEventStreamDecoder(model="phi-4", is_messages_api=True)
+    texts = [
+        _content_of(chunk)
+        for chunk in decoder.iter_bytes(iter(chunks))
+        if chunk is not None and _content_of(chunk) is not None
+    ]
+
+    assert texts == [f"token{i} " for i in range(len(frames))]

tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py

@@ -0,0 +1,107 @@
+"""
+Regression tests for LIT-4313: the native `sagemaker/` streaming path must
+forward each AWS event-stream frame as it arrives instead of buffering to a
+fixed 1024-byte threshold and then draining a burst of tokens.
+
+The buffering came from `response.aiter_bytes(chunk_size=1024)`: httpx's
+ByteChunker withholds bytes until `chunk_size` accumulates, so the first token
+could not be produced until enough later frames had arrived to cross 1024 bytes,
+inflating TTFT and turning a steady provider stream into gap-then-burst delivery.
+"""
+
+import binascii
+import json
+import struct
+from typing import AsyncIterator
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+from litellm.llms.sagemaker.completion.handler import SagemakerLLM
+
+
+def _encode_header(name: str, value: str) -> bytes:
+    name_b = name.encode("utf-8")
+    value_b = value.encode("utf-8")
+    return struct.pack("B", len(name_b)) + name_b + struct.pack("B", 7) + struct.pack(">H", len(value_b)) + value_b
+
+
+def _encode_event_frame(payload: bytes) -> bytes:
+    """Encode one AWS event-stream message that botocore's EventStreamBuffer decodes."""
+    headers = {
+        ":event-type": "PayloadPart",
+        ":content-type": "application/json",
+        ":message-type": "event",
+    }
+    headers_b = b"".join(_encode_header(k, v) for k, v in headers.items())
+    total_len = 16 + len(headers_b) + len(payload)
+    prelude = struct.pack(">I", total_len) + struct.pack(">I", len(headers_b))
+    prelude_crc = struct.pack(">I", binascii.crc32(prelude) & 0xFFFFFFFF)
+    message = prelude + prelude_crc + headers_b + payload
+    message_crc = struct.pack(">I", binascii.crc32(message) & 0xFFFFFFFF)
+    return message + message_crc
+
+
+def _token_frame(text: str) -> bytes:
+    # SageMaker HF TGI streaming payloads are `{"token": {"text": ...}}` blobs.
+    sse = "data: " + json.dumps({"token": {"text": text}}) + "\n\n"
+    return _encode_event_frame(sse.encode("utf-8"))
+
+
+def _make_frames(n: int) -> list[bytes]:
+    frames = [_token_frame(f"token{i} ") for i in range(n)]
+    assert all(len(f) < 1024 for f in frames)
+    return frames
+
+
+class _CountingAsyncStream(httpx.AsyncByteStream):
+    """Yields provider frames one at a time and records how many have been pulled."""
+
+    def __init__(self, frames: list[bytes]) -> None:
+        self._frames = frames
+        self.consumed = 0
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        for frame in self._frames:
+            self.consumed += 1
+            yield frame
+
+
+class _FakeAsyncClient:
+    def __init__(self, response: httpx.Response) -> None:
+        self._response = response
+
+    async def post(self, *args, **kwargs) -> httpx.Response:
+        return self._response
+
+
+@pytest.mark.asyncio
+async def test_async_native_streaming_forwards_each_frame_incrementally():
+    """Each token must be emitted after exactly one newly-pulled source frame.
+
+    With the old `chunk_size=1024` the httpx chunker would swallow several small
+    frames before yielding, so the first token would arrive only after `consumed`
+    had already crossed multiple frames, and tokens would then replay in a burst.
+    """
+    frames = _make_frames(24)
+    stream = _CountingAsyncStream(frames)
+    response = httpx.Response(200, stream=stream)
+
+    completion_stream = await SagemakerLLM().make_async_call(
+        api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream",
+        headers={},
+        data="",
+        logging_obj=MagicMock(),
+        client=_FakeAsyncClient(response),
+    )
+
+    consumed_at_token = []
+    texts = []
+    async for chunk in completion_stream:
+        if chunk is not None and chunk["text"]:
+            consumed_at_token.append(stream.consumed)
+            texts.append(chunk["text"])
+
+    assert texts == [f"token{i} " for i in range(len(frames))]
+    assert consumed_at_token == list(range(1, len(frames) + 1))