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

universal_mcp/client/cli.py

@@ -0,0 +1,181 @@
+from typing import Any
+
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.prompt import Prompt
+from rich.table import Table
+from typer import Typer
+
+from universal_mcp.client.agents import AgentType, ReActAgent
+from universal_mcp.logger import setup_logger
+from universal_mcp.tools.adapters import ToolFormat
+from universal_mcp.tools.manager import ToolManager
+
+app = Typer(name="client")
+
+
+class RichCLI:
+    def __init__(self):
+        self.console = Console()
+
+    def display_welcome(self, agent_name: str):
+        """Display welcome message"""
+        welcome_text = f"""
+# Welcome to {agent_name}!
+
+Available commands:
+- Type your questions naturally
+- `/help` - Show help
+- `/tools` - List available tools
+- `/switch <agent_type>` - Switch agent type (react/codeact)
+- `/exit` - Exit the application
+        """
+
+        self.console.print(Panel(Markdown(welcome_text), title="🤖 AI Agent CLI", border_style="blue"))
+
+    def display_agent_response(self, response: str, agent_name: str):
+        """Display agent response with formatting"""
+        self.console.print(Panel(Markdown(response), title=f"🤖 {agent_name}", border_style="green", padding=(1, 2)))
+
+    def display_thinking(self, thought: str):
+        """Display agent's thinking process"""
+        if thought:
+            self.console.print(Panel(thought, title="💭 Thinking", border_style="yellow", padding=(1, 2)))
+
+    def display_tools(self, tools: list):
+        """Display available tools in a table"""
+        table = Table(title="🛠️ Available Tools")
+        table.add_column("Tool Name", style="cyan")
+        table.add_column("Description", style="white")
+
+        for tool in tools:
+            func_info = tool["function"]
+            table.add_row(func_info["name"], func_info["description"])
+
+        self.console.print(table)
+
+    def display_error(self, error: str):
+        """Display error message"""
+        self.console.print(Panel(error, title="❌ Error", border_style="red"))
+
+    def get_user_input(self) -> str:
+        """Get user input with rich prompt"""
+        return Prompt.ask("[bold blue]You[/bold blue]", console=self.console)
+
+    def display_info(self, message: str):
+        """Display info message"""
+        self.console.print(f"[bold cyan]ℹ️ {message}[/bold cyan]")
+
+
+class AgentCLI:
+    def __init__(self):
+        self.cli = RichCLI()
+        self.tool_manager = ToolManager(default_format=ToolFormat.OPENAI)
+        self.current_agent = None
+        self.agent_type = AgentType.REACT
+        # self.tool_manager.register_tools_from_app(SampleToolApp(), tags=["all"])
+
+    def create_agent(self, agent_type: AgentType, model: str = "gpt-4o") -> Any:
+        """Create an agent based on type"""
+        instructions = """You are a helpful AI assistant. Use the available tools to help users with their requests.
+        Think step by step and provide clear explanations of your reasoning."""
+
+        if agent_type == AgentType.REACT:
+            return ReActAgent("ReAct Agent", instructions, model)
+        else:
+            raise ValueError("Unknown agent type")
+
+    def switch_agent(self, agent_type: str):
+        """Switch to a different agent type"""
+        try:
+            self.agent_type = AgentType(agent_type.lower())
+            self.current_agent = self.create_agent(self.agent_type)
+            self.cli.display_info(f"Switched to {self.agent_type.value} agent")
+        except ValueError:
+            self.cli.display_error(f"Unknown agent type: {agent_type}")
+
+    async def process_command(self, user_input: str) -> bool | None:
+        """Process special commands, return True if command was processed"""
+        if user_input.startswith("/"):
+            command_parts = user_input[1:].split()
+            command = command_parts[0].lower()
+
+            if command == "help":
+                self.show_help()
+                return True
+            elif command == "tools":
+                self.cli.display_tools(self.tool_manager.list_tools())
+                return True
+            elif command == "switch" and len(command_parts) > 1:
+                self.switch_agent(command_parts[1])
+                return True
+            elif command == "exit" or command == "quit" or command == "q":
+                self.cli.display_info("Goodbye! 👋")
+                return False
+            else:
+                self.cli.display_error(f"Unknown command: {command}")
+                return True
+
+        return None  # Not a command
+
+    def show_help(self):
+        """Show help information"""
+        help_text = """
+# Available Commands
+
+- `/help` - Show this help message
+- `/tools` - List all available tools
+- `/exit` - Exit the application
+"
+        """
+        self.cli.console.print(help_text)
+
+    async def run(self, model):
+        """Main application loop"""
+
+        # Initialize agent
+        self.current_agent = self.create_agent(self.agent_type, model)
+
+        # Display welcome
+        self.cli.display_welcome(self.current_agent.name)
+
+        # Main loop
+        while True:
+            try:
+                user_input = self.cli.get_user_input()
+
+                if not user_input.strip():
+                    continue
+
+                # Process commands
+                command_result = await self.process_command(user_input)
+                if command_result is False:  # Exit command
+                    break
+                elif command_result is True:  # Command processed
+                    continue
+
+                # Process with agent
+                response = await self.current_agent.process_step(user_input, self.tool_manager)
+
+                self.cli.display_agent_response(response, self.current_agent.name)
+
+            except KeyboardInterrupt:
+                self.cli.display_info("\nGoodbye! 👋")
+                break
+            except Exception as e:
+                self.cli.display_error(f"An error occurred: {str(e)}")
+
+
+@app.command()
+def run(model: str = "openrouter/auto"):
+    """Run the agent CLI"""
+    import asyncio
+
+    setup_logger(log_file=None, level="WARNING")
+    agent_cli = AgentCLI()
+    asyncio.run(agent_cli.run(model=model))
+
+
+if __name__ == "__main__":
+    app()

universal_mcp/client/oauth.py

@@ -1,21 +1,45 @@
 import threading
 import time
 from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any
 from urllib.parse import parse_qs, urlparse
 
 from universal_mcp.utils.singleton import Singleton
 
 
 class CallbackHandler(BaseHTTPRequestHandler):
-    """Simple HTTP handler to capture OAuth callback."""
-
-    def __init__(self, request, client_address, server, callback_data):
-        """Initialize with callback data storage."""
+    """Handles the HTTP GET request for an OAuth 2.0 callback.
+
+    This handler is designed to capture the authorization code and state
+    (or an error) returned by an OAuth 2.0 authorization server as query
+    parameters in the redirect URI. It stores these values in a shared
+    `callback_data` dictionary.
+
+    It sends a simple HTML response to the user's browser indicating
+    success or failure of the authorization attempt.
+    """
+
+    def __init__(self, request, client_address, server, callback_data: dict):
+        """Initializes the CallbackHandler.
+
+        Args:
+            request: The HTTP request.
+            client_address: The client's address.
+            server: The server instance.
+            callback_data (dict): A dictionary shared with the `CallbackServer`
+                to store the captured OAuth parameters (e.g.,
+                `authorization_code`, `state`, `error`).
+        """
         self.callback_data = callback_data
         super().__init__(request, client_address, server)
 
     def do_GET(self):
-        """Handle GET request from OAuth redirect."""
+        """Handles the GET request from the OAuth authorization server's redirect.
+
+        Parses the URL query parameters to find 'code' and 'state', or 'error'.
+        Stores these values into the `self.callback_data` dictionary.
+        Responds to the browser with a success or failure HTML page.
+        """
         parsed = urlparse(self.path)
         query_params = parse_qs(parsed.query)
 
@@ -44,7 +68,7 @@ class CallbackHandler(BaseHTTPRequestHandler):
             <html>
             <body>
                 <h1>Authorization Failed</h1>
-                <p>Error: {query_params['error'][0]}</p>
+                <p>Error: {query_params["error"][0]}</p>
                 <p>You can close this window and return to the terminal.</p>
             </body>
             </html>
@@ -54,23 +78,68 @@ class CallbackHandler(BaseHTTPRequestHandler):
             self.send_response(404)
             self.end_headers()
 
-    def log_message(self, format, *args):
-        """Suppress default logging."""
+    def log_message(self, format: str, *args: Any):
+        """Suppresses the default logging of HTTP requests.
+
+        Overrides the base class method to prevent request logs from being
+        printed to stderr, keeping the console cleaner during the OAuth flow.
+        """
         pass
 
 
 class CallbackServer(metaclass=Singleton):
-    """Simple server to handle OAuth callbacks."""
-
-    def __init__(self, port=3000):
+    """A singleton HTTP server to manage OAuth 2.0 redirect callbacks.
+
+    This server runs in a background thread, listening on a specified
+    localhost port. It uses the `CallbackHandler` to capture the
+    authorization code or error returned by an OAuth 2.0 provider
+    after user authentication.
+
+    Being a Singleton, only one instance of this server will run per
+    application, even if instantiated multiple times.
+
+    Attributes:
+        port (int): The port number on localhost where the server listens.
+        server (HTTPServer | None): The underlying `HTTPServer` instance.
+            None if the server is not running.
+        thread (threading.Thread | None): The background thread in which
+            the server runs. None if the server is not running.
+        callback_data (dict): A dictionary to store data received from the
+            OAuth callback (e.g., `authorization_code`, `state`, `error`).
+            This is shared with the `CallbackHandler`.
+        _running (bool): A flag indicating whether the server is currently
+            started and listening.
+    """
+
+    def __init__(self, port: int = 3000):
+        """Initializes the CallbackServer.
+
+        Args:
+            port (int, optional): The port number on localhost for the server
+                to listen on. Defaults to 3000.
+        """
         self.port = port
         self.server = None
         self.thread = None
         self.callback_data = {"authorization_code": None, "state": None, "error": None}
         self._running = False
 
+    @property
+    def is_running(self) -> bool:
+        return self._running
+
     def _create_handler_with_data(self):
-        """Create a handler class with access to callback data."""
+        """Creates a `CallbackHandler` subclass with shared `callback_data`.
+
+        This method dynamically defines a new handler class that inherits from
+        `CallbackHandler`. The purpose is to allow the handler instances
+        to access and modify the `self.callback_data` dictionary of this
+        `CallbackServer` instance, enabling communication of OAuth parameters
+        from the handler back to the server logic.
+
+        Returns:
+            type: A new class, subclass of `CallbackHandler`.
+        """
         callback_data = self.callback_data
 
         class DataCallbackHandler(CallbackHandler):
@@ -80,7 +149,13 @@ class CallbackServer(metaclass=Singleton):
         return DataCallbackHandler
 
     def start(self):
-        """Start the callback server in a background thread."""
+        """Starts the HTTP callback server in a background daemon thread.
+
+        If the server is not already running, it initializes an `HTTPServer`
+        with a specialized `CallbackHandler` and starts it in a new
+        daemon thread. This allows the main application flow to continue
+        while waiting for the OAuth callback.
+        """
         if self._running:
             return
         handler_class = self._create_handler_with_data()
@@ -91,15 +166,37 @@ class CallbackServer(metaclass=Singleton):
         self._running = True
 
     def stop(self):
-        """Stop the callback server."""
+        """Stops the HTTP callback server and cleans up resources.
+
+        Shuts down the `HTTPServer` and waits for its background thread
+        to complete.
+        """
         if self.server:
             self.server.shutdown()
             self.server.server_close()
         if self.thread:
             self.thread.join(timeout=1)
 
-    def wait_for_callback(self, timeout=300):
-        """Wait for OAuth callback with timeout."""
+    def wait_for_callback(self, timeout: int = 300) -> str:
+        """Waits for the OAuth callback to provide an authorization code.
+
+        This method polls the `self.callback_data` dictionary until an
+        authorization code is received or an error is reported by the
+        `CallbackHandler`, or until the timeout is reached.
+
+        Args:
+            timeout (int, optional): The maximum time in seconds to wait
+                for the callback. Defaults to 300 seconds (5 minutes).
+
+        Returns:
+            str: The received authorization code.
+
+        Raises:
+            Exception: If an error is reported in the callback
+                       (e.g., "OAuth error: <error_message>").
+            Exception: If the timeout is reached before a code or error
+                       is received (e.g., "Timeout waiting for OAuth callback").
+        """
         start_time = time.time()
         while time.time() - start_time < timeout:
             if self.callback_data["authorization_code"]:
@@ -109,6 +206,13 @@ class CallbackServer(metaclass=Singleton):
             time.sleep(0.1)
         raise Exception("Timeout waiting for OAuth callback")
 
-    def get_state(self):
-        """Get the received state parameter."""
+    def get_state(self) -> str | None:
+        """Retrieves the 'state' parameter received during the OAuth callback.
+
+        The state parameter is often used to prevent cross-site request forgery (CSRF)
+        attacks by matching its value with one sent in the initial authorization request.
+
+        Returns:
+            str | None: The 'state' parameter value if received, otherwise None.
+        """
         return self.callback_data["state"]

universal_mcp/client/token_store.py

@@ -6,27 +6,86 @@ from universal_mcp.stores.store import KeyringStore
 
 
 class TokenStore(MCPTokenStorage):
-    """Simple in-memory token storage implementation."""
+    """Persistent storage for OAuth tokens and client information using KeyringStore.
+
+    This class implements the `mcp.client.auth.TokenStorage` interface,
+    providing a mechanism to securely store and retrieve OAuth 2.0 tokens
+    (as `OAuthToken` objects) and OAuth client registration details
+    (as `OAuthClientInformationFull` objects).
+
+    It utilizes an underlying `KeyringStore` instance, which typically
+    delegates to the operating system's secure credential management
+    system (e.g., macOS Keychain, Windows Credential Manager, Linux KWallet).
+    This ensures that sensitive token data is stored securely and persistently.
+
+    Attributes:
+        store (KeyringStore): The `KeyringStore` instance used for actually
+            storing and retrieving the serialized token and client info data.
+    """
 
     def __init__(self, store: KeyringStore):
+        """Initializes the TokenStore.
+
+        Args:
+            store (KeyringStore): An instance of `KeyringStore` that will be
+                used for the actual persistence of tokens and client information.
+        """
         self.store = store
-        self._tokens: OAuthToken | None = None
-        self._client_info: OAuthClientInformationFull | None = None
+        # These are not meant to be persistent caches in this implementation
+        # self._tokens: OAuthToken | None = None
+        # self._client_info: OAuthClientInformationFull | None = None
 
     async def get_tokens(self) -> OAuthToken | None:
+        """Retrieves OAuth tokens from the persistent KeyringStore.
+
+        Fetches the JSON string representation of tokens from the store using
+        the key "tokens" and deserializes it into an `OAuthToken` object.
+
+        Returns:
+            OAuthToken | None: The deserialized `OAuthToken` object if found
+                               and successfully parsed, otherwise None.
+        """
         try:
             return OAuthToken.model_validate_json(self.store.get("tokens"))
         except KeyNotFoundError:
             return None
 
     async def set_tokens(self, tokens: OAuthToken) -> None:
+        """Serializes OAuth tokens to JSON and saves them to the KeyringStore.
+
+        The provided `OAuthToken` object is converted to its JSON string
+        representation and stored in the `KeyringStore` under the key "tokens".
+
+        Args:
+            tokens (OAuthToken): The `OAuthToken` object to store.
+        """
         self.store.set("tokens", tokens.model_dump_json())
 
     async def get_client_info(self) -> OAuthClientInformationFull | None:
+        """Retrieves OAuth client information from the persistent KeyringStore.
+
+        Fetches the JSON string representation of client information from the
+        store using the key "client_info" and deserializes it into an
+        `OAuthClientInformationFull` object.
+
+        Returns:
+            OAuthClientInformationFull | None: The deserialized object if found
+                                              and successfully parsed, otherwise None.
+        """
         try:
             return OAuthClientInformationFull.model_validate_json(self.store.get("client_info"))
         except KeyNotFoundError:
             return None
 
     async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        """Serializes OAuth client information to JSON and saves it to KeyringStore.
+
+        The provided `OAuthClientInformationFull` object is converted to its
+        JSON string representation and stored in the `KeyringStore` under the
+        key "client_info".
+
+        Args:
+            client_info (OAuthClientInformationFull): The client information object
+                to store.
+        """
         self.store.set("client_info", client_info.model_dump_json())