py-mcpdock-cli 1.0.13__py3-none-any.whl → 1.0.18__py3-none-any.whl

cli/runners/stdio_runner.py

@@ -7,327 +7,28 @@ through standard input/output pipes.
 import json
 import asyncio
 from logging import error
+import mcp.types as types
 import signal
 import sys
 import os
 import anyio
+import threading
 from typing import Dict, Any, Optional, Awaitable
-from rich import print as rprint
 from contextlib import AsyncExitStack
-from pydantic import BaseModel
+
 from ..utils.logger import verbose
 from ..types.registry import RegistryServer
 from ..utils.runtime import get_runtime_environment
-from mcp import StdioServerParameters
+from mcp import StdioServerParameters, ClientSession
 from mcp.client.stdio import stdio_client
-from mcp.types import ClientRequest, ServerRequest, JSONRPCRequest, JSONRPCMessage
-
-
-# 定义一个通用的 Model，用于接收任何 JSON 响应
-class DictModel(BaseModel):
-    @classmethod
-    def model_validate(cls, value):
-        if isinstance(value, dict):
-            return value
-        return dict(value)
-
-
-# 处理客户端请求
-async def process_client_request(message, session):
-    """处理从客户端接收的请求并转发到服务器"""
-    original_id = message.get("id")
-    method = message.get("method")
-    params = message.get("params", {})
-
-    # 添加特殊处理，记录初始化请求的详细信息
-    if method == "initialize":
-        verbose(f"[Runner] 收到初始化请求 ID: {original_id}, 详细内容: {json.dumps(message)}")
-
-        # 对初始化请求使用SDK的initialize方法，而不是简单转发
-        try:
-            verbose("[Runner] 使用SDK的initialize方法处理初始化请求")
-            # 直接调用session.initialize()获取规范的响应
-            init_result = await session.initialize()
-            verbose(f"[Runner] SDK初始化原始响应类型: {type(init_result)}")
-            verbose(f"[Runner] SDK初始化原始响应属性: {dir(init_result)}")
-
-            # 检查响应内容的具体格式
-            if hasattr(init_result, "model_dump"):
-                result_content = init_result.model_dump()
-                verbose(f"[Runner] 使用model_dump()解析结果: {json.dumps(result_content)[:200]}...")
-            elif hasattr(init_result, "dict"):
-                result_content = init_result.dict()
-                verbose(f"[Runner] 使用dict()解析结果: {json.dumps(result_content)[:200]}...")
-            else:
-                result_content = init_result
-                verbose(
-                    f"[Runner] 直接使用结果: {json.dumps(result_content)[:200] if isinstance(result_content, dict) else str(result_content)[:200]}...")
-
-            # 构造完整的JSON-RPC响应
-            response = json.dumps({
-                "jsonrpc": "2.0",
-                "id": original_id,
-                "result": result_content
-            })
-
-            verbose(f"[Runner] 通过SDK获取到规范的初始化响应: {response[:200]}...")
-
-            # 再额外发送一个initialized通知给上游客户端吗？
-            # 通常客户端接收到initialize响应后，会自己发送initialized通知
-            # 所以这里不需要主动发送
-
-            return response
-        except Exception as e:
-            error(f"[Runner] SDK初始化失败: {str(e)}")
-            import traceback
-            error(f"[Runner] 异常堆栈: {traceback.format_exc()}")
-            # 如果SDK初始化失败，回退到常规请求处理
-            verbose("[Runner] 回退到常规请求处理方式")
-
-    verbose(f"[stdin] Processing request with id: {original_id}, method: {method}")
-
-    # 常规请求处理：确定请求类型和构建请求对象
-    req_obj = create_request_object(message, method)
-
-    # 发送请求并处理响应
-    try:
-        verbose(f"[Runner] 向下游服务器发送请求，method: {method}, id: {original_id}")
-        result = await send_request_with_timeout(session, req_obj, original_id)
-
-        if method == "initialize":
-            verbose(f"[Runner] 收到初始化响应: {result}")
-
-        return result
-    except Exception as e:
-        error(f"[Runner] 请求处理异常 ({method}): {str(e)}")
-        raise
-
-
-# 创建请求对象
-def create_request_object(message, method):
-    """根据方法类型创建适当的请求对象"""
-    # 作为代理，我们不需要严格验证方法是否符合标准列表
-    # 直接创建ClientRequest对象，透明转发所有请求
-    msg = dict(message)
-    msg.pop("jsonrpc", None)  # 移除 jsonrpc 字段，SDK会自动添加
-    msg.pop("id", None)       # 移除 id 字段，我们会在响应中重新添加
-
-    try:
-        # 尝试创建 ClientRequest
-        return ClientRequest(**msg)
-    except Exception as e:
-        # 如果创建失败，回退到 ServerRequest
-        verbose(f"[Runner] 创建 ClientRequest 失败，回退到 ServerRequest: {str(e)}")
-        return ServerRequest(method=method, params=message.get("params", {}))
-
-
-async def send_request_with_timeout(session, req_obj, original_id, timeout_seconds=60):
-    """发送请求并处理超时和错误情况"""
-    try:
-        # 初始化 resp 为 None，防止超时时未定义
-        resp = None
-        # 使用超时机制
-        with anyio.move_on_after(timeout_seconds):
-            # 记录请求信息
-            verbose(f"[Runner] 发送请求，原始ID={original_id}, 方法={req_obj.method if hasattr(req_obj, 'method') else '未知'}")
-
-            # 发送请求并等待响应
-            resp = await session.send_request(req_obj, DictModel)
-
-        if resp:
-            # 直接使用原始ID构造响应
-            return json.dumps({
-                "id": original_id,
-                "jsonrpc": "2.0",
-                "result": resp
-            })
-        else:
-            return json.dumps({
-                "id": original_id,
-                "jsonrpc": "2.0",
-                "error": {
-                    "code": -32000,
-                    "message": "Request timed out or empty response"
-                }
-            })
-
-    except Exception as e:
-        # 处理请求错误
-        error_msg = str(e)
-        error_code = -32603
-
-        # 分类错误类型
-        if "timed out" in error_msg.lower():
-            error_code = -32001
-        elif "connection" in error_msg.lower():
-            error_code = -32002
-
-        return json.dumps({
-            "id": original_id,
-            "jsonrpc": "2.0",
-            "error": {
-                "code": error_code,
-                "message": f"Request failed: {error_msg}"
-            }
-        })
-
-
-# 初始化 MCP 会话
-async def initialize_session(session):
-    """初始化 MCP 协议会话，转发上游客户端的初始化请求到下游服务器"""
-    try:
-        # 关键点: 作为代理，我们不应该主动调用 session.initialize()
-        # 上游客户端会发送初始化请求，我们应该在 handle_stdin 函数中处理
-        verbose("[Runner] MCP 代理准备就绪，等待上游客户端的初始化请求...")
-        return True
-    except Exception as init_error:
-        error_msg = f"代理初始化失败: {str(init_error)}"
-        error(f"[Runner] {error_msg}")
-        return False
-
-
-# 处理来自标准输入的消息
-async def handle_stdin(session, is_shutting_down):
-    """处理从标准输入接收的消息并转发到服务器"""
-    loop = asyncio.get_event_loop()
-
-    while not is_shutting_down:
-        line = await loop.run_in_executor(None, sys.stdin.readline)
-        if not line:
-            break
-
-        try:
-            message = json.loads(line)
-            verbose(f"[stdin] Received message: {line.strip()}")
-
-            method = message.get("method", "")
-            # 根据消息类型处理
-            if "id" in message:  # 这是请求，需要响应
-                response = await process_client_request(message, session)
-                sys.stdout.write(response + "\n")
-                sys.stdout.flush()
-                verbose(f"[stdin] Response sent for method: {method}")
-            else:  # 这是通知，不需要响应
-                await session.send(message)
-                verbose(f"[stdin] Notification sent for method: {method}")
-
-            verbose(f"[stdin] Processed: {line.strip()}")
-
-        except json.JSONDecodeError as e:
-            error(f"[stdin] JSON decode error: {e}")
-        except Exception as e:
-            error(f"[stdin] Error processing input: {e}")
-            # 如果是请求(有ID)，才需要发送错误响应
-            try:
-                if 'message' in locals() and isinstance(message, dict) and "id" in message:
-                    error_resp = json.dumps({
-                        "jsonrpc": "2.0",
-                        "id": message.get("id"),
-                        "error": {
-                            "code": -32700,
-                            "message": f"Parse error: {str(e)}"
-                        }
-                    })
-                    sys.stdout.write(error_resp + "\n")
-                    sys.stdout.flush()
-                    verbose(f"[stdin] Sent error response for parse error")
-            except Exception as err:
-                error(f"[stdin] Failed to send error response: {err}")
-
-
-# 处理单个服务器消息
-async def handle_single_server_message(data):
-    """处理单个从服务器接收的消息并输出到标准输出"""
-    try:
-        # 记录接收到的原始数据类型，帮助调试
-        verbose(f"[server_raw] Received data type: {type(data)}")
-
-        # 尝试获取原始数据的字符串表示用于调试
-        raw_data_str = str(data)
-        if len(raw_data_str) > 500:
-            raw_data_str = raw_data_str[:500] + "..."
-        verbose(f"[server_raw] 原始数据: {raw_data_str}")
-
-        # 根据数据类型进行处理
-        if hasattr(data, "model_dump"):
-            content = data.model_dump()
-            verbose(f"[server_raw] Processed pydantic v2 model: {type(data)}")
-        elif hasattr(data, "dict"):
-            content = data.dict()
-            verbose(f"[server_raw] Processed pydantic v1 model: {type(data)}")
-        elif isinstance(data, dict):
-            content = data
-            verbose(f"[server_raw] Processed dict with {len(data)} keys")
-        else:
-            # 尝试转换为字符串，然后解析为JSON
-            try:
-                content = json.loads(str(data))
-                verbose(f"[server_raw] Converted to JSON: {type(data)}")
-            except:
-                content = {"data": str(data)}
-                verbose(f"[server_raw] Used raw string for unknown type: {type(data)}")
-
-        # 检查是否是初始化响应
-        is_init_response = False
-        if isinstance(content, dict):
-            if "result" in content and isinstance(content["result"], dict):
-                result = content["result"]
-                if "protocolVersion" in result or "serverInfo" in result:
-                    is_init_response = True
-                    verbose(f"[server] 检测到初始化响应: {json.dumps(content)[:200]}...")
-
-                # 特别检查是否包含tools字段，这对于VSCode非常重要
-                if "tools" in result:
-                    verbose(f"[server] 检测到tools字段，工具数量: {len(result['tools'])}")
-
-        # 检查数据是否已经是标准的 JSON-RPC 消息
-        if isinstance(content, dict):
-            if "jsonrpc" in content and ("id" in content or "method" in content):
-                # 已经是标准格式，直接输出
-                output = json.dumps(content)
-                verbose(f"[server] Standard JSON-RPC message detected, id: {content.get('id')}")
-
-            elif "result" in content and not "jsonrpc" in content:
-                # 是结果但缺少 jsonrpc 字段，构造标准响应
-                output = json.dumps({
-                    "jsonrpc": "2.0",
-                    "id": 1,  # 默认ID，应该不会被用到
-                    "result": content["result"] if "result" in content else content
-                })
-                verbose(f"[server] Fixed response format by adding jsonrpc")
 
-            else:
-                # 其他类型的消息，包装为通知
-                output = json.dumps(content)
-                verbose("[server] Passing through data as-is")
-        else:
-            # 非字典类型，直接序列化
-            output = json.dumps(content)
-
-        # 写入 stdout 并立即刷新，确保 VS Code 能收到
-        sys.stdout.write(output + "\n")
-        sys.stdout.flush()
-        verbose(f"[server] Response sent to stdout: {output}")
-
-    except Exception as e:
-        error(f"[server] Error processing server message: {e}")
-        import traceback
-        error(f"[server] 异常堆栈: {traceback.format_exc()}")
-        # 尝试发送错误响应
-        try:
-            error_resp = json.dumps({
-                "jsonrpc": "2.0",
-                "id": 1,  # 使用默认ID
-                "error": {
-                    "code": -32603,
-                    "message": f"Internal error: {str(e)}"
-                }
-            })
-            sys.stdout.write(error_resp + "\n")
-            sys.stdout.flush()
-            verbose(f"[server] Sent error response due to: {e}")
-        except:
-            error("[server] Failed to send error response")
+# 导入工具模块
+from ..utils.process_utils import find_server_process_by_command, save_pid_to_file
+from ..utils.mcp_utils import (
+    process_client_request,
+    handle_single_server_message,
+    initialize_session
+)
 
 
 async def create_stdio_runner(
@@ -353,23 +54,54 @@ async def create_stdio_runner(
         verbose("[Runner] Starting cleanup...")
         is_shutting_down = True
         try:
+            # Remove PID by server name and client name
+            from ..utils.process_utils import remove_pid_by_server_name
+
+            # 尝试获取第三级进程作为客户端标识 - 与保存时使用相同的逻辑
+            client_name = "unknown_client"
+            try:
+                import psutil
+                current_process = psutil.Process()
+
+                # 收集完整的进程链
+                process_chain = []
+                process = current_process
+                while process:
+                    process_chain.append(process.name())
+                    process = process.parent()
+
+                # 打印完整进程链用于调试
+                verbose(f"[Runner] 清理时的完整进程链: {' -> '.join(reversed(process_chain))}")
+
+                # 尝试获取第四级进程（如果存在）
+                if len(process_chain) >= 3:
+                    client_name = process_chain[2]
+                    verbose(f"[Runner] 清理时使用第四级进程作为客户端: {client_name}")
+                else:
+                    # 如果进程链不够长，使用最后一个进程（顶层进程）
+                    client_name = process_chain[-1] if process_chain else "unknown_client"
+                    verbose(f"[Runner] 清理时进程链不够长，使用顶层进程作为客户端: {client_name}")
+            except Exception as e:
+                verbose(f"[Runner] 清理时获取客户端进程名称失败: {e}，使用默认值: {client_name}")
+
+            remove_pid_by_server_name(server_details.qualifiedName, client_name)
+            verbose(f"[Runner] 已从PID文件中移除服务器 '{server_details.qualifiedName}' 的客户端 '{client_name}' 记录")
+
             await exit_stack.aclose()
             verbose("[Runner] Resources closed successfully")
         except Exception as error:
             handle_error(error, "Error during cleanup")
         verbose("[Runner] Cleanup completed")
 
-    def handle_sigint(sig, frame):
+    async def handle_sigint():
         verbose("[Runner] Received interrupt signal, shutting down...")
-        asyncio.create_task(cleanup())
+        await cleanup()
         # 立即打印一条确认消息，让用户知道CTRL+C已被捕获
-        print("\n[CTRL+C] 正在关闭服务，请稍候...", flush=True)
+        verbose("\n[CTRL+C] 正在关闭服务，请稍候...")
         # 可选：设置一个短暂的超时，然后强制退出
         import threading
         threading.Timer(2.0, lambda: os._exit(0)).start()
 
-    signal.signal(signal.SIGINT, handle_sigint)
-
     # 获取连接配置
     stdio_connection = next((conn for conn in server_details.connections if conn.type == "stdio"), None)
     if not stdio_connection:
@@ -404,13 +136,64 @@ async def create_stdio_runner(
         async with stdio_client(server_params, errlog=sys.stderr) as (read_stream, write_stream):
             verbose("Stdio proxy client connection established")
 
+            # 查找服务器进程ID - 优先使用父子进程关系
+            from ..utils.process_utils import find_server_process_by_command, save_pid_to_file, find_server_process_from_current
+
+            # 首先尝试通过父子进程关系查找
+            is_found_from_current, current_pid, cmd_str = find_server_process_from_current()
+            if is_found_from_current and current_pid:
+                verbose(f"[Runner] 通过父子进程关系找到服务器进程，PID: {current_pid}")
+                server_pid = current_pid
+            else:
+                # 如果父子进程关系查找失败，尝试通过命令行查找
+                server_pid = find_server_process_by_command(command, args)
+
+            if server_pid:
+                verbose(f"[Runner] 已找到MCP服务器进程，PID: {server_pid}")
+
+                # 将进程ID写入PID文件
+                # 尝试获取第四级进程作为客户端标识
+                client_name = "unknown_client"
+                try:
+                    import psutil
+                    current_process = psutil.Process()
+
+                    # 收集完整的进程链
+                    process_chain = []
+                    process = current_process
+                    while process:
+                        process_chain.append(process.name())
+                        process = process.parent()
+
+                    # 打印完整进程链用于调试
+                    verbose(f"[Runner] 完整进程链: {' -> '.join(reversed(process_chain))}")
+
+                    # 尝试获取第四级进程（如果存在）
+                    if len(process_chain) >= 3:
+                        client_name = process_chain[2]
+                        verbose(f"[Runner] 使用第三级进程作为客户端: {client_name}")
+                    else:
+                        # 如果进程链不够长，使用最后一个进程（顶层进程）
+                        client_name = process_chain[-1] if process_chain else "unknown_client"
+                        verbose(f"[Runner] 进程链不够长，使用顶层进程作为客户端: {client_name}")
+                except Exception as e:
+                    verbose(f"[Runner] 获取客户端进程名称失败: {e}，使用默认值: {client_name}")
+
+                verbose(f"[Runner] 使用 '{client_name}' 作为客户端标识")
+                pid_file_path = save_pid_to_file(server_pid, server_details.qualifiedName,
+                                                 client_name, command, args)
+                if pid_file_path:
+                    verbose(f"[Runner] 已将服务器PID信息安全地保存到: {pid_file_path}")
+            else:
+                verbose(f"[Runner] 未能找到MCP服务器进程，将无法获取其PID")
+
             # 创建 MCP 客户端会话
             from mcp import ClientSession
             session = await exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
 
             # 注册消息处理回调
             def handle_server_message(msg):
-                rprint(f"[magenta][server][/magenta] {json.dumps(msg, ensure_ascii=False)}")
+                verbose(f"[magenta][server][/magenta] {json.dumps(msg, ensure_ascii=False)}")
             session.on_message = handle_server_message
 
             # 初始化 MCP 协议
@@ -421,7 +204,11 @@ async def create_stdio_runner(
             verbose("[Runner] 开始处理循环，使用同步阻塞模式")
 
             # 打印启动消息
-            rprint("[cyan]MCP client running. Press Ctrl+C to stop.[/cyan]")
+            verbose("[cyan]MCP client running. Press Ctrl+C to stop.[/cyan]")
+
+            # 获取事件循环并在主任务中添加信号处理程序
+            loop = asyncio.get_event_loop()
+            loop.add_signal_handler(signal.SIGINT, lambda: asyncio.create_task(handle_sigint()))
 
             # 循环处理客户端请求，直到关闭
             while not is_shutting_down:
@@ -445,8 +232,13 @@ async def create_stdio_runner(
                         verbose(f"[stdin] Response sent for method: {method}")
                     else:  # 这是通知，不需要响应
                         # 创建通知对象并发送
-                        notification_obj = create_request_object(message, method)
-                        await session.send_notification(notification_obj)
+                        # notification_obj = create_request_object(message, method)
+                        # await session.send_notification(notification_obj)
+                        await session.send_notification(
+                            types.ClientNotification(
+                                types.InitializedNotification(method=method)
+                            )
+                        )
                         verbose(f"[stdin] Notification sent for method: {method}")
 
                     verbose(f"[stdin] Processed: {line.strip()}")
@@ -488,7 +280,7 @@ async def create_stdio_runner(
             verbose("[Runner] 处理循环结束")
 
     except Exception as e:
-        rprint(f"[red]Error running stdio proxy: {e}[/red]")
+        verbose(f"[red]Error running stdio proxy: {e}[/red]")
         raise
     finally:
         await cleanup()