From 567bf9513ce6612355282f2f197ed0ed04328e61 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Thu, 28 May 2026 10:27:14 +0800
Subject: [PATCH] Python: fix: serialize MCP loader reloads

---
 python/packages/core/agent_framework/_mcp.py | 12 +++++
 python/packages/core/tests/core/test_mcp.py  | 54 ++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py
index b2942de2a0..13a2760f41 100644
--- a/python/packages/core/agent_framework/_mcp.py
+++ b/python/packages/core/agent_framework/_mcp.py
@@ -255,6 +255,8 @@ def __init__(
         self._exit_stack = AsyncExitStack()
         self._lifecycle_lock = asyncio.Lock()
         self._lifecycle_request_lock = asyncio.Lock()
+        self._load_tools_lock = asyncio.Lock()
+        self._load_prompts_lock = asyncio.Lock()
         self._lifecycle_queue: asyncio.Queue[tuple[str, bool, bool, asyncio.Future[None]]] | None = None
         self._lifecycle_owner_task: asyncio.Task[None] | None = None
         self.session = session
@@ -1010,6 +1012,11 @@ def _determine_approval_mode(
         return self.approval_mode  # type: ignore[return-value]
 
     async def load_prompts(self) -> None:
+        """Load prompts from the MCP server."""
+        async with self._load_prompts_lock:
+            await self._load_prompts_locked()
+
+    async def _load_prompts_locked(self) -> None:
         """Load prompts from the MCP server.
 
         Retrieves available prompts from the connected MCP server and converts
@@ -1092,6 +1099,11 @@ async def load_prompts(self) -> None:
             params = types.PaginatedRequestParams(cursor=prompt_list.nextCursor)
 
     async def load_tools(self) -> None:
+        """Load tools from the MCP server."""
+        async with self._load_tools_lock:
+            await self._load_tools_locked()
+
+    async def _load_tools_locked(self) -> None:
         """Load tools from the MCP server.
 
         Retrieves available tools from the connected MCP server and converts
diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py
index 6273eb76e6..cbe12e8454 100644
--- a/python/packages/core/tests/core/test_mcp.py
+++ b/python/packages/core/tests/core/test_mcp.py
@@ -3202,6 +3202,33 @@ async def mock_list_tools(params=None):
     assert [f.name for f in tool._functions] == ["tool_1", "tool_2", "tool_3"]
 
 
+async def test_load_tools_serializes_concurrent_loads():
+    tool = MCPTool(name="test_tool")
+    tool.session = AsyncMock()
+    tool.load_tools_flag = True
+
+    page = Mock()
+    page.tools = [
+        types.Tool(
+            name="tool_1",
+            description="First tool",
+            inputSchema={"type": "object", "properties": {"param": {"type": "string"}}},
+        ),
+    ]
+    page.nextCursor = None
+
+    async def list_tools(params=None):
+        await asyncio.sleep(0)
+        return page
+
+    tool.session.list_tools = AsyncMock(side_effect=list_tools)
+
+    await asyncio.gather(tool.load_tools(), tool.load_tools())
+
+    assert tool.session.list_tools.call_count == 2
+    assert [f.name for f in tool._functions] == ["tool_1"]
+
+
 async def test_load_prompts_pagination_with_duplicates():
     """Test that load_prompts prevents duplicates across paginated results."""
     from unittest.mock import AsyncMock, MagicMock
@@ -3260,6 +3287,33 @@ async def mock_list_prompts(params=None):
     assert [f.name for f in tool._functions] == ["prompt_1", "prompt_2"]
 
 
+async def test_load_prompts_serializes_concurrent_loads():
+    tool = MCPTool(name="test_tool")
+    tool.session = AsyncMock()
+    tool.load_prompts_flag = True
+
+    page = Mock()
+    page.prompts = [
+        types.Prompt(
+            name="prompt_1",
+            description="First prompt",
+            arguments=[types.PromptArgument(name="arg1", description="Arg 1", required=True)],
+        ),
+    ]
+    page.nextCursor = None
+
+    async def list_prompts(params=None):
+        await asyncio.sleep(0)
+        return page
+
+    tool.session.list_prompts = AsyncMock(side_effect=list_prompts)
+
+    await asyncio.gather(tool.load_prompts(), tool.load_prompts())
+
+    assert tool.session.list_prompts.call_count == 2
+    assert [f.name for f in tool._functions] == ["prompt_1"]
+
+
 async def test_load_tools_pagination_exception_handling():
     """Test that load_tools handles exceptions during pagination gracefully."""
     from unittest.mock import AsyncMock