[feature] Add response-level input/output token ids to chat completions via SglExt (#34488)

test/registered/unit/entrypoints/openai/test_serving_chat.py

@@ -127,6 +127,9 @@ class _MockTokenizerManager:
             reasoning_parser=None,
             stream_response_default_include_usage=False,
             default_chat_template_kwargs=None,
+            return_input_ids=False,
+            return_output_ids=False,
+            incremental_streaming_output=False,
         )
         self.model_path = self.server_args.model_path
         # The manager tracks the served name itself; a weight update rewrites it.
@@ -2525,6 +2528,162 @@ class ServingChatTestCase(unittest.TestCase):
         self.assertEqual(len(chunks), 2)
         self.assertIn("error", chunks[0])
 
+    def test_streaming_abort_with_ids_enabled(self):
+        """Test that a terminal abort with input/output ids enabled yields only an error and [DONE]."""
+        err_msg = "Aborted by scheduler"
+        err_code = HTTPStatus.INTERNAL_SERVER_ERROR
+
+        async def _mock_generate_abort():
+            yield {
+                "text": "Partial ",
+                "prompt_token_ids": [4, 5, 6],
+                "output_ids": [1, 2, 3],
+                "meta_info": {
+                    "id": "chatcmpl-test",
+                    "prompt_tokens": 10,
+                    "completion_tokens": 2,
+                    "cached_tokens": 0,
+                    "finish_reason": {
+                        "type": "abort",
+                        "status_code": err_code,
+                        "message": err_msg,
+                    },
+                    "output_token_logprobs": None,
+                    "output_top_logprobs": None,
+                },
+                "index": 0,
+            }
+
+        self.tm.generate_request.return_value = _mock_generate_abort()
+
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+            temperature=0.7,
+            max_tokens=100,
+            stream=True,
+            return_input_ids_in_sglext=True,
+            return_output_ids_in_sglext=True,
+        )
+
+        with patch(
+            "sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
+        ) as conv_mock:
+            conv_ins = Mock()
+            conv_ins.get_prompt.return_value = "Test prompt"
+            conv_mock.return_value = conv_ins
+
+            adapted_request, _ = self.chat._convert_to_internal_request(
+                req, self.fastapi_request
+            )
+
+            async def run_stream():
+                chunks = []
+                try:
+                    async for chunk in self.chat._generate_chat_stream(
+                        adapted_request, req, self.fastapi_request
+                    ):
+                        chunks.append(chunk)
+                except Exception as e:
+                    print(f"Error during stream iteration: {e}")
+                return chunks
+
+        loop = get_or_create_event_loop()
+        chunks = loop.run_until_complete(run_stream())
+
+        # Exactly one error chunk followed by [DONE]; no sglext ids leak.
+        self.assertIn("error", chunks[0])
+        self.assertEqual(chunks[1], "data: [DONE]\n\n")
+        self.assertFalse(
+            any("input_ids" in c or "output_ids" in c for c in chunks),
+            "sglext ids event leaked after abort error",
+        )
+
+        error_chunk_data = json.loads(chunks[0][len("data: ") :])
+        self.assertEqual(error_chunk_data["error"]["message"], err_msg)
+        self.assertEqual(error_chunk_data["error"]["code"], err_code.value)
+
+    def test_streaming_error_abort_still_finalizes_other_choices(self):
+        """An error abort ends generation but still runs finalization, so a
+        sibling choice's finish chunk and the usage chunk follow the error."""
+        err_code = HTTPStatus.SERVICE_UNAVAILABLE
+
+        async def _mock_generate():
+            yield {
+                "text": "hello",
+                "prompt_token_ids": [4, 5, 6],
+                "output_ids": [1, 2],
+                "meta_info": {
+                    "id": "chatcmpl-multi-abort",
+                    "prompt_tokens": 3,
+                    "completion_tokens": 2,
+                    "cached_tokens": 0,
+                    "finish_reason": {"type": "stop"},
+                    "output_token_logprobs": None,
+                    "output_top_logprobs": None,
+                },
+                "index": 0,
+            }
+            yield {
+                "text": "partial",
+                "prompt_token_ids": [4, 5, 6],
+                "output_ids": [7],
+                "meta_info": {
+                    "id": "chatcmpl-multi-abort",
+                    "prompt_tokens": 3,
+                    "completion_tokens": 1,
+                    "cached_tokens": 0,
+                    "finish_reason": {
+                        "type": "abort",
+                        "status_code": err_code,
+                        "message": "Aborted by scheduler",
+                    },
+                    "output_token_logprobs": None,
+                    "output_top_logprobs": None,
+                },
+                "index": 1,
+            }
+
+        self.tm.generate_request.return_value = _mock_generate()
+
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+            max_tokens=100,
+            n=2,
+            stream=True,
+            stream_options={"include_usage": True},
+            return_input_ids_in_sglext=True,
+            return_output_ids_in_sglext=True,
+        )
+
+        with patch(
+            "sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
+        ) as conv_mock:
+            conv_ins = Mock()
+            conv_ins.get_prompt.return_value = "Test prompt"
+            conv_mock.return_value = conv_ins
+
+            adapted_request, _ = self.chat._convert_to_internal_request(
+                req, self.fastapi_request
+            )
+            chunks = self._run_chat_stream(adapted_request, req)
+
+        error_idx = next(i for i, c in enumerate(chunks) if "error" in c)
+        after_error = self._parse_chunks(chunks[error_idx + 1 :])
+
+        finish_reasons = [
+            choice["finish_reason"]
+            for c in after_error
+            for choice in c.get("choices", [])
+            if choice.get("finish_reason") is not None
+        ]
+        self.assertEqual(finish_reasons, ["stop"])
+        self.assertTrue(
+            any(c.get("usage") is not None for c in after_error),
+            "usage chunk dropped after error abort",
+        )
+
     def _run_chat_stream(self, adapted_request, req):
         async def run_stream():
             chunks = []
@@ -3008,6 +3167,331 @@ class ServingChatTestCase(unittest.TestCase):
             },
         )
 
+    def _output_ids_ret(self, *output_ids_by_choice):
+        return [
+            {
+                "text": "Answer",
+                "prompt_token_ids": [1, 2, 3],
+                "output_ids": list(output_ids),
+                "meta_info": {
+                    "id": "chatcmpl-output-ids",
+                    "prompt_tokens": 3,
+                    "completion_tokens": len(output_ids),
+                    "cached_tokens": 0,
+                    "finish_reason": {"type": "stop", "matched": None},
+                    "weight_version": "default",
+                },
+            }
+            for output_ids in output_ids_by_choice
+        ]
+
+    def test_non_streaming_ids_emit_sglext(self):
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+            n=2,
+            return_input_ids_in_sglext=True,
+            return_output_ids_in_sglext=True,
+        )
+
+        response = self.chat._build_chat_response(
+            req, self._output_ids_ret([5, 6, 7], [8, 9]), 123
+        )
+
+        self.assertIsNotNone(response.sglext)
+        self.assertEqual(response.sglext.input_ids, [1, 2, 3])
+        self.assertEqual(response.sglext.output_ids, [[5, 6, 7], [8, 9]])
+        dumped = json.loads(response.model_dump_json())
+        self.assertEqual(dumped["sglext"]["input_ids"], [1, 2, 3])
+        self.assertEqual(dumped["sglext"]["output_ids"], [[5, 6, 7], [8, 9]])
+        # None sglext fields stay stripped from the serialized response.
+        self.assertNotIn("routed_experts", dumped["sglext"])
+
+    def test_non_streaming_ids_not_returned_by_default(self):
+        req = ChatCompletionRequest(
+            model="x", messages=[{"role": "user", "content": "Hi?"}]
+        )
+
+        response = self.chat._build_chat_response(
+            req, self._output_ids_ret([5, 6, 7]), 123
+        )
+
+        self.assertIsNone(response.sglext)
+
+    def test_non_streaming_ids_server_default_enables_flag(self):
+        self.tm.server_args.return_input_ids = True
+        self.tm.server_args.return_output_ids = True
+        req = ChatCompletionRequest(
+            model="x", messages=[{"role": "user", "content": "Hi?"}]
+        )
+
+        response = self.chat._build_chat_response(
+            req, self._output_ids_ret([5, 6, 7]), 123
+        )
+
+        self.assertIsNotNone(response.sglext)
+        self.assertEqual(response.sglext.input_ids, [1, 2, 3])
+        self.assertEqual(response.sglext.output_ids, [[5, 6, 7]])
+
+    def test_ids_headers_enable_flags(self):
+        self.fastapi_request.headers = {
+            "x-sglext-return-input-ids": "1",
+            "x-sglext-return-output-ids": "1",
+        }
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+        )
+        processed_messages = MessageProcessingResult(
+            "Test prompt",
+            [1, 2, 3],
+            None,
+            None,
+            [],
+            ["</s>"],
+            None,
+        )
+
+        with patch.object(
+            self.chat, "_process_messages", return_value=processed_messages
+        ):
+            _, processed_request = self.chat._convert_to_internal_request(
+                req, self.fastapi_request
+            )
+
+        self.assertTrue(processed_request.return_input_ids_in_sglext)
+        self.assertTrue(processed_request.return_output_ids_in_sglext)
+
+    def _run_output_ids_stream(
+        self,
+        chunk_output_ids,
+        incremental,
+        framed=False,
+        return_raw=False,
+        cached_tokens_details=None,
+    ):
+        """Stream chunks with incremental or cumulative output_ids;
+        return parsed sglext chunks, or raw SSE strings when return_raw."""
+        self.tm.server_args.incremental_streaming_output = incremental
+        if framed:
+            self.fastapi_request.headers["x-sglext-ids-framed"] = "1"
+
+        async def _mock_generate():
+            for i, ids in enumerate(chunk_output_ids):
+                finished = i == len(chunk_output_ids) - 1
+                yield {
+                    "text": "chunk",
+                    "prompt_token_ids": [1, 2, 3],
+                    "output_ids": list(ids),
+                    "meta_info": {
+                        "id": "chatcmpl-output-ids-stream",
+                        "prompt_tokens": 3,
+                        "completion_tokens": 1 + i,
+                        "cached_tokens": 0,
+                        "cached_tokens_details": cached_tokens_details,
+                        "finish_reason": (
+                            {"type": "stop", "matched": None} if finished else None
+                        ),
+                        "output_token_logprobs": None,
+                        "output_top_logprobs": None,
+                    },
+                    "index": 0,
+                }
+
+        self.tm.generate_request.return_value = _mock_generate()
+
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+            max_tokens=100,
+            stream=True,
+            return_input_ids_in_sglext=True,
+            return_output_ids_in_sglext=True,
+            return_cached_tokens_details=cached_tokens_details is not None,
+        )
+
+        with patch(
+            "sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
+        ) as conv_mock:
+            conv_ins = Mock()
+            conv_ins.get_prompt.return_value = "Test prompt"
+            conv_mock.return_value = conv_ins
+
+            adapted_request, _ = self.chat._convert_to_internal_request(
+                req, self.fastapi_request
+            )
+            chunks = self._run_chat_stream(adapted_request, req)
+
+        if return_raw:
+            return chunks
+        return [c for c in self._parse_chunks(chunks) if "sglext" in c]
+
+    def test_streaming_output_ids_incremental_accumulates_deltas(self):
+        sglext_chunks = self._run_output_ids_stream([[5, 6], [7]], incremental=True)
+
+        self.assertEqual(len(sglext_chunks), 1)
+        self.assertEqual(sglext_chunks[0]["choices"], [])
+        # input_ids is captured once as a flat shared prompt, not accumulated.
+        self.assertEqual(sglext_chunks[0]["sglext"]["input_ids"], [1, 2, 3])
+        self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
+
+    def test_streaming_output_ids_non_incremental_keeps_latest_full_list(self):
+        sglext_chunks = self._run_output_ids_stream(
+            [[5, 6], [5, 6, 7]], incremental=False
+        )
+
+        self.assertEqual(len(sglext_chunks), 1)
+        self.assertEqual(sglext_chunks[0]["sglext"]["input_ids"], [1, 2, 3])
+        self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
+
+    _SGLEXT_IDS_EVENT_PREFIX = "event: sglext_ids\ndata: "
+
+    def test_streaming_framed_ids_emits_named_event(self):
+        raw = self._run_output_ids_stream(
+            [[5, 6], [5, 6, 7]], incremental=False, framed=True, return_raw=True
+        )
+
+        named = [c for c in raw if c.startswith(self._SGLEXT_IDS_EVENT_PREFIX)]
+        self.assertEqual(len(named), 1)
+        payload = json.loads(named[0][len(self._SGLEXT_IDS_EVENT_PREFIX) :])
+        self.assertEqual(payload["choices"], [])
+        # The named event carries ONLY the id fields.
+        self.assertEqual(set(payload["sglext"]), {"input_ids", "output_ids"})
+        self.assertEqual(payload["sglext"]["input_ids"], [1, 2, 3])
+        self.assertEqual(payload["sglext"]["output_ids"], [[5, 6, 7]])
+        # All requested sglext fields were ids, so no plain sglext chunk.
+        self.assertEqual([c for c in self._parse_chunks(raw) if "sglext" in c], [])
+
+    def test_streaming_framed_splits_non_id_fields_into_plain_chunk(self):
+        raw = self._run_output_ids_stream(
+            [[5, 6], [5, 6, 7]],
+            incremental=False,
+            framed=True,
+            return_raw=True,
+            cached_tokens_details={"device": 2, "host": 1},
+        )
+
+        named = [c for c in raw if c.startswith(self._SGLEXT_IDS_EVENT_PREFIX)]
+        self.assertEqual(len(named), 1)
+        named_payload = json.loads(named[0][len(self._SGLEXT_IDS_EVENT_PREFIX) :])
+        self.assertEqual(set(named_payload["sglext"]), {"input_ids", "output_ids"})
+
+        plain_sglext = [c for c in self._parse_chunks(raw) if "sglext" in c]
+        self.assertEqual(len(plain_sglext), 1)
+        self.assertEqual(
+            plain_sglext[0]["sglext"]["cached_tokens_details"],
+            {"device": 2, "host": 1},
+        )
+        self.assertNotIn("input_ids", plain_sglext[0]["sglext"])
+        self.assertNotIn("output_ids", plain_sglext[0]["sglext"])
+        # The plain non-id chunk precedes the named ids event.
+        plain_raw_idx = next(
+            i
+            for i, c in enumerate(raw)
+            if c.startswith("data: ") and "cached_tokens_details" in c
+        )
+        self.assertLess(plain_raw_idx, raw.index(named[0]))
+
+    def _run_output_ids_stream_with_graceful_abort(
+        self, normal_chunks, abort_output_ids, incremental, abort_completion_tokens
+    ):
+        """Stream normal chunks then a graceful-abort chunk; return parsed sglext chunks."""
+        self.tm.server_args.incremental_streaming_output = incremental
+
+        async def _mock_generate():
+            generated = 0
+            for ids in normal_chunks:
+                generated = generated + len(ids) if incremental else len(ids)
+                yield {
+                    "text": "chunk",
+                    "output_ids": list(ids),
+                    "meta_info": {
+                        "id": "chatcmpl-abort-ids-stream",
+                        "prompt_tokens": 3,
+                        "completion_tokens": generated,
+                        "cached_tokens": 0,
+                        "finish_reason": None,
+                        "output_token_logprobs": None,
+                        "output_top_logprobs": None,
+                    },
+                    "index": 0,
+                }
+            # Graceful abort terminal chunk (no status_code): falls through to
+            # the normal finalization path.
+            yield {
+                "text": "chunk",
+                "output_ids": list(abort_output_ids),
+                "meta_info": {
+                    "id": "chatcmpl-abort-ids-stream",
+                    "prompt_tokens": 3,
+                    "completion_tokens": abort_completion_tokens,
+                    "cached_tokens": 0,
+                    "finish_reason": {"type": "abort", "message": "Aborted."},
+                    "output_token_logprobs": None,
+                    "output_top_logprobs": None,
+                },
+                "index": 0,
+            }
+
+        self.tm.generate_request.return_value = _mock_generate()
+
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+            max_tokens=100,
+            stream=True,
+            return_output_ids_in_sglext=True,
+        )
+
+        with patch(
+            "sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
+        ) as conv_mock:
+            conv_ins = Mock()
+            conv_ins.get_prompt.return_value = "Test prompt"
+            conv_mock.return_value = conv_ins
+
+            adapted_request, _ = self.chat._convert_to_internal_request(
+                req, self.fastapi_request
+            )
+            chunks = self._run_chat_stream(adapted_request, req)
+
+        return [c for c in self._parse_chunks(chunks) if "sglext" in c]
+
+    def test_streaming_output_ids_incremental_ignores_abort_chunk(self):
+        # Abort re-sends the last token; it must not be appended again.
+        sglext_chunks = self._run_output_ids_stream_with_graceful_abort(
+            normal_chunks=[[5, 6], [7]],
+            abort_output_ids=[7],
+            incremental=True,
+            abort_completion_tokens=3,
+        )
+
+        self.assertEqual(len(sglext_chunks), 1)
+        self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
+
+    def test_streaming_output_ids_incremental_keeps_coalesced_abort_deltas(self):
+        # Coalesced abort chunk: keep real deltas, drop the repeated last token.
+        sglext_chunks = self._run_output_ids_stream_with_graceful_abort(
+            normal_chunks=[[5, 6]],
+            abort_output_ids=[7, 8, 8],
+            incremental=True,
+            abort_completion_tokens=4,
+        )
+
+        self.assertEqual(len(sglext_chunks), 1)
+        self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7, 8]])
+
+    def test_streaming_output_ids_non_incremental_abort_supersedes_earlier(self):
+        sglext_chunks = self._run_output_ids_stream_with_graceful_abort(
+            normal_chunks=[[5, 6]],
+            abort_output_ids=[5, 6, 7],
+            incremental=False,
+            abort_completion_tokens=3,
+        )
+
+        self.assertEqual(len(sglext_chunks), 1)
+        self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
+
     def test_streaming_parallel_sampling_orders_spec_details_by_choice(self):
         async def mock_generate():
             for index in (1, 0):