A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Feature][Responses API]Support MCP tools with streaming mode + background mode (#23927)
Signed-off-by: wuhang <wuhang6@huawei.com>
wuhang committed
Sep 4, 2025 at 12:05 UTC
a38f8bd54c861d37acd1bf6497b86edf664a6ab7
3 files changed
+138
-26
tests/entrypoints/openai/test_response_api_with_harmony.py
+18
-1
@@ -275,7 +275,8 @@ async def test_stateful_multi_turn(client: OpenAI, model_name: str):
275
276
@pytest.mark.asyncio
277
@pytest.mark.parametrize("model_name", [MODEL_NAME])
278
-async def test_streaming(client: OpenAI, model_name: str):
278
+@pytest.mark.parametrize("background", [True, False])
279
+async def test_streaming(client: OpenAI, model_name: str, background: bool):
280
# TODO: Add back when web search and code interpreter are available in CI
281
prompts = [
282
"tell me a story about a cat in 20 words",
@@ -300,11 +301,16 @@ async def test_streaming(client: OpenAI, model_name: str):
301
# },
302
],
303
stream=True,
304
+ background=background,
305
)
306
307
events = []
308
current_event_mode = None
309
+ resp_id = None
310
async for event in response:
311
+ if event.type == "response.created":
312
+ resp_id = event.response.id
313
+
314
if current_event_mode != event.type:
315
current_event_mode = event.type
316
print(f"\n[{event.type}] ", end="", flush=True)
@@ -322,6 +328,17 @@ async def test_streaming(client: OpenAI, model_name: str):
328
329
assert len(events) > 0
330
331
+ if background:
332
+ starting_after = 5
333
+ async with await client.responses.retrieve(
334
+ response_id=resp_id,
335
+ stream=True,
336
+ starting_after=starting_after) as stream:
337
+ counter = starting_after
338
+ async for event in stream:
339
+ counter += 1
340
+ assert event == events[counter]
341
+
342
343
@pytest.mark.asyncio
344
@pytest.mark.parametrize("model_name", [MODEL_NAME])
vllm/entrypoints/openai/api_server.py
+14
-2
@@ -616,14 +616,23 @@ async def create_responses(request: ResponsesRequest, raw_request: Request):
616
617
618
@router.get("/v1/responses/{response_id}")
619
-async def retrieve_responses(response_id: str, raw_request: Request):
619
+async def retrieve_responses(
620
+ response_id: str,
621
+ raw_request: Request,
622
+ starting_after: Optional[int] = None,
623
+ stream: Optional[bool] = False,
624
+):
625
handler = responses(raw_request)
626
if handler is None:
627
return base(raw_request).create_error_response(
628
message="The model does not support Responses API")
629
630
try:
626
- response = await handler.retrieve_responses(response_id)
631
+ response = await handler.retrieve_responses(
632
+ response_id,
633
+ starting_after=starting_after,
634
+ stream=stream,
635
+ )
636
except Exception as e:
637
raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
638
detail=str(e)) from e
@@ -631,6 +640,9 @@ async def retrieve_responses(response_id: str, raw_request: Request):
640
if isinstance(response, ErrorResponse):
641
return JSONResponse(content=response.model_dump(),
642
status_code=response.error.code)
643
+ elif stream:
644
+ return StreamingResponse(content=response,
645
+ media_type="text/event-stream")
646
return JSONResponse(content=response.model_dump())
647
648
vllm/entrypoints/openai/serving_responses.py
+106
-23
@@ -4,6 +4,7 @@
4
import asyncio
5
import json
6
import time
7
+from collections import deque
8
from collections.abc import AsyncGenerator, AsyncIterator, Sequence
9
from contextlib import AsyncExitStack
10
from copy import copy
@@ -55,7 +56,7 @@ from vllm.entrypoints.openai.protocol import (ErrorResponse,
56
# yapf: enable
57
from vllm.entrypoints.openai.serving_engine import OpenAIServing
58
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
58
-from vllm.entrypoints.tool_server import MCPToolServer, ToolServer
59
+from vllm.entrypoints.tool_server import ToolServer
60
from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
61
from vllm.logger import init_logger
62
from vllm.logprobs import Logprob as SampleLogprob
@@ -168,6 +169,11 @@ class OpenAIServingResponses(OpenAIServing):
169
# never remove messages from the store.
170
self.msg_store: dict[str, list[ChatCompletionMessageParam]] = {}
171
172
+ # HACK(wuhang): This is a hack. We should use a better store.
173
+ # FIXME: If enable_store=True, this may cause a memory leak since we
174
+ # never remove events from the store.
175
+ self.event_store: dict[str, tuple[deque[str], asyncio.Event]] = {}
176
+
177
self.background_tasks: dict[str, asyncio.Task] = {}
178
179
self.tool_server = tool_server
@@ -249,15 +255,6 @@ class OpenAIServingResponses(OpenAIServing):
255
if raw_request:
256
raw_request.state.request_metadata = request_metadata
257
252
- if self.tool_server is not None and isinstance(
253
- self.tool_server,
254
- MCPToolServer) and request.stream and request.tools and any(
255
- tool.type in ["web_search_preview", "code_interpreter"]
256
- for tool in request.tools):
257
- return self.create_error_response(
258
- "MCP tool server is not supported in background mode and "
259
- "streaming mode")
260
-
258
# Schedule the request and get the result generator.
259
generators: list[AsyncGenerator[ConversationContext, None]] = []
260
@@ -329,25 +326,44 @@ class OpenAIServingResponses(OpenAIServing):
326
self.response_store[response.id] = response
327
328
# Run the request in the background.
332
- task = asyncio.create_task(
333
- self._run_background_request(
334
- request,
335
- sampling_params,
336
- result_generator,
337
- context,
338
- model_name,
339
- tokenizer,
340
- request_metadata,
341
- created_time,
342
- ),
343
- name=f"create_{response.id}",
344
- )
329
+ if request.stream:
330
+ task = asyncio.create_task(
331
+ self._run_background_request_stream(
332
+ request,
333
+ sampling_params,
334
+ result_generator,
335
+ context,
336
+ model_name,
337
+ tokenizer,
338
+ request_metadata,
339
+ created_time,
340
+ ),
341
+ name=f"create_{request.request_id}",
342
+ )
343
+ else:
344
+ task = asyncio.create_task(
345
+ self._run_background_request(
346
+ request,
347
+ sampling_params,
348
+ result_generator,
349
+ context,
350
+ model_name,
351
+ tokenizer,
352
+ request_metadata,
353
+ created_time,
354
+ ),
355
+ name=f"create_{response.id}",
356
+ )
357
358
# For cleanup.
359
response_id = response.id
360
self.background_tasks[response_id] = task
361
task.add_done_callback(
362
lambda _: self.background_tasks.pop(response_id, None))
363
+
364
+ if request.stream:
365
+ return self.responses_background_stream_generator(
366
+ request.request_id)
367
return response
368
369
if request.stream:
@@ -736,6 +752,40 @@ class OpenAIServingResponses(OpenAIServing):
752
prev_outputs.append(response_msg)
753
return messages
754
755
+ async def _run_background_request_stream(
756
+ self,
757
+ request: ResponsesRequest,
758
+ *args,
759
+ **kwargs,
760
+ ):
761
+ event_deque: deque[str] = deque()
762
+ new_event_signal = asyncio.Event()
763
+ self.event_store[request.request_id] = (event_deque, new_event_signal)
764
+ response = None
765
+ try:
766
+ generator = self.responses_stream_generator(
767
+ request, *args, **kwargs)
768
+ async for event in generator:
769
+ event_deque.append(event)
770
+ new_event_signal.set() # Signal new event available
771
+ except Exception as e:
772
+ logger.exception("Background request failed for %s",
773
+ request.request_id)
774
+ response = self.create_error_response(str(e))
775
+ finally:
776
+ # Mark as finished with a special marker
777
+ event_deque.append("__STREAM_END__")
778
+ new_event_signal.set()
779
+
780
+ if response is not None and isinstance(response, ErrorResponse):
781
+ # If the request has failed, update the status to "failed".
782
+ response_id = request.request_id
783
+ async with self.response_store_lock:
784
+ stored_response = self.response_store.get(response_id)
785
+ assert stored_response is not None
786
+ if stored_response.status not in ("completed", "cancelled"):
787
+ stored_response.status = "failed"
788
+
789
async def _run_background_request(
790
self,
791
request: ResponsesRequest,
@@ -759,9 +809,36 @@ class OpenAIServingResponses(OpenAIServing):
809
if stored_response.status not in ("completed", "cancelled"):
810
stored_response.status = "failed"
811
812
+ async def responses_background_stream_generator(
813
+ self,
814
+ response_id: str,
815
+ starting_after: Optional[int] = None,
816
+ ):
817
+ if response_id not in self.event_store:
818
+ raise ValueError(f"Unknown response_id: {response_id}")
819
+
820
+ event_deque, new_event_signal = self.event_store[response_id]
821
+ start_index = 0 if starting_after is None else starting_after + 1
822
+ current_index = start_index
823
+
824
+ while True:
825
+ new_event_signal.clear()
826
+
827
+ # Yield existing events from start_index
828
+ while current_index < len(event_deque):
829
+ event = event_deque[current_index]
830
+ if event == "__STREAM_END__":
831
+ return
832
+ yield event
833
+ current_index += 1
834
+
835
+ await new_event_signal.wait()
836
+
837
async def retrieve_responses(
838
self,
839
response_id: str,
840
+ starting_after: Optional[int],
841
+ stream: Optional[bool],
842
) -> Union[ErrorResponse, ResponsesResponse]:
843
if not response_id.startswith("resp_"):
844
return self._make_invalid_id_error(response_id)
@@ -771,6 +848,12 @@ class OpenAIServingResponses(OpenAIServing):
848
849
if response is None:
850
return self._make_not_found_error(response_id)
851
+
852
+ if stream:
853
+ return self.responses_background_stream_generator(
854
+ response_id,
855
+ starting_after,
856
+ )
857
return response
858
859
async def cancel_responses(