From 61ee657a681c40bce5c4b2a93de24c81a0d5d81b Mon Sep 17 00:00:00 2001 From: longcoding Date: Thu, 23 Jul 2026 17:31:13 +0800 Subject: [PATCH] fix(mcp): reuse initialized client sessions --- src/kimi_cli/soul/toolset.py | 29 +++++-- tests/core/test_mcp_stdio_lifecycle.py | 105 +++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 7 deletions(-) create mode 100644 tests/core/test_mcp_stdio_lifecycle.py diff --git a/src/kimi_cli/soul/toolset.py b/src/kimi_cli/soul/toolset.py index 5d66344aaa..8b65aa904f 100644 --- a/src/kimi_cli/soul/toolset.py +++ b/src/kimi_cli/soul/toolset.py @@ -778,7 +778,7 @@ def _mark_oauth_unauthorized(server_name: str) -> None: server_name=server_name, ) self._mcp_servers[server_name] = MCPServerInfo( - status="unauthorized", client=None, tools=[] + status="unauthorized", client=None, tools=[], connection_stack=None ) async def _connect_server( @@ -790,11 +790,13 @@ async def _connect_server( server_info.status = "connecting" try: assert server_info.client is not None - async with server_info.client as client: - for tool in await client.list_tools(): - server_info.tools.append( - MCPTool(server_name, tool, client, runtime=runtime) - ) + # Keep the initialized session alive for the toolset lifetime. Re-entering a + # keep-alive stdio transport would send initialize again to the same process. + connection_stack = contextlib.AsyncExitStack() + server_info.connection_stack = connection_stack + client = await connection_stack.enter_async_context(server_info.client) + for tool in await client.list_tools(): + server_info.tools.append(MCPTool(server_name, tool, client, runtime=runtime)) for tool in server_info.tools: self.add(tool) @@ -803,6 +805,9 @@ async def _connect_server( logger.info("Connected MCP server: {server_name}", server_name=server_name) return server_name, None except Exception as e: + if server_info.connection_stack is not None: + await server_info.connection_stack.aclose() + server_info.connection_stack = None logger.error( "Failed to connect MCP server: {server_name}, error: {error}", server_name=server_name, @@ -859,7 +864,10 @@ async def _connect(): client = fastmcp.Client(MCPConfig(mcpServers={server_name: server_config})) self._mcp_servers[server_name] = MCPServerInfo( - status="pending", client=client, tools=[] + status="pending", + client=client, + tools=[], + connection_stack=None, ) if in_background: @@ -890,6 +898,12 @@ async def cleanup(self) -> None: with contextlib.suppress(Exception, asyncio.CancelledError): await self._mcp_loading_task for server_info in self._mcp_servers.values(): + if server_info.connection_stack is not None: + try: + await server_info.connection_stack.aclose() + except Exception: + logger.warning("Failed to close MCP connection", exc_info=True) + server_info.connection_stack = None if server_info.client is not None: try: await server_info.client.close() @@ -902,6 +916,7 @@ class MCPServerInfo: status: Literal["pending", "connecting", "connected", "failed", "unauthorized"] client: fastmcp.Client[Any] | None tools: list[MCPTool[Any]] + connection_stack: contextlib.AsyncExitStack | None class MCPTool[T: ClientTransport](CallableTool): diff --git a/tests/core/test_mcp_stdio_lifecycle.py b/tests/core/test_mcp_stdio_lifecycle.py new file mode 100644 index 0000000000..4bf86d5ec9 --- /dev/null +++ b/tests/core/test_mcp_stdio_lifecycle.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock + +from fastmcp.mcp_config import MCPConfig, StdioMCPServer + +from kimi_cli.soul.toolset import KimiToolset + +_STRICT_MCP_SERVER = r""" +import json +import pathlib +import sys + +counter_path = pathlib.Path(sys.argv[1]) +initialize_count = 0 + +for line in sys.stdin: + request = json.loads(line) + method = request.get("method") + request_id = request.get("id") + if method == "initialize": + initialize_count += 1 + counter_path.write_text(str(initialize_count), encoding="utf-8") + if initialize_count > 1: + response = { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32600, "message": "Server is already initialized"}, + } + else: + response = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "protocolVersion": request["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "strict-test-server", "version": "1.0"}, + }, + } + elif method == "tools/list": + response = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "tools": [{ + "name": "echo", + "description": "Echo text", + "inputSchema": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + }, + }], + }, + } + elif method == "tools/call": + response = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "content": [{"type": "text", "text": request["params"]["arguments"]["text"]}], + "isError": False, + }, + } + else: + continue + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() +""" + + +async def test_stdio_mcp_client_reuses_its_initialized_session(tmp_path: Path) -> None: + server = tmp_path / "strict_mcp_server.py" + server.write_text(_STRICT_MCP_SERVER, encoding="utf-8") + initialize_counter = tmp_path / "initialize-count.txt" + config = MCPConfig( + mcpServers={ + "strict": StdioMCPServer( + command=sys.executable, + args=[str(server), str(initialize_counter)], + ) + } + ) + + runtime = MagicMock() + runtime.config.mcp.client.tool_call_timeout_ms = 5_000 + runtime.approval.request = AsyncMock(return_value=MagicMock()) + toolset = KimiToolset() + + try: + await toolset.load_mcp_tools([config], runtime, in_background=False) + tool = toolset.find("echo") + assert tool is not None + + first = await cast(Any, tool)(text="first") + second = await cast(Any, tool)(text="second") + + assert first.is_error is False + assert second.is_error is False + assert initialize_counter.read_text(encoding="utf-8") == "1" + finally: + await toolset.cleanup()