[Bugfix][Responses] Clean up MCP tool sessions once before closing them (#56988)
Signed-off-by: 子华 <huaxi.shx@alibaba-inc.com> Co-authored-by: Codex <noreply@openai.com>
shaohuaxi committed
Sep 15, 2026 at 23:07 UTC
e0301afb2a1364fe594a52d9ce649264a62fa243
2 files changed
+90
-2
tests/entrypoints/openai/responses/test_serving_responses.py
+83
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from contextlib import AsyncExitStack
4
+from contextlib import AsyncExitStack, asynccontextmanager
5
from types import SimpleNamespace
6
from unittest.mock import AsyncMock, MagicMock
7
@@ -41,6 +41,7 @@ from vllm.entrypoints.mcp.tool_server import ToolServer
41
from vllm.entrypoints.openai.responses.context import (
42
ConversationContext,
43
HarmonyContext,
44
+ ParsableContext,
45
SimpleContext,
46
)
47
from vllm.entrypoints.openai.responses.protocol import (
@@ -551,6 +552,87 @@ def test_responses_render_result_rejects_non_single_prompt(engine_inputs):
552
class TestInitializeToolSessions:
553
"""Test class for _initialize_tool_sessions method"""
554
555
+ @pytest.fixture(params=[ParsableContext, HarmonyContext])
556
+ def tool_context(self, request):
557
+ context = request.param.__new__(request.param)
558
+ context.available_tools = ["browser", "python"]
559
+ context._tool_sessions = {}
560
+ context.called_tools = set()
561
+ return context
562
+
563
+ @pytest.mark.asyncio
564
+ @pytest.mark.parametrize(
565
+ ("available_tools", "called_tools"),
566
+ [
567
+ ([], []),
568
+ (["browser"], ["browser"]),
569
+ (["browser", "python"], []),
570
+ (["browser", "python"], ["browser"]),
571
+ (["browser", "python"], ["python"]),
572
+ (["browser", "python"], ["browser", "python"]),
573
+ ],
574
+ ids=["empty", "single", "unused", "first", "last", "both"],
575
+ )
576
+ async def test_mcp_sessions_cleaned_once_before_close(
577
+ self, tool_context, available_tools, called_tools
578
+ ):
579
+ events = []
580
+ sessions = {}
581
+
582
+ @asynccontextmanager
583
+ async def new_session(name, *_args):
584
+ session = MagicMock()
585
+ session.call_tool = AsyncMock(
586
+ side_effect=lambda *_: events.append(("cleanup", name))
587
+ )
588
+ sessions[name] = session
589
+ try:
590
+ yield session
591
+ finally:
592
+ events.append(("close", name))
593
+
594
+ tool_server = MagicMock(spec=ToolServer)
595
+ tool_server.new_session.side_effect = new_session
596
+ tool_context.available_tools = available_tools
597
+ async with AsyncExitStack() as stack:
598
+ await tool_context.init_tool_sessions(tool_server, stack, "req", {})
599
+ # Reinitializing existing sessions must not duplicate cleanup either.
600
+ await tool_context.init_tool_sessions(tool_server, stack, "req", {})
601
+ tool_context.called_tools.update(called_tools)
602
+
603
+ for name, session in sessions.items():
604
+ if name in called_tools:
605
+ session.call_tool.assert_awaited_once_with("cleanup_session", {})
606
+ assert events.index(("cleanup", name)) < events.index(("close", name))
607
+ else:
608
+ session.call_tool.assert_not_awaited()
609
+ assert [name for event, name in events if event == "close"] == list(
610
+ reversed(available_tools)
611
+ )
612
+
613
+ @pytest.mark.asyncio
614
+ async def test_mcp_partial_initialization_closes_opened_session(self, tool_context):
615
+ closed = []
616
+ session = MagicMock(call_tool=AsyncMock())
617
+
618
+ @asynccontextmanager
619
+ async def new_session(name, *_args):
620
+ if name == "python":
621
+ raise RuntimeError("session initialization failed")
622
+ try:
623
+ yield session
624
+ finally:
625
+ closed.append(name)
626
+
627
+ tool_server = MagicMock(spec=ToolServer)
628
+ tool_server.new_session.side_effect = new_session
629
+ with pytest.raises(RuntimeError, match="session initialization failed"):
630
+ async with AsyncExitStack() as stack:
631
+ await tool_context.init_tool_sessions(tool_server, stack, "req", {})
632
+
633
+ assert closed == ["browser"]
634
+ session.call_tool.assert_not_awaited()
635
+
636
@pytest_asyncio.fixture
637
async def serving_responses_instance(self):
638
"""Create a real OpenAIServingResponses instance for testing"""
vllm/entrypoints/openai/responses/context.py
+7
-1
@@ -565,6 +565,7 @@ class ParsableContext(ConversationContext):
565
mcp_tools: dict[str, Mcp],
566
):
567
if tool_server:
568
+ initialized_session = False
569
for tool_name in self.available_tools:
570
if tool_name in self._tool_sessions:
571
continue
@@ -577,6 +578,8 @@ class ParsableContext(ConversationContext):
578
tool_server.new_session(tool_name, request_id, headers)
579
)
580
self._tool_sessions[tool_name] = tool_session
581
+ initialized_session = True
582
+ if initialized_session:
583
exit_stack.push_async_exit(self.cleanup_session)
584
585
async def cleanup_session(self, *args, **kwargs) -> None:
@@ -863,6 +866,7 @@ class HarmonyContext(ConversationContext):
866
mcp_tools: dict[str, Mcp],
867
):
868
if tool_server:
869
+ initialized_session = False
870
for tool_name in self.available_tools:
871
if tool_name not in self._tool_sessions:
872
tool_type = _map_tool_name_to_tool_type(tool_name)
@@ -873,7 +877,9 @@ class HarmonyContext(ConversationContext):
877
tool_server.new_session(tool_name, request_id, headers)
878
)
879
self._tool_sessions[tool_name] = tool_session
876
- exit_stack.push_async_exit(self.cleanup_session)
880
+ initialized_session = True
881
+ if initialized_session:
882
+ exit_stack.push_async_exit(self.cleanup_session)
883
884
async def call_container_tool(
885
self, tool_session: Union["ClientSession", Tool], last_msg: Message