universal-mcp 0.1.7rc1__py3-none-any.whl → 0.1.8__py3-none-any.whl

universal_mcp/servers/server.py

@@ -1,181 +1,236 @@
 import os
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from typing import Any
+from urllib.parse import urlparse
 
 import httpx
 from loguru import logger
 from mcp.server.fastmcp import FastMCP
-from mcp.server.fastmcp.exceptions import ToolError
 from mcp.types import TextContent
 
-from universal_mcp.applications import app_from_slug
-from universal_mcp.config import AppConfig, IntegrationConfig, StoreConfig
-from universal_mcp.exceptions import NotAuthorizedError
-from universal_mcp.integrations import AgentRIntegration, ApiKeyIntegration
-from universal_mcp.stores import store_from_config
+from universal_mcp.applications import Application, app_from_slug
+from universal_mcp.config import AppConfig, ServerConfig, StoreConfig
+from universal_mcp.integrations import AgentRIntegration, integration_from_config
+from universal_mcp.stores import BaseStore, store_from_config
+from universal_mcp.tools.tools import ToolManager
 
 
-class Server(FastMCP, ABC):
-    """
-    Server is responsible for managing the applications and the store
-    It also acts as a router for the applications, and exposed to the client
+class BaseServer(FastMCP, ABC):
+    """Base server class with common functionality.
+
+    This class provides core server functionality including store setup,
+    tool management, and application loading.
+
+    Args:
+        config: Server configuration
+        **kwargs: Additional keyword arguments passed to FastMCP
     """
 
-    def __init__(
-        self, name: str, description: str, store: StoreConfig | None = None, **kwargs
-    ):
-        super().__init__(name, description, **kwargs)
-        logger.info(f"Initializing server: {name} with store: {store}")
-        self.store = store_from_config(store) if store else None
-        self._setup_store(store)
-        self._load_apps()
+    def __init__(self, config: ServerConfig, **kwargs):
+        super().__init__(config.name, config.description, **kwargs)
+        logger.info(
+            f"Initializing server: {config.name} ({config.type}) with store: {config.store}"
+        )
 
-    def _setup_store(self, store_config: StoreConfig | None):
-        """
-        Setup the store for the server.
-        """
-        if store_config is None:
-            return
-        self.store = store_from_config(store_config)
-        self.add_tool(self.store.set)
-        self.add_tool(self.store.delete)
-        # self.add_tool(self.store.get)
+        self.config = config  # Store config at base level for consistency
+        self._tool_manager = ToolManager(warn_on_duplicate_tools=True)
 
     @abstractmethod
-    def _load_apps(self):
+    def _load_apps(self) -> None:
+        """Load and register applications."""
         pass
 
-    async def call_tool(self, name: str, arguments: dict[str, Any]):
-        """Call a tool by name with arguments."""
-        logger.info(f"Calling tool: {name} with arguments: {arguments}")
-        try:
-            result = await super().call_tool(name, arguments)
-            logger.info(f"Tool {name} completed successfully")
+    def add_tool(self, tool: Callable) -> None:
+        """Add a tool to the server.
+
+        Args:
+            tool: Tool to add
+        """
+        self._tool_manager.add_tool(tool)
+
+    async def list_tools(self) -> list[dict]:
+        """List all available tools in MCP format.
+
+        Returns:
+            List of tool definitions
+        """
+        return self._tool_manager.list_tools(format="mcp")
+
+    def _format_tool_result(self, result: Any) -> list[TextContent]:
+        """Format tool result into TextContent list.
+
+        Args:
+            result: Raw tool result
+
+        Returns:
+            List of TextContent objects
+        """
+        if isinstance(result, str):
+            return [TextContent(type="text", text=result)]
+        elif isinstance(result, list) and all(
+            isinstance(item, TextContent) for item in result
+        ):
             return result
-        except ToolError as e:
-            raised_error = e.__cause__
-            if isinstance(raised_error, NotAuthorizedError):
-                logger.warning(
-                    f"Not authorized to call tool {name}: {raised_error.message}"
-                )
-                return [TextContent(type="text", text=raised_error.message)]
-            else:
-                logger.error(f"Error calling tool {name}: {str(e)}")
-                raise e
+        else:
+            logger.warning(
+                f"Tool returned unexpected type: {type(result)}. Wrapping in TextContent."
+            )
+            return [TextContent(type="text", text=str(result))]
 
+    async def call_tool(
+        self, name: str, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Call a tool with comprehensive error handling.
 
-class LocalServer(Server):
-    """
-    Local server for development purposes
+        Args:
+            name: Tool name
+            arguments: Tool arguments
+
+        Returns:
+            List of TextContent results
+
+        Raises:
+            ToolError: If tool execution fails
+        """
+        logger.info(f"Calling tool: {name} with arguments: {arguments}")
+        result = await self._tool_manager.call_tool(name, arguments)
+        logger.info(f"Tool '{name}' completed successfully")
+        return self._format_tool_result(result)
+
+
+class LocalServer(BaseServer):
+    """Local development server implementation.
+
+    Args:
+        config: Server configuration
+        **kwargs: Additional keyword arguments passed to FastMCP
     """
 
-    def __init__(
-        self,
-        apps_list: list[AppConfig] = None,
-        **kwargs,
-    ):
-        if not apps_list:
-            self.apps_list = []
-        else:
-            self.apps_list = apps_list
-        super().__init__(**kwargs)
-
-    def _get_store(self, store_config: StoreConfig | None):
-        logger.info(f"Getting store: {store_config}")
-        # No store override, use the one from the server
-        if store_config is None:
-            return self.store
-        return store_from_config(store_config)
-
-    def _get_integration(self, integration_config: IntegrationConfig | None):
-        if not integration_config:
+    def __init__(self, config: ServerConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.store = self._setup_store(config.store)
+        self._load_apps()
+
+    def _setup_store(self, store_config: StoreConfig | None) -> BaseStore | None:
+        """Setup and configure the store.
+
+        Args:
+            store_config: Store configuration
+
+        Returns:
+            Configured store instance or None if no config provided
+        """
+        if not store_config:
             return None
-        if integration_config.type == "api_key":
-            store = self._get_store(integration_config.store)
-            integration = ApiKeyIntegration(integration_config.name, store=store)
-            if integration_config.credentials:
-                integration.set_credentials(integration_config.credentials)
-            return integration
-        return None
-
-    def _load_app(self, app_config: AppConfig):
-        name = app_config.name
-        integration = self._get_integration(app_config.integration)
-        app = app_from_slug(name)(integration=integration)
-        return app
-
-    def _load_apps(self):
-        logger.info(f"Loading apps: {self.apps_list}")
-        for app_config in self.apps_list:
-            try:
-                app = self._load_app(app_config)
-                if app:
-                    tools = app.list_tools()
-                for tool in tools:
-                    tool_name = tool.__name__
-                    name = app.name + "_" + tool_name
-                    description = tool.__doc__
-                    if (
-                        app_config.actions is None
-                        or len(app_config.actions) == 0
-                        or name in app_config.actions
-                    ):
-                        self.add_tool(tool, name=name, description=description)
-            except Exception as e:
-                logger.error(f"Error loading app {app_config.name}: {e}")
-
-
-class AgentRServer(Server):
-    """
-    AgentR server. Connects to the AgentR API to get the apps and tools. Only supports agentr integrations.
+
+        store = store_from_config(store_config)
+        self.add_tool(store.set)
+        self.add_tool(store.delete)
+        return store
+
+    def _load_app(self, app_config: AppConfig) -> Application | None:
+        """Load a single application with its integration.
+
+        Args:
+            app_config: Application configuration
+
+        Returns:
+            Configured application instance or None if loading fails
+        """
+        try:
+            integration = (
+                integration_from_config(app_config.integration, store=self.store)
+                if app_config.integration
+                else None
+            )
+            return app_from_slug(app_config.name)(integration=integration)
+        except Exception as e:
+            logger.error(f"Failed to load app {app_config.name}: {e}", exc_info=True)
+            return None
+
+    def _load_apps(self) -> None:
+        """Load all configured applications."""
+        logger.info(f"Loading apps: {self.config.apps}")
+        for app_config in self.config.apps:
+            app = self._load_app(app_config)
+            if app:
+                self._tool_manager.register_tools_from_app(app, app_config.actions)
+
+
+class AgentRServer(BaseServer):
+    """AgentR API-connected server implementation.
+
+    Args:
+        config: Server configuration
+        api_key: Optional API key for AgentR authentication. If not provided,
+                will attempt to read from AGENTR_API_KEY environment variable.
+        **kwargs: Additional keyword arguments passed to FastMCP
     """
 
-    def __init__(
-        self, name: str, description: str, api_key: str | None = None, **kwargs
-    ):
+    def __init__(self, config: ServerConfig, api_key: str | None = None, **kwargs):
         self.api_key = api_key or os.getenv("AGENTR_API_KEY")
         self.base_url = os.getenv("AGENTR_BASE_URL", "https://api.agentr.dev")
+
         if not self.api_key:
             raise ValueError("API key required - get one at https://agentr.dev")
-        super().__init__(name, description=description, **kwargs)
+        parsed = urlparse(self.base_url)
+        if not all([parsed.scheme, parsed.netloc]):
+            raise ValueError(f"Invalid base URL format: {self.base_url}")
+        super().__init__(config, **kwargs)
+        self.integration = AgentRIntegration(name="agentr", api_key=self.api_key)
+        self._load_apps()
 
-    def _load_app(self, app_config: AppConfig):
-        name = app_config.name
-        if app_config.integration:
-            integration_name = app_config.integration.name
-            integration = AgentRIntegration(integration_name, api_key=self.api_key)
-        else:
-            integration = None
-        app = app_from_slug(name)(integration=integration)
-        return app
-
-    def _list_apps_with_integrations(self) -> list[AppConfig]:
-        # TODO: get this from the API
-        response = httpx.get(
-            f"{self.base_url}/api/apps/", headers={"X-API-KEY": self.api_key}
-        )
-        response.raise_for_status()
-        apps = response.json()
+    def _fetch_apps(self) -> list[AppConfig]:
+        """Fetch available apps from AgentR API.
 
-        logger.info(f"Apps: {apps}")
-        return [AppConfig.model_validate(app) for app in apps]
+        Returns:
+            List of application configurations
 
-    def _load_apps(self):
-        apps = self._list_apps_with_integrations()
-        for app_config in apps:
-            try:
+        Raises:
+            httpx.HTTPError: If API request fails
+        """
+        try:
+            response = httpx.get(
+                f"{self.base_url}/api/apps/",
+                headers={"X-API-KEY": self.api_key},
+                timeout=10,
+            )
+            response.raise_for_status()
+            return [AppConfig.model_validate(app) for app in response.json()]
+        except httpx.HTTPError as e:
+            logger.error(f"Failed to fetch apps from AgentR: {e}", exc_info=True)
+            raise
+
+    def _load_app(self, app_config: AppConfig) -> Application | None:
+        """Load a single application with AgentR integration.
+
+        Args:
+            app_config: Application configuration
+
+        Returns:
+            Configured application instance or None if loading fails
+        """
+        try:
+            integration = (
+                AgentRIntegration(
+                    name=app_config.integration.name, api_key=self.api_key
+                )
+                if app_config.integration
+                else None
+            )
+            return app_from_slug(app_config.name)(integration=integration)
+        except Exception as e:
+            logger.error(f"Failed to load app {app_config.name}: {e}", exc_info=True)
+            return None
+
+    def _load_apps(self) -> None:
+        """Load all apps available from AgentR."""
+        try:
+            for app_config in self._fetch_apps():
                 app = self._load_app(app_config)
                 if app:
-                    tools = app.list_tools()
-                for tool in tools:
-                    tool_name = tool.__name__
-                    name = app.name + "_" + tool_name
-                    description = tool.__doc__
-                    if (
-                        app_config.actions is None
-                        or len(app_config.actions) == 0
-                        or name in app_config.actions
-                    ):
-                        self.add_tool(tool, name=name, description=description)
-            except Exception as e:
-                logger.error(f"Error loading app {app_config.name}: {e}")
+                    self._tool_manager.register_tools_from_app(app, app_config.actions)
+        except Exception:
+            logger.error("Failed to load apps", exc_info=True)
+            raise

universal_mcp/stores/__init__.py

@@ -1,5 +1,10 @@
 from universal_mcp.config import StoreConfig
-from universal_mcp.stores.store import EnvironmentStore, KeyringStore, MemoryStore
+from universal_mcp.stores.store import (
+    BaseStore,
+    EnvironmentStore,
+    KeyringStore,
+    MemoryStore,
+)
 
 
 def store_from_config(store_config: StoreConfig):
@@ -13,4 +18,4 @@ def store_from_config(store_config: StoreConfig):
         raise ValueError(f"Invalid store type: {store_config.type}")
 
 
-__all__ = [MemoryStore, EnvironmentStore, KeyringStore]
+__all__ = [BaseStore, MemoryStore, EnvironmentStore, KeyringStore]

universal_mcp/stores/store.py

@@ -1,18 +1,31 @@
 import os
 from abc import ABC, abstractmethod
+from typing import Any
 
 import keyring
 from loguru import logger
 
 
-class Store(ABC):
+class StoreError(Exception):
+    """Base exception class for store-related errors."""
+
+    pass
+
+
+class KeyNotFoundError(StoreError):
+    """Exception raised when a key is not found in the store."""
+
+    pass
+
+
+class BaseStore(ABC):
     """
     Abstract base class defining the interface for credential stores.
     All credential stores must implement get, set and delete methods.
     """
 
     @abstractmethod
-    def get(self, key: str):
+    def get(self, key: str) -> Any:
         """
         Retrieve a value from the store by key.
 
@@ -20,44 +33,54 @@ class Store(ABC):
             key (str): The key to look up
 
         Returns:
-            The stored value if found, None otherwise
+            Any: The stored value
+
+        Raises:
+            KeyNotFoundError: If the key is not found in the store
+            StoreError: If there is an error accessing the store
         """
         pass
 
     @abstractmethod
-    def set(self, key: str, value: str):
+    def set(self, key: str, value: str) -> None:
         """
         Store a value in the store with the given key.
 
         Args:
             key (str): The key to store the value under
             value (str): The value to store
+
+        Raises:
+            StoreError: If there is an error storing the value
         """
         pass
 
     @abstractmethod
-    def delete(self, key: str):
+    def delete(self, key: str) -> None:
         """
         Delete a value from the store by key.
 
         Args:
             key (str): The key to delete
+
+        Raises:
+            KeyNotFoundError: If the key is not found in the store
+            StoreError: If there is an error deleting the value
         """
         pass
 
 
-class MemoryStore:
+class MemoryStore(BaseStore):
     """
-    Acts as credential store for the applications.
-    Responsible for storing and retrieving credentials.
-    Ideally should be a key value store that keeps data in memory.
+    In-memory credential store implementation.
+    Stores credentials in a dictionary that persists only for the duration of the program execution.
     """
 
     def __init__(self):
         """Initialize an empty dictionary to store the data."""
-        self.data = {}
+        self.data: dict[str, str] = {}
 
-    def get(self, key: str):
+    def get(self, key: str) -> Any:
         """
         Retrieve a value from the in-memory store by key.
 
@@ -65,11 +88,16 @@ class MemoryStore:
             key (str): The key to look up
 
         Returns:
-            The stored value if found, None otherwise
+            Any: The stored value
+
+        Raises:
+            KeyNotFoundError: If the key is not found in the store
         """
-        return self.data.get(key)
+        if key not in self.data:
+            raise KeyNotFoundError(f"Key '{key}' not found in memory store")
+        return self.data[key]
 
-    def set(self, key: str, value: str):
+    def set(self, key: str, value: str) -> None:
         """
         Store a value in the in-memory store with the given key.
 
@@ -79,27 +107,28 @@ class MemoryStore:
         """
         self.data[key] = value
 
-    def delete(self, key: str):
+    def delete(self, key: str) -> None:
         """
         Delete a value from the in-memory store by key.
 
         Args:
             key (str): The key to delete
+
+        Raises:
+            KeyNotFoundError: If the key is not found in the store
         """
+        if key not in self.data:
+            raise KeyNotFoundError(f"Key '{key}' not found in memory store")
         del self.data[key]
 
 
-class EnvironmentStore(Store):
+class EnvironmentStore(BaseStore):
     """
-    Store that uses environment variables to store credentials.
-    Implements the Store interface using OS environment variables as the backend.
+    Environment variable-based credential store implementation.
+    Uses OS environment variables to store and retrieve credentials.
     """
 
-    def __init__(self):
-        """Initialize the environment store."""
-        pass
-
-    def get(self, key: str):
+    def get(self, key: str) -> Any:
         """
         Retrieve a value from environment variables by key.
 
@@ -107,11 +136,17 @@ class EnvironmentStore(Store):
             key (str): The environment variable name to look up
 
         Returns:
-            dict: Dictionary containing the api_key from environment variable
+            Any: The stored value
+
+        Raises:
+            KeyNotFoundError: If the environment variable is not found
         """
-        return {"api_key": os.getenv(key)}
+        value = os.getenv(key)
+        if value is None:
+            raise KeyNotFoundError(f"Environment variable '{key}' not found")
+        return value
 
-    def set(self, key: str, value: str):
+    def set(self, key: str, value: str) -> None:
         """
         Set an environment variable.
 
@@ -121,20 +156,25 @@ class EnvironmentStore(Store):
         """
         os.environ[key] = value
 
-    def delete(self, key: str):
+    def delete(self, key: str) -> None:
         """
         Delete an environment variable.
 
         Args:
             key (str): The environment variable name to delete
+
+        Raises:
+            KeyNotFoundError: If the environment variable is not found
         """
+        if key not in os.environ:
+            raise KeyNotFoundError(f"Environment variable '{key}' not found")
         del os.environ[key]
 
 
-class KeyringStore(Store):
+class KeyringStore(BaseStore):
     """
-    Store that uses keyring to store credentials.
-    Implements the Store interface using system keyring as the backend.
+    System keyring-based credential store implementation.
+    Uses the system's secure credential storage facility via the keyring library.
     """
 
     def __init__(self, app_name: str = "universal_mcp"):
@@ -146,7 +186,7 @@ class KeyringStore(Store):
         """
         self.app_name = app_name
 
-    def get(self, key: str):
+    def get(self, key: str) -> Any:
         """
         Retrieve a password from the system keyring.
 
@@ -154,28 +194,51 @@ class KeyringStore(Store):
             key (str): The key to look up
 
         Returns:
-            The stored password if found, None otherwise
+            Any: The stored value
+
+        Raises:
+            KeyNotFoundError: If the key is not found in the keyring
+            StoreError: If there is an error accessing the keyring
         """
-        logger.info(f"Getting password for {key} from keyring")
-        return keyring.get_password(self.app_name, key)
+        try:
+            logger.info(f"Getting password for {key} from keyring")
+            value = keyring.get_password(self.app_name, key)
+            if value is None:
+                raise KeyNotFoundError(f"Key '{key}' not found in keyring")
+            return value
+        except Exception as e:
+            raise KeyNotFoundError(f"Key '{key}' not found in keyring") from e
 
-    def set(self, key: str, value: str):
+    def set(self, key: str, value: str) -> None:
         """
         Store a password in the system keyring.
 
         Args:
             key (str): The key to store the password under
             value (str): The password to store
+
+        Raises:
+            StoreError: If there is an error storing in the keyring
         """
-        logger.info(f"Setting password for {key} in keyring")
-        keyring.set_password(self.app_name, key, value)
+        try:
+            logger.info(f"Setting password for {key} in keyring")
+            keyring.set_password(self.app_name, key, value)
+        except Exception as e:
+            raise StoreError(f"Error storing in keyring: {str(e)}") from e
 
-    def delete(self, key: str):
+    def delete(self, key: str) -> None:
         """
         Delete a password from the system keyring.
 
         Args:
             key (str): The key to delete
+
+        Raises:
+            KeyNotFoundError: If the key is not found in the keyring
+            StoreError: If there is an error deleting from the keyring
         """
-        logger.info(f"Deleting password for {key} from keyring")
-        keyring.delete_password(self.app_name, key)
+        try:
+            logger.info(f"Deleting password for {key} from keyring")
+            keyring.delete_password(self.app_name, key)
+        except Exception as e:
+            raise KeyNotFoundError(f"Key '{key}' not found in keyring") from e