fix(sagemaker)!: parse every streamed record however SageMaker cuts the payload parts (#47495)

* fix(sagemaker): parse every streamed record however SageMaker cuts the payload parts

SageMaker cuts a container's response stream into PayloadParts at arbitrary
byte offsets, so one part can hold several SSE events or JSON lines, or only
part of one. The decoder treated each part as one record, so vLLM's [DONE]
was logged at ERROR on every stream and a split at an unlucky byte dropped
the rest of the answer.

Peel records off a running buffer with JSONDecoder.raw_decode, skipping SSE
framing and [DONE] between them.

* test(sagemaker): cover streamed records cut or merged across payload parts

* fix(sagemaker): decode UTF-8 across parts, skip SSE comment lines, and log container error records

A UTF-8 character split across two PayloadParts now decodes whole through one incremental decoder per stream.
SSE comment lines (TGI's keep-alive) and event, id and retry field lines are skipped as framing.
A container error record is logged at ERROR with its message and skipped, and any other record that is not a
chunk is logged and skipped, the same way on the sync and async paths.

* fix(sagemaker): raise a container's error record and ModelStreamError on both stream paths

A record with a non-null error now raises a 500 SagemakerError carrying the record after the chunks before it, on sync and async. Async no longer swallows a non-200 event, so a ModelStreamError raises as it does on sync. A chunk with a null error field is delivered, and any other record that is not a chunk is logged and skipped on both paths.

* fix(sagemaker): answer a streamed error with the container's status and fail a stream cut mid-record

A container's error record now raises with the status the container gave it (vLLM's error.code, the older flat vLLM shape, TGI's http_status_code or error_type) instead of a fixed 500 that the router retried as a 503. A ModelStreamError raises a SagemakerError naming its ErrorCode and Message, a stream that ends inside a record raises instead of ending as a stop, and a skipped record is logged by its keys only

* fix(sagemaker): trust only 422 and 429 from a streamed error record and fail LMI generation errors

* test(sagemaker): drive an LMI failed-generation stream through the proxy on a sagemaker/ deployment

* test(sagemaker): fall back to the next model when a stream fails before any content

litellm/llms/sagemaker/common_utils.py

@@ -1,13 +1,15 @@
+import codecs
 import functools
 import json
+import re
 from collections.abc import AsyncIterator, Iterator, Mapping
-from typing import Final
+from types import MappingProxyType
+from typing import Final, Protocol
 
 import httpx
 from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, ValidationError, with_config
 from typing_extensions import ReadOnly, Required, TypedDict
 
-import litellm
 from litellm import verbose_logger
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.types.utils import GenericStreamingChunk as GChunk
@@ -40,12 +42,90 @@ class _SageMakerMessagesChunkDataRoot(RootModel[_SageMakerMessagesChunkData]):
     model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
 
 
-def _parse_sagemaker_chunk_data(data: str) -> _SageMakerChunkData:
-    return _SageMakerChunkDataRoot.model_validate(json.loads(data)).root
+_STREAM_FRAMING: Final = re.compile(r"(?:\s|data:|\[DONE\]|(?::|event:|id:|retry:)[^\n]*\n)*")
+_JSON_DECODER: Final = json.JSONDecoder()
+_RAW_DECODED: Final = TypeAdapter(tuple[object, int])
+_NO_RECORD: Final = object()
+_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
+_CLIENT_ERROR_TYPE_STATUS: Final = MappingProxyType(
+    {
+        "overloaded": 429,
+        "validation": 422,
+        "template_error": 422,
+        "missing_template_variable": 422,
+        "tool_error": 422,
+    }
+)
+_CLIENT_ERROR_STATUSES: Final = frozenset(_CLIENT_ERROR_TYPE_STATUS.values())
 
 
-def _parse_sagemaker_messages_chunk_data(data: str) -> _SageMakerMessagesChunkData:
-    return _SageMakerMessagesChunkDataRoot.model_validate(json.loads(data)).root
+def _framing_end(text: str) -> int:
+    framing: Final = _STREAM_FRAMING.match(text)
+    return framing.end() if framing else 0
+
+
+class _StreamRecordBuffer:
+    """
+    SageMaker cuts a container's stream into PayloadParts at arbitrary points, so one part can hold
+    several SSE events or JSON lines, or only part of one, down to half of a UTF-8 character
+    """
+
+    def __init__(self) -> None:
+        self._decoder: Final = codecs.getincrementaldecoder("utf-8")(errors="replace")
+        self._pending: str = ""
+
+    def feed(self, data: bytes) -> tuple[object, ...]:
+        self._pending += self._decoder.decode(data)
+        return tuple(iter(self._pop_record, _NO_RECORD))
+
+    def close(self) -> "SagemakerError | str":
+        unparsed: Final = self._pending + self._decoder.decode(b"", final=True)
+        record_start: Final = unparsed[_framing_end(unparsed) :]
+        if record_start.startswith(("{", "[")):
+            return SagemakerError(
+                status_code=500,
+                message=f"SageMaker stream ended inside a record, {len(record_start)} characters into it",
+            )
+        return unparsed
+
+    def _pop_record(self) -> object:
+        start: Final = _framing_end(self._pending)
+        try:
+            record, end = _RAW_DECODED.validate_python(_JSON_DECODER.raw_decode(self._pending, start))
+        except json.JSONDecodeError:
+            self._pending = self._pending[start:]
+            return _NO_RECORD
+        self._pending = self._pending[end:]
+        return record
+
+
+def _fields(value: object) -> Mapping[str, object]:
+    return STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else _NO_FIELDS
+
+
+def _client_error_status(value: object) -> int | None:
+    return value if isinstance(value, int) and value in _CLIENT_ERROR_STATUSES else None
+
+
+def _streamed_error_status(record: Mapping[str, object]) -> int:
+    detail: Final = _fields(record.get("error"))
+    candidates: Final = (
+        detail.get("code"),
+        detail.get("http_status_code"),
+        record.get("code"),
+        _CLIENT_ERROR_TYPE_STATUS.get(str(record.get("error_type"))),
+    )
+    return next(filter(None, map(_client_error_status, candidates)), 500)
+
+
+def _failed_generation(record: Mapping[str, object]) -> "SagemakerError | None":
+    error_message: Final = _fields(record.get("token")).get("error_msg")
+    if not error_message and _fields(record.get("details")).get("finish_reason") != "error":
+        return None
+    return SagemakerError(
+        status_code=500,
+        message=f"SageMaker endpoint streamed a failed generation: {error_message or 'no error message'}",
+    )
 
 
 def _load_sagemaker_response_stream_shape():
@@ -110,6 +190,10 @@ def client_error_details(error: Exception) -> tuple[int, str]:
     return response.response_metadata.http_status_code, response.error.message or str(error)
 
 
+class _EventStreamMessage(Protocol):
+    def to_response_dict(self) -> Mapping[str, object]: ...
+
+
 class SagemakerError(BaseLLMException):
     def __init__(
         self,
@@ -120,6 +204,24 @@ class SagemakerError(BaseLLMException):
         super().__init__(status_code=status_code, message=message, headers=headers)
 
 
+class _StreamExceptionBody(BaseModel):
+    model_config = ConfigDict(extra="ignore", frozen=True, hide_input_in_errors=True)
+    message: str = Field(default="", alias="Message")
+    error_code: str | None = Field(default=None, alias="ErrorCode")
+
+
+def _stream_exception_error(headers: Mapping[str, object], body: object) -> SagemakerError:
+    name: Final = headers.get(":exception-type") or headers.get(":error-code") or "an error"
+    try:
+        detail: Final = _StreamExceptionBody.model_validate_json(body if isinstance(body, bytes) else b"")
+    except ValidationError:
+        return SagemakerError(
+            status_code=500, message=f"SageMaker stream failed with {name}: {headers.get(':error-message', '')}"
+        )
+    code: Final = f" ({detail.error_code})" if detail.error_code else ""
+    return SagemakerError(status_code=500, message=f"SageMaker stream failed with {name}{code}: {detail.message}")
+
+
 class AWSEventStreamDecoder:
     def __init__(self, model: str, is_messages_api: bool | None = None) -> None:
         from botocore.parsers import EventStreamJSONParser
@@ -159,52 +261,55 @@ class AWSEventStreamDecoder:
             usage=None,
         )
 
+    def _record_to_chunk(self, record: object) -> GChunk | StreamingChatCompletionChunk | SagemakerError | None:
+        fields: Final = _fields(record)
+        if fields.get("error") is not None or fields.get("object") == "error":
+            return SagemakerError(
+                status_code=_streamed_error_status(fields),
+                message=f"SageMaker endpoint streamed an error: {json.dumps(record)}",
+            )
+        if failed_generation := _failed_generation(fields):
+            return failed_generation
+        try:
+            if self.is_messages_api:
+                return self._chunk_parser_messages_api(
+                    chunk_data=_SageMakerMessagesChunkDataRoot.model_validate(record).root
+                )
+            return self._chunk_parser(chunk_data=_SageMakerChunkDataRoot.model_validate(record).root)
+        except ValidationError:
+            verbose_logger.error(
+                "Skipping a SageMaker stream record that is not a chunk: %s with keys %s",
+                type(record).__name__,
+                sorted(fields),
+            )
+            return None
+
+    def _chunks(self, records: tuple[object, ...]) -> Iterator[GChunk | StreamingChatCompletionChunk]:
+        for outcome in map(self._record_to_chunk, records):
+            if isinstance(outcome, SagemakerError):
+                raise outcome
+            if outcome is not None:
+                yield outcome
+
     def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[GChunk | StreamingChatCompletionChunk | None]:
         """Given an iterator that yields lines, iterate over it & yield every event encountered"""
         from botocore.eventstream import EventStreamBuffer
 
         event_stream_buffer: Final = EventStreamBuffer()
-        accumulated_json = ""
+        records: Final = _StreamRecordBuffer()
 
         for chunk in iterator:
             event_stream_buffer.add_data(chunk)
             for event in event_stream_buffer:
-                message = self._parse_message_from_event(event)
-                if message:
-                    # remove data: prefix and "\n\n" at the end
-                    message = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(message) or ""
-                    message = message.replace("\n\n", "")
+                if payload := self._parse_message_from_event(event):
+                    yield from self._chunks(records.feed(payload))
 
-                    # Accumulate JSON data
-                    accumulated_json += message
-
-                    # Try to parse the accumulated JSON
-                    try:
-                        if self.is_messages_api:
-                            yield self._chunk_parser_messages_api(
-                                chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
-                            )
-                        else:
-                            yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
-                        # Reset accumulated_json after successful parsing
-                        accumulated_json = ""
-                    except json.JSONDecodeError:
-                        # If it's not valid JSON yet, continue to the next event
-                        continue
-
-        # Handle any remaining data after the iterator is exhausted
-        if accumulated_json:
-            try:
-                if self.is_messages_api:
-                    yield self._chunk_parser_messages_api(
-                        chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
-                    )
-                else:
-                    yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
-            except json.JSONDecodeError:
-                # Handle or log any unparseable data at the end
-                verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json)
-                yield None
+        unparsed: Final = records.close()
+        if isinstance(unparsed, SagemakerError):
+            raise unparsed
+        if unparsed:
+            verbose_logger.error("Warning: Unparseable JSON data remained, %d characters of it", len(unparsed))
+            yield None
 
     async def aiter_bytes(
         self, iterator: AsyncIterator[bytes]
@@ -213,58 +318,23 @@ class AWSEventStreamDecoder:
         from botocore.eventstream import EventStreamBuffer
 
         event_stream_buffer: Final = EventStreamBuffer()
-        accumulated_json = ""
+        records: Final = _StreamRecordBuffer()
 
         async for chunk in iterator:
             event_stream_buffer.add_data(chunk)
             for event in event_stream_buffer:
-                try:
-                    message = self._parse_message_from_event(event)
-                    if message:
-                        verbose_logger.debug("sagemaker  parsed chunk bytes %s", message)
-                        # remove data: prefix and "\n\n" at the end
-                        message = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(message) or ""
-                        message = message.replace("\n\n", "")
+                if payload := self._parse_message_from_event(event):
+                    for parsed_chunk in self._chunks(records.feed(payload)):
+                        yield parsed_chunk
 
-                        # Accumulate JSON data
-                        accumulated_json += message
+        unparsed: Final = records.close()
+        if isinstance(unparsed, SagemakerError):
+            raise unparsed
+        if unparsed:
+            verbose_logger.error("Warning: Unparseable JSON data remained, %d characters of it", len(unparsed))
+            yield None
 
-                        # Try to parse the accumulated JSON
-                        if self.is_messages_api:
-                            yield self._chunk_parser_messages_api(
-                                chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
-                            )
-                        else:
-                            yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
-                        # Reset accumulated_json after successful parsing
-                        accumulated_json = ""
-                except json.JSONDecodeError:
-                    # If it's not valid JSON yet, continue to the next event
-                    continue
-                except UnicodeDecodeError as e:
-                    verbose_logger.warning("UnicodeDecodeError: %s. Attempting to combine with next event.", e)
-                    continue
-                except Exception as e:
-                    verbose_logger.error("Error parsing message: %s. Attempting to combine with next event.", e)
-                    continue
-
-        # Handle any remaining data after the iterator is exhausted
-        if accumulated_json:
-            try:
-                if self.is_messages_api:
-                    yield self._chunk_parser_messages_api(
-                        chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
-                    )
-                else:
-                    yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
-            except json.JSONDecodeError:
-                # Handle or log any unparseable data at the end
-                verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json)
-                yield None
-            except Exception as e:
-                verbose_logger.error("Final error parsing accumulated JSON: %s", e)
-
-    def _parse_message_from_event(self, event) -> str | None:
+    def _parse_message_from_event(self, event: _EventStreamMessage) -> bytes | None:
         response_stream_shape: Final = get_sagemaker_response_stream_shape()
         if response_stream_shape is None:
             raise SagemakerError(
@@ -276,18 +346,14 @@ class AWSEventStreamDecoder:
             )
         response_dict: Final = event.to_response_dict()
         parsed_response: Final = self.parser.parse(response_dict, response_stream_shape)
+        body: Final = response_dict.get("body")
 
         if response_dict["status_code"] != 200:
-            raise ValueError(f"Bad response code, expected 200: {response_dict}")
+            raise _stream_exception_error(_fields(response_dict.get("headers")), body)
 
         if "chunk" in parsed_response:
             chunk = parsed_response.get("chunk")
             if not chunk:
                 return None
-            return chunk.get("bytes").decode()
-        else:
-            chunk = response_dict.get("body")
-            if not chunk:
-                return None
-
-            return chunk.decode()
+            return chunk.get("bytes")
+        return body if isinstance(body, bytes) and body else None

tests/integration/providers/_sagemaker_chat_container.py

@@ -16,6 +16,7 @@ SECRET_KEY: Final = "synthetic-secret-key-for-testing"
 REGION: Final = "us-east-1"
 PROMPT: Final = "synthetic inference component request"
 CONTENT: Final = ("sagemaker ", "wire control")
+USAGE: Final[Mapping[str, int]] = MappingProxyType({"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15})
 NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}})
 JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
 _MARKER_PREFIX: Final = "marker-"
@@ -69,27 +70,61 @@ def _event_header(name: str, value: str) -> bytes:
     return struct.pack("!B", len(name_bytes)) + name_bytes + b"\x07" + struct.pack("!H", len(value_bytes)) + value_bytes
 
 
+def _frame(headers: bytes, body: bytes) -> bytes:
+    prelude: Final = struct.pack("!II", 12 + len(headers) + len(body) + 4, len(headers))
+    message: Final = prelude + struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + headers + body
+    return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF)
+
+
 def _payload_part(body: bytes) -> bytes:
     headers: Final = (
         _event_header(":event-type", "PayloadPart")
         + _event_header(":content-type", "application/octet-stream")
         + _event_header(":message-type", "event")
     )
-    prelude: Final = struct.pack("!II", 12 + len(headers) + len(body) + 4, len(headers))
-    message: Final = prelude + struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + headers + body
-    return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF)
+    return _frame(headers, body)
+
+
+def model_stream_error(message: str, error_code: str) -> bytes:
+    headers: Final = (
+        _event_header(":exception-type", "ModelStreamError")
+        + _event_header(":content-type", "application/json")
+        + _event_header(":message-type", "exception")
+    )
+    return _frame(headers, json.dumps({"Message": message, "ErrorCode": error_code}).encode())
 
 
 def container_reply(identity: str, stream: bool) -> Reply:
     if not stream:
         return Reply(body=_completion(identity))
-    chunks: Final = (
+    return event_stream_reply(
+        (
+            _chunk(identity, {"role": "assistant", "content": CONTENT[0]}, None),
+            _chunk(identity, {"content": CONTENT[1]}, None),
+            _chunk(identity, {}, "stop", usage=dict(USAGE)),
+        )
+    )
+
+
+def event_stream_reply(payloads: tuple[bytes, ...], *, then: tuple[bytes, ...] = ()) -> Reply:
+    return Reply(chunks=(*map(_payload_part, payloads), *then), content_type="application/vnd.amazon.eventstream")
+
+
+def vllm_stream_events(identity: str) -> tuple[bytes, ...]:
+    usage_event: Final = {
+        "id": identity,
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": SERVED_MODEL,
+        "choices": [],
+        "usage": dict(USAGE),
+    }
+    return (
         _chunk(identity, {"role": "assistant", "content": CONTENT[0]}, None),
         _chunk(identity, {"content": CONTENT[1]}, None),
-        _chunk(identity, {}, "stop", usage={"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}),
-    )
-    return Reply(
-        chunks=tuple(_payload_part(chunk) for chunk in chunks), content_type="application/vnd.amazon.eventstream"
+        _chunk(identity, {}, "stop"),
+        f"data: {json.dumps(usage_event, separators=(',', ':'))}\n\n".encode(),
+        b"data: [DONE]\n\n",
     )
 
 

tests/integration/providers/test_sagemaker_chat_wire.py

@@ -1,11 +1,14 @@
 import asyncio
+import itertools
+import json
 import uuid
-from collections.abc import Callable, Mapping
+from collections.abc import AsyncIterator, Callable, Mapping
 from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final, Literal, TypeAlias
 
 import httpx
+import openai
 import pytest
 from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
 from integration._support.database import read_rows
@@ -26,14 +29,18 @@ from integration.providers._sagemaker_chat_container import (
     REGION,
     SECRET_KEY,
     SERVED_MODEL,
+    USAGE,
     answered_choices,
     assert_signed,
     by_marker,
     container_body,
     container_reply,
+    event_stream_reply,
     marked_prompt,
     marker,
+    model_stream_error,
     only_request,
+    vllm_stream_events,
 )
 from pydantic import JsonValue
 
@@ -260,6 +267,77 @@ def test_sagemaker_chat_keeps_the_deployment_aws_credentials_out_of_the_containe
     assert SECRET_KEY not in string_value(row["stored"]), row
 
 
+def _seven_byte_parts(stream: bytes) -> tuple[bytes, ...]:
+    return tuple(stream[start : start + 7] for start in range(0, len(stream), 7))
+
+
+def _event_pairs(events: tuple[bytes, ...]) -> tuple[bytes, ...]:
+    return tuple(b"".join(events[start : start + 2]) for start in range(0, len(events), 2))
+
+
+_PART_CUTTINGS: Final[Mapping[str, Callable[[tuple[bytes, ...]], tuple[bytes, ...]]]] = MappingProxyType(
+    {
+        "one-event-per-part": lambda events: events,
+        "usage-and-done-share-a-part": lambda events: (*events[:-2], events[-2] + events[-1]),
+        "event-pairs-share-a-part": _event_pairs,
+        "seven-byte-parts": lambda events: _seven_byte_parts(b"".join(events)),
+        "keep-alive-comments-between-events": lambda events: tuple(
+            itertools.chain.from_iterable((b":\n\n", event) for event in events)
+        ),
+    }
+)
+
+
+def _sse(record: Mapping[str, JsonValue]) -> bytes:
+    return f"data: {json.dumps(record)}\n\n".encode()
+
+
+@dataclass(frozen=True, slots=True)
+class _StreamedAnswer:
+    call_id: str
+    text: str
+    usages: tuple[tuple[int, int, int], ...]
+
+
+async def _chat_openai_stream_with_usage(gateway: Gateway, model: str, prompt: str) -> _StreamedAnswer:
+    async with async_openai_client(gateway) as client:
+        stream: Final = await client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=16,
+            stream=True,
+            stream_options={"include_usage": True},
+            extra_body=dict(NO_CACHE),
+        )
+        chunks: Final = [chunk async for chunk in stream]
+    return _StreamedAnswer(
+        stream.response.headers["x-litellm-call-id"],
+        "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices),
+        tuple(
+            (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens)
+            for chunk in chunks
+            if chunk.usage
+        ),
+    )
+
+
+@pytest.mark.parametrize("cutting", tuple(_PART_CUTTINGS))
+def test_sagemaker_chat_vllm_stream_reaches_the_caller_whole_however_sagemaker_cuts_its_payload_parts(
+    gateway: Gateway, cutting: str
+) -> None:
+    identity: Final = f"sagemaker-chat-{uuid.uuid4().hex}"
+    parts: Final = _PART_CUTTINGS[cutting](vllm_stream_events(identity))
+    expected_usage: Final = (USAGE["prompt_tokens"], USAGE["completion_tokens"], USAGE["total_tokens"])
+
+    with wire_server(lambda _request: event_stream_reply(parts)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, {})
+        answer: Final = asyncio.run(_chat_openai_stream_with_usage(gateway, model, marked_prompt()))
+        only_request(wire)
+    assert (answer.text, answer.usages) == ("".join(CONTENT), (expected_usage,)), answer
+    stored: Final = JSON_OBJECT.validate_json(string_value(_spend_row(answer.call_id)["stored"]))
+    assert (stored["prompt_tokens"], stored["completion_tokens"], stored["total_tokens"]) == expected_usage, stored
+
+
 @dataclass(frozen=True, slots=True)
 class _Shape:
     provider: str
@@ -388,6 +466,148 @@ def test_sagemaker_chat_container_error_reaches_the_caller_after_one_attempt_wit
     assert_signed(request, ACCESS_KEY, REGION, None)
 
 
+def test_sagemaker_chat_container_validation_error_record_answers_422_after_one_attempt(gateway: Gateway) -> None:
+    failure: Final = _sse({"error": {"message": "synthetic prompt too long", "http_status_code": 422}})
+
+    with wire_server(lambda _request: event_stream_reply((failure,))) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, {})
+        response: Final = _chat(gateway, model, marked_prompt(), {"stream": True})
+        only_request(wire)
+    assert response.status_code == 422, response.text
+    assert "synthetic prompt too long" in response.text, response.text
+
+
+async def _texts_until_error(
+    stream: openai.AsyncStream[openai.types.chat.ChatCompletionChunk],
+) -> AsyncIterator[str | openai.APIError]:
+    try:
+        async for chunk in stream:
+            yield "".join(choice.delta.content or "" for choice in chunk.choices)
+    except openai.APIError as error:
+        yield error
+
+
+async def _chat_openai_stream_until_error(
+    gateway: Gateway, model: str, prompt: str, extra: Mapping[str, JsonValue] = MappingProxyType({})
+) -> tuple[str | openai.APIError, ...]:
+    async with async_openai_client(gateway) as client:
+        stream: Final = await client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=16,
+            stream=True,
+            extra_body={**NO_CACHE, **extra},
+        )
+        return tuple([item async for item in _texts_until_error(stream)])
+
+
+_MID_STREAM_FAILURES: Final[Mapping[str, tuple[Callable[[tuple[bytes, ...]], Reply], str]]] = MappingProxyType(
+    {
+        "error-record": (
+            lambda events: event_stream_reply(
+                (*events[:2], _sse({"error": {"message": "synthetic decode failure", "code": 500}}))
+            ),
+            "synthetic decode failure",
+        ),
+        "model-stream-error": (
+            lambda events: event_stream_reply(
+                events[:2], then=(model_stream_error("synthetic stream timeout", "ModelInvocationTimeExceeded"),)
+            ),
+            "SageMaker stream failed with ModelStreamError (ModelInvocationTimeExceeded): synthetic stream timeout",
+        ),
+        "stream-ends-inside-a-record": (
+            lambda events: event_stream_reply((*events[:2], events[2][:30])),
+            "SageMaker stream ended inside a record",
+        ),
+    }
+)
+
+
+@pytest.mark.parametrize("failure", tuple(_MID_STREAM_FAILURES))
+def test_sagemaker_chat_stream_failure_after_the_first_chunks_reaches_the_caller_as_an_error(
+    gateway: Gateway, failure: str
+) -> None:
+    reply, expected = _MID_STREAM_FAILURES[failure]
+    events: Final = vllm_stream_events(f"sagemaker-chat-{uuid.uuid4().hex}")
+
+    with wire_server(lambda _request: reply(events)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, {})
+        outcome: Final = asyncio.run(_chat_openai_stream_until_error(gateway, model, marked_prompt()))
+        only_request(wire)
+    errors: Final = tuple(item for item in outcome if isinstance(item, openai.APIError))
+    assert "".join(item for item in outcome if isinstance(item, str)) == "".join(CONTENT), outcome
+    assert len(errors) == 1 and expected in errors[0].message, outcome
+
+
+_PRE_CONTENT_FAILURES: Final[Mapping[str, Reply]] = MappingProxyType(
+    {
+        "error-record": event_stream_reply((_sse({"error": {"message": "synthetic decode failure", "code": 500}}),)),
+        "model-stream-error": event_stream_reply(
+            (), then=(model_stream_error("synthetic stream timeout", "ModelInvocationTimeExceeded"),)
+        ),
+    }
+)
+
+
+@pytest.mark.parametrize("failure", tuple(_PRE_CONTENT_FAILURES))
+def test_sagemaker_chat_stream_failure_before_any_content_falls_back_to_the_next_model(
+    gateway: Gateway, failure: str
+) -> None:
+    events: Final = vllm_stream_events(f"sagemaker-chat-{uuid.uuid4().hex}")
+
+    with (
+        wire_server(lambda _request: _PRE_CONTENT_FAILURES[failure]) as primary_wire,
+        wire_server(lambda _request: event_stream_reply(events)) as fallback_wire,
+        gateway.scenario() as scenario,
+    ):
+        primary: Final = _deployment(scenario, primary_wire, {})
+        fallback: Final = _deployment(scenario, fallback_wire, {})
+        outcome: Final = asyncio.run(
+            _chat_openai_stream_until_error(gateway, primary, marked_prompt(), {"fallbacks": [fallback]})
+        )
+        primary_attempts: Final = primary_wire.drain()
+        only_request(fallback_wire)
+    assert outcome and all(isinstance(item, str) for item in outcome), outcome
+    assert "".join(item for item in outcome if isinstance(item, str)) == "".join(CONTENT), outcome
+    assert len(primary_attempts) == 1, primary_attempts
+
+
+def _lmi_line(record: Mapping[str, JsonValue]) -> bytes:
+    return f"{json.dumps(dict(record), ensure_ascii=False)}\n".encode()
+
+
+_LMI_TOKENS: Final = ("sagemaker", " lmi tokens")
+_LMI_ERROR: Final = "exception occurred during rolling batch inference"
+# djl-serving 8a99e6e5: rolling_batch.py's error Token rendered by output_formatter._jsonlines_output_formatter
+_LMI_FAILED_GENERATION: Final[Mapping[str, JsonValue]] = MappingProxyType(
+    {
+        "token": {"id": -1, "text": "", "log_prob": -1, "special_token": True, "error_msg": _LMI_ERROR},
+        "generated_text": "".join(_LMI_TOKENS),
+        "details": {"finish_reason": "error"},
+    }
+)
+
+
+def test_sagemaker_lmi_failed_generation_after_the_first_tokens_reaches_the_caller_as_an_error(
+    gateway: Gateway,
+) -> None:
+    tokens: Final = tuple(
+        _lmi_line({"token": {"id": index, "text": text, "log_prob": -0.1}}) for index, text in enumerate(_LMI_TOKENS, 1)
+    )
+    reply: Final = event_stream_reply((*tokens, _lmi_line(_LMI_FAILED_GENERATION)))
+
+    with wire_server(lambda _request: reply) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(
+            scenario, wire, {"hf_model_name": "meta-llama/Llama-2-7b-chat-hf"}, provider="sagemaker"
+        )
+        outcome: Final = asyncio.run(_chat_openai_stream_until_error(gateway, model, marked_prompt()))
+        only_request(wire)
+    errors: Final = tuple(item for item in outcome if isinstance(item, openai.APIError))
+    assert "".join(item for item in outcome if isinstance(item, str)) == "".join(_LMI_TOKENS), outcome
+    assert len(errors) == 1, outcome
+    assert f"SageMaker endpoint streamed a failed generation: {_LMI_ERROR}" in errors[0].message, outcome
+
+
 @dataclass(frozen=True, slots=True)
 class _Call:
     endpoint: Endpoint

tests/unit/llms/sagemaker/test_sagemaker_common_utils.py

@@ -1,4 +1,9 @@
 import json
+import logging
+import struct
+import zlib
+from collections.abc import AsyncIterator, Mapping
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -6,8 +11,10 @@ import pytest
 
 from botocore.exceptions import ClientError, ConnectionClosedError
 
-from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder, client_error_details
+from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder, SagemakerError, client_error_details
 from litellm.llms.sagemaker.completion.transformation import SagemakerConfig
+from litellm.types.utils import GenericStreamingChunk as GChunk
+from litellm.types.utils import StreamingChatCompletionChunk
 
 
 # --------------------------------------------------------------------------- #
@@ -232,6 +239,423 @@ async def test_aiter_bytes_valid_chunk_followed_by_unicode_error():
         assert chunks[0]["text"] == "hello"  # Verify the content of the valid chunk
 
 
+def _event_header(name: str, value: str) -> bytes:
+    return struct.pack("!B", len(name)) + name.encode() + b"\x07" + struct.pack("!H", len(value)) + value.encode()
+
+
+_PAYLOAD_PART_HEADERS: Final = (
+    _event_header(":event-type", "PayloadPart")
+    + _event_header(":content-type", "application/octet-stream")
+    + _event_header(":message-type", "event")
+)
+_MODEL_STREAM_ERROR_HEADERS: Final = (
+    _event_header(":exception-type", "ModelStreamError")
+    + _event_header(":content-type", "application/json")
+    + _event_header(":message-type", "exception")
+)
+
+
+def _frame(headers: bytes, body: str | bytes) -> bytes:
+    encoded: Final = body.encode() if isinstance(body, str) else body
+    prelude: Final = struct.pack("!II", 16 + len(headers) + len(encoded), len(headers))
+    message: Final = prelude + struct.pack("!I", zlib.crc32(prelude)) + headers + encoded
+    return message + struct.pack("!I", zlib.crc32(message))
+
+
+def _payload_part(body: str | bytes) -> bytes:
+    return _frame(_PAYLOAD_PART_HEADERS, body)
+
+
+def _cut(stream: str, size: int) -> tuple[bytes, ...]:
+    encoded: Final = stream.encode()
+    return tuple(encoded[start : start + size] for start in range(0, len(encoded), size))
+
+
+def _sse(record: Mapping[str, object]) -> str:
+    return f"data: {json.dumps(record, ensure_ascii=False)}\n\n"
+
+
+def _vllm_chunk(choices: list[Mapping[str, object]], **extra: object) -> str:
+    return _sse(
+        {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": choices, **extra}
+    )
+
+
+_VLLM_EVENTS: Final = (
+    _vllm_chunk([{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}]),
+    _vllm_chunk([{"index": 0, "delta": {"content": "Hél"}, "finish_reason": None}]),
+    _vllm_chunk([{"index": 0, "delta": {"content": "lo 👋"}, "finish_reason": "stop"}]),
+    _vllm_chunk([], usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}),
+    "data: [DONE]\n\n",
+)
+_VLLM_STREAM: Final = "".join(_VLLM_EVENTS)
+_TGI_EVENTS: Final = tuple(
+    f"data:{json.dumps({'index': 0, 'token': {'text': text}}, ensure_ascii=False)}\n\n" for text in ("Hél", "lo 👋")
+)
+_TGI_STREAM: Final = "".join(_TGI_EVENTS)
+_TGI_KEEP_ALIVE_STREAM: Final = (
+    f"{_TGI_EVENTS[0]}:\n\n: keep-alive\n\nevent: message\nid: 2\nretry: 1000\n{_TGI_EVENTS[1]}"
+)
+_LMI_STREAM: Final = "".join(
+    f"{json.dumps({'token': {'text': text}}, ensure_ascii=False)}\n" for text in ("Hél", "lo 👋")
+)
+
+
+async def _stream_frames(
+    frames: tuple[bytes, ...], is_messages_api: bool, mode: str
+) -> AsyncIterator[GChunk | StreamingChatCompletionChunk | None]:
+    decoder: Final = AWSEventStreamDecoder(model="m", is_messages_api=is_messages_api)
+    if mode == "sync":
+        for chunk in decoder.iter_bytes(iter(frames)):
+            yield chunk
+        return
+
+    async def stream():
+        for frame in frames:
+            yield frame
+
+    async for chunk in decoder.aiter_bytes(stream()):
+        yield chunk
+
+
+async def _decode(
+    parts: tuple[str | bytes, ...], is_messages_api: bool, mode: str
+) -> list[GChunk | StreamingChatCompletionChunk | None]:
+    return [chunk async for chunk in _stream_frames(tuple(map(_payload_part, parts)), is_messages_api, mode)]
+
+
+def _text(chunk: GChunk | StreamingChatCompletionChunk | None, is_messages_api: bool) -> str | None:
+    if chunk is None:
+        return None
+    return chunk.choices[0].delta.content if is_messages_api else chunk["text"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+@pytest.mark.parametrize(
+    "parts",
+    [
+        pytest.param(_VLLM_EVENTS, id="one-event-per-part"),
+        pytest.param((*_VLLM_EVENTS[:3], _VLLM_EVENTS[3] + _VLLM_EVENTS[4]), id="usage-and-done-share-a-part"),
+        pytest.param(
+            (_VLLM_EVENTS[0] + _VLLM_EVENTS[1], _VLLM_EVENTS[2] + _VLLM_EVENTS[3], _VLLM_EVENTS[4]),
+            id="event-pairs-share-a-part",
+        ),
+        pytest.param((*_VLLM_EVENTS[:4], "data: [DO", "NE]\n\n"), id="done-split-across-parts"),
+        pytest.param((_VLLM_STREAM,), id="whole-stream-in-one-part"),
+        pytest.param(_cut(_VLLM_STREAM, 7), id="seven-byte-parts"),
+        pytest.param(_cut(_VLLM_STREAM, 1), id="one-byte-parts-split-every-utf8-character"),
+    ],
+)
+async def test_messages_api_stream_keeps_every_chunk_however_sagemaker_cuts_the_parts(
+    parts: tuple[str | bytes, ...], mode: str, caplog: pytest.LogCaptureFixture
+):
+    with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+        chunks: Final = await _decode(parts, is_messages_api=True, mode=mode)
+
+    assert None not in chunks
+    assert "".join(chunk.choices[0].delta.content for chunk in chunks if chunk.choices) == "Héllo 👋"
+    assert [(chunk.usage.prompt_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage] == [(7, 10)]
+    assert caplog.records == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+@pytest.mark.parametrize(
+    "parts",
+    [
+        pytest.param(_TGI_EVENTS, id="tgi-sse-one-event-per-part"),
+        pytest.param((_TGI_STREAM,), id="tgi-sse-whole-stream-in-one-part"),
+        pytest.param(_cut(_TGI_STREAM, 5), id="tgi-sse-five-byte-parts"),
+        pytest.param(_cut(_TGI_STREAM, 1), id="tgi-sse-one-byte-parts-split-every-utf8-character"),
+        pytest.param((_TGI_KEEP_ALIVE_STREAM,), id="tgi-sse-comments-and-fields-whole-stream-in-one-part"),
+        pytest.param(_cut(_TGI_KEEP_ALIVE_STREAM, 1), id="tgi-sse-comments-and-fields-one-byte-parts"),
+        pytest.param((_LMI_STREAM,), id="json-lines-whole-stream-in-one-part"),
+        pytest.param(_cut(_LMI_STREAM, 5), id="json-lines-five-byte-parts"),
+        pytest.param(tuple(_LMI_STREAM.splitlines()), id="undelimited-objects-one-per-part"),
+    ],
+)
+async def test_completion_stream_keeps_every_token_however_sagemaker_cuts_the_parts(
+    parts: tuple[str | bytes, ...], mode: str, caplog: pytest.LogCaptureFixture
+):
+    with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+        chunks: Final = await _decode(parts, is_messages_api=False, mode=mode)
+
+    assert None not in chunks
+    assert "".join(chunk["text"] for chunk in chunks) == "Héllo 👋"
+    assert caplog.records == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+@pytest.mark.parametrize(
+    "last_part",
+    [
+        pytest.param(_VLLM_EVENTS[2][:20], id="cut-between-characters"),
+        pytest.param(
+            _VLLM_EVENTS[2].encode()[: _VLLM_EVENTS[2].encode().index("👋".encode()) + 2], id="cut-inside-a-character"
+        ),
+    ],
+)
+async def test_stream_that_ends_inside_a_record_raises_after_the_complete_ones(mode: str, last_part: str | bytes):
+    stream: Final = _stream_frames(tuple(map(_payload_part, (*_VLLM_EVENTS[:2], last_part))), True, mode)
+
+    assert [_text(await anext(stream), True) for _ in range(2)] == ["", "Hél"]
+    with pytest.raises(SagemakerError, match="ended inside a record") as exc_info:
+        await anext(stream)
+    assert exc_info.value.status_code == 500
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+async def test_text_left_over_that_is_not_a_record_is_logged_once(mode: str, caplog: pytest.LogCaptureFixture):
+    with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+        chunks: Final = await _decode((*_VLLM_EVENTS[:3], "data: upstream hiccup\n\n"), is_messages_api=True, mode=mode)
+
+    assert chunks[-1] is None
+    assert "".join(chunk.choices[0].delta.content for chunk in chunks[:-1]) == "Héllo 👋"
+    assert [record.levelno for record in caplog.records] == [logging.ERROR]
+    assert "hiccup" not in caplog.text
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+@pytest.mark.parametrize(
+    ("parts", "is_messages_api", "status_code"),
+    [
+        pytest.param(
+            (
+                _VLLM_EVENTS[1],
+                _sse(
+                    {
+                        "error": {
+                            "message": "CUDA out of memory while decoding",
+                            "type": "InternalServerError",
+                            "code": 500,
+                        }
+                    }
+                )
+                + "data: [DONE]\n\n",
+            ),
+            True,
+            500,
+            id="vllm-500-error-event-shares-a-part-with-done",
+        ),
+        pytest.param(
+            (
+                _VLLM_EVENTS[1],
+                _sse(
+                    {"error": {"message": "CUDA out of memory while decoding", "type": "BadRequestError", "code": 400}}
+                ),
+                "data: [DONE]\n\n",
+            ),
+            True,
+            500,
+            id="vllm-400-error-event-in-its-own-part",
+        ),
+        pytest.param(
+            (
+                _VLLM_EVENTS[1],
+                _sse(
+                    {
+                        "object": "error",
+                        "message": "CUDA out of memory while decoding",
+                        "type": "BadRequestError",
+                        "code": 400,
+                    }
+                ),
+            ),
+            True,
+            500,
+            id="legacy-vllm-flat-error-event",
+        ),
+        pytest.param(
+            (
+                _VLLM_EVENTS[1],
+                _sse({"error": {"message": "CUDA out of memory while decoding", "http_status_code": 422}}),
+            ),
+            True,
+            422,
+            id="tgi-messages-api-error-event",
+        ),
+        pytest.param(
+            (
+                _VLLM_EVENTS[1],
+                _sse({"error": {"message": "CUDA out of memory while decoding", "http_status_code": 424}}),
+            ),
+            True,
+            500,
+            id="tgi-messages-api-generation-error-event",
+        ),
+        pytest.param(
+            (
+                _VLLM_EVENTS[1],
+                _sse(
+                    {"error": {"message": "CUDA out of memory while decoding", "type": "RateLimitError", "code": 429}}
+                ),
+            ),
+            True,
+            429,
+            id="vllm-429-error-event",
+        ),
+        pytest.param(
+            (_VLLM_EVENTS[1], _sse({"error": {"message": "CUDA out of memory while decoding", "code": "oom"}})),
+            True,
+            500,
+            id="error-event-without-a-status",
+        ),
+        pytest.param(
+            (
+                _TGI_EVENTS[0]
+                + f"data:{json.dumps({'error': 'CUDA out of memory while decoding', 'error_type': 'generation'})}\n\n",
+            ),
+            False,
+            500,
+            id="tgi-generation-error-event-shares-a-part-with-a-token",
+        ),
+        pytest.param(
+            (
+                _TGI_EVENTS[0],
+                f"data:{json.dumps({'error': 'CUDA out of memory while decoding', 'error_type': 'overloaded'})}\n\n",
+            ),
+            False,
+            429,
+            id="tgi-overloaded-error-event",
+        ),
+        pytest.param(
+            (
+                _TGI_EVENTS[0],
+                f"data:{json.dumps({'error': 'CUDA out of memory while decoding', 'error_type': 'incomplete_generation'})}\n\n",
+            ),
+            False,
+            500,
+            id="tgi-incomplete-generation-error-event",
+        ),
+    ],
+)
+async def test_container_error_event_raises_with_its_message_and_status_after_the_chunks_before_it(
+    parts: tuple[str, ...], is_messages_api: bool, status_code: int, mode: str
+):
+    stream: Final = _stream_frames(tuple(map(_payload_part, parts)), is_messages_api, mode)
+
+    assert _text(await anext(stream), is_messages_api) == "Hél"
+    with pytest.raises(SagemakerError, match="CUDA out of memory while decoding") as exc_info:
+        await anext(stream)
+    assert exc_info.value.status_code == status_code
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+@pytest.mark.parametrize(
+    ("last_record", "message"),
+    [
+        pytest.param(
+            {
+                "token": {"id": 2, "text": "", "log_prob": 0.0, "error_msg": "CUDA out of memory while decoding"},
+                "generated_text": "Hél",
+                "details": {"finish_reason": "error"},
+            },
+            "SageMaker endpoint streamed a failed generation: CUDA out of memory while decoding",
+            id="token-error-message",
+        ),
+        pytest.param(
+            {
+                "token": {"id": 2, "text": "", "log_prob": 0.0},
+                "generated_text": "Hél",
+                "details": {"finish_reason": "error"},
+            },
+            "SageMaker endpoint streamed a failed generation: no error message",
+            id="error-finish-without-a-message",
+        ),
+    ],
+)
+async def test_lmi_failed_generation_raises_after_the_chunks_before_it(
+    last_record: Mapping[str, object], message: str, mode: str
+):
+    lines: Final = "".join(
+        f"{json.dumps(record, ensure_ascii=False)}\n"
+        for record in ({"token": {"id": 1, "text": "Hél", "log_prob": -0.1}}, last_record)
+    )
+    stream: Final = _stream_frames((_payload_part(lines),), False, mode)
+
+    assert _text(await anext(stream), False) == "Hél"
+    with pytest.raises(SagemakerError) as exc_info:
+        await anext(stream)
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.message == message
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+async def test_model_stream_error_raises_after_the_chunks_before_it(mode: str):
+    frames: Final = (
+        _payload_part(_VLLM_EVENTS[1]),
+        _frame(
+            _MODEL_STREAM_ERROR_HEADERS,
+            json.dumps({"Message": "container exited", "ErrorCode": "ModelInvocationTimeExceeded"}),
+        ),
+        _payload_part(_VLLM_EVENTS[2]),
+    )
+    stream: Final = _stream_frames(frames, True, mode)
+
+    assert _text(await anext(stream), True) == "Hél"
+    with pytest.raises(SagemakerError) as exc_info:
+        await anext(stream)
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.message == (
+        "SageMaker stream failed with ModelStreamError (ModelInvocationTimeExceeded): container exited"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+@pytest.mark.parametrize(
+    ("parts", "is_messages_api"),
+    [
+        pytest.param(
+            tuple(
+                _vllm_chunk([{"index": 0, "delta": {"content": text}, "finish_reason": None}], error=None)
+                for text in ("Hél", "lo 👋")
+            ),
+            True,
+            id="vllm-chunks-with-null-error",
+        ),
+        pytest.param(
+            tuple(
+                f"data:{json.dumps({'index': 0, 'token': {'text': text}, 'error': None}, ensure_ascii=False)}\n\n"
+                for text in ("Hél", "lo 👋")
+            ),
+            False,
+            id="tgi-tokens-with-null-error",
+        ),
+    ],
+)
+async def test_chunk_with_a_null_error_field_is_delivered(
+    parts: tuple[str, ...], is_messages_api: bool, mode: str, caplog: pytest.LogCaptureFixture
+):
+    with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+        chunks: Final = await _decode(parts, is_messages_api=is_messages_api, mode=mode)
+
+    assert "".join(_text(chunk, is_messages_api) or "" for chunk in chunks) == "Héllo 👋"
+    assert caplog.records == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sync", "async"])
+async def test_record_that_is_not_a_chunk_is_logged_and_skipped(mode: str, caplog: pytest.LogCaptureFixture):
+    with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+        chunks: Final = await _decode(
+            (_VLLM_EVENTS[1], _sse({"object": "keepalive", "detail": "still decoding"}), _VLLM_EVENTS[2]),
+            is_messages_api=True,
+            mode=mode,
+        )
+
+    assert "".join(chunk.choices[0].delta.content for chunk in chunks) == "Héllo 👋"
+    assert [record.levelno for record in caplog.records] == [logging.ERROR]
+    assert "['detail', 'object']" in caplog.records[0].getMessage()
+    assert "still decoding" not in caplog.records[0].getMessage()
+
+
 class TestSagemakerTransform:
     def setup_method(self):
         self.config = SagemakerConfig()