fast-agent-mcp 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

fast_agent/mcp/mcp_aggregator.py

@@ -1,4 +1,7 @@
 from asyncio import Lock, gather
+from collections import Counter
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -17,11 +20,12 @@ from mcp.types import (
     CallToolResult,
     ListToolsResult,
     Prompt,
+    ServerCapabilities,
     TextContent,
     Tool,
 )
 from opentelemetry import trace
-from pydantic import AnyUrl, BaseModel, ConfigDict
+from pydantic import AnyUrl, BaseModel, ConfigDict, Field
 
 from fast_agent.context_dependent import ContextDependent
 from fast_agent.core.logging.logger import get_logger
@@ -30,6 +34,7 @@ from fast_agent.mcp.common import SEP, create_namespaced_name, is_namespaced_nam
 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
+from fast_agent.mcp.transport_tracking import TransportSnapshot
 
 if TYPE_CHECKING:
     from fast_agent.context import Context
@@ -52,6 +57,49 @@ class NamespacedTool(BaseModel):
     namespaced_tool_name: str
 
 
+@dataclass
+class ServerStats:
+    call_counts: Counter = field(default_factory=Counter)
+    last_call_at: datetime | None = None
+    last_error_at: datetime | None = None
+
+    def record(self, operation_type: str, success: bool) -> None:
+        self.call_counts[operation_type] += 1
+        now = datetime.now(timezone.utc)
+        self.last_call_at = now
+        if not success:
+            self.last_error_at = now
+
+
+class ServerStatus(BaseModel):
+    server_name: str
+    implementation_name: str | None = None
+    implementation_version: str | None = None
+    server_capabilities: ServerCapabilities | None = None
+    client_capabilities: Mapping[str, Any] | None = None
+    client_info_name: str | None = None
+    client_info_version: str | None = None
+    transport: str | None = None
+    is_connected: bool | None = None
+    last_call_at: datetime | None = None
+    last_error_at: datetime | None = None
+    staleness_seconds: float | None = None
+    call_counts: Dict[str, int] = Field(default_factory=dict)
+    error_message: str | None = None
+    instructions_available: bool | None = None
+    instructions_enabled: bool | None = None
+    instructions_included: bool | None = None
+    roots_configured: bool | None = None
+    roots_count: int | None = None
+    elicitation_mode: str | None = None
+    sampling_mode: str | None = None
+    spoofing_enabled: bool | None = None
+    session_id: str | None = None
+    transport_channels: TransportSnapshot | None = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
 class MCPAggregator(ContextDependent):
     """
     Aggregates multiple MCP servers. When a developer calls, e.g. call_tool(...),
@@ -140,6 +188,10 @@ class MCPAggregator(ContextDependent):
         # Lock for refreshing tools from a server
         self._refresh_lock = Lock()
 
+        # Track runtime stats per server
+        self._server_stats: Dict[str, ServerStats] = {}
+        self._stats_lock = Lock()
+
     def _create_progress_callback(self, server_name: str, tool_name: str) -> "ProgressFnT":
         """Create a progress callback function for tool execution."""
 
@@ -282,6 +334,9 @@ class MCPAggregator(ContextDependent):
                     server_name, client_session_factory=self._create_session_factory(server_name)
                 )
 
+                # Record the initialize call that happened during connection setup
+                await self._record_server_call(server_name, "initialize", True)
+
             logger.info(
                 f"MCP Servers initialized for agent '{self.agent_name}'",
                 data={
@@ -290,27 +345,39 @@ class MCPAggregator(ContextDependent):
                 },
             )
 
-        async def fetch_tools(client: ClientSession, server_name: str) -> List[Tool]:
+        async def fetch_tools(server_name: str) -> List[Tool]:
             # Only fetch tools if the server supports them
             if not await self.server_supports_feature(server_name, "tools"):
                 logger.debug(f"Server '{server_name}' does not support tools")
                 return []
 
             try:
-                result: ListToolsResult = await client.list_tools()
+                result: ListToolsResult = await self._execute_on_server(
+                    server_name=server_name,
+                    operation_type="tools/list",
+                    operation_name="",
+                    method_name="list_tools",
+                    method_args={},
+                )
                 return result.tools or []
             except Exception as e:
                 logger.error(f"Error loading tools from server '{server_name}'", data=e)
                 return []
 
-        async def fetch_prompts(client: ClientSession, server_name: str) -> List[Prompt]:
+        async def fetch_prompts(server_name: str) -> List[Prompt]:
             # Only fetch prompts if the server supports them
             if not await self.server_supports_feature(server_name, "prompts"):
                 logger.debug(f"Server '{server_name}' does not support prompts")
                 return []
 
             try:
-                result = await client.list_prompts()
+                result = await self._execute_on_server(
+                    server_name=server_name,
+                    operation_type="prompts/list",
+                    operation_name="",
+                    method_name="list_prompts",
+                    method_args={},
+                )
                 return getattr(result, "prompts", [])
             except Exception as e:
                 logger.debug(f"Error loading prompts from server '{server_name}': {e}")
@@ -320,20 +387,9 @@ class MCPAggregator(ContextDependent):
             tools: List[Tool] = []
             prompts: List[Prompt] = []
 
-            if self.connection_persistence:
-                server_connection = await self._persistent_connection_manager.get_server(
-                    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:
-                async with gen_client(
-                    server_name,
-                    server_registry=self.context.server_registry,
-                    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)
+            # Use _execute_on_server for consistent tracking regardless of connection mode
+            tools = await fetch_tools(server_name)
+            prompts = await fetch_prompts(server_name)
 
             return server_name, tools, prompts
 
@@ -461,6 +517,50 @@ class MCPAggregator(ContextDependent):
         for server_name in self.server_names:
             await self._refresh_server_tools(server_name)
 
+    async def _record_server_call(
+        self, server_name: str, operation_type: str, success: bool
+    ) -> None:
+        async with self._stats_lock:
+            stats = self._server_stats.setdefault(server_name, ServerStats())
+            stats.record(operation_type, success)
+
+            # For stdio servers, also emit synthetic transport events to create activity timeline
+            await self._notify_stdio_transport_activity(server_name, operation_type, success)
+
+    async def _notify_stdio_transport_activity(
+        self, server_name: str, operation_type: str, success: bool
+    ) -> None:
+        """Notify transport metrics of activity for stdio servers to create activity timeline."""
+        if not self._persistent_connection_manager:
+            return
+
+        try:
+            # Get the server connection and check if it's stdio transport
+            server_conn = self._persistent_connection_manager.running_servers.get(server_name)
+            if not server_conn:
+                return
+
+            server_config = getattr(server_conn, "server_config", None)
+            if not server_config or server_config.transport != "stdio":
+                return
+
+            # Get transport metrics and emit synthetic message event
+            transport_metrics = getattr(server_conn, "transport_metrics", None)
+            if transport_metrics:
+                # Import here to avoid circular imports
+                from fast_agent.mcp.transport_tracking import ChannelEvent
+
+                # Create a synthetic message event to represent the MCP operation
+                event = ChannelEvent(
+                    channel="stdio",
+                    event_type="message",
+                    detail=f"{operation_type} ({'success' if success else 'error'})"
+                )
+                transport_metrics.record_event(event)
+        except Exception:
+            # Don't let transport tracking errors break normal operation
+            logger.debug("Failed to notify stdio transport activity for %s", server_name, exc_info=True)
+
     async def get_server_instructions(self) -> Dict[str, tuple[str, List[str]]]:
         """
         Get instructions from all connected servers along with their tool names.
@@ -492,6 +592,174 @@ class MCPAggregator(ContextDependent):
 
         return instructions
 
+    async def collect_server_status(self) -> Dict[str, ServerStatus]:
+        """Return aggregated status information for each configured server."""
+        if not self.initialized:
+            await self.load_servers()
+
+        now = datetime.now(timezone.utc)
+        status_map: Dict[str, ServerStatus] = {}
+
+        for server_name in self.server_names:
+            stats = self._server_stats.get(server_name)
+            last_call = stats.last_call_at if stats else None
+            last_error = stats.last_error_at if stats else None
+            staleness = (now - last_call).total_seconds() if last_call else None
+            call_counts = dict(stats.call_counts) if stats else {}
+
+            implementation_name = None
+            implementation_version = None
+            capabilities: ServerCapabilities | None = None
+            client_capabilities: Mapping[str, Any] | None = None
+            client_info_name = None
+            client_info_version = None
+            is_connected = None
+            error_message = None
+            instructions_available = None
+            instructions_enabled = None
+            instructions_included = None
+            roots_configured = None
+            roots_count = None
+            elicitation_mode = None
+            sampling_mode = None
+            spoofing_enabled = None
+            server_cfg = None
+            session_id = None
+            server_conn = None
+            transport: str | None = None
+            transport_snapshot: TransportSnapshot | None = None
+
+            manager = getattr(self, "_persistent_connection_manager", None)
+            if self.connection_persistence and manager is not None:
+                try:
+                    server_conn = await manager.get_server(
+                        server_name,
+                        client_session_factory=self._create_session_factory(server_name),
+                    )
+                    implementation = getattr(server_conn, "server_implementation", None)
+                    if implementation:
+                        implementation_name = getattr(implementation, "name", None)
+                        implementation_version = getattr(implementation, "version", None)
+                    capabilities = getattr(server_conn, "server_capabilities", None)
+                    client_capabilities = getattr(server_conn, "client_capabilities", None)
+                    session = server_conn.session
+                    client_info = getattr(session, "client_info", None) if session else None
+                    if client_info:
+                        client_info_name = getattr(client_info, "name", None)
+                        client_info_version = getattr(client_info, "version", None)
+                    is_connected = server_conn.is_healthy()
+                    error_message = getattr(server_conn, "_error_message", None)
+                    instructions_available = getattr(
+                        server_conn, "server_instructions_available", None
+                    )
+                    instructions_enabled = getattr(
+                        server_conn, "server_instructions_enabled", None
+                    )
+                    instructions_included = bool(getattr(server_conn, "server_instructions", None))
+                    server_cfg = getattr(server_conn, "server_config", None)
+                    if session:
+                        elicitation_mode = getattr(session, "effective_elicitation_mode", elicitation_mode)
+                        session_id = getattr(server_conn, "session_id", None)
+                        if not session_id and getattr(server_conn, "_get_session_id_cb", None):
+                            try:
+                                session_id = server_conn._get_session_id_cb()  # type: ignore[attr-defined]
+                            except Exception:
+                                session_id = None
+                    metrics = getattr(server_conn, "transport_metrics", None)
+                    if metrics is not None:
+                        try:
+                            transport_snapshot = metrics.snapshot()
+                        except Exception:
+                            logger.debug(
+                                "Failed to snapshot transport metrics for server '%s'",
+                                server_name,
+                                exc_info=True,
+                            )
+                except Exception as exc:
+                    logger.debug(
+                        f"Failed to collect status for server '{server_name}'",
+                        data={"error": str(exc)},
+                    )
+
+            if server_cfg is None and self.context and getattr(self.context, "server_registry", None):
+                try:
+                    server_cfg = self.context.server_registry.get_server_config(server_name)
+                except Exception:
+                    server_cfg = None
+
+            if server_cfg is not None:
+                instructions_enabled = (
+                    instructions_enabled
+                    if instructions_enabled is not None
+                    else server_cfg.include_instructions
+                )
+                roots = getattr(server_cfg, "roots", None)
+                roots_configured = bool(roots)
+                roots_count = len(roots) if roots else 0
+                transport = getattr(server_cfg, "transport", transport)
+                elicitation = getattr(server_cfg, "elicitation", None)
+                elicitation_mode = (
+                    getattr(elicitation, "mode", None)
+                    if elicitation
+                    else elicitation_mode
+                )
+                sampling_cfg = getattr(server_cfg, "sampling", None)
+                spoofing_enabled = bool(getattr(server_cfg, "implementation", None))
+                if implementation_name is None and getattr(server_cfg, "implementation", None):
+                    implementation_name = server_cfg.implementation.name
+                    implementation_version = getattr(server_cfg.implementation, "version", None)
+                if session_id is None:
+                    if server_cfg.transport == "stdio":
+                        session_id = "local"
+                    elif server_conn and getattr(server_conn, "_get_session_id_cb", None):
+                        try:
+                            session_id = server_conn._get_session_id_cb()  # type: ignore[attr-defined]
+                        except Exception:
+                            session_id = None
+
+                if sampling_cfg is not None:
+                    sampling_mode = "configured"
+                else:
+                    auto_sampling = True
+                    if self.context and getattr(self.context, "config", None):
+                        auto_sampling = getattr(self.context.config, "auto_sampling", True)
+                    sampling_mode = "auto" if auto_sampling else "off"
+            else:
+                # Fall back to defaults when config missing
+                auto_sampling = True
+                if self.context and getattr(self.context, "config", None):
+                    auto_sampling = getattr(self.context.config, "auto_sampling", True)
+                sampling_mode = sampling_mode or ("auto" if auto_sampling else "off")
+
+            status_map[server_name] = ServerStatus(
+                server_name=server_name,
+                implementation_name=implementation_name,
+                implementation_version=implementation_version,
+                server_capabilities=capabilities,
+                client_capabilities=client_capabilities,
+                client_info_name=client_info_name,
+                client_info_version=client_info_version,
+                transport=transport,
+                is_connected=is_connected,
+                last_call_at=last_call,
+                last_error_at=last_error,
+                staleness_seconds=staleness,
+                call_counts=call_counts,
+                error_message=error_message,
+                instructions_available=instructions_available,
+                instructions_enabled=instructions_enabled,
+                instructions_included=instructions_included,
+                roots_configured=roots_configured,
+                roots_count=roots_count,
+                elicitation_mode=elicitation_mode,
+                sampling_mode=sampling_mode,
+                spoofing_enabled=spoofing_enabled,
+                session_id=session_id,
+                transport_channels=transport_snapshot,
+            )
+
+        return status_map
+
     async def _execute_on_server(
         self,
         server_name: str,
@@ -554,13 +822,17 @@ class MCPAggregator(ContextDependent):
                     # Re-raise the original exception to propagate it
                     raise e
 
+        success_flag: bool | None = None
+        result: R | None = None
+
         # Try initial execution
         try:
             if self.connection_persistence:
                 server_connection = await self._persistent_connection_manager.get_server(
                     server_name, client_session_factory=self._create_session_factory(server_name)
                 )
-                return await try_execute(server_connection.session)
+                result = await try_execute(server_connection.session)
+                success_flag = True
             else:
                 logger.debug(
                     f"Creating temporary connection to server: {server_name}",
@@ -582,7 +854,7 @@ class MCPAggregator(ContextDependent):
                             "agent_name": self.agent_name,
                         },
                     )
-                    return result
+                    success_flag = True
         except ConnectionError:
             # Server offline - attempt reconnection
             from fast_agent.ui import console
@@ -613,7 +885,7 @@ class MCPAggregator(ContextDependent):
 
                 # Success!
                 console.console.print(f"[dim green]MCP server {server_name} online[/dim green]")
-                return result
+                success_flag = True
 
             except Exception:
                 # Reconnection failed
@@ -621,10 +893,19 @@ class MCPAggregator(ContextDependent):
                     f"[dim red]MCP server {server_name} offline - failed to reconnect[/dim red]"
                 )
                 error_msg = f"MCP server {server_name} offline - failed to reconnect"
+                success_flag = False
                 if error_factory:
-                    return error_factory(error_msg)
+                    result = error_factory(error_msg)
                 else:
                     raise Exception(error_msg)
+        except Exception:
+            success_flag = False
+            raise
+        finally:
+            if success_flag is not None:
+                await self._record_server_call(server_name, operation_type, success_flag)
+
+        return result
 
     async def _parse_resource_name(self, name: str, resource_type: str) -> tuple[str, str]:
         """
@@ -701,7 +982,7 @@ class MCPAggregator(ContextDependent):
 
             return await self._execute_on_server(
                 server_name=server_name,
-                operation_type="tool",
+                operation_type="tools/call",
                 operation_name=local_tool_name,
                 method_name="call_tool",
                 method_args={
@@ -794,7 +1075,7 @@ class MCPAggregator(ContextDependent):
 
             result = await self._execute_on_server(
                 server_name=server_name,
-                operation_type="prompt",
+                operation_type="prompts/get",
                 operation_name=local_prompt_name or "default",
                 method_name="get_prompt",
                 method_args=method_args,
@@ -842,7 +1123,7 @@ class MCPAggregator(ContextDependent):
 
                     result = await self._execute_on_server(
                         server_name=s_name,
-                        operation_type="prompt",
+                        operation_type="prompts/get",
                         operation_name=local_prompt_name,
                         method_name="get_prompt",
                         method_args=method_args,
@@ -890,7 +1171,7 @@ class MCPAggregator(ContextDependent):
 
                     result = await self._execute_on_server(
                         server_name=s_name,
-                        operation_type="prompt",
+                        operation_type="prompts/get",
                         operation_name=local_prompt_name,
                         method_name="get_prompt",
                         method_args=method_args,
@@ -913,7 +1194,7 @@ class MCPAggregator(ContextDependent):
                         try:
                             prompt_list_result = await self._execute_on_server(
                                 server_name=s_name,
-                                operation_type="prompts-list",
+                                operation_type="prompts/list",
                                 operation_name="",
                                 method_name="list_prompts",
                                 error_factory=lambda _: None,
@@ -987,7 +1268,7 @@ class MCPAggregator(ContextDependent):
             # Fetch from server
             result = await self._execute_on_server(
                 server_name=server_name,
-                operation_type="prompts-list",
+                operation_type="prompts/list",
                 operation_name="",
                 method_name="list_prompts",
                 error_factory=lambda _: None,
@@ -1026,7 +1307,7 @@ class MCPAggregator(ContextDependent):
             try:
                 result = await self._execute_on_server(
                     server_name=s_name,
-                    operation_type="prompts-list",
+                    operation_type="prompts/list",
                     operation_name="",
                     method_name="list_prompts",
                     error_factory=lambda _: None,
@@ -1212,7 +1493,7 @@ class MCPAggregator(ContextDependent):
         # Use the _execute_on_server method to call read_resource on the server
         result = await self._execute_on_server(
             server_name=server_name,
-            operation_type="resource",
+            operation_type="resources/read",
             operation_name=resource_uri,
             method_name="read_resource",
             method_args={"uri": uri},
@@ -1263,7 +1544,7 @@ class MCPAggregator(ContextDependent):
                 # Use the _execute_on_server method to call list_resources on the server
                 result = await self._execute_on_server(
                     server_name=s_name,
-                    operation_type="resources-list",
+                    operation_type="resources/list",
                     operation_name="",
                     method_name="list_resources",
                     method_args={},  # Empty dictionary instead of None
@@ -1316,7 +1597,7 @@ class MCPAggregator(ContextDependent):
                 # Use the _execute_on_server method to call list_tools on the server
                 result = await self._execute_on_server(
                     server_name=s_name,
-                    operation_type="tools-list",
+                    operation_type="tools/list",
                     operation_name="",
                     method_name="list_tools",
                     method_args={},