fast-agent-mcp 0.2.12__py3-none-any.whl → 0.2.13__py3-none-any.whl

{fast_agent_mcp-0.2.12.dist-info → fast_agent_mcp-0.2.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fast-agent-mcp
-Version: 0.2.12
+Version: 0.2.13
 Summary: Define, Prompt and Test MCP enabled Agents and Workflows
 Author-email: Shaun Smith <fastagent@llmindset.co.uk>, Sarmad Qadri <sarmad@lastmileai.dev>
 License:                                  Apache License

{fast_agent_mcp-0.2.12.dist-info → fast_agent_mcp-0.2.13.dist-info}/RECORD

@@ -102,7 +102,7 @@ mcp_agent/mcp/prompts/prompt_load.py,sha256=Zo0FogqWFEG5FtF1d9ZH-RWsCSSMsi5FIEQH
 mcp_agent/mcp/prompts/prompt_server.py,sha256=SiUR2xYfd3vEpghnYRdzz2rFEMtAbCKx2xzUXgvz1g8,18501
 mcp_agent/mcp/prompts/prompt_template.py,sha256=EejiqGkau8OizORNyKTUwUjrPof5V-hH1H_MBQoQfXw,15732
 mcp_agent/mcp_server/__init__.py,sha256=zBU51ITHIEPScd9nRafnhEddsWqXRPAAvHhkrbRI2_4,155
-mcp_agent/mcp_server/agent_server.py,sha256=Nzws9h-itEqx9ml87Wx_mKhyoHehAvcAgsve0PZgghI,5905
+mcp_agent/mcp_server/agent_server.py,sha256=eMPUHcHgXNCldwOAh648D0TLNMPC_zSQtqShsT8s6Eg,16010
 mcp_agent/resources/examples/data-analysis/analysis-campaign.py,sha256=QdNdo0-7LR4Uzw61hEU_jVKmWyk6A9YpGo81kMwVobM,7267
 mcp_agent/resources/examples/data-analysis/analysis.py,sha256=M9z8Q4YC5OGuqSa5uefYmmfmctqMn-WqCSfg5LI407o,2609
 mcp_agent/resources/examples/data-analysis/fastagent.config.yaml,sha256=ini94PHyJCfgpjcjHKMMbGuHs6LIj46F1NwY0ll5HVk,1609
@@ -143,8 +143,8 @@ mcp_agent/resources/examples/workflows/parallel.py,sha256=n0dFN26QvYd2wjgohcaUBf
 mcp_agent/resources/examples/workflows/router.py,sha256=E4x_-c3l4YW9w1i4ARcDtkdeqIdbWEGfsMzwLYpdbVc,1677
 mcp_agent/resources/examples/workflows/short_story.txt,sha256=X3y_1AyhLFN2AKzCKvucJtDgAFIJfnlbsbGZO5bBWu0,1187
 mcp_agent/ui/console_display.py,sha256=TVGDtJ37hc6UG0ei9g7ZPZZfFNeS1MYozt-Mx8HsPCk,9752
-fast_agent_mcp-0.2.12.dist-info/METADATA,sha256=9PC3h7NaQW5XCnjmsldIGvAQ450FZ_GuT24UwDyL_R0,29940
-fast_agent_mcp-0.2.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-fast_agent_mcp-0.2.12.dist-info/entry_points.txt,sha256=bRniFM5zk3Kix5z7scX0gf9VnmGQ2Cz_Q1Gh7Ir4W00,186
-fast_agent_mcp-0.2.12.dist-info/licenses/LICENSE,sha256=cN3FxDURL9XuzE5mhK9L2paZo82LTfjwCYVT7e3j0e4,10939
-fast_agent_mcp-0.2.12.dist-info/RECORD,,
+fast_agent_mcp-0.2.13.dist-info/METADATA,sha256=aPMXpckeF8JsbXqQ8gGFWXpN208nlwUqkXmeG4yKd-w,29940
+fast_agent_mcp-0.2.13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+fast_agent_mcp-0.2.13.dist-info/entry_points.txt,sha256=bRniFM5zk3Kix5z7scX0gf9VnmGQ2Cz_Q1Gh7Ir4W00,186
+fast_agent_mcp-0.2.13.dist-info/licenses/LICENSE,sha256=cN3FxDURL9XuzE5mhK9L2paZo82LTfjwCYVT7e3j0e4,10939
+fast_agent_mcp-0.2.13.dist-info/RECORD,,

mcp_agent/mcp_server/agent_server.py

@@ -1,6 +1,12 @@
-# src/mcp_agent/mcp_server/agent_server.py
+"""
+Enhanced AgentMCPServer with robust shutdown handling for SSE transport.
+"""
 
 import asyncio
+import os
+import signal
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import Set
 
 from mcp.server.fastmcp import Context as MCPContext
 from mcp.server.fastmcp import FastMCP
@@ -9,6 +15,9 @@ import mcp_agent
 import mcp_agent.core
 import mcp_agent.core.prompt
 from mcp_agent.core.agent_app import AgentApp
+from mcp_agent.logging.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class AgentMCPServer:
@@ -20,14 +29,30 @@ class AgentMCPServer:
         server_name: str = "FastAgent-MCP-Server",
         server_description: str | None = None,
     ) -> None:
+        """Initialize the server with the provided agent app."""
         self.agent_app = agent_app
         self.mcp_server: FastMCP = FastMCP(
             name=server_name,
             instructions=server_description
             or f"This server provides access to {len(agent_app._agents)} agents",
         )
+        # Shutdown coordination
+        self._graceful_shutdown_event = asyncio.Event()
+        self._force_shutdown_event = asyncio.Event()
+        self._shutdown_timeout = 5.0  # Seconds to wait for graceful shutdown
+
+        # Resource management
+        self._exit_stack = AsyncExitStack()
+        self._active_connections: Set[any] = set()
+
+        # Server state
+        self._server_task = None
+
+        # Set up agent tools
         self.setup_tools()
 
+        logger.info(f"AgentMCPServer initialized with {len(agent_app._agents)} agents")
+
     def setup_tools(self) -> None:
         """Register all agents as MCP tools."""
         for agent_name, agent in self.agent_app._agents.items():
@@ -43,7 +68,6 @@ class AgentMCPServer:
         )
         async def send_message(message: str, ctx: MCPContext) -> str:
             """Send a message to the agent and return its response."""
-
             # Get the agent's context
             agent_context = getattr(agent, "context", None)
 
@@ -76,34 +100,212 @@ class AgentMCPServer:
             # that matches the structure that FastMCP expects (list of dicts with role/content)
             return [{"role": msg.role, "content": msg.content} for msg in prompt_messages]
 
+    def _setup_signal_handlers(self):
+        """Set up signal handlers for graceful and forced shutdown."""
+        loop = asyncio.get_running_loop()
+
+        def handle_signal(is_term=False):
+            # Use asyncio.create_task to handle the signal in an async-friendly way
+            asyncio.create_task(self._handle_shutdown_signal(is_term))
+
+        # Register handlers for SIGINT (Ctrl+C) and SIGTERM
+        for sig, is_term in [(signal.SIGINT, False), (signal.SIGTERM, True)]:
+            loop.add_signal_handler(sig, lambda term=is_term: handle_signal(term))
+
+        logger.debug("Signal handlers installed")
+
+    async def _handle_shutdown_signal(self, is_term=False):
+        """Handle shutdown signals with proper escalation."""
+        signal_name = "SIGTERM" if is_term else "SIGINT (Ctrl+C)"
+
+        # If force shutdown already requested, exit immediately
+        if self._force_shutdown_event.is_set():
+            logger.info("Force shutdown already in progress, exiting immediately...")
+            os._exit(1)
+
+        # If graceful shutdown already in progress, escalate to forced
+        if self._graceful_shutdown_event.is_set():
+            logger.info(f"Second {signal_name} received. Forcing shutdown...")
+            self._force_shutdown_event.set()
+            # Allow a brief moment for final cleanup, then force exit
+            await asyncio.sleep(0.5)
+            os._exit(1)
+
+        # First signal - initiate graceful shutdown
+        logger.info(f"{signal_name} received. Starting graceful shutdown...")
+        print(f"\n{signal_name} received. Starting graceful shutdown...")
+        print("Press Ctrl+C again to force exit.")
+        self._graceful_shutdown_event.set()
+
     def run(self, transport: str = "sse", host: str = "0.0.0.0", port: int = 8000) -> None:
-        """Run the MCP server."""
+        """Run the MCP server synchronously."""
         if transport == "sse":
-            # For running as a web server
             self.mcp_server.settings.host = host
             self.mcp_server.settings.port = port
 
-        self.mcp_server.run(transport=transport)
+        try:
+            self.mcp_server.run(transport=transport)
+        except KeyboardInterrupt:
+            print("\nServer stopped by user (CTRL+C)")
+        finally:
+            # Run an async cleanup in a new event loop
+            asyncio.run(self.shutdown())
 
     async def run_async(
         self, transport: str = "sse", host: str = "0.0.0.0", port: int = 8000
     ) -> None:
-        """Run the MCP server asynchronously."""
+        """Run the MCP server asynchronously with improved shutdown handling."""
+        # Use different handling strategies based on transport type
         if transport == "sse":
+            # For SSE, use our enhanced shutdown handling
+            self._setup_signal_handlers()
+
             self.mcp_server.settings.host = host
             self.mcp_server.settings.port = port
+
+            # Start the server in a separate task so we can monitor it
+            self._server_task = asyncio.create_task(self._run_server_with_shutdown(transport))
+
             try:
-                await self.mcp_server.run_sse_async()
-            except (asyncio.CancelledError, KeyboardInterrupt):
-                print("Server Stopped (CTRL+C)")
-                return
+                # Wait for the server task to complete
+                await self._server_task
+            except asyncio.CancelledError:
+                logger.info("Server task cancelled.")
+                print("\nServer task cancelled.")
+            except Exception as e:
+                logger.error(f"Server error: {e}", exc_info=True)
+                print(f"\nServer error: {e}")
+            finally:
+                # Ensure cleanup happens
+                await self.shutdown()
+                logger.info("Server shutdown complete.")
+                print("\nServer shutdown complete.")
         else:  # stdio
+            # For STDIO, use simpler approach that respects STDIO lifecycle
+            # STDIO will naturally terminate when streams close
             try:
+                # Run directly without extra monitoring or signal handlers
+                # This preserves the natural lifecycle of STDIO connections
                 await self.mcp_server.run_stdio_async()
             except (asyncio.CancelledError, KeyboardInterrupt):
-                # Gracefully handle cancellation during shutdown
-                print("Server Stopped (CTRL+C)")
-                return
+                logger.info("Server stopped (CTRL+C)")
+                print("\nServer stopped (CTRL+C)")
+
+            # Only perform minimal cleanup needed for STDIO
+            # Don't use our full shutdown procedure which could keep process alive
+            await self._cleanup_stdio()
+
+    async def _run_server_with_shutdown(self, transport: str):
+        """Run the server with proper shutdown handling."""
+        # This method is only used for SSE transport
+        if transport != "sse":
+            raise ValueError("This method should only be used with SSE transport")
+
+        # Start a monitor task for shutdown
+        shutdown_monitor = asyncio.create_task(self._monitor_shutdown())
+
+        try:
+            # Patch SSE server to track connections if needed
+            if hasattr(self.mcp_server, "_sse_transport") and self.mcp_server._sse_transport:
+                # Store the original connect_sse method
+                original_connect = self.mcp_server._sse_transport.connect_sse
+
+                # Create a wrapper that tracks connections
+                @asynccontextmanager
+                async def tracked_connect_sse(*args, **kwargs):
+                    async with original_connect(*args, **kwargs) as streams:
+                        self._active_connections.add(streams)
+                        try:
+                            yield streams
+                        finally:
+                            self._active_connections.discard(streams)
+
+                # Replace with our tracking version
+                self.mcp_server._sse_transport.connect_sse = tracked_connect_sse
+
+            # Run the server (SSE only)
+            await self.mcp_server.run_sse_async()
+        finally:
+            # Cancel the monitor when the server exits
+            shutdown_monitor.cancel()
+            try:
+                await shutdown_monitor
+            except asyncio.CancelledError:
+                pass
+
+    async def _monitor_shutdown(self):
+        """Monitor for shutdown signals and coordinate proper shutdown sequence."""
+        try:
+            # Wait for graceful shutdown request
+            await self._graceful_shutdown_event.wait()
+            logger.info("Graceful shutdown initiated")
+
+            # Two possible paths:
+            # 1. Wait for force shutdown
+            # 2. Wait for shutdown timeout
+            force_shutdown_task = asyncio.create_task(self._force_shutdown_event.wait())
+            timeout_task = asyncio.create_task(asyncio.sleep(self._shutdown_timeout))
+
+            # Wait for either force shutdown or timeout
+            done, pending = await asyncio.wait(
+                [force_shutdown_task, timeout_task], return_when=asyncio.FIRST_COMPLETED
+            )
+
+            # Cancel the remaining task
+            for task in pending:
+                task.cancel()
+
+            # Determine the shutdown reason
+            if force_shutdown_task in done:
+                logger.info("Force shutdown requested")
+                print("\nForced shutdown initiated...")
+            else:
+                logger.info(f"Graceful shutdown timed out after {self._shutdown_timeout} seconds")
+                print(f"\nGraceful shutdown timed out after {self._shutdown_timeout} seconds")
+
+            # Force close any remaining SSE connections
+            await self._close_sse_connections()
+
+            # Cancel the server task if running
+            if self._server_task and not self._server_task.done():
+                logger.info("Cancelling server task")
+                self._server_task.cancel()
+        except asyncio.CancelledError:
+            # Monitor was cancelled - clean exit
+            pass
+        except Exception as e:
+            logger.error(f"Error in shutdown monitor: {e}", exc_info=True)
+
+    async def _close_sse_connections(self):
+        """Force close all SSE connections."""
+        # Close tracked connections
+        for conn in list(self._active_connections):
+            try:
+                if hasattr(conn, "close"):
+                    await conn.close()
+                elif hasattr(conn, "aclose"):
+                    await conn.aclose()
+            except Exception as e:
+                logger.error(f"Error closing connection: {e}")
+            self._active_connections.discard(conn)
+
+        # Access the SSE transport if it exists to close stream writers
+        if (
+            hasattr(self.mcp_server, "_sse_transport")
+            and self.mcp_server._sse_transport is not None
+        ):
+            sse = self.mcp_server._sse_transport
+
+            # Close all read stream writers
+            if hasattr(sse, "_read_stream_writers"):
+                writers = list(sse._read_stream_writers.items())
+                for session_id, writer in writers:
+                    try:
+                        logger.debug(f"Closing SSE connection: {session_id}")
+                        await writer.aclose()
+                        sse._read_stream_writers.pop(session_id, None)
+                    except Exception as e:
+                        logger.error(f"Error closing SSE connection {session_id}: {e}")
 
     async def with_bridged_context(self, agent_context, mcp_context, func, *args, **kwargs):
         """
@@ -146,7 +348,44 @@ class AgentMCPServer:
             if hasattr(agent_context, "mcp_context"):
                 delattr(agent_context, "mcp_context")
 
+    async def _cleanup_stdio(self):
+        """Minimal cleanup for STDIO transport to avoid keeping process alive."""
+        logger.info("Performing minimal STDIO cleanup")
+
+        # Just clean up agent resources directly without the full shutdown sequence
+        # This preserves the natural exit process for STDIO
+        for agent_name, agent in self.agent_app._agents.items():
+            try:
+                if hasattr(agent, "shutdown"):
+                    await agent.shutdown()
+            except Exception as e:
+                logger.error(f"Error shutting down agent {agent_name}: {e}")
+
+        logger.info("STDIO cleanup complete")
+
     async def shutdown(self):
         """Gracefully shutdown the MCP server and its resources."""
-        # Your MCP server may have additional cleanup code here
-        pass
+        logger.info("Running full shutdown procedure")
+
+        # Skip if already in shutdown
+        if self._graceful_shutdown_event.is_set():
+            return
+
+        # Signal shutdown
+        self._graceful_shutdown_event.set()
+
+        # Close SSE connections
+        await self._close_sse_connections()
+
+        # Close any resources in the exit stack
+        await self._exit_stack.aclose()
+
+        # Shutdown any agent resources
+        for agent_name, agent in self.agent_app._agents.items():
+            try:
+                if hasattr(agent, "shutdown"):
+                    await agent.shutdown()
+            except Exception as e:
+                logger.error(f"Error shutting down agent {agent_name}: {e}")
+
+        logger.info("Full shutdown complete")