mcp-use 1.3.6__py3-none-any.whl → 1.3.8__py3-none-any.whl

mcp_use/agents/remote.py

@@ -0,0 +1,239 @@
+"""
+Remote agent implementation for executing agents via API.
+"""
+
+import json
+import os
+from typing import Any, TypeVar
+
+import httpx
+from langchain.schema import BaseMessage
+from pydantic import BaseModel
+
+from ..logging import logger
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class RemoteAgent:
+    """Agent that executes remotely via API."""
+
+    def __init__(self, agent_id: str, api_key: str | None = None, base_url: str = "https://cloud.mcp-use.com"):
+        """Initialize remote agent.
+
+        Args:
+            agent_id: The ID of the remote agent to execute
+            api_key: API key for authentication. If None, will check MCP_USE_API_KEY env var
+            base_url: Base URL for the remote API
+        """
+        self.agent_id = agent_id
+        self.base_url = base_url
+
+        # Handle API key validation
+        if api_key is None:
+            api_key = os.getenv("MCP_USE_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "API key is required for remote execution. "
+                "Please provide it as a parameter or set the MCP_USE_API_KEY environment variable. "
+                "You can get an API key from https://cloud.mcp-use.com"
+            )
+
+        self.api_key = api_key
+        # Configure client with reasonable timeouts for agent execution
+        self._client = httpx.AsyncClient(
+            timeout=httpx.Timeout(
+                connect=10.0,  # 10 seconds to establish connection
+                read=300.0,  # 5 minutes to read response (agents can take time)
+                write=10.0,  # 10 seconds to send request
+                pool=10.0,  # 10 seconds to get connection from pool
+            )
+        )
+
+    def _pydantic_to_json_schema(self, model_class: type[T]) -> dict[str, Any]:
+        """Convert a Pydantic model to JSON schema for API transmission.
+
+        Args:
+            model_class: The Pydantic model class to convert
+
+        Returns:
+            JSON schema representation of the model
+        """
+        return model_class.model_json_schema()
+
+    def _parse_structured_response(self, response_data: Any, output_schema: type[T]) -> T:
+        """Parse the API response into the structured output format.
+
+        Args:
+            response_data: Raw response data from the API
+            output_schema: The Pydantic model to parse into
+
+        Returns:
+            Parsed structured output
+        """
+        # Handle different response formats
+        if isinstance(response_data, dict):
+            if "result" in response_data:
+                outer_result = response_data["result"]
+                # Check if this is a nested result structure (agent execution response)
+                if isinstance(outer_result, dict) and "result" in outer_result:
+                    # Extract the actual structured output from the nested result
+                    result_data = outer_result["result"]
+                else:
+                    # Use the outer result directly
+                    result_data = outer_result
+            else:
+                result_data = response_data
+        elif isinstance(response_data, str):
+            try:
+                result_data = json.loads(response_data)
+            except json.JSONDecodeError:
+                # If it's not valid JSON, try to create the model from the string content
+                result_data = {"content": response_data}
+        else:
+            result_data = response_data
+
+        # Parse into the Pydantic model
+        try:
+            return output_schema.model_validate(result_data)
+        except Exception as e:
+            logger.warning(f"Failed to parse structured output: {e}")
+            # Fallback: try to parse it as raw content if the model has a content field
+            if hasattr(output_schema, "model_fields") and "content" in output_schema.model_fields:
+                return output_schema.model_validate({"content": str(result_data)})
+            raise
+
+    async def run(
+        self,
+        query: str,
+        max_steps: int | None = None,
+        manage_connector: bool = True,
+        external_history: list[BaseMessage] | None = None,
+        output_schema: type[T] | None = None,
+    ) -> str | T:
+        """Run a query on the remote agent.
+
+        Args:
+            query: The query to execute
+            max_steps: Maximum number of steps (default: 10)
+            manage_connector: Ignored for remote execution
+            external_history: Ignored for remote execution (not supported yet)
+            output_schema: Optional Pydantic model for structured output
+
+        Returns:
+            The result from the remote agent execution (string or structured output)
+        """
+        if external_history is not None:
+            logger.warning("External history is not yet supported for remote execution")
+
+        payload = {"query": query, "max_steps": max_steps or 10}
+
+        # Add structured output schema if provided
+        if output_schema is not None:
+            payload["output_schema"] = self._pydantic_to_json_schema(output_schema)
+            logger.info(f"🔧 Using structured output with schema: {output_schema.__name__}")
+
+        headers = {"Content-Type": "application/json", "x-api-key": self.api_key}
+
+        url = f"{self.base_url}/api/v1/agents/{self.agent_id}/run"
+
+        try:
+            logger.info(f"🌐 Executing query on remote agent {self.agent_id}")
+            response = await self._client.post(url, json=payload, headers=headers)
+            response.raise_for_status()
+
+            result = response.json()
+            logger.info(f"🔧 Response: {result}")
+            logger.info("✅ Remote execution completed successfully")
+
+            # Check for error responses (even with 200 status)
+            if isinstance(result, dict):
+                # Check for actual error conditions (not just presence of error field)
+                if result.get("status") == "error" or (result.get("error") is not None):
+                    error_msg = result.get("error", str(result))
+                    logger.error(f"❌ Remote agent execution failed: {error_msg}")
+                    raise RuntimeError(f"Remote agent execution failed: {error_msg}")
+
+                # Check if the response indicates agent initialization failure
+                if "failed to initialize" in str(result):
+                    logger.error(f"❌ Agent initialization failed: {result}")
+                    raise RuntimeError(
+                        f"Agent initialization failed on remote server. "
+                        f"This usually indicates:\n"
+                        f"• Invalid agent configuration (LLM model, system prompt)\n"
+                        f"• Missing or invalid MCP server configurations\n"
+                        f"• Network connectivity issues with MCP servers\n"
+                        f"• Missing environment variables or credentials\n"
+                        f"Raw error: {result}"
+                    )
+
+            # Handle structured output
+            if output_schema is not None:
+                return self._parse_structured_response(result, output_schema)
+
+            # Regular string output
+            if isinstance(result, dict) and "result" in result:
+                return result["result"]
+            elif isinstance(result, str):
+                return result
+            else:
+                return str(result)
+
+        except httpx.HTTPStatusError as e:
+            status_code = e.response.status_code
+            response_text = e.response.text
+
+            # Provide specific error messages based on status code
+            if status_code == 401:
+                logger.error(f"❌ Authentication failed: {response_text}")
+                raise RuntimeError(
+                    "Authentication failed: Invalid or missing API key. "
+                    "Please check your API key and ensure the MCP_USE_API_KEY environment variable is set correctly."
+                ) from e
+            elif status_code == 403:
+                logger.error(f"❌ Access forbidden: {response_text}")
+                raise RuntimeError(
+                    f"Access denied: You don't have permission to execute agent '{self.agent_id}'. "
+                    "Check if the agent exists and you have the necessary permissions."
+                ) from e
+            elif status_code == 404:
+                logger.error(f"❌ Agent not found: {response_text}")
+                raise RuntimeError(
+                    f"Agent not found: Agent '{self.agent_id}' does not exist or you don't have access to it. "
+                    "Please verify the agent ID and ensure it exists in your account."
+                ) from e
+            elif status_code == 422:
+                logger.error(f"❌ Validation error: {response_text}")
+                raise RuntimeError(
+                    f"Request validation failed: {response_text}. "
+                    "Please check your query parameters and output schema format."
+                ) from e
+            elif status_code == 500:
+                logger.error(f"❌ Server error: {response_text}")
+                raise RuntimeError(
+                    "Internal server error occurred during agent execution. "
+                    "Please try again later or contact support if the issue persists."
+                ) from e
+            else:
+                logger.error(f"❌ Remote execution failed with status {status_code}: {response_text}")
+                raise RuntimeError(f"Remote agent execution failed: {status_code} - {response_text}") from e
+        except httpx.TimeoutException as e:
+            logger.error(f"❌ Remote execution timed out: {e}")
+            raise RuntimeError(
+                "Remote agent execution timed out. The server may be overloaded or the query is taking too long to "
+                "process. Try again or use a simpler query."
+            ) from e
+        except httpx.ConnectError as e:
+            logger.error(f"❌ Remote execution connection error: {e}")
+            raise RuntimeError(
+                f"Remote agent connection failed: Cannot connect to {self.base_url}. "
+                f"Check if the server is running and the URL is correct."
+            ) from e
+        except Exception as e:
+            logger.error(f"❌ Remote execution error: {e}")
+            raise RuntimeError(f"Remote agent execution failed: {str(e)}") from e
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        await self._client.aclose()
+        logger.info("🔌 Remote agent client closed")

mcp_use/client.py

@@ -9,6 +9,8 @@ import json
 import warnings
 from typing import Any
 
+from mcp.client.session import ElicitationFnT, SamplingFnT
+
 from mcp_use.types.sandbox import SandboxOptions
 
 from .config import create_connector_from_config, load_config_file
@@ -26,8 +28,11 @@ class MCPClient:
     def __init__(
         self,
         config: str | dict[str, Any] | None = None,
+        allowed_servers: list[str] | None = None,
         sandbox: bool = False,
         sandbox_options: SandboxOptions | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
     ) -> None:
         """Initialize a new MCP client.
 
@@ -36,13 +41,16 @@ class MCPClient:
                    If None, an empty configuration is used.
             sandbox: Whether to use sandboxed execution mode for running MCP servers.
             sandbox_options: Optional sandbox configuration options.
+            sampling_callback: Optional sampling callback function.
         """
         self.config: dict[str, Any] = {}
+        self.allowed_servers: list[str] = allowed_servers
         self.sandbox = sandbox
         self.sandbox_options = sandbox_options
         self.sessions: dict[str, MCPSession] = {}
         self.active_sessions: list[str] = []
-
+        self.sampling_callback = sampling_callback
+        self.elicitation_callback = elicitation_callback
         # Load configuration if provided
         if config is not None:
             if isinstance(config, str):
@@ -56,6 +64,8 @@ class MCPClient:
         config: dict[str, Any],
         sandbox: bool = False,
         sandbox_options: SandboxOptions | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
     ) -> "MCPClient":
         """Create a MCPClient from a dictionary.
 
@@ -63,12 +73,25 @@ class MCPClient:
             config: The configuration dictionary.
             sandbox: Whether to use sandboxed execution mode for running MCP servers.
             sandbox_options: Optional sandbox configuration options.
+            sampling_callback: Optional sampling callback function.
+            elicitation_callback: Optional elicitation callback function.
         """
-        return cls(config=config, sandbox=sandbox, sandbox_options=sandbox_options)
+        return cls(
+            config=config,
+            sandbox=sandbox,
+            sandbox_options=sandbox_options,
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+        )
 
     @classmethod
     def from_config_file(
-        cls, filepath: str, sandbox: bool = False, sandbox_options: SandboxOptions | None = None
+        cls,
+        filepath: str,
+        sandbox: bool = False,
+        sandbox_options: SandboxOptions | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
     ) -> "MCPClient":
         """Create a MCPClient from a configuration file.
 
@@ -76,8 +99,16 @@ class MCPClient:
             filepath: The path to the configuration file.
             sandbox: Whether to use sandboxed execution mode for running MCP servers.
             sandbox_options: Optional sandbox configuration options.
+            sampling_callback: Optional sampling callback function.
+            elicitation_callback: Optional elicitation callback function.
         """
-        return cls(config=load_config_file(filepath), sandbox=sandbox, sandbox_options=sandbox_options)
+        return cls(
+            config=load_config_file(filepath),
+            sandbox=sandbox,
+            sandbox_options=sandbox_options,
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+        )
 
     def add_server(
         self,
@@ -151,7 +182,11 @@ class MCPClient:
 
         # Create connector with options
         connector = create_connector_from_config(
-            server_config, sandbox=self.sandbox, sandbox_options=self.sandbox_options
+            server_config,
+            sandbox=self.sandbox,
+            sandbox_options=self.sandbox_options,
+            sampling_callback=self.sampling_callback,
+            elicitation_callback=self.elicitation_callback,
         )
 
         # Create the session
@@ -187,9 +222,10 @@ class MCPClient:
             warnings.warn("No MCP servers defined in config", UserWarning, stacklevel=2)
             return {}
 
-        # Create sessions for all servers
+        # Create sessions only for allowed servers if applicable else create for all servers
         for name in servers:
-            await self.create_session(name, auto_initialize)
+            if self.allowed_servers is None or name in self.allowed_servers:
+                await self.create_session(name, auto_initialize)
 
         return self.sessions
 

mcp_use/config.py

@@ -7,6 +7,8 @@ This module provides functionality to load MCP configuration from JSON files.
 import json
 from typing import Any
 
+from mcp.client.session import ElicitationFnT, SamplingFnT
+
 from mcp_use.types.sandbox import SandboxOptions
 
 from .connectors import (
@@ -36,6 +38,8 @@ def create_connector_from_config(
     server_config: dict[str, Any],
     sandbox: bool = False,
     sandbox_options: SandboxOptions | None = None,
+    sampling_callback: SamplingFnT | None = None,
+    elicitation_callback: ElicitationFnT | None = None,
 ) -> BaseConnector:
     """Create a connector based on server configuration.
     This function can be called with just the server_config parameter:
@@ -44,7 +48,7 @@ def create_connector_from_config(
         server_config: The server configuration section
         sandbox: Whether to use sandboxed execution mode for running MCP servers.
         sandbox_options: Optional sandbox configuration options.
-
+        sampling_callback: Optional sampling callback function.
     Returns:
         A configured connector instance
     """
@@ -55,6 +59,8 @@ def create_connector_from_config(
             command=server_config["command"],
             args=server_config["args"],
             env=server_config.get("env", None),
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
         )
 
     # Sandboxed connector
@@ -64,6 +70,8 @@ def create_connector_from_config(
             args=server_config["args"],
             env=server_config.get("env", None),
             e2b_options=sandbox_options,
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
         )
 
     # HTTP connector
@@ -72,6 +80,10 @@ def create_connector_from_config(
             base_url=server_config["url"],
             headers=server_config.get("headers", None),
             auth_token=server_config.get("auth_token", None),
+            timeout=server_config.get("timeout", 5),
+            sse_read_timeout=server_config.get("sse_read_timeout", 60 * 5),
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
         )
 
     # WebSocket connector

mcp_use/connectors/base.py

@@ -6,13 +6,17 @@ must implement.
 """
 
 from abc import ABC, abstractmethod
+from datetime import timedelta
 from typing import Any
 
-from mcp import ClientSession
+from mcp import ClientSession, Implementation
+from mcp.client.session import ElicitationFnT, SamplingFnT
 from mcp.shared.exceptions import McpError
 from mcp.types import CallToolResult, GetPromptResult, Prompt, ReadResourceResult, Resource, Tool
 from pydantic import AnyUrl
 
+import mcp_use
+
 from ..logging import logger
 from ..task_managers import ConnectionManager
 
@@ -23,7 +27,11 @@ class BaseConnector(ABC):
     This class defines the interface that all MCP connectors must implement.
     """
 
-    def __init__(self):
+    def __init__(
+        self,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
+    ):
         """Initialize base connector with common attributes."""
         self.client_session: ClientSession | None = None
         self._connection_manager: ConnectionManager | None = None
@@ -33,6 +41,17 @@ class BaseConnector(ABC):
         self._connected = False
         self._initialized = False  # Track if client_session.initialize() has been called
         self.auto_reconnect = True  # Whether to automatically reconnect on connection loss (not configurable for now)
+        self.sampling_callback = sampling_callback
+        self.elicitation_callback = elicitation_callback
+
+    @property
+    def client_info(self) -> Implementation:
+        """Get the client info for the connector."""
+        return Implementation(
+            name="mcp-use",
+            version=mcp_use.__version__,
+            url="https://github.com/mcp-use/mcp-use",
+        )
 
     @abstractmethod
     async def connect(self) -> None:
@@ -110,28 +129,41 @@ class BaseConnector(ABC):
 
         if server_capabilities.tools:
             # Get available tools directly from client session
-            tools_result = await self.client_session.list_tools()
-            self._tools = tools_result.tools if tools_result else []
+            try:
+                tools_result = await self.client_session.list_tools()
+                self._tools = tools_result.tools if tools_result else []
+            except Exception as e:
+                logger.error(f"Error listing tools: {e}")
+                self._tools = []
         else:
             self._tools = []
 
         if server_capabilities.resources:
             # Get available resources directly from client session
-            resources_result = await self.client_session.list_resources()
-            self._resources = resources_result.resources if resources_result else []
+            try:
+                resources_result = await self.client_session.list_resources()
+                self._resources = resources_result.resources if resources_result else []
+            except Exception as e:
+                logger.error(f"Error listing resources: {e}")
+                self._resources = []
         else:
             self._resources = []
 
         if server_capabilities.prompts:
             # Get available prompts directly from client session
-            prompts_result = await self.client_session.list_prompts()
-            self._prompts = prompts_result.prompts if prompts_result else []
+            try:
+                prompts_result = await self.client_session.list_prompts()
+                self._prompts = prompts_result.prompts if prompts_result else []
+            except Exception as e:
+                logger.error(f"Error listing prompts: {e}")
+                self._prompts = []
         else:
             self._prompts = []
 
         logger.debug(
             f"MCP session initialized with {len(self._tools)} tools, "
-            "{len(self._resources)} resources, and {len(self._prompts)} prompts"
+            f"{len(self._resources)} resources, "
+            f"and {len(self._prompts)} prompts"
         )
 
         return result
@@ -235,12 +267,15 @@ class BaseConnector(ABC):
                     "Connection to MCP server has been lost. Auto-reconnection is disabled. Please reconnect manually."
                 )
 
-    async def call_tool(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
+    async def call_tool(
+        self, name: str, arguments: dict[str, Any], read_timeout_seconds: timedelta | None = None
+    ) -> CallToolResult:
         """Call an MCP tool with automatic reconnection handling.
 
         Args:
             name: The name of the tool to call.
             arguments: The arguments to pass to the tool.
+            read_timeout_seconds: timeout seconds when calling tool
 
         Returns:
             The result of the tool call.
@@ -254,7 +289,7 @@ class BaseConnector(ABC):
 
         logger.debug(f"Calling tool '{name}' with arguments: {arguments}")
         try:
-            result = await self.client_session.call_tool(name, arguments)
+            result = await self.client_session.call_tool(name, arguments, read_timeout_seconds)
             logger.debug(f"Tool '{name}' called with result: {result}")
             return result
         except Exception as e:

mcp_use/connectors/http.py

@@ -7,9 +7,10 @@ through HTTP APIs with SSE or Streamable HTTP for transport.
 
 import httpx
 from mcp import ClientSession
+from mcp.client.session import ElicitationFnT, SamplingFnT
 
 from ..logging import logger
-from ..task_managers import ConnectionManager, SseConnectionManager, StreamableHttpConnectionManager
+from ..task_managers import SseConnectionManager, StreamableHttpConnectionManager
 from .base import BaseConnector
 
 
@@ -27,6 +28,8 @@ class HttpConnector(BaseConnector):
         headers: dict[str, str] | None = None,
         timeout: float = 5,
         sse_read_timeout: float = 60 * 5,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
     ):
         """Initialize a new HTTP connector.
 
@@ -36,8 +39,10 @@ class HttpConnector(BaseConnector):
             headers: Optional additional headers.
             timeout: Timeout for HTTP operations in seconds.
             sse_read_timeout: Timeout for SSE read operations in seconds.
+            sampling_callback: Optional sampling callback.
+            elicitation_callback: Optional elicitation callback.
         """
-        super().__init__()
+        super().__init__(sampling_callback=sampling_callback, elicitation_callback=elicitation_callback)
         self.base_url = base_url.rstrip("/")
         self.auth_token = auth_token
         self.headers = headers or {}
@@ -46,14 +51,6 @@ class HttpConnector(BaseConnector):
         self.timeout = timeout
         self.sse_read_timeout = sse_read_timeout
 
-    async def _setup_client(self, connection_manager: ConnectionManager) -> None:
-        """Set up the client session with the provided connection manager."""
-
-        self._connection_manager = connection_manager
-        read_stream, write_stream = await self._connection_manager.start()
-        self.client_session = ClientSession(read_stream, write_stream, sampling_callback=None)
-        await self.client_session.__aenter__()
-
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
         if self._connected:
@@ -76,7 +73,13 @@ class HttpConnector(BaseConnector):
             read_stream, write_stream = await connection_manager.start()
 
             # Test if this actually works by trying to create a client session and initialize it
-            test_client = ClientSession(read_stream, write_stream, sampling_callback=None)
+            test_client = ClientSession(
+                read_stream,
+                write_stream,
+                sampling_callback=self.sampling_callback,
+                elicitation_callback=self.elicitation_callback,
+                client_info=self.client_info,
+            )
             await test_client.__aenter__()
 
             try:
@@ -154,7 +157,13 @@ class HttpConnector(BaseConnector):
                     read_stream, write_stream = await connection_manager.start()
 
                     # Create the client session for SSE
-                    self.client_session = ClientSession(read_stream, write_stream, sampling_callback=None)
+                    self.client_session = ClientSession(
+                        read_stream,
+                        write_stream,
+                        sampling_callback=self.sampling_callback,
+                        elicitation_callback=self.elicitation_callback,
+                        client_info=self.client_info,
+                    )
                     await self.client_session.__aenter__()
                     self.transport_type = "SSE"
 

mcp_use/connectors/sandbox.py

@@ -12,6 +12,7 @@ import time
 
 import aiohttp
 from mcp import ClientSession
+from mcp.client.session import ElicitationFnT, SamplingFnT
 
 from ..logging import logger
 from ..task_managers import SseConnectionManager
@@ -50,6 +51,8 @@ class SandboxConnector(BaseConnector):
         e2b_options: SandboxOptions | None = None,
         timeout: float = 5,
         sse_read_timeout: float = 60 * 5,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
     ):
         """Initialize a new sandbox connector.
 
@@ -61,8 +64,10 @@ class SandboxConnector(BaseConnector):
                         See SandboxOptions for available options and defaults.
             timeout: Timeout for the sandbox process in seconds.
             sse_read_timeout: Timeout for the SSE connection in seconds.
+            sampling_callback: Optional sampling callback.
+            elicitation_callback: Optional elicitation callback.
         """
-        super().__init__()
+        super().__init__(sampling_callback=sampling_callback, elicitation_callback=elicitation_callback)
         if Sandbox is None:
             raise ImportError(
                 "E2B SDK (e2b-code-interpreter) not found. Please install it with "
@@ -217,7 +222,13 @@ class SandboxConnector(BaseConnector):
             read_stream, write_stream = await self._connection_manager.start()
 
             # Create the client session
-            self.client_session = ClientSession(read_stream, write_stream, sampling_callback=None)
+            self.client_session = ClientSession(
+                read_stream,
+                write_stream,
+                sampling_callback=self.sampling_callback,
+                elicitation_callback=self.elicitation_callback,
+                client_info=self.client_info,
+            )
             await self.client_session.__aenter__()
 
             # Mark as connected