utim-cli 1.0.0__py3-none-any.whl

utim_cli/mcp_client.py

@@ -0,0 +1,198 @@
+import os
+import json
+import asyncio
+import threading
+from typing import Dict, List, Any, Optional
+
+try:
+    from mcp import ClientSession, StdioServerParameters
+    from mcp.client.stdio import stdio_client
+except ImportError:
+    pass
+
+class MCPManager:
+    def __init__(self):
+        self.sessions: Dict[str, ClientSession] = {}
+        self._exit_stacks = {}
+        self._tool_to_server: Dict[str, str] = {}
+        self._started = False
+        self.cached_tools: List[Dict[str, Any]] = []
+        self.server_tools: Dict[str, List[str]] = {}  # server_name -> list of tool names
+        self._loop = None
+        self._thread = None
+
+    def start_loop(self):
+        if self._loop is None:
+            self._loop = asyncio.new_event_loop()
+            self._thread = threading.Thread(target=self._run_loop, daemon=True)
+            self._thread.start()
+
+    def _run_loop(self):
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def run_coro(self, coro):
+        self.start_loop()
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        return future.result()
+
+    def start(self):
+        if self._started:
+            return
+        
+        # Check if mcp library is available
+        try:
+            from mcp import ClientSession, StdioServerParameters
+            from mcp.client.stdio import stdio_client
+        except ImportError:
+            # Create config if it doesn't exist, just in case
+            config_path = os.path.abspath(os.path.join(".utim", "mcp.json"))
+            if not os.path.exists(config_path):
+                os.makedirs(os.path.dirname(config_path), exist_ok=True)
+                with open(config_path, "w") as f:
+                    json.dump({"mcpServers": {}}, f, indent=2)
+            self._started = True
+            return
+
+        self.start_loop()
+        self.run_coro(self._start_async())
+
+    async def _start_async(self):
+        config_path = os.path.abspath(os.path.join(".utim", "mcp.json"))
+        if not os.path.exists(config_path):
+            # Create dummy config
+            os.makedirs(os.path.dirname(config_path), exist_ok=True)
+            with open(config_path, "w") as f:
+                json.dump({"mcpServers": {}}, f, indent=2)
+            self._started = True
+            return
+
+        try:
+            with open(config_path, "r") as f:
+                config = json.load(f)
+        except Exception:
+            self._started = True
+            return
+
+        servers = config.get("mcpServers", {})
+        
+        for name, srv_config in servers.items():
+            command = srv_config.get("command")
+            args = srv_config.get("args", [])
+            env = srv_config.get("env", {})
+            
+            if not command:
+                continue
+
+            # Strip outer quotes from argument strings if they exist (common config error for space-containing paths)
+            cleaned_args = []
+            for arg in args:
+                if isinstance(arg, str):
+                    if (arg.startswith('"') and arg.endswith('"')) or (arg.startswith("'") and arg.endswith("'")):
+                        arg = arg[1:-1]
+                cleaned_args.append(arg)
+
+            try:
+                # Merge current env
+                merged_env = os.environ.copy()
+                merged_env.update(env)
+                
+                # Setup StdioServerParameters with stdout wrapper to prevent JSON-RPC pollution
+                import sys
+                import shutil
+                current_dir = os.path.dirname(os.path.abspath(__file__))
+                wrapper_path = os.path.join(current_dir, "mcp_clean_wrapper.py")
+                resolved_command = shutil.which(command) or command
+                
+                server_params = StdioServerParameters(
+                    command=sys.executable,
+                    args=["-u", wrapper_path, resolved_command] + cleaned_args,
+                    env=merged_env
+                )
+                
+                from contextlib import AsyncExitStack
+                stack = AsyncExitStack()
+                self._exit_stacks[name] = stack
+                
+                read, write = await stack.enter_async_context(stdio_client(server_params))
+                session = await stack.enter_async_context(ClientSession(read, write))
+                await session.initialize()
+                
+                self.sessions[name] = session
+                
+                # Retrieve tools from the server and cache them
+                try:
+                    res = await session.list_tools()
+                    self.server_tools[name] = []
+                    for t in res.tools:
+                        full_name = f"{name}__{t.name}"
+                        self._tool_to_server[full_name] = name
+                        self.server_tools[name].append(t.name)
+                        
+                        # Convert MCP tool schema to OpenAI/OpenRouter format
+                        self.cached_tools.append({
+                            "type": "function",
+                            "function": {
+                                "name": full_name,
+                                "description": f"[{name}] {t.description}",
+                                "parameters": t.inputSchema
+                            }
+                        })
+                except Exception as e:
+                    print(f"Error fetching tools from MCP server {name}: {e}")
+                    
+            except Exception as e:
+                import traceback
+                print(f"Error starting MCP server {name}: {e}")
+                
+        self._started = True
+
+    def get_tools(self) -> List[Dict[str, Any]]:
+        """Synchronously returns the cached list of MCP tools."""
+        return self.cached_tools
+
+    def get_notification_context(self) -> str:
+        """Returns the context notification string for connected MCP servers."""
+        notifications = []
+        for server_name, tool_names in self.server_tools.items():
+            if tool_names:
+                tools_str = ", ".join(tool_names)
+                notifications.append(f"this mcp server {server_name} is connected and a new set of tools are available for u: {tools_str}")
+        return "\n".join(notifications)
+
+    def call_tool(self, server_name: str, tool_name: str, arguments: Dict[str, Any]) -> str:
+        """Synchronously routes tool call execution to the background thread."""
+        return self.run_coro(self._call_tool_async(server_name, tool_name, arguments))
+
+    async def _call_tool_async(self, server_name: str, tool_name: str, arguments: Dict[str, Any]) -> str:
+        if server_name not in self.sessions:
+            return f"Error: MCP server '{server_name}' not running."
+            
+        session = self.sessions[server_name]
+        try:
+            res = await session.call_tool(tool_name, arguments)
+            if getattr(res, "isError", False):
+                return f"Error from {server_name}: " + "\n".join(c.text for c in res.content)
+            return "\n".join(c.text for c in res.content)
+        except Exception as e:
+            return f"Error calling {tool_name} on {server_name}: {str(e)}"
+
+    def restart(self):
+        """Synchronously restarts and reloads all MCP sessions."""
+        self.run_coro(self._restart_async())
+
+    async def _restart_async(self):
+        for name, stack in list(self._exit_stacks.items()):
+            try:
+                await stack.aclose()
+            except Exception:
+                pass
+        self.sessions.clear()
+        self._exit_stacks.clear()
+        self._tool_to_server.clear()
+        self.cached_tools.clear()
+        self.server_tools.clear()
+        self._started = False
+        await self._start_async()
+
+mcp_manager = MCPManager()