fast-agent-mcp 0.4.7__py3-none-any.whl

fast_agent/mcp/logger_textio.py

@@ -0,0 +1,108 @@
+"""
+Utilities for MCP stdio client integration with our logging system.
+"""
+
+import io
+import os
+from typing import TextIO
+
+from fast_agent.core.logging.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class LoggerTextIO(TextIO):
+    """
+    A TextIO implementation that logs to our application logger.
+    This implements the full TextIO interface as specified by Python.
+
+    Args:
+        server_name: The name of the server to include in logs
+    """
+
+    def __init__(self, server_name: str) -> None:
+        super().__init__()
+        self.server_name = server_name
+        # Use a StringIO for buffering
+        self._buffer = io.StringIO()
+        # Keep track of complete and partial lines
+        self._line_buffer = ""
+
+    def write(self, s: str) -> int:
+        """
+        Write data to our buffer and log any complete lines.
+        """
+        if not s:
+            return 0
+
+        # Handle line buffering for clean log output
+        text = self._line_buffer + s
+        lines = text.split("\n")
+
+        # If the text ends with a newline, the last line is complete
+        if text.endswith("\n"):
+            complete_lines = lines
+            self._line_buffer = ""
+        else:
+            # Otherwise, the last line is incomplete
+            complete_lines = lines[:-1]
+            self._line_buffer = lines[-1]
+
+        # Log complete lines but at debug level instead of info to prevent console spam
+        for line in complete_lines:
+            if line.strip():  # Only log non-empty lines
+                logger.debug(f"{self.server_name} (stderr): {line}")
+
+        # Always write to the underlying buffer
+        return self._buffer.write(s)
+
+    def flush(self) -> None:
+        """Flush the internal buffer."""
+        self._buffer.flush()
+
+    def close(self) -> None:
+        """Close the stream."""
+        # Log any remaining content in the line buffer
+        if self._line_buffer and self._line_buffer.strip():
+            logger.debug(f"{self.server_name} (stderr): {self._line_buffer}")
+        self._buffer.close()
+
+    def readable(self) -> bool:
+        return False
+
+    def writable(self) -> bool:
+        return True
+
+    def seekable(self) -> bool:
+        return False
+
+    def fileno(self) -> int:
+        """
+        Return a file descriptor for /dev/null.
+        This prevents output from showing on the terminal
+        while still allowing our write() method to capture it for logging.
+        """
+        if not hasattr(self, "_devnull_fd"):
+            self._devnull_fd = os.open(os.devnull, os.O_WRONLY)
+        return self._devnull_fd
+
+    def __del__(self):
+        """Clean up the devnull file descriptor."""
+        if hasattr(self, "_devnull_fd"):
+            try:
+                os.close(self._devnull_fd)
+            except (OSError, AttributeError):
+                pass
+
+
+def get_stderr_handler(server_name: str) -> TextIO:
+    """
+    Get a stderr handler that routes MCP server errors to our logger.
+
+    Args:
+        server_name: The name of the server to include in logs
+
+    Returns:
+        A TextIO object that can be used as stderr by MCP
+    """
+    return LoggerTextIO(server_name)

fast_agent/mcp/mcp_agent_client_session.py

@@ -0,0 +1,411 @@
+"""
+A derived client session for the MCP Agent framework.
+It adds logging and supports sampling requests.
+"""
+
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from mcp import ClientSession, ServerNotification
+from mcp.shared.message import MessageMetadata
+from mcp.shared.session import (
+    ProgressFnT,
+    ReceiveResultT,
+    SendRequestT,
+)
+from mcp.types import (
+    CallToolRequest,
+    CallToolRequestParams,
+    CallToolResult,
+    GetPromptRequest,
+    GetPromptRequestParams,
+    GetPromptResult,
+    Implementation,
+    ListRootsResult,
+    ReadResourceRequest,
+    ReadResourceRequestParams,
+    ReadResourceResult,
+    Root,
+    ToolListChangedNotification,
+)
+from pydantic import FileUrl
+
+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 fast_agent.config import MCPServerSettings
+    from fast_agent.mcp.transport_tracking import TransportChannelMetrics
+
+logger = get_logger(__name__)
+
+
+async def list_roots(ctx: ClientSession) -> ListRootsResult:
+    """List roots callback that will be called by the MCP library."""
+
+    if server_config := get_server_config(ctx):
+        if server_config.roots:
+            roots = [
+                Root(
+                    uri=FileUrl(
+                        root.server_uri_alias or root.uri,
+                    ),
+                    name=root.name,
+                )
+                for root in server_config.roots
+            ]
+            return ListRootsResult(roots=roots)
+
+    return ListRootsResult(roots=[])
+
+
+class MCPAgentClientSession(ClientSession, ContextDependent):
+    """
+    MCP Agent framework acts as a client to the servers providing tools/resources/prompts for the agent workloads.
+    This is a simple client session for those server connections, and supports
+        - handling sampling requests
+        - notifications
+        - MCP root configuration
+
+    Developers can extend this class to add more custom functionality as needed
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        # Extract server_name if provided in kwargs
+        from importlib.metadata import version
+
+        self.session_server_name = kwargs.pop("server_name", None)
+        # Extract the notification callbacks if provided
+        self._tool_list_changed_callback = kwargs.pop("tool_list_changed_callback", None)
+        # Extract server_config if provided
+        self.server_config: MCPServerSettings | None = kwargs.pop("server_config", None)
+        # Extract agent_model if provided (for auto_sampling fallback)
+        self.agent_model: str | None = kwargs.pop("agent_model", None)
+        # Extract agent_name if provided
+        self.agent_name: str | None = kwargs.pop("agent_name", None)
+        # Extract api_key if provided
+        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)
+        # Extract transport metrics tracker if provided
+        self._transport_metrics: TransportChannelMetrics | None = kwargs.pop(
+            "transport_metrics", None
+        )
+
+        # Track the effective elicitation mode for diagnostics
+        self.effective_elicitation_mode: str | None = "none"
+
+        version = version("fast-agent-mcp") or "dev"
+        fast_agent: Implementation = Implementation(name="fast-agent-mcp", version=version)
+        if self.server_config and self.server_config.implementation:
+            fast_agent = self.server_config.implementation
+
+        # Only register callbacks if the server_config has the relevant settings
+        list_roots_cb = list_roots if (self.server_config and self.server_config.roots) else None
+
+        # Register sampling callback if either:
+        # 1. Sampling is explicitly configured, OR
+        # 2. Application-level auto_sampling is enabled
+        sampling_cb = None
+        if self.server_config and self.server_config.sampling:
+            # Explicit sampling configuration
+            sampling_cb = sample
+        elif self._should_enable_auto_sampling():
+            # Auto-sampling enabled at application level
+            sampling_cb = sample
+
+        # Use custom elicitation handler if provided, otherwise resolve using factory
+        if custom_elicitation_handler is not None:
+            elicitation_handler = custom_elicitation_handler
+        else:
+            # Try to resolve using factory
+            elicitation_handler = None
+            try:
+                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:
+                    # Create a minimal agent config for the factory
+                    agent_config = AgentConfig(
+                        name=self.agent_name or "unknown",
+                        model=self.agent_model or "unknown",
+                        elicitation_handler=None,
+                    )
+                    elicitation_handler = resolve_elicitation_handler(
+                        agent_config, context.config, self.server_config
+                    )
+            except Exception:
+                # If factory resolution fails, we'll use default fallback
+                pass
+
+            # Fallback to forms handler only if factory resolution wasn't attempted
+            if elicitation_handler is None and not self.server_config:
+                from fast_agent.mcp.elicitation_handlers import forms_elicitation_handler
+
+                elicitation_handler = forms_elicitation_handler
+
+        # Determine effective elicitation mode for diagnostics
+        if self.server_config and getattr(self.server_config, "elicitation", None):
+            self.effective_elicitation_mode = self.server_config.elicitation.mode or "forms"
+        elif elicitation_handler is not None:
+            # Use global config if available to distinguish auto-cancel
+            try:
+                from fast_agent.context import get_current_context
+
+                context = get_current_context()
+                mode = None
+                if context and getattr(context, "config", None):
+                    elicitation_cfg = getattr(context.config, "elicitation", None)
+                    if isinstance(elicitation_cfg, dict):
+                        mode = elicitation_cfg.get("mode")
+                    else:
+                        mode = getattr(elicitation_cfg, "mode", None)
+                self.effective_elicitation_mode = (mode or "forms").lower()
+            except Exception:
+                self.effective_elicitation_mode = "forms"
+        else:
+            self.effective_elicitation_mode = "none"
+
+        super().__init__(
+            *args,
+            **kwargs,
+            list_roots_callback=list_roots_cb,
+            sampling_callback=sampling_cb,
+            client_info=fast_agent,
+            elicitation_callback=elicitation_handler,
+        )
+
+    def _should_enable_auto_sampling(self) -> bool:
+        """Check if auto_sampling is enabled at the application level."""
+        try:
+            from fast_agent.context import get_current_context
+
+            context = get_current_context()
+            if context and context.config:
+                return getattr(context.config, "auto_sampling", True)
+        except Exception:
+            pass
+        return True  # Default to True if can't access config
+
+    async def send_request(
+        self,
+        request: SendRequestT,
+        result_type: type[ReceiveResultT],
+        request_read_timeout_seconds: timedelta | None = None,
+        metadata: MessageMetadata | None = None,
+        progress_callback: ProgressFnT | None = None,
+    ) -> ReceiveResultT:
+        logger.debug("send_request: request=", data=request.model_dump())
+        request_id = getattr(self, "_request_id", None)
+        try:
+            result = await super().send_request(
+                request=request,
+                result_type=result_type,
+                request_read_timeout_seconds=request_read_timeout_seconds,
+                metadata=metadata,
+                progress_callback=progress_callback,
+            )
+            logger.debug(
+                "send_request: response=",
+                data=result.model_dump() if result is not None else "no response returned",
+            )
+            self._attach_transport_channel(request_id, result)
+            return result
+        except Exception as e:
+            from anyio import ClosedResourceError
+
+            from fast_agent.core.exceptions import ServerSessionTerminatedError
+
+            # Check for session terminated error (404 from server)
+            if self._is_session_terminated_error(e):
+                raise ServerSessionTerminatedError(
+                    server_name=self.session_server_name or "unknown",
+                    details="Server returned 404 - session may have expired due to server restart",
+                ) from e
+
+            # Handle connection closure errors (transport closed)
+            if isinstance(e, ClosedResourceError):
+                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
+
+            logger.error(f"send_request failed: {str(e)}")
+            raise
+
+    def _is_session_terminated_error(self, exc: Exception) -> bool:
+        """Check if exception is a session terminated error (code 32600 from 404)."""
+        from mcp.shared.exceptions import McpError
+
+        from fast_agent.core.exceptions import ServerSessionTerminatedError
+
+        if isinstance(exc, McpError):
+            error_data = getattr(exc, "error", None)
+            if error_data:
+                code = getattr(error_data, "code", None)
+                if code == ServerSessionTerminatedError.SESSION_TERMINATED_CODE:
+                    return True
+        return False
+
+    def _attach_transport_channel(self, request_id, result) -> None:
+        if self._transport_metrics is None or request_id is None or result is None:
+            return
+        channel = self._transport_metrics.consume_response_channel(request_id)
+        if not channel:
+            return
+        try:
+            setattr(result, "transport_channel", channel)
+        except Exception:
+            # If result cannot be mutated, ignore silently
+            pass
+
+    async def _received_notification(self, notification: ServerNotification) -> None:
+        """
+        Can be overridden by subclasses to handle a notification without needing
+        to listen on the message stream.
+        """
+        logger.debug(
+            "_received_notification: notification=",
+            data=notification.model_dump(),
+        )
+
+        # Call parent notification handler first
+        await super()._received_notification(notification)
+
+        # Then process our specific notification types
+        match notification.root:
+            case ToolListChangedNotification():
+                # Simple notification handling - just call the callback if it exists
+                if self._tool_list_changed_callback and self.session_server_name:
+                    logger.info(
+                        f"Tool list changed for server '{self.session_server_name}', triggering callback"
+                    )
+                    # Use asyncio.create_task to prevent blocking the notification handler
+                    import asyncio
+
+                    asyncio.create_task(
+                        self._handle_tool_list_change_callback(self.session_server_name)
+                    )
+                else:
+                    logger.debug(
+                        f"Tool list changed for server '{self.session_server_name}' but no callback registered"
+                    )
+
+        return None
+
+    async def _handle_tool_list_change_callback(self, server_name: str) -> None:
+        """
+        Helper method to handle tool list change callback in a separate task
+        to prevent blocking the notification handler
+        """
+        try:
+            await self._tool_list_changed_callback(server_name)
+        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,
+        progress_callback: ProgressFnT | None = None,
+        **kwargs,
+    ) -> CallToolResult:
+        """Call a tool with optional metadata and progress callback support.
+
+        Always uses our overridden send_request to ensure session terminated errors
+        are properly detected and converted to ServerSessionTerminatedError.
+        """
+        from mcp.types import RequestParams
+
+        # Always create request ourselves to ensure we go through our send_request override
+        # This is critical for session terminated detection to work
+        params = CallToolRequestParams(name=name, arguments=arguments)
+
+        if _meta:
+            # Safe merge - preserve existing meta fields like progressToken
+            existing_meta = kwargs.get("meta")
+            if existing_meta:
+                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)
+
+            params_dict = params.model_dump(by_alias=True)
+            params_dict["_meta"] = meta_obj.model_dump()
+            params = CallToolRequestParams.model_validate(params_dict)
+
+        request = CallToolRequest(method="tools/call", params=params)
+        return await self.send_request(
+            request, CallToolResult, progress_callback=progress_callback
+        )
+
+    async def read_resource(
+        self, uri: str, _meta: dict | None = None, **kwargs
+    ) -> ReadResourceResult:
+        """Read a resource with optional metadata support.
+
+        Always uses our overridden send_request to ensure session terminated errors
+        are properly detected and converted to ServerSessionTerminatedError.
+        """
+        from mcp.types import RequestParams
+
+        # Always create request ourselves to ensure we go through our send_request override
+        params = ReadResourceRequestParams(uri=uri)
+
+        if _meta:
+            # Safe merge - preserve existing meta fields like progressToken
+            existing_meta = kwargs.get("meta")
+            if existing_meta:
+                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)
+            params = ReadResourceRequestParams(uri=uri, meta=meta_obj)
+
+        request = ReadResourceRequest(method="resources/read", params=params)
+        return await self.send_request(request, ReadResourceResult)
+
+    async def get_prompt(
+        self, name: str, arguments: dict | None = None, _meta: dict | None = None, **kwargs
+    ) -> GetPromptResult:
+        """Get a prompt with optional metadata support.
+
+        Always uses our overridden send_request to ensure session terminated errors
+        are properly detected and converted to ServerSessionTerminatedError.
+        """
+        from mcp.types import RequestParams
+
+        # Always create request ourselves to ensure we go through our send_request override
+        params = GetPromptRequestParams(name=name, arguments=arguments)
+
+        if _meta:
+            # Safe merge - preserve existing meta fields like progressToken
+            existing_meta = kwargs.get("meta")
+            if existing_meta:
+                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)
+            params = GetPromptRequestParams(name=name, arguments=arguments, meta=meta_obj)
+
+        request = GetPromptRequest(method="prompts/get", params=params)
+        return await self.send_request(request, GetPromptResult)