[Bugfix][Mooncake] Fix MultiConnector stats and proxy SSE responses (#61218)

Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>

.buildkite/test_areas/misc.yaml

@@ -214,6 +214,7 @@ steps:
     - vllm/v1/
     - tests/v1
     - tests/watermarking
+    - examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
   device: cpu-small
   commands:
     # split the test to avoid interference

examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py

@@ -340,7 +340,10 @@ async def _handle_completions(api: str, request: Request):
             ):
                 yield chunk
 
-        return StreamingResponse(generate_stream(), media_type="application/json")
+        media_type = (
+            "text/event-stream" if req_data.get("stream") else "application/json"
+        )
+        return StreamingResponse(generate_stream(), media_type=media_type)
 
     except Exception as e:
         import sys

tests/v1/kv_connector/unit/test_mooncake_proxy.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import itertools
+import runpy
+from pathlib import Path
+from unittest.mock import AsyncMock
+
+import httpx
+import pytest
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("endpoint", ["/v1/completions", "/v1/chat/completions"])
+@pytest.mark.parametrize("stream", [False, True])
+async def test_response_media_type_matches_payload(monkeypatch, endpoint, stream):
+    """Clients must select SSE parsing for streaming completion responses."""
+    path = (
+        Path(__file__).parents[4]
+        / "examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py"
+    )
+    proxy = runpy.run_path(str(path))["_handle_completions"].__globals__
+    app = proxy["app"]
+    app.state.ready = asyncio.Event()
+    app.state.ready.set()
+    app.state.prefill_iterator = itertools.cycle([({}, 0)])
+    app.state.decode_iterator = itertools.cycle([0])
+    app.state.decode_clients = [{}]
+    payload = b'{"choices":[{"text":"hello"}]}'
+    body = b"data: " + payload + b"\n\ndata: [DONE]\n\n" if stream else payload
+
+    async def decode_response(*args, **kwargs):
+        yield body
+
+    monkeypatch.setitem(proxy, "send_request_to_service", AsyncMock())
+    monkeypatch.setitem(proxy, "stream_service_response", decode_response)
+    async with httpx.AsyncClient(
+        transport=httpx.ASGITransport(app=app), base_url="http://proxy"
+    ) as client:
+        response = await client.post(endpoint, json={"stream": stream})
+
+    expected_type = "text/event-stream" if stream else "application/json"
+    assert response.headers["content-type"].split(";")[0] == expected_type
+    assert response.content == body

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -805,6 +805,23 @@ class TestMultiConnectorStats:
 
         child_metrics.observe.assert_called_once_with({"transfer_duration": [1.5]}, 2)
 
+    @pytest.mark.parametrize("with_exporter", [False, True])
+    def test_prom_metrics_skips_stats_only_connectors(self, with_exporter):
+        """Stats-only children must not prevent supported metrics from exporting."""
+        child_metrics = MagicMock()
+        metrics = object.__new__(MultiKVConnectorPromMetrics)
+        metrics._prom_metrics = (
+            {"MooncakeStoreConnector": child_metrics} if with_exporter else {}
+        )
+        payload = {"MooncakeConnector": {"bytes_transferred": [1024]}}
+        if with_exporter:
+            payload["MooncakeStoreConnector"] = {"load_get_bytes": [2048]}
+
+        metrics.observe(payload, engine_idx=2)
+
+        if with_exporter:
+            child_metrics.observe.assert_called_once_with({"load_get_bytes": [2048]}, 2)
+
     def test_aggregate_same_connector(self):
         """Test aggregating stats from the same connector type."""
         stats1 = MultiKVConnectorStats(

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -125,11 +125,9 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
 
     def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
         for connector_id, stats_data in transfer_stats_data.items():
-            assert connector_id in self._prom_metrics, (
-                f"{connector_id} is not contained in the list of registered connectors "
-                f"with Prometheus metrics support: {self._prom_metrics.keys()}"
-            )
-            self._prom_metrics[connector_id].observe(stats_data, engine_idx)
+            prom_metrics = self._prom_metrics.get(connector_id)
+            if prom_metrics is not None:
+                prom_metrics.observe(stats_data, engine_idx)
 
 
 class MultiConnector(KVConnectorBase_V1, SupportsHMA):