fix(mcp): reuse initialized client sessions

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):

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()