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

mcp_agent/agents/base_agent.py → fast_agent/agents/mcp_agent.py

@@ -7,22 +7,23 @@ and delegates operations to an attached AugmentedLLMProtocol instance.
 
 import asyncio
 import fnmatch
-import uuid
+from abc import ABC
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
     Dict,
     List,
     Mapping,
     Optional,
+    Sequence,
     Tuple,
     Type,
     TypeVar,
     Union,
 )
 
-from a2a.types import AgentCapabilities, AgentCard, AgentSkill
+import mcp
+from a2a.types import AgentCard, AgentSkill
 from mcp.types import (
     CallToolResult,
     EmbeddedResource,
@@ -33,42 +34,38 @@ from mcp.types import (
     TextContent,
     Tool,
 )
-from opentelemetry import trace
 from pydantic import BaseModel
 
-from mcp_agent.core.agent_types import AgentConfig, AgentType
-from mcp_agent.core.exceptions import PromptExitError
-from mcp_agent.core.prompt import Prompt
-from mcp_agent.core.request_params import RequestParams
-from mcp_agent.human_input.types import (
-    HUMAN_INPUT_SIGNAL_NAME,
-    HumanInputCallback,
-    HumanInputRequest,
-    HumanInputResponse,
+from fast_agent.agents.agent_types import AgentConfig, AgentType
+from fast_agent.agents.llm_agent import DEFAULT_CAPABILITIES
+from fast_agent.agents.tool_agent import ToolAgent
+from fast_agent.constants import HUMAN_INPUT_TOOL_NAME
+from fast_agent.core.exceptions import PromptExitError
+from fast_agent.core.logging.logger import get_logger
+from fast_agent.interfaces import FastAgentLLMProtocol
+from fast_agent.mcp.helpers.content_helpers import normalize_to_extended_list
+from fast_agent.mcp.mcp_aggregator import MCPAggregator
+from fast_agent.tools.elicitation import (
+    get_elicitation_tool,
+    run_elicitation_form,
+    set_elicitation_input_callback,
 )
-from mcp_agent.logging.logger import get_logger
-from mcp_agent.mcp.interfaces import AgentProtocol, AugmentedLLMProtocol
-from mcp_agent.mcp.mcp_aggregator import MCPAggregator
-from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
+from fast_agent.types import PromptMessageExtended, RequestParams
 
 # Define a TypeVar for models
 ModelT = TypeVar("ModelT", bound=BaseModel)
 
 # Define a TypeVar for AugmentedLLM and its subclasses
-LLM = TypeVar("LLM", bound=AugmentedLLMProtocol)
+LLM = TypeVar("LLM", bound=FastAgentLLMProtocol)
 
-HUMAN_INPUT_TOOL_NAME = "__human_input__"
 if TYPE_CHECKING:
-    from mcp_agent.context import Context
-    from mcp_agent.llm.usage_tracking import UsageAccumulator
+    from rich.text import Text
 
-
-DEFAULT_CAPABILITIES = AgentCapabilities(
-    streaming=False, pushNotifications=False, stateTransitionHistory=False
-)
+    from fast_agent.context import Context
+    from fast_agent.llm.usage_tracking import UsageAccumulator
 
 
-class BaseAgent(MCPAggregator, AgentProtocol):
+class McpAgent(ABC, ToolAgent):
     """
     A base Agent class that implements the AgentProtocol interface.
 
@@ -79,258 +76,141 @@ class BaseAgent(MCPAggregator, AgentProtocol):
     def __init__(
         self,
         config: AgentConfig,
-        functions: Optional[List[Callable]] = None,
         connection_persistence: bool = True,
-        human_input_callback: Optional[HumanInputCallback] = None,
-        context: Optional["Context"] = None,
-        **kwargs: Dict[str, Any],
+        # legacy human_input_callback removed
+        context: "Context | None" = None,
+        **kwargs,
     ) -> None:
-        self.config = config
-
         super().__init__(
+            config=config,
             context=context,
+            **kwargs,
+        )
+
+        # Create aggregator with composition
+        self._aggregator = MCPAggregator(
             server_names=self.config.servers,
             connection_persistence=connection_persistence,
             name=self.config.name,
+            context=context,
+            config=self.config,  # Pass the full config for access to elicitation_handler
             **kwargs,
         )
 
-        self._context = context
-        self.tracer = trace.get_tracer(__name__)
-        self.name = self.config.name
         self.instruction = self.config.instruction
-        self.functions = functions or []
-        self.executor = self.context.executor if context and hasattr(context, "executor") else None
-        self.logger = get_logger(f"{__name__}.{self.name}")
+        self.executor = context.executor if context else None
+        self.logger = get_logger(f"{__name__}.{self._name}")
 
         # Store the default request params from config
         self._default_request_params = self.config.default_request_params
 
-        # Initialize the LLM to None (will be set by attach_llm)
-        self._llm: Optional[AugmentedLLMProtocol] = None
+        # set with the "attach" method
+        self._llm: FastAgentLLMProtocol | None = None
 
-        # Map function names to tools
-        self._function_tool_map: Dict[str, Any] = {}
+        # Instantiate human input tool once if enabled in config
+        self._human_input_tool: Tool | None = None
+        if self.config.human_input:
+            try:
+                self._human_input_tool = get_elicitation_tool()
+            except Exception:
+                self._human_input_tool = None
 
-        if not self.config.human_input:
-            self.human_input_callback = None
-        else:
-            self.human_input_callback: Optional[HumanInputCallback] = human_input_callback
-            if not human_input_callback and context and hasattr(context, "human_input_handler"):
-                self.human_input_callback = context.human_input_handler
+        # Register the MCP UI handler as the elicitation callback so fast_agent.tools can call it
+        # without importing MCP types. This avoids circular imports and ensures the callback is ready.
+        try:
+            from fast_agent.human_input.elicitation_handler import elicitation_input_callback
+            from fast_agent.human_input.types import HumanInputRequest
+
+            async def _mcp_elicitation_adapter(
+                request_payload: dict,
+                agent_name: str | None = None,
+                server_name: str | None = None,
+                server_info: dict | None = None,
+            ) -> str:
+                req = HumanInputRequest(**request_payload)
+                resp = await elicitation_input_callback(
+                    request=req,
+                    agent_name=agent_name,
+                    server_name=server_name,
+                    server_info=server_info,
+                )
+                return resp.response if isinstance(resp.response, str) else str(resp.response)
+
+            set_elicitation_input_callback(_mcp_elicitation_adapter)
+        except Exception:
+            # If UI handler import fails, leave callback unset; tool will error with a clear message
+            pass
+
+    async def __aenter__(self):
+        """Initialize the agent and its MCP aggregator."""
+        await self._aggregator.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Clean up the agent and its MCP aggregator."""
+        await self._aggregator.__aexit__(exc_type, exc_val, exc_tb)
 
     async def initialize(self) -> None:
         """
         Initialize the agent and connect to the MCP servers.
         NOTE: This method is called automatically when the agent is used as an async context manager.
         """
-        await self.__aenter__()  # This initializes the connection manager and loads the servers
-
-    async def attach_llm(
-        self,
-        llm_factory: Union[Type[AugmentedLLMProtocol], Callable[..., AugmentedLLMProtocol]],
-        model: Optional[str] = None,
-        request_params: Optional[RequestParams] = None,
-        **additional_kwargs,
-    ) -> AugmentedLLMProtocol:
-        """
-        Create and attach an LLM instance to this agent.
-
-        Parameters have the following precedence (highest to lowest):
-        1. Explicitly passed parameters to this method
-        2. Agent's default_request_params
-        3. LLM's default values
-
-        Args:
-            llm_factory: A class or callable that constructs an AugmentedLLM
-            model: Optional model name override
-            request_params: Optional request parameters override
-            **additional_kwargs: Additional parameters passed to the LLM constructor
-
-        Returns:
-            The created LLM instance
-        """
-        # Start with agent's default params
-        effective_params = (
-            self._default_request_params.model_copy() if self._default_request_params else None
-        )
-
-        # Override with explicitly passed request_params
-        if request_params:
-            if effective_params:
-                # Update non-None values
-                for k, v in request_params.model_dump(exclude_unset=True).items():
-                    if v is not None:
-                        setattr(effective_params, k, v)
-            else:
-                effective_params = request_params
-
-        # Override model if explicitly specified
-        if model and effective_params:
-            effective_params.model = model
-
-        # Create the LLM instance
-        self._llm = llm_factory(
-            agent=self, request_params=effective_params, context=self._context, **additional_kwargs
-        )
-
-        return self._llm
+        await self.__aenter__()
 
     async def shutdown(self) -> None:
         """
         Shutdown the agent and close all MCP server connections.
         NOTE: This method is called automatically when the agent is used as an async context manager.
         """
-        await super().close()
+        await self._aggregator.close()
+
+    @property
+    def initialized(self) -> bool:
+        """Check if both the agent and aggregator are initialized."""
+        return self._initialized and self._aggregator.initialized
+
+    @initialized.setter
+    def initialized(self, value: bool) -> None:
+        """Set the initialized state of both agent and aggregator."""
+        self._initialized = value
+        self._aggregator.initialized = value
 
     async def __call__(
         self,
-        message: Union[str, PromptMessageMultipart] | None = None,
-        agent_name: Optional[str] = None,
-        default_prompt: str = "",
+        message: Union[
+            str,
+            PromptMessage,
+            PromptMessageExtended,
+            Sequence[Union[str, PromptMessage, PromptMessageExtended]],
+        ],
     ) -> str:
-        """
-        Make the agent callable to send messages or start an interactive prompt.
-
-        Args:
-            message: Optional message to send to the agent
-            agent_name: Optional name of the agent (for consistency with DirectAgentApp)
-            default: Default message to use in interactive prompt mode
-
-        Returns:
-            The agent's response as a string or the result of the interactive session
-        """
-        if message:
-            return await self.send(message)
-        return await self.prompt(default_prompt=default_prompt)
-
-    async def generate_str(self, message: str, request_params: RequestParams | None) -> str:
-        result: PromptMessageMultipart = await self.generate([Prompt.user(message)], request_params)
-        return result.first_text()
+        return await self.send(message)
 
     async def send(
-        self, 
-        message: Union[str, PromptMessage, PromptMessageMultipart],
-        request_params: RequestParams | None = None
+        self,
+        message: Union[
+            str,
+            PromptMessage,
+            PromptMessageExtended,
+            Sequence[Union[str, PromptMessage, PromptMessageExtended]],
+        ],
+        request_params: RequestParams | None = None,
     ) -> str:
         """
         Send a message to the agent and get a response.
 
         Args:
             message: Message content in various formats:
-                - String: Converted to a user PromptMessageMultipart
-                - PromptMessage: Converted to PromptMessageMultipart
-                - PromptMessageMultipart: Used directly
+                - String: Converted to a user PromptMessageExtended
+                - PromptMessage: Converted to PromptMessageExtended
+                - PromptMessageExtended: Used directly
                 - request_params: Optional request parameters
 
         Returns:
             The agent's response as a string
         """
-        # Convert the input to a PromptMessageMultipart
-        prompt = self._normalize_message_input(message)
-
-        # Use the LLM to generate a response
-        response = await self.generate([prompt], request_params)
-        return response.all_text()
-
-    def _normalize_message_input(
-        self, message: Union[str, PromptMessage, PromptMessageMultipart]
-    ) -> PromptMessageMultipart:
-        """
-        Convert a message of any supported type to PromptMessageMultipart.
-
-        Args:
-            message: Message in various formats (string, PromptMessage, or PromptMessageMultipart)
-
-        Returns:
-            A PromptMessageMultipart object
-        """
-        # Handle single message
-        if isinstance(message, str):
-            return Prompt.user(message)
-        elif isinstance(message, PromptMessage):
-            return PromptMessageMultipart(role=message.role, content=[message.content])
-        elif isinstance(message, PromptMessageMultipart):
-            return message
-        else:
-            # Try to convert to string as fallback
-            return Prompt.user(str(message))
-
-    async def prompt(self, default_prompt: str = "") -> str:
-        """
-        Start an interactive prompt session with the agent.
-
-        Args:
-            default_prompt: The initial prompt to send to the agent
-
-        Returns:
-            The result of the interactive session
-        """
-        ...
-
-    async def request_human_input(self, request: HumanInputRequest) -> str:
-        """
-        Request input from a human user. Pauses the workflow until input is received.
-
-        Args:
-            request: The human input request
-
-        Returns:
-            The input provided by the human
-
-        Raises:
-            TimeoutError: If the timeout is exceeded
-        """
-        if not self.human_input_callback:
-            raise ValueError("Human input callback not set")
-
-        # Generate a unique ID for this request to avoid signal collisions
-        request_id = f"{HUMAN_INPUT_SIGNAL_NAME}_{self.name}_{uuid.uuid4()}"
-        request.request_id = request_id
-        # Use metadata as a dictionary to pass agent name
-        request.metadata = {"agent_name": self.name}
-        self.logger.debug("Requesting human input:", data=request)
-
-        if not self.executor:
-            raise ValueError("No executor available")
-
-        async def call_callback_and_signal() -> None:
-            try:
-                assert self.human_input_callback is not None
-                user_input = await self.human_input_callback(request)
-
-                self.logger.debug("Received human input:", data=user_input)
-                await self.executor.signal(signal_name=request_id, payload=user_input)
-            except PromptExitError as e:
-                # Propagate the exit error through the signal system
-                self.logger.info("User requested to exit session")
-                await self.executor.signal(
-                    signal_name=request_id,
-                    payload={"exit_requested": True, "error": str(e)},
-                )
-            except Exception as e:
-                await self.executor.signal(
-                    request_id, payload=f"Error getting human input: {str(e)}"
-                )
-
-        asyncio.create_task(call_callback_and_signal())
-
-        self.logger.debug("Waiting for human input signal")
-
-        # Wait for signal (workflow is paused here)
-        result = await self.executor.wait_for_signal(
-            signal_name=request_id,
-            request_id=request_id,
-            workflow_id=request.workflow_id,
-            signal_description=request.description or request.prompt,
-            timeout_seconds=request.timeout_seconds,
-            signal_type=HumanInputResponse,
-        )
-
-        if isinstance(result, dict) and result.get("exit_requested", False):
-            raise PromptExitError(result.get("error", "User requested to exit FastAgent session"))
-        self.logger.debug("Received human input signal", data=result)
-        return result
+        response = await self.generate(message, request_params)
+        return response.last_text() or ""
 
     def _matches_pattern(self, name: str, pattern: str, server_name: str) -> bool:
         """
@@ -359,11 +239,8 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         Returns:
             ListToolsResult with available tools
         """
-        if not self.initialized:
-            await self.initialize()
-
-        # Get all tools from the parent class
-        result = await super().list_tools()
+        # Get all tools from the aggregator
+        result = await self._aggregator.list_tools()
 
         # Apply filtering if tools are specified in config
         if self.config.tools is not None:
@@ -378,27 +255,16 @@ class BaseAgent(MCPAggregator, AgentProtocol):
 
                 # Check if this server has tool filters
                 if server_name and server_name in self.config.tools:
-                        # Check if tool matches any pattern for this server
-                        for pattern in self.config.tools[server_name]:
-                            if self._matches_pattern(tool.name, pattern, server_name):
-                                filtered_tools.append(tool)
-                                break
+                    # Check if tool matches any pattern for this server
+                    for pattern in self.config.tools[server_name]:
+                        if self._matches_pattern(tool.name, pattern, server_name):
+                            filtered_tools.append(tool)
+                            break
             result.tools = filtered_tools
 
-        if not self.human_input_callback:
-            return result
-
-        # Add a human_input_callback as a tool
-        from mcp.server.fastmcp.tools import Tool as FastTool
-
-        human_input_tool: FastTool = FastTool.from_function(self.request_human_input)
-        result.tools.append(
-            Tool(
-                name=HUMAN_INPUT_TOOL_NAME,
-                description=human_input_tool.description,
-                inputSchema=human_input_tool.parameters,
-            )
-        )
+        # Append human input tool if enabled and available
+        if self.config.human_input and getattr(self, "_human_input_tool", None):
+            result.tools.append(self._human_input_tool)
 
         return result
 
@@ -414,50 +280,43 @@ class BaseAgent(MCPAggregator, AgentProtocol):
             Result of the tool call
         """
         if name == HUMAN_INPUT_TOOL_NAME:
-            # Call the human input tool
+            # Call the elicitation-backed human input tool
             return await self._call_human_input_tool(arguments)
         else:
-            return await super().call_tool(name, arguments)
+            return await self._aggregator.call_tool(name, arguments)
 
     async def _call_human_input_tool(
         self, arguments: Dict[str, Any] | None = None
     ) -> CallToolResult:
         """
-        Handle human input request via tool calling.
+        Handle human input via an elicitation form.
 
-        Args:
-            arguments: Tool arguments
+        Expected inputs:
+        - Either an object with optional 'message' and a 'schema' JSON Schema (object), or
+        - The JSON Schema (object) itself as the arguments.
 
-        Returns:
-            Result of the human input request
+        Constraints:
+        - No more than 7 top-level properties are allowed in the schema.
         """
-        # Handle human input request
         try:
-            # Make sure arguments is not None
-            if arguments is None:
-                arguments = {}
-
-            # Extract request data
-            request_data = arguments.get("request")
-
-            # Handle both string and dict request formats
-            if isinstance(request_data, str):
-                request = HumanInputRequest(prompt=request_data)
-            elif isinstance(request_data, dict):
-                request = HumanInputRequest(**request_data)
-            else:
-                # Fallback for invalid or missing request data
-                request = HumanInputRequest(prompt="Please provide input:")
-
-            result = await self.request_human_input(request=request)
-
-            # Use response attribute if available, otherwise use the result directly
-            response_text = (
-                result.response if isinstance(result, HumanInputResponse) else str(result)
-            )
-
+            # Run via shared tool runner
+            resp_text = await run_elicitation_form(arguments, agent_name=self._name)
+            if resp_text == "__DECLINED__":
+                return CallToolResult(
+                    isError=False,
+                    content=[TextContent(type="text", text="The Human declined the input request")],
+                )
+            if resp_text in ("__CANCELLED__", "__DISABLE_SERVER__"):
+                return CallToolResult(
+                    isError=False,
+                    content=[
+                        TextContent(type="text", text="The Human cancelled the input request")
+                    ],
+                )
+            # Success path: return the (JSON) response as-is
             return CallToolResult(
-                content=[TextContent(type="text", text=f"Human response: {response_text}")]
+                isError=False,
+                content=[TextContent(type="text", text=resp_text)],
             )
 
         except PromptExitError:
@@ -476,7 +335,6 @@ class BaseAgent(MCPAggregator, AgentProtocol):
             import traceback
 
             print(f"Error in _call_human_input_tool: {traceback.format_exc()}")
-
             return CallToolResult(
                 isError=True,
                 content=[TextContent(type="text", text=f"Error requesting human input: {str(e)}")],
@@ -486,6 +344,7 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         self,
         prompt_name: str,
         arguments: Dict[str, str] | None = None,
+        namespace: str | None = None,
         server_name: str | None = None,
     ) -> GetPromptResult:
         """
@@ -494,20 +353,21 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         Args:
             prompt_name: Name of the prompt, optionally namespaced
             arguments: Optional dictionary of arguments to pass to the prompt template
-            server_name: Optional name of the server to get the prompt from
+            namespace: Optional namespace (server) to get the prompt from
 
         Returns:
             GetPromptResult containing the prompt information
         """
-        return await super().get_prompt(prompt_name, arguments, server_name)
+        target = namespace if namespace is not None else server_name
+        return await self._aggregator.get_prompt(prompt_name, arguments, target)
 
     async def apply_prompt(
         self,
         prompt: Union[str, GetPromptResult],
         arguments: Dict[str, str] | None = None,
-        agent_name: str | None = None,
-        server_name: str | None = None,
         as_template: bool = False,
+        namespace: str | None = None,
+        **_: Any,
     ) -> str:
         """
         Apply an MCP Server Prompt by name or GetPromptResult and return the assistant's response.
@@ -519,9 +379,8 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         Args:
             prompt: The name of the prompt to apply OR a GetPromptResult object
             arguments: Optional dictionary of string arguments to pass to the prompt template
-            agent_name: Optional agent name (ignored at this level, used by multi-agent apps)
-            server_name: Optional name of the server to get the prompt from
             as_template: If True, store as persistent template (always included in context)
+            namespace: Optional namespace/server to resolve the prompt from
 
         Returns:
             The assistant's response or error message
@@ -533,7 +392,7 @@ class BaseAgent(MCPAggregator, AgentProtocol):
             # Get the prompt - this will search all servers if needed
             self.logger.debug(f"Loading prompt '{prompt_name}'")
             prompt_result: GetPromptResult = await self.get_prompt(
-                prompt_name, arguments, server_name
+                prompt_name, arguments, namespace
             )
 
             if not prompt_result or not prompt_result.messages:
@@ -557,7 +416,7 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         self.logger.debug(f"Using prompt '{namespaced_name}'")
 
         # Convert prompt messages to multipart format using the safer method
-        multipart_messages = PromptMessageMultipart.from_get_prompt_result(prompt_result)
+        multipart_messages = PromptMessageExtended.from_get_prompt_result(prompt_result)
 
         if as_template:
             # Use apply_prompt_template to store as persistent prompt messages
@@ -579,13 +438,13 @@ class BaseAgent(MCPAggregator, AgentProtocol):
             server_name: Optional name of the MCP server to retrieve the resource from
 
         Returns:
-            List of EmbeddedResource objects ready to use in a PromptMessageMultipart
+            List of EmbeddedResource objects ready to use in a PromptMessageExtended
 
         Raises:
             ValueError: If the server doesn't exist or the resource couldn't be found
         """
         # Get the raw resource result
-        result: ReadResourceResult = await self.get_resource(resource_uri, server_name)
+        result: ReadResourceResult = await self._aggregator.get_resource(resource_uri, server_name)
 
         # Convert each resource content to an EmbeddedResource
         embedded_resources: List[EmbeddedResource] = []
@@ -597,10 +456,32 @@ class BaseAgent(MCPAggregator, AgentProtocol):
 
         return embedded_resources
 
+    async def get_resource(
+        self, resource_uri: str, namespace: str | None = None, server_name: str | None = None
+    ) -> ReadResourceResult:
+        """
+        Get a resource from an MCP server.
+
+        Args:
+            resource_uri: URI of the resource to retrieve
+            namespace: Optional namespace (server) to retrieve the resource from
+
+        Returns:
+            ReadResourceResult containing the resource data
+
+        Raises:
+            ValueError: If the server doesn't exist or the resource couldn't be found
+        """
+        # Get the raw resource result
+        target = namespace if namespace is not None else server_name
+        result: ReadResourceResult = await self._aggregator.get_resource(resource_uri, target)
+        return result
+
     async def with_resource(
         self,
-        prompt_content: Union[str, PromptMessage, PromptMessageMultipart],
+        prompt_content: Union[str, PromptMessage, PromptMessageExtended],
         resource_uri: str,
+        namespace: str | None = None,
         server_name: str | None = None,
     ) -> str:
         """
@@ -609,62 +490,96 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         Args:
             prompt_content: Content in various formats:
                 - String: Converted to a user message with the text
-                - PromptMessage: Converted to PromptMessageMultipart
-                - PromptMessageMultipart: Used directly
+                - PromptMessage: Converted to PromptMessageExtended
+                - PromptMessageExtended: Used directly
             resource_uri: URI of the resource to retrieve
-            server_name: Optional name of the MCP server to retrieve the resource from
+            namespace: Optional namespace (server) to retrieve the resource from
 
         Returns:
             The agent's response as a string
         """
         # Get the embedded resources
         embedded_resources: List[EmbeddedResource] = await self.get_embedded_resources(
-            resource_uri, server_name
+            resource_uri, namespace if namespace is not None else server_name
         )
 
         # Create or update the prompt message
-        prompt: PromptMessageMultipart
+        prompt: PromptMessageExtended
         if isinstance(prompt_content, str):
             # Create a new prompt with the text and resources
             content = [TextContent(type="text", text=prompt_content)]
             content.extend(embedded_resources)
-            prompt = PromptMessageMultipart(role="user", content=content)
+            prompt = PromptMessageExtended(role="user", content=content)
         elif isinstance(prompt_content, PromptMessage):
-            # Convert PromptMessage to PromptMessageMultipart and add resources
+            # Convert PromptMessage to PromptMessageExtended and add resources
             content = [prompt_content.content]
             content.extend(embedded_resources)
-            prompt = PromptMessageMultipart(role=prompt_content.role, content=content)
-        elif isinstance(prompt_content, PromptMessageMultipart):
+            prompt = PromptMessageExtended(role=prompt_content.role, content=content)
+        elif isinstance(prompt_content, PromptMessageExtended):
             # Add resources to the existing prompt
             prompt = prompt_content
             prompt.content.extend(embedded_resources)
         else:
             raise TypeError(
-                "prompt_content must be a string, PromptMessage, or PromptMessageMultipart"
+                "prompt_content must be a string, PromptMessage, or PromptMessageExtended"
             )
 
-        response: PromptMessageMultipart = await self.generate([prompt], None)
+        response: PromptMessageExtended = await self.generate([prompt], None)
         return response.first_text()
 
-    async def generate(
-        self,
-        multipart_messages: List[PromptMessageMultipart],
-        request_params: RequestParams | None = None,
-    ) -> PromptMessageMultipart:
-        """
-        Create a completion with the LLM using the provided messages.
-        Delegates to the attached LLM.
+    async def run_tools(self, request: PromptMessageExtended) -> PromptMessageExtended:
+        """Override ToolAgent's run_tools to use MCP tools via aggregator."""
+        if not request.tool_calls:
+            self.logger.warning("No tool calls found in request", data=request)
+            return PromptMessageExtended(role="user", tool_results={})
+
+        tool_results: dict[str, CallToolResult] = {}
+
+        # Cache available tool names (original, not namespaced) for display
+        available_tools = [
+            namespaced_tool.tool.name
+            for namespaced_tool in self._aggregator._namespaced_tool_map.values()
+        ]
+
+        # Process each tool call using our aggregator
+        for correlation_id, tool_request in request.tool_calls.items():
+            tool_name = tool_request.params.name
+            tool_args = tool_request.params.arguments or {}
+
+            # Get the original tool name for display (not namespaced)
+            namespaced_tool = self._aggregator._namespaced_tool_map.get(tool_name)
+            display_tool_name = namespaced_tool.tool.name if namespaced_tool else tool_name
+
+            self.display.show_tool_call(
+                name=self._name,
+                tool_args=tool_args,
+                bottom_items=available_tools,
+                tool_name=display_tool_name,
+                highlight_items=tool_name,
+                max_item_length=12,
+            )
 
-        Args:
-            multipart_messages: List of multipart messages to send to the LLM
-            request_params: Optional parameters to configure the request
+            try:
+                # Use our aggregator to call the MCP tool
+                result = await self.call_tool(tool_name, tool_args)
+                tool_results[correlation_id] = result
 
-        Returns:
-            The LLM's response as a PromptMessageMultipart
-        """
-        assert self._llm
-        with self.tracer.start_as_current_span(f"Agent: '{self.name}' generate"):
-            return await self._llm.generate(multipart_messages, request_params)
+                # Show tool result (like ToolAgent does)
+                self.display.show_tool_result(name=self._name, result=result)
+
+                self.logger.debug(f"MCP tool {display_tool_name} executed successfully")
+            except Exception as e:
+                self.logger.error(f"MCP tool {display_tool_name} failed: {e}")
+                error_result = CallToolResult(
+                    content=[TextContent(type="text", text=f"Error: {str(e)}")],
+                    isError=True,
+                )
+                tool_results[correlation_id] = error_result
+
+                # Show error result too
+                self.display.show_tool_result(name=self._name, result=error_result)
+
+        return PromptMessageExtended(role="user", tool_results=tool_results)
 
     async def apply_prompt_template(self, prompt_result: GetPromptResult, prompt_name: str) -> str:
         """
@@ -679,21 +594,30 @@ class BaseAgent(MCPAggregator, AgentProtocol):
             String representation of the assistant's response if generated
         """
         assert self._llm
-        with self.tracer.start_as_current_span(f"Agent: '{self.name}' apply_prompt_template"):
+        with self._tracer.start_as_current_span(f"Agent: '{self._name}' apply_prompt_template"):
             return await self._llm.apply_prompt_template(prompt_result, prompt_name)
 
     async def structured(
         self,
-        multipart_messages: List[PromptMessageMultipart],
+        messages: Union[
+            str,
+            PromptMessage,
+            PromptMessageExtended,
+            List[Union[str, PromptMessage, PromptMessageExtended]],
+        ],
         model: Type[ModelT],
         request_params: RequestParams | None = None,
-    ) -> Tuple[ModelT | None, PromptMessageMultipart]:
+    ) -> Tuple[ModelT | None, PromptMessageExtended]:
         """
         Apply the prompt and return the result as a Pydantic model.
-        Delegates to the attached LLM.
+        Normalizes input messages and delegates to the attached LLM.
 
         Args:
-            prompt: List of PromptMessageMultipart objects
+            messages: Message(s) in various formats:
+                - String: Converted to a user PromptMessageExtended
+                - PromptMessage: Converted to PromptMessageExtended
+                - PromptMessageExtended: Used directly
+                - List of any combination of the above
             model: The Pydantic model class to parse the result into
             request_params: Optional parameters to configure the LLM request
 
@@ -701,17 +625,20 @@ class BaseAgent(MCPAggregator, AgentProtocol):
             An instance of the specified model, or None if coercion fails
         """
         assert self._llm
-        with self.tracer.start_as_current_span(f"Agent: '{self.name}' structured"):
-            return await self._llm.structured(multipart_messages, model, request_params)
+        # Normalize all input types to a list of PromptMessageExtended
+        normalized_messages = normalize_to_extended_list(messages)
+
+        with self._tracer.start_as_current_span(f"Agent: '{self._name}' structured"):
+            return await self._llm.structured(normalized_messages, model, request_params)
 
     async def apply_prompt_messages(
-        self, prompts: List[PromptMessageMultipart], request_params: RequestParams | None = None
+        self, prompts: List[PromptMessageExtended], request_params: RequestParams | None = None
     ) -> str:
         """
         Apply a list of prompt messages and return the result.
 
         Args:
-            prompts: List of PromptMessageMultipart messages
+            prompts: List of PromptMessageExtended messages
             request_params: Optional request parameters
 
         Returns:
@@ -721,21 +648,21 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         response = await self.generate(prompts, request_params)
         return response.first_text()
 
-    async def list_prompts(self, server_name: str | None = None) -> Mapping[str, List[Prompt]]:
+    async def list_prompts(
+        self, namespace: str | None = None, server_name: str | None = None
+    ) -> Mapping[str, List[mcp.types.Prompt]]:
         """
         List all prompts available to this agent, filtered by configuration.
 
         Args:
-            server_name: Optional server name to list prompts from
+            namespace: Optional namespace (server) to list prompts from
 
         Returns:
             Dictionary mapping server names to lists of Prompt objects
         """
-        if not self.initialized:
-            await self.initialize()
-
-        # Get all prompts from the parent class
-        result = await super().list_prompts(server_name)
+        # Get all prompts from the aggregator
+        target = namespace if namespace is not None else server_name
+        result = await self._aggregator.list_prompts(target)
 
         # Apply filtering if prompts are specified in config
         if self.config.prompts is not None:
@@ -756,21 +683,21 @@ class BaseAgent(MCPAggregator, AgentProtocol):
 
         return result
 
-    async def list_resources(self, server_name: str | None = None) -> Dict[str, List[str]]:
+    async def list_resources(
+        self, namespace: str | None = None, server_name: str | None = None
+    ) -> Dict[str, List[str]]:
         """
         List all resources available to this agent, filtered by configuration.
 
         Args:
-            server_name: Optional server name to list resources from
+            namespace: Optional namespace (server) to list resources from
 
         Returns:
             Dictionary mapping server names to lists of resource URIs
         """
-        if not self.initialized:
-            await self.initialize()
-
-        # Get all resources from the parent class
-        result = await super().list_resources(server_name)
+        # Get all resources from the aggregator
+        target = namespace if namespace is not None else server_name
+        result = await self._aggregator.list_resources(target)
 
         # Apply filtering if resources are specified in config
         if self.config.resources is not None:
@@ -791,21 +718,21 @@ class BaseAgent(MCPAggregator, AgentProtocol):
 
         return result
 
-    async def list_mcp_tools(self, server_name: str | None = None) -> Mapping[str, List[Tool]]:
+    async def list_mcp_tools(
+        self, namespace: str | None = None, server_name: str | None = None
+    ) -> Mapping[str, List[Tool]]:
         """
         List all tools available to this agent, grouped by server and filtered by configuration.
 
         Args:
-            server_name: Optional server name to list tools from
+            namespace: Optional namespace (server) to list tools from
 
         Returns:
             Dictionary mapping server names to lists of Tool objects (with original names, not namespaced)
         """
-        if not self.initialized:
-            await self.initialize()
-
-        # Get all tools from the parent class
-        result = await super().list_mcp_tools(server_name)
+        # Get all tools from the aggregator
+        target = namespace if namespace is not None else server_name
+        result = await self._aggregator.list_mcp_tools(target)
 
         # Apply filtering if tools are specified in config
         if self.config.tools is not None:
@@ -824,24 +751,15 @@ class BaseAgent(MCPAggregator, AgentProtocol):
                         filtered_result[server] = filtered_tools
             result = filtered_result
 
-        # Add human input tool to a special server if human input is configured
-        if self.human_input_callback:
-            from mcp.server.fastmcp.tools import Tool as FastTool
-
-            human_input_tool: FastTool = FastTool.from_function(self.request_human_input)
+        # Add elicitation-backed human input tool to a special server if enabled and available
+        if self.config.human_input and getattr(self, "_human_input_tool", None):
             special_server_name = "__human_input__"
-            
+
             # If the special server doesn't exist in result, create it
             if special_server_name not in result:
                 result[special_server_name] = []
-            
-            result[special_server_name].append(
-                Tool(
-                    name=HUMAN_INPUT_TOOL_NAME,
-                    description=human_input_tool.description,
-                    inputSchema=human_input_tool.parameters,
-                )
-            )
+
+            result[special_server_name].append(self._human_input_tool)
 
         return result
 
@@ -863,19 +781,93 @@ class BaseAgent(MCPAggregator, AgentProtocol):
             skills.append(await self.convert(tool))
 
         return AgentCard(
-            name=self.name,
+            skills=skills,
+            name=self._name,
             description=self.instruction,
-            url=f"fast-agent://agents/{self.name}/",
+            url=f"fast-agent://agents/{self._name}/",
             version="0.1",
             capabilities=DEFAULT_CAPABILITIES,
-            defaultInputModes=["text/plain"],
-            defaultOutputModes=["text/plain"],
+            default_input_modes=["text/plain"],
+            default_output_modes=["text/plain"],
             provider=None,
-            documentationUrl=None,
-            authentication=None,
-            skills=skills,
+            documentation_url=None,
+        )
+
+    async def show_assistant_message(
+        self,
+        message: PromptMessageExtended,
+        bottom_items: List[str] | None = None,
+        highlight_items: str | List[str] | None = None,
+        max_item_length: int | None = None,
+        name: str | None = None,
+        model: str | None = None,
+        additional_message: Optional["Text"] = None,
+    ) -> None:
+        """
+        Display an assistant message with MCP servers in the bottom bar.
+
+        This override adds the list of connected MCP servers to the bottom bar
+        and highlights servers that were used for tool calls in this message.
+        """
+        # Get the list of MCP servers (if not provided)
+        if bottom_items is None:
+            if self._aggregator and self._aggregator.server_names:
+                server_names = self._aggregator.server_names
+            else:
+                server_names = []
+        else:
+            server_names = bottom_items
+
+        # Extract servers from tool calls in the message for highlighting
+        if highlight_items is None:
+            highlight_servers = self._extract_servers_from_message(message)
+        else:
+            # Convert to list if needed
+            if isinstance(highlight_items, str):
+                highlight_servers = [highlight_items]
+            else:
+                highlight_servers = highlight_items
+
+        # Call parent's implementation with server information
+        await super().show_assistant_message(
+            message=message,
+            bottom_items=server_names,
+            highlight_items=highlight_servers,
+            max_item_length=max_item_length or 12,
+            name=name,
+            model=model,
+            additional_message=additional_message,
         )
 
+    def _extract_servers_from_message(self, message: PromptMessageExtended) -> List[str]:
+        """
+        Extract server names from tool calls in the message.
+
+        Args:
+            message: The message containing potential tool calls
+
+        Returns:
+            List of server names that were called
+        """
+        servers = []
+
+        # Check if message has tool calls
+        if message.tool_calls:
+            for tool_request in message.tool_calls.values():
+                tool_name = tool_request.params.name
+
+                # Use aggregator's mapping to find the server for this tool
+                if tool_name in self._aggregator._namespaced_tool_map:
+                    namespaced_tool = self._aggregator._namespaced_tool_map[tool_name]
+                    if namespaced_tool.server_name not in servers:
+                        servers.append(namespaced_tool.server_name)
+
+        return servers
+
+    async def _parse_resource_name(self, name: str, resource_type: str) -> tuple[str, str]:
+        """Delegate resource name parsing to the aggregator."""
+        return await self._aggregator._parse_resource_name(name, resource_type)
+
     async def convert(self, tool: Tool) -> AgentSkill:
         """
         Convert a Tool to an AgentSkill.
@@ -885,26 +877,26 @@ class BaseAgent(MCPAggregator, AgentProtocol):
         return AgentSkill(
             id=tool.name,
             name=tool_without_namespace,
-            description=tool.description,
+            description=tool.description or "",
             tags=["tool"],
             examples=None,
-            inputModes=None,  # ["text/plain"],
+            input_modes=None,  # ["text/plain"],
             # cover TextContent | ImageContent ->
             # https://github.com/modelcontextprotocol/modelcontextprotocol/pull/223
             # https://github.com/modelcontextprotocol/modelcontextprotocol/pull/93
-            outputModes=None,  # ,["text/plain", "image/*"],
+            output_modes=None,  # ,["text/plain", "image/*"],
         )
 
     @property
-    def message_history(self) -> List[PromptMessageMultipart]:
+    def message_history(self) -> List[PromptMessageExtended]:
         """
-        Return the agent's message history as PromptMessageMultipart objects.
+        Return the agent's message history as PromptMessageExtended objects.
 
         This history can be used to transfer state between agents or for
         analysis and debugging purposes.
 
         Returns:
-            List of PromptMessageMultipart objects representing the conversation history
+            List of PromptMessageExtended objects representing the conversation history
         """
         if self._llm:
             return self._llm.message_history