fast-agent-mcp 0.2.27__py3-none-any.whl → 0.2.29__py3-none-any.whl

mcp_agent/mcp/mcp_aggregator.py

@@ -215,13 +215,20 @@ class MCPAggregator(ContextDependent):
                 )
 
                 # Create a wrapper to capture the parameters for the client session
-                def session_factory(read_stream, write_stream, read_timeout):
+                def session_factory(read_stream, write_stream, read_timeout, **kwargs):
+                    # Get agent's model if this aggregator is part of an agent
+                    agent_model = None
+                    if hasattr(self, 'config') and self.config and hasattr(self.config, 'model'):
+                        agent_model = self.config.model
+                    
                     return MCPAgentClientSession(
                         read_stream,
                         write_stream,
                         read_timeout,
                         server_name=server_name,
+                        agent_model=agent_model,
                         tool_list_changed_callback=self._handle_tool_list_changed,
+                        **kwargs  # Pass through any additional kwargs like server_config
                     )
 
                 await self._persistent_connection_manager.get_server(
@@ -269,13 +276,20 @@ class MCPAggregator(ContextDependent):
                 prompts = await fetch_prompts(server_connection.session, server_name)
             else:
                 # Create a factory function for the client session
-                def create_session(read_stream, write_stream, read_timeout):
+                def create_session(read_stream, write_stream, read_timeout, **kwargs):
+                    # Get agent's model if this aggregator is part of an agent
+                    agent_model = None
+                    if hasattr(self, 'config') and self.config and hasattr(self.config, 'model'):
+                        agent_model = self.config.model
+                    
                     return MCPAgentClientSession(
                         read_stream,
                         write_stream,
                         read_timeout,
                         server_name=server_name,
+                        agent_model=agent_model,
                         tool_list_changed_callback=self._handle_tool_list_changed,
+                        **kwargs  # Pass through any additional kwargs like server_config
                     )
 
                 async with gen_client(
@@ -797,12 +811,13 @@ class MCPAggregator(ContextDependent):
             messages=[],
         )
 
-    async def list_prompts(self, server_name: str | None = None) -> Mapping[str, List[Prompt]]:
+    async def list_prompts(self, server_name: str | None = None, agent_name: str | None = None) -> Mapping[str, List[Prompt]]:
         """
         List available prompts from one or all servers.
 
         :param server_name: Optional server name to list prompts from. If not provided,
                            lists prompts from all servers.
+        :param agent_name: Optional agent name (ignored at this level, used by multi-agent apps)
         :return: Dictionary mapping server names to lists of Prompt objects
         """
         if not self.initialized:

mcp_agent/mcp/mcp_connection_manager.py

@@ -165,11 +165,12 @@ class ServerConnection:
             else None
         )
 
-        session = self._client_session_factory(read_stream, send_stream, read_timeout)
-
-        # Make the server config available to the session for initialization
-        if hasattr(session, "server_config"):
-            session.server_config = self.server_config
+        session = self._client_session_factory(
+            read_stream, 
+            send_stream, 
+            read_timeout,
+            server_config=self.server_config
+        )
 
         self.session = session
 

mcp_agent/mcp/sampling.py

@@ -10,6 +10,7 @@ from mcp.types import CreateMessageRequestParams, CreateMessageResult, TextConte
 from mcp_agent.core.agent_types import AgentConfig
 from mcp_agent.llm.sampling_converter import SamplingConverter
 from mcp_agent.logging.logger import get_logger
+from mcp_agent.mcp.helpers.server_config_helpers import get_server_config
 from mcp_agent.mcp.interfaces import AugmentedLLMProtocol
 
 if TYPE_CHECKING:
@@ -78,18 +79,47 @@ async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) ->
     """
     model = None
     try:
-        # Extract model from server config
-        if (
-            hasattr(mcp_ctx, "session")
-            and hasattr(mcp_ctx.session, "server_config")
-            and mcp_ctx.session.server_config
-            and hasattr(mcp_ctx.session.server_config, "sampling")
-            and mcp_ctx.session.server_config.sampling.model
-        ):
-            model = mcp_ctx.session.server_config.sampling.model
+        # Extract model from server config using type-safe helper
+        server_config = get_server_config(mcp_ctx)
+        
+        # First priority: explicitly configured sampling model
+        if server_config and hasattr(server_config, "sampling") and server_config.sampling:
+            model = server_config.sampling.model
+        
+        # Second priority: auto_sampling fallback (if enabled at application level)
+        if model is None:
+            # Check if auto_sampling is enabled
+            auto_sampling_enabled = False
+            try:
+                from mcp_agent.context import get_current_context
+                app_context = get_current_context()
+                if app_context and app_context.config:
+                    auto_sampling_enabled = getattr(app_context.config, 'auto_sampling', True)
+            except Exception as e:
+                logger.debug(f"Could not get application config: {e}")
+                auto_sampling_enabled = True  # Default to enabled
+            
+            if auto_sampling_enabled:
+                # Import here to avoid circular import
+                from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
+                
+                # Try agent's model first (from the session)
+                if (hasattr(mcp_ctx, 'session') and 
+                    isinstance(mcp_ctx.session, MCPAgentClientSession) and
+                    mcp_ctx.session.agent_model):
+                    model = mcp_ctx.session.agent_model
+                    logger.debug(f"Using agent's model for sampling: {model}")
+                else:
+                    # Fall back to system default model
+                    try:
+                        if app_context and app_context.config and app_context.config.default_model:
+                            model = app_context.config.default_model
+                            logger.debug(f"Using system default model for sampling: {model}")
+                    except Exception as e:
+                        logger.debug(f"Could not get system default model: {e}")
 
         if model is None:
-            raise ValueError("No model configured")
+            raise ValueError("No model configured for sampling (server config, agent model, or system default)")
 
         # Create an LLM instance
         llm = create_sampling_llm(params, model)

mcp_agent/mcp_server_registry.py

@@ -27,6 +27,7 @@ from mcp_agent.config import (
     get_settings,
 )
 from mcp_agent.logging.logger import get_logger
+from mcp_agent.mcp.hf_auth import add_hf_auth_header
 from mcp_agent.mcp.logger_textio import get_stderr_handler
 from mcp_agent.mcp.mcp_connection_manager import (
     MCPConnectionManager,
@@ -70,9 +71,15 @@ class ServerRegistry:
             config (Settings): The Settings object containing the server configurations.
             config_path (str): Path to the YAML configuration file.
         """
-        self.registry = (
-            self.load_registry_from_file(config_path) if config is None else config.mcp.servers
-        )
+        if config is None:
+            self.registry = self.load_registry_from_file(config_path)
+        elif config.mcp is not None and hasattr(config.mcp, 'servers') and config.mcp.servers is not None:
+            # Ensure config.mcp exists, has a 'servers' attribute, and it's not None
+            self.registry = config.mcp.servers
+        else:
+            # Default to an empty dictionary if config.mcp is None or has no 'servers'
+            self.registry = {}
+
         self.init_hooks: Dict[str, InitHookCallable] = {}
         self.connection_manager = MCPConnectionManager(self)
 
@@ -88,8 +95,13 @@ class ServerRegistry:
         Raises:
             ValueError: If the configuration is invalid.
         """
+        servers = {} 
 
-        servers = get_settings(config_path).mcp.servers or {}
+        settings = get_settings(config_path)
+        
+        if settings.mcp is not None and hasattr(settings.mcp, 'servers') and settings.mcp.servers is not None:
+            return settings.mcp.servers
+        
         return servers
 
     @asynccontextmanager
@@ -165,11 +177,14 @@ class ServerRegistry:
             if not config.url:
                 raise ValueError(f"URL is required for SSE transport: {server_name}")
 
+            # Apply HuggingFace authentication if appropriate
+            headers = add_hf_auth_header(config.url, config.headers)
+
             # Use sse_client to get the read and write streams
             async with _add_none_to_context(
                 sse_client(
                     config.url,
-                    config.headers,
+                    headers,
                     sse_read_timeout=config.read_transport_sse_timeout_seconds,
                 )
             ) as (read_stream, write_stream, _):
@@ -187,9 +202,12 @@ class ServerRegistry:
                         logger.debug(f"{server_name}: Closed session to server")
         elif config.transport == "http":
             if not config.url:
-                raise ValueError(f"URL is required for SSE transport: {server_name}")
+                raise ValueError(f"URL is required for HTTP transport: {server_name}")
+
+            # Apply HuggingFace authentication if appropriate
+            headers = add_hf_auth_header(config.url, config.headers)
 
-            async with streamablehttp_client(config.url, config.headers) as (
+            async with streamablehttp_client(config.url, headers) as (
                 read_stream,
                 write_stream,
                 _,

mcp_agent/tools/tool_definition.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class ToolDefinition:
+    """
+    Represents a definition of a tool available to the agent.
+    """
+
+    name: str
+    description: Optional[str] = None
+    inputSchema: Dict[str, Any] = field(default_factory=dict)
+    # Add other relevant fields if necessary based on how tools are defined in fast-agent