fast-agent-mcp 0.2.57__py3-none-any.whl → 0.3.0__py3-none-any.whl

{mcp_agent → fast_agent}/mcp/mcp_agent_client_session.py

@@ -30,13 +30,13 @@ from mcp.types import (
 )
 from pydantic import FileUrl
 
-from mcp_agent.context_dependent import ContextDependent
-from mcp_agent.logging.logger import get_logger
-from mcp_agent.mcp.helpers.server_config_helpers import get_server_config
-from mcp_agent.mcp.sampling import sample
+from fast_agent.context_dependent import ContextDependent
+from fast_agent.core.logging.logger import get_logger
+from fast_agent.mcp.helpers.server_config_helpers import get_server_config
+from fast_agent.mcp.sampling import sample
 
 if TYPE_CHECKING:
-    from mcp_agent.config import MCPServerSettings
+    from fast_agent.config import MCPServerSettings
 
 logger = get_logger(__name__)
 
@@ -88,6 +88,8 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
         self.api_key: str | None = kwargs.pop("api_key", None)
         # Extract custom elicitation handler if provided
         custom_elicitation_handler = kwargs.pop("elicitation_handler", None)
+        # Extract optional context for ContextDependent mixin without passing it to ClientSession
+        self._context = kwargs.pop("context", None)
 
         version = version("fast-agent-mcp") or "dev"
         fast_agent: Implementation = Implementation(name="fast-agent-mcp", version=version)
@@ -119,9 +121,9 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
             # Try to resolve using factory
             elicitation_handler = None
             try:
-                from mcp_agent.context import get_current_context
-                from mcp_agent.core.agent_types import AgentConfig
-                from mcp_agent.mcp.elicitation_factory import resolve_elicitation_handler
+                from fast_agent.agents.agent_types import AgentConfig
+                from fast_agent.context import get_current_context
+                from fast_agent.mcp.elicitation_factory import resolve_elicitation_handler
 
                 context = get_current_context()
                 if context and context.config:
@@ -141,7 +143,7 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
             # Fallback to forms handler only if factory resolution wasn't attempted
             # If factory was attempted and returned None, respect that (means no elicitation capability)
             if elicitation_handler is None and not self.server_config:
-                from mcp_agent.mcp.elicitation_handlers import forms_elicitation_handler
+                from fast_agent.mcp.elicitation_handlers import forms_elicitation_handler
 
                 elicitation_handler = forms_elicitation_handler
 
@@ -157,7 +159,7 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
     def _should_enable_auto_sampling(self) -> bool:
         """Check if auto_sampling is enabled at the application level."""
         try:
-            from mcp_agent.context import get_current_context
+            from fast_agent.context import get_current_context
 
             context = get_current_context()
             if context and context.config:
@@ -190,12 +192,16 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
             return result
         except Exception as e:
             # Handle connection errors cleanly
+            # Looking at the MCP SDK, this should probably handle MCPError
             from anyio import ClosedResourceError
-            
+
             if isinstance(e, ClosedResourceError):
                 # Show clean offline message and convert to ConnectionError
-                from mcp_agent import console
-                console.console.print(f"[dim red]MCP server {self.session_server_name} offline[/dim red]")
+                from fast_agent.ui import console
+
+                console.console.print(
+                    f"[dim red]MCP server {self.session_server_name} offline[/dim red]"
+                )
                 raise ConnectionError(f"MCP server {self.session_server_name} offline") from e
             else:
                 logger.error(f"send_request failed: {str(e)}")
@@ -206,7 +212,7 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
         Can be overridden by subclasses to handle a notification without needing
         to listen on the message stream.
         """
-        logger.info(
+        logger.debug(
             "_received_notification: notification=",
             data=notification.model_dump(),
         )
@@ -245,86 +251,85 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
         except Exception as e:
             logger.error(f"Error in tool list changed callback: {e}")
 
+    # TODO -- decide whether to make this override type safe or not (modify SDK)
     async def call_tool(
-        self, 
-        name: str, 
-        arguments: dict | None = None, 
-        _meta: dict | None = None,
-        **kwargs
+        self, name: str, arguments: dict | None = None, _meta: dict | None = None, **kwargs
     ) -> CallToolResult:
         """Call a tool with optional metadata support."""
         if _meta:
             from mcp.types import RequestParams
-            
+
             # Safe merge - preserve existing meta fields like progressToken
-            existing_meta = kwargs.get('meta')
+            existing_meta = kwargs.get("meta")
             if existing_meta:
-                meta_dict = existing_meta.model_dump() if hasattr(existing_meta, 'model_dump') else {}
+                meta_dict = (
+                    existing_meta.model_dump() if hasattr(existing_meta, "model_dump") else {}
+                )
                 meta_dict.update(_meta)
                 meta_obj = RequestParams.Meta(**meta_dict)
             else:
                 meta_obj = RequestParams.Meta(**_meta)
-                
+
             # Create CallToolRequestParams without meta, then add _meta via model_dump
             params = CallToolRequestParams(name=name, arguments=arguments)
             params_dict = params.model_dump(by_alias=True)
             params_dict["_meta"] = meta_obj.model_dump()
-            
+
             # Create request with proper types
             request = CallToolRequest(
-                method="tools/call",
-                params=CallToolRequestParams.model_validate(params_dict)
+                method="tools/call", params=CallToolRequestParams.model_validate(params_dict)
             )
-            
+
             return await self.send_request(request, CallToolResult)
         else:
             return await super().call_tool(name, arguments, **kwargs)
 
-    async def read_resource(self, uri: str, _meta: dict | None = None, **kwargs) -> ReadResourceResult:
+    async def read_resource(
+        self, uri: str, _meta: dict | None = None, **kwargs
+    ) -> ReadResourceResult:
         """Read a resource with optional metadata support."""
         if _meta:
             from mcp.types import RequestParams
-            
+
             # Safe merge - preserve existing meta fields like progressToken
-            existing_meta = kwargs.get('meta')
+            existing_meta = kwargs.get("meta")
             if existing_meta:
-                meta_dict = existing_meta.model_dump() if hasattr(existing_meta, 'model_dump') else {}
+                meta_dict = (
+                    existing_meta.model_dump() if hasattr(existing_meta, "model_dump") else {}
+                )
                 meta_dict.update(_meta)
                 meta_obj = RequestParams.Meta(**meta_dict)
             else:
                 meta_obj = RequestParams.Meta(**_meta)
-                
+
             request = ReadResourceRequest(
-                method="resources/read",
-                params=ReadResourceRequestParams(uri=uri, meta=meta_obj)
+                method="resources/read", params=ReadResourceRequestParams(uri=uri, meta=meta_obj)
             )
             return await self.send_request(request, ReadResourceResult)
         else:
             return await super().read_resource(uri, **kwargs)
 
     async def get_prompt(
-        self, 
-        name: str, 
-        arguments: dict | None = None, 
-        _meta: dict | None = None,
-        **kwargs
+        self, name: str, arguments: dict | None = None, _meta: dict | None = None, **kwargs
     ) -> GetPromptResult:
         """Get a prompt with optional metadata support."""
         if _meta:
             from mcp.types import RequestParams
-            
+
             # Safe merge - preserve existing meta fields like progressToken
-            existing_meta = kwargs.get('meta')
+            existing_meta = kwargs.get("meta")
             if existing_meta:
-                meta_dict = existing_meta.model_dump() if hasattr(existing_meta, 'model_dump') else {}
+                meta_dict = (
+                    existing_meta.model_dump() if hasattr(existing_meta, "model_dump") else {}
+                )
                 meta_dict.update(_meta)
                 meta_obj = RequestParams.Meta(**meta_dict)
             else:
                 meta_obj = RequestParams.Meta(**_meta)
-                
+
             request = GetPromptRequest(
                 method="prompts/get",
-                params=GetPromptRequestParams(name=name, arguments=arguments, meta=meta_obj)
+                params=GetPromptRequestParams(name=name, arguments=arguments, meta=meta_obj),
             )
             return await self.send_request(request, GetPromptResult)
         else:

{mcp_agent → fast_agent}/mcp/mcp_aggregator.py

@@ -23,16 +23,16 @@ from mcp.types import (
 from opentelemetry import trace
 from pydantic import AnyUrl, BaseModel, ConfigDict
 
-from mcp_agent.context_dependent import ContextDependent
-from mcp_agent.event_progress import ProgressAction
-from mcp_agent.logging.logger import get_logger
-from mcp_agent.mcp.common import SEP, create_namespaced_name, is_namespaced_name
-from mcp_agent.mcp.gen_client import gen_client
-from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
-from mcp_agent.mcp.mcp_connection_manager import MCPConnectionManager
+from fast_agent.context_dependent import ContextDependent
+from fast_agent.core.logging.logger import get_logger
+from fast_agent.event_progress import ProgressAction
+from fast_agent.mcp.common import SEP, create_namespaced_name, is_namespaced_name
+from fast_agent.mcp.gen_client import gen_client
+from fast_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
+from fast_agent.mcp.mcp_connection_manager import MCPConnectionManager
 
 if TYPE_CHECKING:
-    from mcp_agent.context import Context
+    from fast_agent.context import Context
 
 
 logger = get_logger(__name__)  # This will be replaced per-instance when agent_name is available
@@ -84,7 +84,7 @@ class MCPAggregator(ContextDependent):
             self._persistent_connection_manager = self.context._connection_manager
 
         # Import the display component here to avoid circular imports
-        from mcp_agent.ui.console_display import ConsoleDisplay
+        from fast_agent.ui.console_display import ConsoleDisplay
 
         # Initialize the display component
         self.display = ConsoleDisplay(config=self.context.config)
@@ -99,14 +99,16 @@ class MCPAggregator(ContextDependent):
     def __init__(
         self,
         server_names: List[str],
-        connection_persistence: bool = True,  # Default to True for better stability
+        connection_persistence: bool = True,
         context: Optional["Context"] = None,
         name: str = None,
+        config: Optional[Any] = None,  # Accept the agent config for elicitation_handler access
         **kwargs,
     ) -> None:
         """
         :param server_names: A list of server names to connect to.
         :param connection_persistence: Whether to maintain persistent connections to servers (default: True).
+        :param config: Optional agent config containing elicitation_handler and other settings.
         Note: The server names must be resolvable by the gen_client function, and specified in the server registry.
         """
         super().__init__(
@@ -117,6 +119,7 @@ class MCPAggregator(ContextDependent):
         self.server_names = server_names
         self.connection_persistence = connection_persistence
         self.agent_name = name
+        self.config = config  # Store the config for access in session factory
         self._persistent_connection_manager: MCPConnectionManager = None
 
         # Set up logger with agent name in namespace if available
@@ -186,10 +189,6 @@ class MCPAggregator(ContextDependent):
     ) -> "MCPAggregator":
         """
         Factory method to create and initialize an MCPAggregator.
-        Use this instead of constructor since we need async initialization.
-        If connection_persistence is True, the aggregator will maintain a
-        persistent connection to the servers for as long as this aggregator is around.
-        By default we do not maintain a persistent connection.
         """
 
         logger.info(f"Creating MCPAggregator with servers: {server_names}")
@@ -211,6 +210,47 @@ class MCPAggregator(ContextDependent):
             logger.error(f"Error creating MCPAggregator: {e}")
             await instance.__aexit__(None, None, None)
 
+    def _create_session_factory(self, server_name: str):
+        """
+        Create a session factory function for the given server.
+        This centralizes the logic for creating MCPAgentClientSession instances.
+
+        Args:
+            server_name: The name of the server to create a session for
+
+        Returns:
+            A factory function that creates MCPAgentClientSession instances
+        """
+
+        def session_factory(read_stream, write_stream, read_timeout, **kwargs):
+            # Get agent's model and name from config if available
+            agent_model: str | None = None
+            agent_name: str | None = None
+            elicitation_handler = None
+            api_key: str | None = None
+
+            # Access config directly if it was passed from BaseAgent
+            if self.config:
+                agent_model = self.config.model
+                agent_name = self.config.name
+                elicitation_handler = self.config.elicitation_handler
+                api_key = self.config.api_key
+
+            return MCPAgentClientSession(
+                read_stream,
+                write_stream,
+                read_timeout,
+                server_name=server_name,
+                agent_model=agent_model,
+                agent_name=agent_name,
+                api_key=api_key,
+                elicitation_handler=elicitation_handler,
+                tool_list_changed_callback=self._handle_tool_list_changed,
+                **kwargs,  # Pass through any additional kwargs like server_config
+            )
+
+        return session_factory
+
     async def load_servers(self) -> None:
         """
         Discover tools from each server in parallel and build an index of namespaced tool names.
@@ -238,39 +278,8 @@ class MCPAggregator(ContextDependent):
                     },
                 )
 
-                # Create a wrapper to capture the parameters for the client session
-                def session_factory(read_stream, write_stream, read_timeout, **kwargs):
-                    # Get agent's model and name if this aggregator is part of an agent
-                    agent_model: str | None = None
-                    agent_name: str | None = None
-                    elicitation_handler = None
-                    api_key: str | None = None
-
-                    # Check if this aggregator is part of an Agent (which has config)
-                    # Import here to avoid circular dependency
-                    from mcp_agent.agents.base_agent import BaseAgent
-
-                    if isinstance(self, BaseAgent):
-                        agent_model = self.config.model
-                        agent_name = self.config.name
-                        elicitation_handler = self.config.elicitation_handler
-                        api_key = self.config.api_key
-
-                    return MCPAgentClientSession(
-                        read_stream,
-                        write_stream,
-                        read_timeout,
-                        server_name=server_name,
-                        agent_model=agent_model,
-                        agent_name=agent_name,
-                        api_key=api_key,
-                        elicitation_handler=elicitation_handler,
-                        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(
-                    server_name, client_session_factory=session_factory
+                    server_name, client_session_factory=self._create_session_factory(server_name)
                 )
 
             logger.info(
@@ -313,46 +322,15 @@ class MCPAggregator(ContextDependent):
 
             if self.connection_persistence:
                 server_connection = await self._persistent_connection_manager.get_server(
-                    server_name, client_session_factory=MCPAgentClientSession
+                    server_name, client_session_factory=self._create_session_factory(server_name)
                 )
                 tools = await fetch_tools(server_connection.session, server_name)
                 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, **kwargs):
-                    # Get agent's model and name if this aggregator is part of an agent
-                    agent_model: str | None = None
-                    agent_name: str | None = None
-                    elicitation_handler = None
-                    api_key: str | None = None
-
-                    # Check if this aggregator is part of an Agent (which has config)
-                    # Import here to avoid circular dependency
-                    from mcp_agent.agents.base_agent import BaseAgent
-
-                    if isinstance(self, BaseAgent):
-                        agent_model = self.config.model
-                        agent_name = self.config.name
-                        elicitation_handler = self.config.elicitation_handler
-                        api_key = self.config.api_key
-
-                    return MCPAgentClientSession(
-                        read_stream,
-                        write_stream,
-                        read_timeout,
-                        server_name=server_name,
-                        agent_model=agent_model,
-                        agent_name=agent_name,
-                        api_key=api_key,
-                        elicitation_handler=elicitation_handler,
-                        tool_list_changed_callback=self._handle_tool_list_changed,
-                        **kwargs,  # Pass through any additional kwargs like server_config
-                    )
-
                 async with gen_client(
                     server_name,
                     server_registry=self.context.server_registry,
-                    client_session_factory=create_session,
+                    client_session_factory=self._create_session_factory(server_name),
                 ) as client:
                     tools = await fetch_tools(client, server_name)
                     prompts = await fetch_prompts(client, server_name)
@@ -410,7 +388,7 @@ class MCPAggregator(ContextDependent):
 
         try:
             server_conn = await self._persistent_connection_manager.get_server(
-                server_name, client_session_factory=MCPAgentClientSession
+                server_name, client_session_factory=self._create_session_factory(server_name)
             )
             # server_capabilities is a property, not a coroutine
             return server_conn.server_capabilities
@@ -516,7 +494,7 @@ class MCPAggregator(ContextDependent):
                 # Get metadata from context for tool, resource, and prompt calls
                 metadata = None
                 if method_name in ["call_tool", "read_resource", "get_prompt"]:
-                    from mcp_agent.llm.augmented_llm import _mcp_metadata_var
+                    from fast_agent.llm.fastagent_llm import _mcp_metadata_var
 
                     metadata = _mcp_metadata_var.get()
 
@@ -549,7 +527,7 @@ class MCPAggregator(ContextDependent):
         try:
             if self.connection_persistence:
                 server_connection = await self._persistent_connection_manager.get_server(
-                    server_name, client_session_factory=MCPAgentClientSession
+                    server_name, client_session_factory=self._create_session_factory(server_name)
                 )
                 return await try_execute(server_connection.session)
             else:
@@ -576,7 +554,7 @@ class MCPAggregator(ContextDependent):
                     return result
         except ConnectionError:
             # Server offline - attempt reconnection
-            from mcp_agent import console
+            from fast_agent.ui import console
 
             console.console.print(
                 f"[dim yellow]MCP server {server_name} reconnecting...[/dim yellow]"
@@ -591,7 +569,8 @@ class MCPAggregator(ContextDependent):
                     await asyncio.sleep(0.1)
 
                     server_connection = await self._persistent_connection_manager.get_server(
-                        server_name, client_session_factory=MCPAgentClientSession
+                        server_name,
+                        client_session_factory=self._create_session_factory(server_name),
                     )
                     result = await try_execute(server_connection.session)
                 else:
@@ -1069,39 +1048,11 @@ class MCPAggregator(ContextDependent):
 
         async with self._refresh_lock:
             try:
-                # Create a factory function that will include our parameters
-                def create_session(read_stream, write_stream, read_timeout):
-                    # Get agent name if available
-                    agent_model: str | None = None
-                    agent_name: str | None = None
-                    elicitation_handler = None
-                    api_key: str | None = None
-
-                    # Import here to avoid circular dependency
-                    from mcp_agent.agents.base_agent import BaseAgent
-
-                    if isinstance(self, BaseAgent):
-                        agent_model = self.config.model
-                        agent_name = self.config.name
-                        elicitation_handler = self.config.elicitation_handler
-                        api_key = self.config.api_key
-
-                    return MCPAgentClientSession(
-                        read_stream,
-                        write_stream,
-                        read_timeout,
-                        server_name=server_name,
-                        agent_model=agent_model,
-                        agent_name=agent_name,
-                        api_key=api_key,
-                        elicitation_handler=elicitation_handler,
-                        tool_list_changed_callback=self._handle_tool_list_changed,
-                    )
-
                 # Fetch new tools from the server
                 if self.connection_persistence:
                     server_connection = await self._persistent_connection_manager.get_server(
-                        server_name, client_session_factory=create_session
+                        server_name,
+                        client_session_factory=self._create_session_factory(server_name),
                     )
                     tools_result = await server_connection.session.list_tools()
                     new_tools = tools_result.tools or []
@@ -1109,7 +1060,7 @@ class MCPAggregator(ContextDependent):
                     async with gen_client(
                         server_name,
                         server_registry=self.context.server_registry,
-                        client_session_factory=create_session,
+                        client_session_factory=self._create_session_factory(server_name),
                     ) as client:
                         tools_result = await client.list_tools()
                         new_tools = tools_result.tools or []

{mcp_agent → fast_agent}/mcp/mcp_connection_manager.py

@@ -26,17 +26,17 @@ from mcp.client.stdio import (
 from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
 from mcp.types import JSONRPCMessage, ServerCapabilities
 
-from mcp_agent.config import MCPServerSettings
-from mcp_agent.context_dependent import ContextDependent
-from mcp_agent.core.exceptions import ServerInitializationError
-from mcp_agent.event_progress import ProgressAction
-from mcp_agent.logging.logger import get_logger
-from mcp_agent.mcp.logger_textio import get_stderr_handler
-from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
+from fast_agent.config import MCPServerSettings
+from fast_agent.context_dependent import ContextDependent
+from fast_agent.core.exceptions import ServerInitializationError
+from fast_agent.core.logging.logger import get_logger
+from fast_agent.event_progress import ProgressAction
+from fast_agent.mcp.logger_textio import get_stderr_handler
+from fast_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
 
 if TYPE_CHECKING:
-    from mcp_agent.context import Context
-    from mcp_agent.mcp_server_registry import InitHookCallable, ServerRegistry
+    from fast_agent.context import Context
+    from fast_agent.mcp_server_registry import ServerRegistry
 
 logger = get_logger(__name__)
 
@@ -88,13 +88,11 @@ class ServerConnection:
             [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None],
             ClientSession,
         ],
-        init_hook: Optional["InitHookCallable"] = None,
     ) -> None:
         self.server_name = server_name
         self.server_config = server_config
         self.session: ClientSession | None = None
         self._client_session_factory = client_session_factory
-        self._init_hook = init_hook
         self._transport_context_factory = transport_context_factory
         # Signal that session is fully up and initialized
         self._initialized_event = Event()
@@ -137,9 +135,6 @@ class ServerConnection:
 
         self.server_capabilities = result.capabilities
         # If there's an init hook, run it
-        if self._init_hook:
-            logger.info(f"{self.server_name}: Executing init hook.")
-            self._init_hook(self.session, self.server_config.auth)
 
         # Now the session is ready for use
         self._initialized_event.set()
@@ -305,13 +300,15 @@ class MCPConnectionManager(ContextDependent):
         """Suppress MCP library's 'Error in sse_reader' messages."""
         if self._mcp_sse_filter_added:
             return
-            
+
         import logging
-        
+
         class MCPSSEErrorFilter(logging.Filter):
             def filter(self, record):
-                return not (record.name == "mcp.client.sse" and "Error in sse_reader" in record.getMessage())
-        
+                return not (
+                    record.name == "mcp.client.sse" and "Error in sse_reader" in record.getMessage()
+                )
+
         mcp_sse_logger = logging.getLogger("mcp.client.sse")
         mcp_sse_logger.addFilter(MCPSSEErrorFilter())
         self._mcp_sse_filter_added = True
@@ -323,7 +320,6 @@ class MCPConnectionManager(ContextDependent):
             [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None],
             ClientSession,
         ],
-        init_hook: Optional["InitHookCallable"] = None,
     ) -> ServerConnection:
         """
         Connect to a server and return a RunningServer instance that will persist
@@ -359,7 +355,7 @@ class MCPConnectionManager(ContextDependent):
             elif config.transport == "sse":
                 # Suppress MCP library error spam
                 self._suppress_mcp_sse_errors()
-                
+
                 return _add_none_to_context(
                     sse_client(
                         config.url,
@@ -377,7 +373,6 @@ class MCPConnectionManager(ContextDependent):
             server_config=config,
             transport_context_factory=transport_context_factory,
             client_session_factory=client_session_factory,
-            init_hook=init_hook or self.server_registry.init_hooks.get(server_name),
         )
 
         async with self._lock:
@@ -395,7 +390,6 @@ class MCPConnectionManager(ContextDependent):
         self,
         server_name: str,
         client_session_factory: Callable,
-        init_hook: Optional["InitHookCallable"] = None,
     ) -> ServerConnection:
         """
         Get a running server instance, launching it if needed.
@@ -416,7 +410,6 @@ class MCPConnectionManager(ContextDependent):
         server_conn = await self.launch_server(
             server_name=server_name,
             client_session_factory=client_session_factory,
-            init_hook=init_hook,
         )
 
         # Wait until it's fully initialized, or an error occurs