universal-mcp 0.1.23rc1__py3-none-any.whl → 0.1.24rc2__py3-none-any.whl

universal_mcp/client/transport.py

@@ -0,0 +1,277 @@
+import os
+import webbrowser
+from contextlib import AsyncExitStack
+from typing import Any, Literal, Self
+
+from loguru import logger
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.auth import OAuthClientProvider
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.auth import OAuthClientMetadata
+from mcp.types import (
+    CallToolResult as MCPCallToolResult,
+)
+from mcp.types import (
+    Tool as MCPTool,
+)
+from openai.types.chat import ChatCompletionToolParam
+
+from universal_mcp.client.oauth import CallbackServer
+from universal_mcp.client.token_store import TokenStore
+from universal_mcp.config import ClientConfig, ClientTransportConfig
+from universal_mcp.stores.store import KeyringStore
+from universal_mcp.tools.adapters import transform_mcp_tool_to_openai_tool
+
+
+class ClientTransport:
+    """
+    Client for connecting to and interacting with a single MCP server.
+
+    Manages the lifecycle of a connection to an MCP server, handles various
+    transport mechanisms (stdio, sse, streamable_http), and facilitates
+    authentication, including OAuth 2.0 client flows. Allows listing tools
+    available on the server and calling them.
+    """
+
+    def __init__(self, name: str, config: ClientTransportConfig) -> None:
+        self.name: str = name
+        self.config: ClientTransportConfig = config
+        self.session: ClientSession | None = None
+        self.server_url: str = config.url
+
+        # Create OAuth authentication handler if needed
+        if self.server_url and not getattr(self.config, "headers", None):
+            # Set up callback server
+            self._callback_server = CallbackServer(port=3000)
+            self.store: KeyringStore | None = KeyringStore(self.name)
+            self.auth: OAuthClientProvider | None = OAuthClientProvider(
+                server_url="/".join(self.server_url.split("/")[:-1]),
+                client_metadata=OAuthClientMetadata.model_validate(self.client_metadata_dict),
+                storage=TokenStore(self.store),
+                redirect_handler=self._default_redirect_handler,
+                callback_handler=self._callback_handler,
+            )
+        else:
+            self._callback_server = None
+            self.store = None
+            self.auth = None
+
+    @property
+    def callback_server(self) -> CallbackServer:
+        if self._callback_server and not self._callback_server.is_running:
+            self._callback_server.start()
+        return self._callback_server
+
+    async def _callback_handler(self) -> tuple[str, str | None]:
+        """Handles the OAuth callback by waiting for and returning auth details."""
+        print("⏳ Waiting for authorization callback...")
+        try:
+            auth_code = self.callback_server.wait_for_callback(timeout=300)
+            return auth_code, self.callback_server.get_state()
+        finally:
+            self.callback_server.stop()
+
+    @property
+    def client_metadata_dict(self) -> dict[str, Any]:
+        """Provides OAuth 2.0 client metadata for registration or authentication."""
+        return {
+            "client_name": self.name,
+            "redirect_uris": [self.callback_server.redirect_uri],  # type: ignore
+            "grant_types": ["authorization_code", "refresh_token"],
+            "response_types": ["code"],
+            "token_endpoint_auth_method": "client_secret_post",
+        }
+
+    async def _default_redirect_handler(self, authorization_url: str) -> None:
+        """Default handler for OAuth redirects; opens URL in a web browser."""
+        print(f"Opening browser for authorization: {authorization_url}")
+        webbrowser.open(authorization_url)
+
+    async def initialize(self, exit_stack: AsyncExitStack) -> None:
+        """
+        Establishes and initializes the connection to the MCP server.
+
+        Raises:
+            ValueError: If the transport type is unknown or if required
+                        configuration for a transport is missing.
+        """
+        transport = getattr(self.config, "transport", None)
+        session = None
+        try:
+            if transport == "stdio":
+                command = self.config.get("command")
+                if not command:
+                    raise ValueError("The command must be a valid string and cannot be None.")
+
+                server_params = StdioServerParameters(
+                    command=command,
+                    args=self.config.get("args", []),
+                    env={**os.environ, **self.config.get("env", {})} if self.config.get("env") else None,
+                )
+                stdio_transport = await exit_stack.enter_async_context(stdio_client(server_params))
+                read, write = stdio_transport
+                session = await exit_stack.enter_async_context(ClientSession(read, write))
+                await session.initialize()
+            elif transport == "streamable_http":
+                url = self.config.get("url")
+                headers = self.config.get("headers", {})
+                if not url:
+                    raise ValueError("'url' must be provided for streamable_http transport.")
+                streamable_http_transport = await exit_stack.enter_async_context(
+                    streamablehttp_client(url=url, headers=headers, auth=self.auth)
+                )
+                read, write, _ = streamable_http_transport
+                session = await exit_stack.enter_async_context(ClientSession(read, write))
+                await session.initialize()
+            elif transport == "sse":
+                url = self.config.get("url")
+                headers = self.config.get("headers", {})
+                if not url:
+                    raise ValueError("'url' must be provided for sse transport.")
+                sse_transport = await exit_stack.enter_async_context(
+                    sse_client(url=url, headers=headers, auth=self.auth)
+                )
+                read, write = sse_transport
+                session = await exit_stack.enter_async_context(ClientSession(read, write))
+                await session.initialize()
+            else:
+                raise ValueError(f"Unknown transport: {transport}")
+            self.session = session
+        except Exception as e:
+            if session:
+                await session.aclose()
+            logger.error(f"Error initializing server {self.name}: {e}")
+            raise
+
+    async def list_tools(self) -> list[MCPTool]:
+        """Lists all tools available on the connected MCP server."""
+        if self.session:
+            try:
+                tools = await self.session.list_tools()
+                return list(tools.tools)
+            except Exception as e:
+                logger.warning(f"Failed to list tools for client {self.name}: {e}")
+        return []
+
+    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> MCPCallToolResult:
+        """Calls a specified tool on the connected MCP server with given arguments."""
+        if self.session:
+            try:
+                return await self.session.call_tool(tool_name, arguments)
+            except Exception as e:
+                logger.error(f"Error calling tool '{tool_name}' on client {self.name}: {e}")
+        return MCPCallToolResult(
+            content=[],
+            isError=True,
+        )
+
+
+class MultiClientTransport:
+    """
+    Aggregates multiple ClientTransport instances to act as a single MCP Server.
+
+    Provides a unified Server interface for a collection of ClientTransport
+    instances, each potentially connected to a different MCP server.
+    Maintains a mapping of tool names to the specific ClientTransport that
+    provides that tool.
+    """
+
+    def __init__(self, clients: dict[str, ClientTransportConfig]):
+        self.clients: list[ClientTransport] = [ClientTransport(name, config) for name, config in clients.items()]
+        self.tool_to_client: dict[str, ClientTransport] = {}
+        self._mcp_tools: list[MCPTool] = []
+        self._exit_stack: AsyncExitStack = AsyncExitStack()
+
+    @classmethod
+    def from_file(cls, path: str) -> Self:
+        mcp_config = ClientConfig.load_json_config(path)
+        return cls(mcp_config.mcpServers)
+
+    def save_to_file(self, path: str) -> None:
+        mcp_config = ClientConfig(mcpServers={name: config.model_dump() for name, config in self.clients.items()})
+        mcp_config.save_json_config(path)
+
+    async def add_client(self, name: str, config: ClientTransportConfig) -> None:
+        if name in self.tool_to_client:
+            logger.warning(f"Client {name} already exists. Skipping.")
+            return
+        self.clients.append(ClientTransport(name, config))
+        self.tool_to_client[name] = self.clients[-1]
+        logger.info(f"Added client: {name}")
+        await self._populate_tool_mapping()
+
+    async def remove_client(self, name: str) -> None:
+        if name not in self.tool_to_client:
+            logger.warning(f"Client {name} not found. Skipping.")
+            return
+        self.clients.remove(self.tool_to_client[name])
+        del self.tool_to_client[name]
+        logger.info(f"Removed client: {name}")
+        await self._populate_tool_mapping()
+
+    async def __aenter__(self):
+        for client in self.clients:
+            await client.initialize(self._exit_stack)
+        await self._populate_tool_mapping()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.clients.clear()
+        self.tool_to_client.clear()
+        self._mcp_tools.clear()
+        await self._exit_stack.aclose()
+
+    async def _populate_tool_mapping(self):
+        self.tool_to_client.clear()
+        self._mcp_tools.clear()
+        for client in self.clients:
+            try:
+                tools = await client.list_tools()
+                for tool in tools:
+                    tool_name = getattr(tool, "name", None)
+                    if tool_name:
+                        if tool_name not in self.tool_to_client:
+                            self._mcp_tools.append(tool)
+                            self.tool_to_client[tool_name] = client
+                            logger.info(f"Found tool: {tool_name} from client: {client.name}")
+                        else:
+                            logger.warning(
+                                f"Duplicate tool name '{tool_name}' found in client '{client.name}'. Skipping."
+                            )
+            except Exception as e:
+                logger.warning(f"Failed to list tools for client {client.name}: {e}")
+
+    async def list_tools(self, format: Literal["mcp", "openai"] = "mcp") -> list[MCPTool | ChatCompletionToolParam]:
+        """
+        Lists all unique tools available from all managed clients.
+
+        Args:
+            format: The desired format for the returned tools.
+
+        Returns:
+            List of tools in the specified format.
+
+        Raises:
+            ValueError: If an unsupported format is requested.
+        """
+        if format == "mcp":
+            return self._mcp_tools
+        elif format == "openai":
+            return [transform_mcp_tool_to_openai_tool(tool) for tool in self._mcp_tools]
+        else:
+            raise ValueError(f"Invalid format: {format}")
+
+    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> MCPCallToolResult:
+        """
+        Calls a tool by routing the request to the appropriate ClientTransport.
+
+        Raises:
+            KeyError: If the tool_name is not found.
+        """
+        client = self.tool_to_client.get(tool_name)
+        if not client:
+            logger.error(f"Tool '{tool_name}' not found in any client.")
+            return MCPCallToolResult(content=[], isError=True)
+        return await client.call_tool(tool_name, arguments)

universal_mcp/config.py

@@ -1,43 +1,89 @@
+import json
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, Self
 
-from pydantic import BaseModel, Field, SecretStr, field_validator
+from pydantic import BaseModel, Field, SecretStr, field_validator, model_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class StoreConfig(BaseModel):
-    """Configuration for credential storage."""
+    """Specifies the configuration for a credential or token store.
 
-    name: str = Field(default="universal_mcp", description="Name of the store")
+    Defines where and how sensitive information like API keys or OAuth tokens
+    should be stored and retrieved.
+    """
+
+    name: str = Field(
+        default="universal_mcp",
+        description="Name of the store service or context (e.g., 'my_app_tokens', 'global_api_keys').",
+    )
     type: Literal["memory", "environment", "keyring", "agentr"] = Field(
-        default="memory", description="Type of credential storage to use"
+        default="memory",
+        description="The type of storage backend to use. 'memory' is transient, 'environment' uses environment variables, 'keyring' uses the system's secure credential manager, 'agentr' delegates to AgentR platform storage.",
+    )
+    path: Path | None = Field(
+        default=None,
+        description="Filesystem path for store types that require it (e.g., a future 'file' store type). Currently not used by memory, environment, or keyring.",
     )
-    path: Path | None = Field(default=None, description="Path to store credentials (if applicable)")
 
 
 class IntegrationConfig(BaseModel):
-    """Configuration for API integrations."""
+    """Defines the authentication and credential management for an application.
 
-    name: str = Field(..., description="Name of the integration")
-    type: Literal["api_key", "oauth", "agentr", "oauth2"] = Field(
-        default="api_key", description="Type of authentication to use"
+    Specifies how a particular application (`AppConfig`) should authenticate
+    with its target service, including the authentication type (e.g., API key,
+    OAuth) and where to find the necessary credentials.
+    """
+
+    name: str = Field(
+        ..., description="A unique name for this integration instance (e.g., 'my_github_oauth', 'tavily_api_key')."
+    )
+    type: Literal["api_key", "oauth", "agentr", "oauth2", "basic_auth"] = Field(
+        default="api_key",
+        description="The authentication mechanism to be used. 'oauth2' is often synonymous with 'oauth'. 'agentr' implies AgentR platform managed authentication.",
+    )
+    credentials: dict[str, Any] | None = Field(
+        default=None,
+        description="Directly provided credentials, if not using a store. Structure depends on the integration type (e.g., {'api_key': 'value'} or {'client_id': 'id', 'client_secret': 'secret'}). Use with caution for sensitive data; prefer using a 'store'.",
+    )
+    store: StoreConfig | None = Field(
+        default=None,
+        description="Configuration for the credential store to be used for this integration, overriding any default server-level store.",
     )
-    credentials: dict[str, Any] | None = Field(default=None, description="Integration-specific credentials")
-    store: StoreConfig | None = Field(default=None, description="Store configuration for credentials")
 
 
 class AppConfig(BaseModel):
-    """Configuration for individual applications."""
+    """Configuration for a single application to be loaded by the MCP server.
 
-    name: str = Field(..., description="Name of the application")
-    integration: IntegrationConfig | None = Field(default=None, description="Integration configuration")
-    actions: list[str] | None = Field(default=None, description="List of available actions")
+    Defines an application's name (slug), its integration settings for
+    authentication, and optionally, a list of specific actions (tools)
+    it provides.
+    """
+
+    name: str = Field(
+        ...,
+        description="The unique name or slug of the application (e.g., 'github', 'google-calendar'). This is often used to dynamically load the application module.",
+    )
+    integration: IntegrationConfig | None = Field(
+        default=None,
+        description="Authentication and credential configuration for this application. If None, the application is assumed not to require authentication or uses a global/default mechanism.",
+    )
+    actions: list[str] | None = Field(
+        default=None,
+        description="A list of specific actions or tools provided by this application that should be exposed. If None or empty, all tools from the application might be exposed by default, depending on the application's implementation.",
+    )
 
 
 class ServerConfig(BaseSettings):
-    """Main server configuration."""
+    """Core configuration settings for the Universal MCP server.
 
-    model_config = SettingsConfigDict(
+    Manages server behavior, including its name, description, connection
+    to AgentR (if applicable), transport protocol, network settings (port/host),
+    applications to load, default credential store, and logging verbosity.
+    Settings can be loaded from environment variables or a .env file.
+    """
+
+    model_config: SettingsConfigDict = SettingsConfigDict(
         env_file=".env",
         env_file_encoding="utf-8",
         case_sensitive=True,
@@ -45,21 +91,50 @@ class ServerConfig(BaseSettings):
     )
 
     name: str = Field(default="Universal MCP", description="Name of the MCP server")
-    description: str = Field(default="Universal MCP", description="Description of the MCP server")
-    api_key: SecretStr | None = Field(default=None, description="API key for authentication", alias="AGENTR_API_KEY")
-    type: Literal["local", "agentr"] = Field(default="agentr", description="Type of server deployment")
+    description: str = Field(
+        default="Universal MCP", description="A brief description of this MCP server's purpose or deployment."
+    )
+    type: Literal["local", "agentr", "other"] = Field(
+        default="agentr",
+        description="Deployment type of the server. 'local' runs apps defined in 'apps' list; 'agentr' dynamically loads apps from the AgentR platform.",
+    )
     transport: Literal["stdio", "sse", "streamable-http"] = Field(
-        default="stdio", description="Transport protocol to use"
+        default="stdio",
+        description="The communication protocol the server will use to interact with clients (e.g., an AI agent).",
+    )
+    port: int = Field(
+        default=8005,
+        description="Network port for 'sse' or 'streamable-http' transports. Must be between 1 and 65535.",
+        ge=1,
+        le=65535,
+    )
+    log_level: str = Field(
+        default="INFO", description="Logging level for the server (e.g., DEBUG, INFO, WARNING, ERROR, CRITICAL)."
+    )
+    # AgentR specific settings
+    base_url: str = Field(
+        default="https://api.agentr.dev",
+        description="The base URL for the AgentR API, used when type is 'agentr' or for AgentR-mediated integrations.",
+        alias="AGENTR_BASE_URL",
+    )
+    api_key: SecretStr | None = Field(
+        default=None,
+        description="The API key for authenticating with the AgentR platform. Stored as a SecretStr for security.",
+        alias="AGENTR_API_KEY",
+    )
+    # Local specific settings
+    apps: list[AppConfig] | None = Field(
+        default=None,
+        description="A list of application configurations to load when server 'type' is 'local'. Ignored if 'type' is 'agentr'.",
+    )
+    store: StoreConfig | None = Field(
+        default=None,
+        description="Default credential store configuration for applications that do not define their own specific store.",
     )
-    port: int = Field(default=8005, description="Port to run the server on (if applicable)", ge=1024, le=65535)
-    host: str = Field(default="localhost", description="Host to bind the server to (if applicable)")
-    apps: list[AppConfig] | None = Field(default=None, description="List of configured applications")
-    store: StoreConfig | None = Field(default=None, description="Default store configuration")
-    debug: bool = Field(default=False, description="Enable debug mode")
-    log_level: str = Field(default="INFO", description="Logging level")
 
     @field_validator("log_level", mode="before")
     def validate_log_level(cls, v: str) -> str:
+        """Validates and normalizes the log_level field."""
         valid_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
         if v.upper() not in valid_levels:
             raise ValueError(f"Invalid log level. Must be one of: {', '.join(valid_levels)}")
@@ -67,6 +142,104 @@ class ServerConfig(BaseSettings):
 
     @field_validator("port", mode="before")
     def validate_port(cls, v: int) -> int:
+        """Validates the port number is within the valid range."""
         if not 1 <= v <= 65535:
             raise ValueError("Port must be between 1 and 65535")
         return v
+
+    @classmethod
+    def load_json_config(cls, path: str = "local_config.json") -> Self:
+        """Loads server configuration from a JSON file.
+
+        Args:
+            path (str, optional): The path to the JSON configuration file.
+                Defaults to "local_config.json".
+
+        Returns:
+            ServerConfig: An instance of ServerConfig populated with data
+                          from the JSON file.
+        """
+        with open(path) as f:
+            data = json.load(f)
+        return cls.model_validate(data)
+
+
+class ClientTransportConfig(BaseModel):
+    """Configuration for how an MCP client connects to an MCP server.
+
+    Specifies the transport protocol and its associated parameters, such as
+    the command for stdio, URL for HTTP-based transports (SSE, streamable_http),
+    and any necessary headers or environment variables.
+    """
+
+    transport: str | None = Field(
+        default=None,
+        description="The transport protocol (e.g., 'stdio', 'sse', 'streamable_http'). Auto-detected in model_validate if not set.",
+    )
+    command: str | None = Field(
+        default=None, description="The command to execute for 'stdio' transport (e.g., 'python -m mcp_server.run')."
+    )
+    args: list[str] = Field(default=[], description="List of arguments for the 'stdio' command.")
+    env: dict[str, str] = Field(default={}, description="Environment variables to set for the 'stdio' command.")
+    url: str | None = Field(default=None, description="The URL for 'sse' or 'streamable_http' transport.")
+    headers: dict[str, str] = Field(
+        default={}, description="HTTP headers to include for 'sse' or 'streamable_http' transport."
+    )
+
+    @model_validator(mode="after")
+    def determine_transport_if_not_set(self) -> Self:
+        """Determines and sets the transport type if not explicitly provided.
+
+        - If `command` is present, transport is set to 'stdio'.
+        - If `url` is present, transport is 'streamable_http' if URL ends with '/mcp',
+          otherwise 'sse' if URL ends with '/sse'.
+        - Raises ValueError if transport cannot be determined or if neither
+          `command` nor `url` is provided.
+        """
+        if self.command:
+            self.transport = "stdio"
+        elif self.url:
+            # Remove search params from url
+            url = self.url.split("?")[0]
+            if url.rstrip("/").endswith("mcp"):
+                self.transport = "streamable_http"
+            elif url.rstrip("/").endswith("sse"):
+                self.transport = "sse"
+            else:
+                raise ValueError(f"Unknown transport: {self.url}")
+        else:
+            raise ValueError("Either command or url must be provided")
+        return self
+
+
+class ClientConfig(BaseSettings):
+    """Configuration for a client application that interacts with MCP servers and an LLM.
+
+    Defines connections to one or more MCP servers (via `mcpServers`) and
+    optionally, settings for an LLM to be used by the client (e.g., by an agent).
+    """
+
+    mcpServers: dict[str, ClientTransportConfig] = Field(
+        ...,
+        description="A dictionary where keys are descriptive names for MCP server connections and values are `ClientTransportConfig` objects defining how to connect to each server.",
+    )
+
+    @classmethod
+    def load_json_config(cls, path: str = "servers.json") -> Self:
+        """Loads client configuration from a JSON file.
+
+        Args:
+            path (str, optional): The path to the JSON configuration file.
+                Defaults to "servers.json".
+
+        Returns:
+            ClientConfig: An instance of ClientConfig populated with data
+                          from the JSON file.
+        """
+        with open(path) as f:
+            data = json.load(f)
+        return cls.model_validate(data)
+
+    def save_json_config(self, path: str) -> None:
+        with open(path, "w") as f:
+            json.dump(self.model_dump(), f, indent=4)

universal_mcp/exceptions.py

@@ -1,25 +1,69 @@
 class NotAuthorizedError(Exception):
-    """Raised when a user is not authorized to access a resource or perform an action."""
+    """Raised when an action is attempted without necessary permissions.
+
+    This typically occurs if a user or process tries to access a protected
+    resource or perform an operation for which they lack the required
+    authorization credentials or roles.
+    """
 
     def __init__(self, message: str):
+        """Initializes the NotAuthorizedError.
+
+        Args:
+            message (str): A descriptive message explaining the authorization failure.
+        """
         self.message = message
+        super().__init__(message)  # Ensure message is passed to base Exception
 
 
 class ToolError(Exception):
-    """Raised when a tool is not found or fails to execute."""
+    """Indicates an issue related to tool discovery, validation, or execution.
+
+    This could be due to a tool not being found, failing during its
+    operation, or having invalid configuration or arguments.
+    """
+
+    pass
 
 
 class InvalidSignature(Exception):
-    """Raised when a signature is invalid."""
+    """Raised when a cryptographic signature verification fails.
+
+    This can occur during webhook validation or any other process that
+    relies on verifying the authenticity and integrity of a message
+    using a digital signature.
+    """
+
+    pass
 
 
 class StoreError(Exception):
-    """Base exception class for store-related errors."""
+    """Base exception for errors related to data or credential stores.
+
+    This serves as a generic error for issues arising from operations
+    on any storage backend (e.g., KeyringStore, EnvironmentStore).
+    Specific store errors should ideally subclass this.
+    """
+
+    pass
 
 
 class KeyNotFoundError(StoreError):
-    """Exception raised when a key is not found in the store."""
+    """Raised when a specified key cannot be found in a data or credential store.
+
+    This is a common error when attempting to retrieve a piece of data
+    (e.g., an API key, token, or client information) that does not exist
+    under the given identifier.
+    """
+
+    pass
 
 
 class ConfigurationError(Exception):
-    """Exception raised when a configuration error occurs."""
+    """Indicates an error was detected in application or server configuration.
+
+    This can be due to missing required settings, invalid values for
+    configuration parameters, or inconsistencies in the provided setup.
+    """
+
+    pass

universal_mcp/integrations/__init__.py

@@ -12,10 +12,7 @@ def integration_from_config(config: IntegrationConfig, store: BaseStore | None =
     if config.type == "api_key":
         return ApiKeyIntegration(config.name, store=store, **kwargs)
     elif config.type == "agentr":
-        api_key = kwargs.get("api_key")
-        if not api_key:
-            raise ValueError("api_key is required for AgentR integration")
-        return AgentRIntegration(config.name, api_key=api_key)
+        return AgentRIntegration(config.name, **kwargs)
     else:
         raise ValueError(f"Unsupported integration type: {config.type}")