fix: cover MCP reconnect edge cases

maintainer edit: handle prompt sessions that report Connection closed outside McpError, and match reconnect registration prefixes with the same sanitization used by MCP wrapper names.
This commit is contained in:
chengyongru
2026-06-04 10:43:09 +08:00
committed by Xubin Ren
parent e9145b7acd
commit d0eba7cd9d
3 changed files with 83 additions and 4 deletions
+26
View File
@@ -464,3 +464,29 @@ async def test_prompt_reconnects_on_session_terminated():
assert output == "fresh prompt"
assert old_session.get_prompt.call_count == 1
assert new_session.get_prompt.call_count == 1
@pytest.mark.asyncio
async def test_prompt_reconnects_on_connection_closed_exception():
"""Prompt should reconnect when the SDK reports a closed session as a generic exception."""
old_session = AsyncMock()
old_session.get_prompt = AsyncMock(side_effect=RuntimeError("Connection closed"))
new_session = AsyncMock()
new_session.get_prompt = AsyncMock(return_value=_make_prompt_result("fresh prompt"))
wrapper = MCPPromptWrapper(old_session, "test_server", _make_prompt_def())
replacement = MCPPromptWrapper(new_session, "test_server", _make_prompt_def())
async def reconnect(server_name: str, tool_name: str, stale_tool):
assert server_name == "test_server"
assert tool_name == "mcp_test_server_prompt_test_prompt"
assert stale_tool is wrapper
return replacement
wrapper.set_reconnect_handler(reconnect)
output = await wrapper.execute()
assert output == "fresh prompt"
assert old_session.get_prompt.call_count == 1
assert new_session.get_prompt.call_count == 1