lazy init mcp session

Signed-off-by: wuhang <wuhang6@huawei.com>

Author: wuhang2014

vllm/entrypoints/context.py

@@ -104,6 +104,14 @@ async def init_tool_sessions(
     ) -> None:
         pass
 
+    @abstractmethod
+    async def __aenter__(self):
+        pass
+
+    @abstractmethod
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
     @abstractmethod
     async def cleanup_session(self) -> None:
         raise NotImplementedError("Should not be called.")
@@ -146,6 +154,12 @@ async def init_tool_sessions(
     ) -> None:
         pass
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
     async def cleanup_session(self) -> None:
         raise NotImplementedError("Should not be called.")
 
@@ -155,12 +169,17 @@ def __init__(
         self,
         messages: list,
         available_tools: list[str],
+        tool_server: Optional[ToolServer],
     ):
         self._messages = messages
         self.finish_reason: str | None = None
         self.available_tools = available_tools
         self._tool_sessions: dict[str, ClientSession | Tool] = {}
         self.called_tools: set[str] = set()
+        self._tool_server = tool_server
+        self._async_exit_stack: Optional[AsyncExitStack] = None
+        self._reference_count = 0
+        self._reference_count_lock = asyncio.Lock()
 
         self.parser = get_streamable_parser_for_assistant()
         self.num_init_messages = len(messages)
@@ -309,6 +328,18 @@ def need_builtin_tool_call(self) -> bool:
             or recipient.startswith("container.")
         )
 
+    async def _get_tool_session(self, tool_name: str) -> Union["ClientSession", Tool]:
+        if tool_name not in self._tool_sessions and self._tool_server is not None:
+            assert self._async_exit_stack is not None, (
+                "Async exit stack not set. Please report this issue."
+            )
+            self._tool_sessions[
+                tool_name
+            ] = await self._async_exit_stack.enter_async_context(
+                self._tool_server.new_session(tool_name)
+            )
+        return self._tool_sessions[tool_name]
+
     async def call_tool(self) -> list[Message]:
         if not self.messages:
             return []
@@ -317,15 +348,15 @@ async def call_tool(self) -> list[Message]:
         if recipient is not None:
             if recipient.startswith("browser."):
                 return await self.call_search_tool(
-                    self._tool_sessions["browser"], last_msg
+                    await self._get_tool_session("browser"), last_msg
                 )
             elif recipient.startswith("python"):
                 return await self.call_python_tool(
-                    self._tool_sessions["python"], last_msg
+                    await self._get_tool_session("python"), last_msg
                 )
             elif recipient.startswith("container."):
                 return await self.call_container_tool(
-                    self._tool_sessions["container"], last_msg
+                    await self._get_tool_session("container"), last_msg
                 )
         raise ValueError("No tool call found")
 
@@ -452,6 +483,25 @@ async def cleanup_tool_session(tool_session):
             )
         )
 
+    async def __aenter__(self):
+        async with self._reference_count_lock:
+            self._reference_count += 1
+            if self._async_exit_stack is None:
+                assert self._reference_count == 1, (
+                    "Reference count of exit stack should be "
+                )
+                "1 when initializing exit stack."
+                self._async_exit_stack = AsyncExitStack()
+                await self._async_exit_stack.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        async with self._reference_count_lock:
+            self._reference_count -= 1
+            if self._reference_count == 0 and self._async_exit_stack is not None:
+                await self._async_exit_stack.__aexit__(exc_type, exc, tb)
+                self._async_exit_stack = None
+
 
 class StreamingHarmonyContext(HarmonyContext):
     def __init__(self, *args, **kwargs):

vllm/entrypoints/openai/serving_responses.py

@@ -352,9 +352,13 @@ async def create_responses(
                 context: ConversationContext
                 if self.use_harmony:
                     if request.stream:
-                        context = StreamingHarmonyContext(messages, available_tools)
+                        context = StreamingHarmonyContext(
+                            messages, available_tools, self.tool_server
+                        )
                     else:
-                        context = HarmonyContext(messages, available_tools)
+                        context = HarmonyContext(
+                            messages, available_tools, self.tool_server
+                        )
                 else:
                     context = SimpleContext()
                 generator = self._generate_with_builtin_tools(
@@ -498,22 +502,6 @@ def _make_request_with_harmony(
 
         return messages, [prompt_token_ids], [engine_prompt]
 
-    async def _initialize_tool_sessions(
-        self,
-        request: ResponsesRequest,
-        context: ConversationContext,
-        exit_stack: AsyncExitStack,
-    ):
-        # we should only initialize the tool session if the request needs tools
-        if len(request.tools) == 0:
-            return
-        mcp_tools = {
-            tool.server_label: tool for tool in request.tools if tool.type == "mcp"
-        }
-        await context.init_tool_sessions(
-            self.tool_server, exit_stack, request.request_id, mcp_tools
-        )
-
     async def responses_full_generator(
         self,
         request: ResponsesRequest,
@@ -528,9 +516,8 @@ async def responses_full_generator(
         if created_time is None:
             created_time = int(time.time())
 
-        async with AsyncExitStack() as exit_stack:
+        async with context:
             try:
-                await self._initialize_tool_sessions(request, context, exit_stack)
                 async for _ in result_generator:
                     pass
             except asyncio.CancelledError:
@@ -1894,12 +1881,9 @@ def _increment_sequence_number_and_return(
             sequence_number += 1
             return event
 
-        async with AsyncExitStack() as exit_stack:
+        async with context:
             processer = None
             if self.use_harmony:
-                # TODO: in streaming, we noticed this bug:
-                # https://github.com/vllm-project/vllm/issues/25697
-                await self._initialize_tool_sessions(request, context, exit_stack)
                 processer = self._process_harmony_streaming_events
             else:
                 processer = self._process_simple_streaming_events