fix: decode SGLang streaming responses

src/inference_endpoint/endpoint_client/adapter_protocol.py

@@ -127,12 +127,11 @@ def parse_sse_chunk(cls, buffer: bytes, end_pos: int) -> list[Any]:
         """
         json_docs = cls.SSE_DATA_PATTERN.findall(buffer[:end_pos])
         parsed: list[Any] = []
-        # Note: if one frame is malformed, remaining frames are skipped
-        try:
-            for json_doc in json_docs:
+        for json_doc in json_docs:
+            try:
                 content = cls.decode_sse_message(json_doc)
                 if content is not None:
                     parsed.append(content)
-        except (msgspec.DecodeError, msgspec.ValidationError) as exc:
-            logger.warning("skipping malformed SSE batch (%s)", type(exc).__name__)
+            except (msgspec.DecodeError, msgspec.ValidationError) as exc:
+                logger.warning("skipping malformed SSE frame (%s)", type(exc).__name__)
         return parsed

src/inference_endpoint/sglang/adapter.py

@@ -122,7 +122,7 @@ def decode_sse_message(cls, json_bytes: bytes) -> SGLangSSEDelta:
         # the output_tokens is the delta from the previous chunk
         token_delta = resp.output_ids
 
-        has_retractions = resp.meta_info.total_retractions > 0
+        has_retractions = resp.meta_info.num_retractions > 0
         return SGLangSSEDelta(
             text=total_text,
             token_delta=token_delta,

src/inference_endpoint/sglang/types.py

@@ -52,13 +52,13 @@ class SGLangGenerateRequest(
 
 class MetaInfo(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False):  # type: ignore[call-arg]
     id: str
-    finish_reason: dict[str, Any]
+    finish_reason: dict[str, Any] | None
     prompt_tokens: int
     weight_version: str
-    total_retractions: int
+    num_retractions: int
     completion_tokens: int
     cached_tokens: int
-    e2e_latency: float
+    e2e_latency: float | None = None
 
 
 class SGLangGenerateResponse(

tests/performance/sglang/test_sglang_adapter.py

@@ -64,7 +64,7 @@ def make_response_bytes(text: str, n_tokens: int) -> bytes:
                 "finish_reason": {"type": "stop"},
                 "prompt_tokens": 10,
                 "weight_version": "v1",
-                "total_retractions": 0,
+                "num_retractions": 0,
                 "completion_tokens": n_tokens,
                 "cached_tokens": 0,
                 "e2e_latency": 0.1,

tests/unit/endpoint_client/test_adapter_protocol.py

@@ -22,10 +22,11 @@
 from inference_endpoint.config.schema import ModelParams
 from inference_endpoint.dataset_manager.transforms import Transform
 from inference_endpoint.endpoint_client.adapter_protocol import HttpRequestAdapter
+from inference_endpoint.sglang.adapter import SGLangGenerateAdapter
 
 
 class _SimpleAdapter(HttpRequestAdapter):
-    """Minimal adapter: raises DecodeError for b'BAD', returns {"x": 1} otherwise."""
+    """Decode JSON objects with integer values."""
 
     @classmethod
     def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]:
@@ -41,15 +42,47 @@ def decode_response(cls, response_bytes: bytes, query_id: str) -> Any:
 
     @classmethod
     def decode_sse_message(cls, json_bytes: bytes) -> Any:
-        if json_bytes.strip() == b"BAD":
-            raise msgspec.DecodeError("bad frame")
-        return {"x": 1}
+        return msgspec.json.decode(json_bytes, type=dict[str, int])
 
 
 @pytest.mark.unit
-def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid():
+@pytest.mark.parametrize("bad_frame", [b'{"ok":}', b'{"ok":"bad"}'])
+def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid(bad_frame):
     """A malformed SSE frame is skipped; surrounding valid frames are preserved."""
-    buffer = b'data: {"ok":1}\n\ndata: BAD\n\ndata: {"ok":1}\n\n'
+    buffer = b'data: {"ok":1}\n\ndata: ' + bad_frame + b'\n\ndata: {"ok":2}\n\n'
     end_pos = len(buffer)
     result = _SimpleAdapter.parse_sse_chunk(buffer, end_pos)
-    assert result == [{"x": 1}, {"x": 1}]
+    assert result == [{"ok": 1}, {"ok": 2}]
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("finish_reason", [None, {"type": "stop"}])
+@pytest.mark.parametrize("latency", [{}, {"e2e_latency": None}, {"e2e_latency": 0.1}])
+@pytest.mark.parametrize("retractions", [0, 2])
+def test_sglang_response_metadata(finish_reason, latency, retractions):
+    response = msgspec.json.encode(
+        {
+            "text": "Hello",
+            "output_ids": [10],
+            "meta_info": {
+                "id": "sample",
+                "finish_reason": finish_reason,
+                "prompt_tokens": 3,
+                "weight_version": "v1",
+                "num_retractions": retractions,
+                "completion_tokens": 1,
+                "cached_tokens": 0,
+                **latency,
+            },
+        }
+    )
+    result = SGLangGenerateAdapter.decode_response(response, "sample")
+    assert result.get_response_output_string() == "Hello"
+    assert result.metadata == {"token_ids": [10], "n_tokens": 1}
+    frame = b"data: " + response + b"\n\n"
+    chunks = SGLangGenerateAdapter.parse_sse_chunk(frame, len(frame))
+    assert len(chunks) == 1
+    assert chunks[0].text == "Hello"
+    assert chunks[0].token_delta == [10]
+    assert chunks[0].total_completion_tokens == 1
+    assert chunks[0].has_retractions is (retractions > 0)