tsugite-cli 0.3.3__py3-none-any.whl

tsugite/mcp_client.py

@@ -0,0 +1,219 @@
+"""MCP client using official Python SDK.
+
+Connects to MCP servers (stdio or HTTP) and loads their tools.
+"""
+
+from typing import List, Optional
+
+from mcp import ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+
+from tsugite.core.tools import Tool
+from tsugite.mcp_config import MCPServerConfig
+
+
+class MCPClient:
+    """Client for connecting to MCP servers.
+
+    Handles both stdio and HTTP transports. Use as async context manager.
+
+    Example:
+        # Create config
+        config = MCPServerConfig(
+            name="basic-memory",
+            command="uvx",
+            args=["basic-memory-mcp"]
+        )
+
+        # Connect and load tools (preferred - automatic cleanup)
+        async with MCPClient(config) as client:
+            tools = await client.get_tools()
+            # Use tools in agent
+            agent = TsugiteAgent(model="...", tools=tools)
+
+        # Or use connect/disconnect manually if needed
+        client = MCPClient(config)
+        await client.connect()
+        tools = await client.get_tools()
+        await client.disconnect()
+    """
+
+    def __init__(self, server_config: MCPServerConfig):
+        """Initialize MCP client.
+
+        Args:
+            server_config: Configuration for MCP server
+        """
+        self.config = server_config
+        self.session = None
+        self.session_ctx = None
+        self.mcp_tools = []  # Raw MCP tool objects
+        self.transport = None
+        self.transport_ctx = None
+
+    async def __aenter__(self):
+        """Enter async context manager - connect to server."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit async context manager - disconnect from server."""
+        await self.disconnect()
+        return False  # Don't suppress exceptions
+
+    async def connect(self):
+        """Connect to MCP server and initialize session.
+
+        Must be called before get_tools() if not using context manager.
+        """
+        if self.config.is_stdio():
+            # Connect via stdio (command-line server)
+            server_params = StdioServerParameters(
+                command=self.config.command,
+                args=self.config.args or [],
+                env=self.config.env,
+            )
+
+            # Start stdio transport using async with
+            self.transport_ctx = stdio_client(server_params)
+            # stdio_client returns a proper async context manager, but pylint can't introspect it
+            read, write = await self.transport_ctx.__aenter__()  # pylint: disable=no-member
+            self.transport = (read, write)
+
+        elif self.config.is_http():
+            # Connect via HTTP using async with
+            self.transport_ctx = streamablehttp_client(self.config.url)
+            # streamablehttp_client returns a proper async context manager, but pylint can't introspect it
+            read, write, _ = await self.transport_ctx.__aenter__()  # pylint: disable=no-member
+            self.transport = (read, write)
+        else:
+            raise ValueError(f"Unknown transport type: {self.config.type}")
+
+        # Create session using async with
+        self.session_ctx = ClientSession(read, write)
+        self.session = await self.session_ctx.__aenter__()
+
+        # Initialize connection
+        await self.session.initialize()
+
+        # Load available tools
+        response = await self.session.list_tools()
+        self.mcp_tools = response.tools
+
+    async def disconnect(self):
+        """Disconnect from MCP server and cleanup resources.
+
+        Closes the session and transport, preventing resource leaks.
+        Should be called when done using the client (unless using context manager).
+        """
+        # Exit session context manager
+        if self.session_ctx:
+            try:
+                await self.session_ctx.__aexit__(None, None, None)
+            except Exception:
+                pass  # Best effort cleanup
+            self.session_ctx = None
+            self.session = None
+
+        # Exit transport context manager
+        if self.transport_ctx:
+            try:
+                # Proper async context manager, but pylint can't introspect it
+                await self.transport_ctx.__aexit__(None, None, None)  # pylint: disable=no-member
+            except Exception:
+                pass  # Best effort cleanup
+            self.transport_ctx = None
+            self.transport = None
+
+    async def get_tools(self, allowed_tools: Optional[List[str]] = None) -> List[Tool]:
+        """Get tools from MCP server as Tool objects.
+
+        Args:
+            allowed_tools: If provided, only return these tools
+
+        Returns:
+            List[Tool]: Tools wrapped in our Tool interface
+        """
+        tools = []
+
+        for mcp_tool in self.mcp_tools:
+            # Filter if needed
+            if allowed_tools and mcp_tool.name not in allowed_tools:
+                continue
+
+            # Convert to our Tool format
+            tool = Tool(
+                name=mcp_tool.name,
+                description=mcp_tool.description or "",
+                parameters=mcp_tool.inputSchema or {},  # Already JSON schema!
+                function=self._create_tool_caller(mcp_tool.name),
+            )
+            tools.append(tool)
+
+        return tools
+
+    def _create_tool_caller(self, tool_name: str):
+        """Create async function that calls MCP tool.
+
+        Returns a function that agents can call.
+        """
+
+        async def call_mcp_tool(**kwargs):
+            """Call MCP tool and return result."""
+            result = await self.session.call_tool(tool_name, arguments=kwargs)
+
+            # Parse result content
+            # MCP returns list of content items
+            if result.content:
+                # Handle different content types
+                content_items = []
+                for item in result.content:
+                    if item.type == "text":
+                        content_items.append(item.text)
+                    elif item.type == "image":
+                        content_items.append(f"[Image: {item.data}]")
+                    # Add other types as needed
+
+                # Return combined content
+                return "\n".join(content_items) if content_items else None
+
+            return None
+
+        return call_mcp_tool
+
+
+async def load_mcp_tools(
+    server_config: MCPServerConfig, allowed_tools: Optional[List[str]] = None
+) -> tuple[MCPClient, List[Tool]]:
+    """Connect to MCP server and load tools.
+
+    Returns both the client and tools. The client must remain alive for tools
+    to work, so caller should call client.disconnect() when done with the tools.
+
+    Args:
+        server_config: MCP server configuration
+        allowed_tools: Optional list of tool names to load
+
+    Returns:
+        tuple[MCPClient, List[Tool]]: Client and tools from MCP server
+
+    Example:
+        from tsugite.mcp_config import load_mcp_config
+
+        # Load config
+        config = load_mcp_config()
+        server_config = config["basic-memory"]
+
+        # Load tools and keep client alive
+        client, tools = await load_mcp_tools(server_config)
+
+        # Use tools...
+
+        # Cleanup when done
+        await client.disconnect()
+    """
+    client = MCPClient(server_config)
+    await client.connect()
+    tools = await client.get_tools(allowed_tools)
+    return client, tools

tsugite/mcp_config.py

@@ -0,0 +1,174 @@
+"""MCP server configuration loading and management."""
+
+import json
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from .xdg import get_xdg_config_path, get_xdg_write_path
+
+
+class MCPServerConfig(BaseModel):
+    """Configuration for a single MCP server."""
+
+    model_config = ConfigDict(
+        extra="forbid",
+        str_strip_whitespace=True,
+    )
+
+    name: str
+    command: Optional[str] = None
+    args: Optional[List[str]] = None
+    env: Optional[Dict[str, str]] = None
+    url: Optional[str] = None
+    type: Optional[str] = None
+
+    @model_validator(mode="after")
+    def validate_server_config(self) -> "MCPServerConfig":
+        """Validate server configuration and auto-detect type."""
+        if self.type is None:
+            if self.command:
+                self.type = "stdio"
+            elif self.url:
+                self.type = "http"
+            else:
+                raise ValueError(f"Server '{self.name}' must have either 'command' or 'url' specified")
+
+        if self.type == "stdio" and not self.command:
+            raise ValueError(f"Stdio server '{self.name}' must have 'command' specified")
+
+        if self.type == "http" and not self.url:
+            raise ValueError(f"HTTP server '{self.name}' must have 'url' specified")
+
+        return self
+
+    def is_stdio(self) -> bool:
+        return self.type == "stdio"
+
+    def is_http(self) -> bool:
+        return self.type == "http"
+
+
+def get_default_config_path() -> Path:
+    return get_xdg_config_path("mcp.json")
+
+
+def get_config_path_for_write() -> Path:
+    return get_xdg_write_path("mcp.json")
+
+
+def load_mcp_config(path: Optional[Path] = None) -> Dict[str, MCPServerConfig]:
+    """Load MCP server configurations from JSON file.
+
+    Args:
+        path: Path to mcp.json file. If None, uses default ~/.tsugite/mcp.json
+
+    Returns:
+        Dictionary mapping server names to MCPServerConfig objects.
+        Returns empty dict if file doesn't exist or is invalid.
+    """
+    if path is None:
+        path = get_default_config_path()
+
+    if not path.exists():
+        return {}
+
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        if "mcpServers" not in data:
+            print(f"Warning: No 'mcpServers' key found in {path}")
+            return {}
+
+        servers = {}
+        for name, config in data["mcpServers"].items():
+            try:
+                servers[name] = MCPServerConfig(
+                    name=name,
+                    command=config.get("command"),
+                    args=config.get("args"),
+                    env=config.get("env"),
+                    url=config.get("url"),
+                    type=config.get("type"),
+                )
+            except ValueError as e:
+                print(f"Warning: Invalid config for server '{name}': {e}")
+                continue
+
+        return servers
+
+    except json.JSONDecodeError as e:
+        print(f"Warning: Failed to parse MCP config at {path}: {e}")
+        return {}
+    except Exception as e:
+        print(f"Warning: Failed to load MCP config from {path}: {e}")
+        return {}
+
+
+def save_mcp_config(servers: Dict[str, MCPServerConfig], path: Optional[Path] = None) -> None:
+    """Save MCP server configurations to JSON file.
+
+    Args:
+        servers: Dictionary mapping server names to MCPServerConfig objects
+        path: Path to mcp.json file. If None, uses appropriate config location
+
+    Raises:
+        IOError: If file cannot be written
+    """
+    if path is None:
+        path = get_config_path_for_write()
+
+    # Ensure directory exists
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Convert MCPServerConfig objects to JSON-serializable dicts
+    config_data = {"mcpServers": {}}
+
+    for name, server in servers.items():
+        # Use Pydantic's model_dump to serialize, excluding None values and the name field
+        server_dict = server.model_dump(exclude_none=True, exclude={"name"})
+
+        # For stdio servers, don't include 'type' field (it's implicit)
+        if server.is_stdio() and "type" in server_dict:
+            del server_dict["type"]
+
+        config_data["mcpServers"][name] = server_dict
+
+    # Write with pretty formatting
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(config_data, f, indent=2)
+
+
+def add_server_to_config(server: MCPServerConfig, path: Optional[Path] = None, overwrite: bool = False) -> bool:
+    """Add or update an MCP server in the configuration file.
+
+    Args:
+        server: MCPServerConfig object to add
+        path: Path to mcp.json file. If None, uses appropriate config location
+        overwrite: If True, overwrite existing server. If False, raise error if exists.
+
+    Returns:
+        True if server was added/updated successfully
+
+    Raises:
+        ValueError: If server already exists and overwrite is False
+    """
+    if path is None:
+        path = get_config_path_for_write()
+
+    # Load existing config
+    servers = load_mcp_config(path)
+
+    # Check if server already exists
+    if server.name in servers and not overwrite:
+        raise ValueError(f"Server '{server.name}' already exists. Use --force to overwrite.")
+
+    # Add/update server
+    servers[server.name] = server
+
+    # Save config
+    save_mcp_config(servers, path)
+
+    return True