diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index dfb7e6cbb1..2ce42d8c33 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -679,6 +679,36 @@ mod tests { assert!(response.mm_placeholders.is_none()); } + #[test] + fn stream_chunk_omits_absent_fields() { + // Contract shared with the Python frontend: absent optional fields + // are omitted from stream chunks, not serialized as null. + let chunk = GenerateStreamResponse { + request_id: "raw-1".to_string(), + choices: vec![GenerateResponseStreamChoice { + index: 0, + logprobs: None, + finish_reason: None, + token_ids: vec![1], + }], + usage: None, + prompt_token_ids: None, + mm_placeholders: None, + }; + + let json = serde_json::to_value(&chunk).expect("serialize chunk"); + + for key in ["usage", "prompt_token_ids", "mm_placeholders"] { + assert!(json.get(key).is_none(), "{key} should be omitted"); + } + for key in ["logprobs", "finish_reason"] { + assert!( + json["choices"][0].get(key).is_none(), + "{key} should be omitted" + ); + } + } + #[test] fn extract_mm_placeholders_returns_none_without_features() { assert!(extract_mm_placeholders(None).is_none()); diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index d9ae424cc8..0cb8a6d397 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -569,6 +569,40 @@ async def test_stream_omits_prompt_metadata_by_default(): assert "mm_placeholders" not in chunk +@pytest.mark.asyncio +async def test_stream_chunks_omit_absent_fields(): + """Absent optional fields are omitted from stream chunks, like Rust.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", token_ids=[20], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert len(data_chunks) == 2 + for chunk in data_chunks: + assert "usage" not in chunk + assert "prompt_token_ids" not in chunk + assert "mm_placeholders" not in chunk + assert "logprobs" not in chunk["choices"][0] + assert "finish_reason" not in data_chunks[0]["choices"][0] + assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" + + @pytest.mark.asyncio async def test_stream_text_only_prompt_omits_mm_placeholders_key(): engine = _mock_engine() diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index a8a2804a06..838fe90c5c 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -486,13 +486,9 @@ class ServingTokens(GenerateBaseServing): total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - # Omit absent prompt metadata (Rust skips None fields). - exclude = { - name - for name in ("prompt_token_ids", "mm_placeholders") - if getattr(chunk, name) is None - } - yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" + # Omit absent fields, like the Rust frontend and the + # final usage chunk below. + yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" total_completion_tokens = sum(num_generated_tokens) final_usage_info = UsageInfo(