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

mcp_agent/mcp_server/agent_server.py

@@ -143,13 +143,38 @@ class AgentMCPServer:
             self.mcp_server.settings.host = host
             self.mcp_server.settings.port = port
 
-        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())
+            # For synchronous run, we can use the simpler approach
+            try:
+                # Add any server attributes that might help with shutdown
+                if not hasattr(self.mcp_server, "_server_should_exit"):
+                    self.mcp_server._server_should_exit = False
+
+                # Run the server
+                self.mcp_server.run(transport=transport)
+            except KeyboardInterrupt:
+                print("\nServer stopped by user (CTRL+C)")
+            except SystemExit as e:
+                # Handle normal exit
+                print(f"\nServer exiting with code {e.code}")
+                # Re-raise to allow normal exit process
+                raise
+            except Exception as e:
+                print(f"\nServer error: {e}")
+            finally:
+                # Run an async cleanup in a new event loop
+                try:
+                    asyncio.run(self.shutdown())
+                except (SystemExit, KeyboardInterrupt):
+                    # These are expected during shutdown
+                    pass
+        else:  # stdio
+            try:
+                self.mcp_server.run(transport=transport)
+            except KeyboardInterrupt:
+                print("\nServer stopped by user (CTRL+C)")
+            finally:
+                # Minimal cleanup for stdio
+                asyncio.run(self._cleanup_stdio())
 
     async def run_async(
         self, transport: str = "sse", host: str = "0.0.0.0", port: int = 8000
@@ -169,20 +194,26 @@ class AgentMCPServer:
             try:
                 # Wait for the server task to complete
                 await self._server_task
-            except asyncio.CancelledError:
-                logger.info("Server task cancelled.")
-                print("\nServer task cancelled.")
+            except (asyncio.CancelledError, KeyboardInterrupt):
+                # Both cancellation and KeyboardInterrupt are expected during shutdown
+                logger.info("Server stopped via cancellation or interrupt")
+                print("\nServer stopped")
+            except SystemExit as e:
+                # Handle normal exit cleanly
+                logger.info(f"Server exiting with code {e.code}")
+                print(f"\nServer exiting with code {e.code}")
+                # If this is exit code 0, let it propagate for normal exit
+                if e.code == 0:
+                    raise
             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.")
+                # Only do minimal cleanup - don't try to be too clever
+                await self._cleanup_stdio()
                 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
@@ -190,9 +221,14 @@ class AgentMCPServer:
             except (asyncio.CancelledError, KeyboardInterrupt):
                 logger.info("Server stopped (CTRL+C)")
                 print("\nServer stopped (CTRL+C)")
-
+            except SystemExit as e:
+                # Handle normal exit cleanly
+                logger.info(f"Server exiting with code {e.code}")
+                print(f"\nServer exiting with code {e.code}")
+                # If this is exit code 0, let it propagate for normal exit
+                if e.code == 0:
+                    raise
             # 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):
@@ -246,7 +282,6 @@ class AgentMCPServer:
             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
             )
@@ -255,21 +290,16 @@ class AgentMCPServer:
             for task in pending:
                 task.cancel()
 
-            # Determine the shutdown reason
+            # Determine shutdown reason
             if force_shutdown_task in done:
-                logger.info("Force shutdown requested")
-                print("\nForced shutdown initiated...")
+                logger.info("Force shutdown requested by user")
+                print("\nForce 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()
+                os._exit(0)
 
-            # 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
@@ -302,11 +332,36 @@ class AgentMCPServer:
                 for session_id, writer in writers:
                     try:
                         logger.debug(f"Closing SSE connection: {session_id}")
+                        # Instead of aclose, try to close more gracefully
+                        # Send a special event to notify client, then close
+                        try:
+                            if hasattr(writer, "send") and not getattr(writer, "_closed", False):
+                                try:
+                                    # Try to send a close event if possible
+                                    await writer.send(Exception("Server shutting down"))
+                                except (AttributeError, asyncio.CancelledError):
+                                    pass
+                        except Exception:
+                            pass
+
+                        # Now close the stream
                         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}")
 
+        # If we have a ASGI lifespan hook, try to signal closure
+        if (
+            hasattr(self.mcp_server, "_lifespan_state")
+            and self.mcp_server._lifespan_state == "started"
+        ):
+            logger.debug("Attempting to signal ASGI lifespan shutdown")
+            try:
+                if hasattr(self.mcp_server, "_on_shutdown"):
+                    await self.mcp_server._on_shutdown()
+            except Exception as e:
+                logger.error(f"Error during ASGI lifespan shutdown: {e}")
+
     async def with_bridged_context(self, agent_context, mcp_context, func, *args, **kwargs):
         """
         Execute a function with bridged context between MCP and agent
@@ -374,18 +429,45 @@ class AgentMCPServer:
         # Signal shutdown
         self._graceful_shutdown_event.set()
 
-        # Close SSE connections
-        await self._close_sse_connections()
+        try:
+            # Close SSE connections
+            await self._close_sse_connections()
 
-        # Close any resources in the exit stack
-        await self._exit_stack.aclose()
+            # 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}")
+            # 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}")
+        except Exception as e:
+            # Log any errors but don't let them prevent shutdown
+            logger.error(f"Error during shutdown: {e}", exc_info=True)
+        finally:
+            logger.info("Full shutdown complete")
+
+    async def _cleanup_minimal(self):
+        """Perform minimal cleanup before simulating a KeyboardInterrupt."""
+        logger.info("Performing minimal cleanup before interrupt")
+
+        # Only close SSE connection writers directly
+        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"):
+                for session_id, writer in list(sse._read_stream_writers.items()):
+                    try:
+                        await writer.aclose()
+                    except Exception:
+                        # Ignore errors during cleanup
+                        pass
 
-        logger.info("Full shutdown complete")
+        # Clear active connections set to prevent further operations
+        self._active_connections.clear()

mcp_agent/cli/commands/config.py

@@ -1,11 +0,0 @@
-from typing import NoReturn
-
-import typer
-
-app = typer.Typer()
-
-
-@app.command()
-def show() -> NoReturn:
-    """Show the configuration."""
-    raise NotImplementedError("The show configuration command has not been implemented yet")

mcp_agent/executor/temporal.py

@@ -1,383 +0,0 @@
-"""
-Temporal based orchestrator for the MCP Agent.
-Temporal provides durable execution and robust workflow orchestration,
-as well as dynamic control flow, making it a good choice for an AI agent orchestrator.
-Read more: https://docs.temporal.io/develop/python/core-application
-"""
-
-import asyncio
-import functools
-import uuid
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    AsyncIterator,
-    Callable,
-    Coroutine,
-    Dict,
-    List,
-    Optional,
-)
-
-from pydantic import ConfigDict
-from temporalio import activity, exceptions, workflow
-from temporalio.client import Client as TemporalClient
-from temporalio.worker import Worker
-
-from mcp_agent.config import TemporalSettings
-from mcp_agent.executor.executor import Executor, ExecutorConfig, R
-from mcp_agent.executor.workflow_signal import (
-    BaseSignalHandler,
-    Signal,
-    SignalHandler,
-    SignalRegistration,
-    SignalValueT,
-)
-
-if TYPE_CHECKING:
-    from mcp_agent.context import Context
-
-
-class TemporalSignalHandler(BaseSignalHandler[SignalValueT]):
-    """Temporal-based signal handling using workflow signals"""
-
-    async def wait_for_signal(self, signal, timeout_seconds=None) -> SignalValueT:
-        if not workflow._Runtime.current():
-            raise RuntimeError(
-                "TemporalSignalHandler.wait_for_signal must be called from within a workflow"
-            )
-
-        unique_signal_name = f"{signal.name}_{uuid.uuid4()}"
-        registration = SignalRegistration(
-            signal_name=signal.name,
-            unique_name=unique_signal_name,
-            workflow_id=workflow.info().workflow_id,
-        )
-
-        # Container for signal value
-        container = {"value": None, "completed": False}
-
-        # Define the signal handler for this specific registration
-        @workflow.signal(name=unique_signal_name)
-        def signal_handler(value: SignalValueT) -> None:
-            container["value"] = value
-            container["completed"] = True
-
-        async with self._lock:
-            # Register both the signal registration and handler atomically
-            self._pending_signals.setdefault(signal.name, []).append(registration)
-            self._handlers.setdefault(signal.name, []).append((unique_signal_name, signal_handler))
-
-        try:
-            # Wait for signal with optional timeout
-            await workflow.wait_condition(lambda: container["completed"], timeout=timeout_seconds)
-
-            return container["value"]
-        except asyncio.TimeoutError as exc:
-            raise TimeoutError(f"Timeout waiting for signal {signal.name}") from exc
-        finally:
-            async with self._lock:
-                # Remove ourselves from _pending_signals
-                if signal.name in self._pending_signals:
-                    self._pending_signals[signal.name] = [
-                        sr
-                        for sr in self._pending_signals[signal.name]
-                        if sr.unique_name != unique_signal_name
-                    ]
-                    if not self._pending_signals[signal.name]:
-                        del self._pending_signals[signal.name]
-
-                # Remove ourselves from _handlers
-                if signal.name in self._handlers:
-                    self._handlers[signal.name] = [
-                        h for h in self._handlers[signal.name] if h[0] != unique_signal_name
-                    ]
-                    if not self._handlers[signal.name]:
-                        del self._handlers[signal.name]
-
-    def on_signal(self, signal_name):
-        """Decorator to register a signal handler."""
-
-        def decorator(func: Callable) -> Callable:
-            # Create unique signal name for this handler
-            unique_signal_name = f"{signal_name}_{uuid.uuid4()}"
-
-            # Create the actual handler that will be registered with Temporal
-            @workflow.signal(name=unique_signal_name)
-            async def wrapped(signal_value: SignalValueT) -> None:
-                # Create a signal object to pass to the handler
-                signal = Signal(
-                    name=signal_name,
-                    payload=signal_value,
-                    workflow_id=workflow.info().workflow_id,
-                )
-                if asyncio.iscoroutinefunction(func):
-                    await func(signal)
-                else:
-                    func(signal)
-
-            # Register the handler under the original signal name
-            self._handlers.setdefault(signal_name, []).append((unique_signal_name, wrapped))
-            return func
-
-        return decorator
-
-    async def signal(self, signal) -> None:
-        self.validate_signal(signal)
-
-        workflow_handle = workflow.get_external_workflow_handle(workflow_id=signal.workflow_id)
-
-        # Send the signal to all registrations of this signal
-        async with self._lock:
-            signal_tasks = []
-
-            if signal.name in self._pending_signals:
-                for pending_signal in self._pending_signals[signal.name]:
-                    registration = pending_signal.registration
-                    if registration.workflow_id == signal.workflow_id:
-                        # Only signal for registrations of that workflow
-                        signal_tasks.append(
-                            workflow_handle.signal(registration.unique_name, signal.payload)
-                        )
-                    else:
-                        continue
-
-            # Notify any registered handler functions
-            if signal.name in self._handlers:
-                for unique_name, _ in self._handlers[signal.name]:
-                    signal_tasks.append(workflow_handle.signal(unique_name, signal.payload))
-
-        await asyncio.gather(*signal_tasks, return_exceptions=True)
-
-    def validate_signal(self, signal) -> None:
-        super().validate_signal(signal)
-        # Add TemporalSignalHandler-specific validation
-        if signal.workflow_id is None:
-            raise ValueError(
-                "No workflow_id provided on Signal. That is required for Temporal signals"
-            )
-
-
-class TemporalExecutorConfig(ExecutorConfig, TemporalSettings):
-    """Configuration for Temporal executors."""
-
-    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
-
-
-class TemporalExecutor(Executor):
-    """Executor that runs @workflows as Temporal workflows, with @workflow_tasks as Temporal activities"""
-
-    def __init__(
-        self,
-        config: TemporalExecutorConfig | None = None,
-        signal_bus: SignalHandler | None = None,
-        client: TemporalClient | None = None,
-        context: Optional["Context"] = None,
-        **kwargs,
-    ) -> None:
-        signal_bus = signal_bus or TemporalSignalHandler()
-        super().__init__(
-            engine="temporal",
-            config=config,
-            signal_bus=signal_bus,
-            context=context,
-            **kwargs,
-        )
-        self.config: TemporalExecutorConfig = (
-            config or self.context.config.temporal or TemporalExecutorConfig()
-        )
-        self.client = client
-        self._worker = None
-        self._activity_semaphore = None
-
-        if config.max_concurrent_activities is not None:
-            self._activity_semaphore = asyncio.Semaphore(self.config.max_concurrent_activities)
-
-    @staticmethod
-    def wrap_as_activity(
-        activity_name: str,
-        func: Callable[..., R] | Coroutine[Any, Any, R],
-        **kwargs: Any,
-    ) -> Coroutine[Any, Any, R]:
-        """
-        Convert a function into a Temporal activity and return its info.
-        """
-
-        @activity.defn(name=activity_name)
-        async def wrapped_activity(*args, **local_kwargs):
-            try:
-                if asyncio.iscoroutinefunction(func):
-                    return await func(*args, **local_kwargs)
-                elif asyncio.iscoroutine(func):
-                    return await func
-                else:
-                    return func(*args, **local_kwargs)
-            except Exception as e:
-                # Handle exceptions gracefully
-                raise e
-
-        return wrapped_activity
-
-    async def _execute_task_as_async(
-        self, task: Callable[..., R] | Coroutine[Any, Any, R], **kwargs: Any
-    ) -> R | BaseException:
-        async def run_task(task: Callable[..., R] | Coroutine[Any, Any, R]) -> R:
-            try:
-                if asyncio.iscoroutine(task):
-                    return await task
-                elif asyncio.iscoroutinefunction(task):
-                    return await task(**kwargs)
-                else:
-                    # Execute the callable and await if it returns a coroutine
-                    loop = asyncio.get_running_loop()
-
-                    # If kwargs are provided, wrap the function with partial
-                    if kwargs:
-                        wrapped_task = functools.partial(task, **kwargs)
-                        result = await loop.run_in_executor(None, wrapped_task)
-                    else:
-                        result = await loop.run_in_executor(None, task)
-
-                    # Handle case where the sync function returns a coroutine
-                    if asyncio.iscoroutine(result):
-                        return await result
-
-                    return result
-            except Exception as e:
-                # TODO: saqadri - adding logging or other error handling here
-                return e
-
-        if self._activity_semaphore:
-            async with self._activity_semaphore:
-                return await run_task(task)
-        else:
-            return await run_task(task)
-
-    async def _execute_task(
-        self, task: Callable[..., R] | Coroutine[Any, Any, R], **kwargs: Any
-    ) -> R | BaseException:
-        func = task.func if isinstance(task, functools.partial) else task
-        is_workflow_task = getattr(func, "is_workflow_task", False)
-        if not is_workflow_task:
-            return await asyncio.create_task(self._execute_task_as_async(task, **kwargs))
-
-        execution_metadata: Dict[str, Any] = getattr(func, "execution_metadata", {})
-
-        # Derive stable activity name, e.g. module + qualname
-        activity_name = execution_metadata.get("activity_name")
-        if not activity_name:
-            activity_name = f"{func.__module__}.{func.__qualname__}"
-
-        schedule_to_close = execution_metadata.get(
-            "schedule_to_close_timeout", self.config.timeout_seconds
-        )
-
-        retry_policy = execution_metadata.get("retry_policy", None)
-
-        _task_activity = self.wrap_as_activity(activity_name=activity_name, func=task)
-
-        # # For partials, we pass the partial's arguments
-        # args = task.args if isinstance(task, functools.partial) else ()
-        try:
-            result = await workflow.execute_activity(
-                activity_name,
-                args=kwargs.get("args", ()),
-                task_queue=self.config.task_queue,
-                schedule_to_close_timeout=schedule_to_close,
-                retry_policy=retry_policy,
-                **kwargs,
-            )
-            return result
-        except Exception as e:
-            # Properly propagate activity errors
-            if isinstance(e, exceptions.ActivityError):
-                raise e.cause if e.cause else e
-            raise
-
-    async def execute(
-        self,
-        *tasks: Callable[..., R] | Coroutine[Any, Any, R],
-        **kwargs: Any,
-    ) -> List[R | BaseException]:
-        # Must be called from within a workflow
-        if not workflow._Runtime.current():
-            raise RuntimeError("TemporalExecutor.execute must be called from within a workflow")
-
-        # TODO: saqadri - validate if async with self.execution_context() is needed here
-        async with self.execution_context():
-            return await asyncio.gather(
-                *(self._execute_task(task, **kwargs) for task in tasks),
-                return_exceptions=True,
-            )
-
-    async def execute_streaming(
-        self,
-        *tasks: Callable[..., R] | Coroutine[Any, Any, R],
-        **kwargs: Any,
-    ) -> AsyncIterator[R | BaseException]:
-        if not workflow._Runtime.current():
-            raise RuntimeError(
-                "TemporalExecutor.execute_streaming must be called from within a workflow"
-            )
-
-        # TODO: saqadri - validate if async with self.execution_context() is needed here
-        async with self.execution_context():
-            # Create futures for all tasks
-            futures = [self._execute_task(task, **kwargs) for task in tasks]
-            pending = set(futures)
-
-            while pending:
-                done, pending = await workflow.wait(pending, return_when=asyncio.FIRST_COMPLETED)
-                for future in done:
-                    try:
-                        result = await future
-                        yield result
-                    except Exception as e:
-                        yield e
-
-    async def ensure_client(self):
-        """Ensure we have a connected Temporal client."""
-        if self.client is None:
-            self.client = await TemporalClient.connect(
-                target_host=self.config.host,
-                namespace=self.config.namespace,
-                api_key=self.config.api_key,
-            )
-
-        return self.client
-
-    async def start_worker(self) -> None:
-        """
-        Start a worker in this process, auto-registering all tasks
-        from the global registry. Also picks up any classes decorated
-        with @workflow_defn as recognized workflows.
-        """
-        await self.ensure_client()
-
-        if self._worker is None:
-            # We'll collect the activities from the global registry
-            # and optionally wrap them with `activity.defn` if we want
-            # (Not strictly required if your code calls `execute_activity("name")` by name)
-            activity_registry = self.context.task_registry
-            activities = []
-            for name in activity_registry.list_activities():
-                activities.append(activity_registry.get_activity(name))
-
-            # Now we attempt to discover any classes that are recognized as workflows
-            # But in this simple example, we rely on the user specifying them or
-            # we might do a dynamic scan.
-            # For demonstration, we'll just assume the user is only using
-            # the workflow classes they decorate with `@workflow_defn`.
-            # We'll rely on them passing the classes or scanning your code.
-
-            self._worker = Worker(
-                client=self.client,
-                task_queue=self.config.task_queue,
-                activities=activities,
-                workflows=[],  # We'll auto-load by Python scanning or let the user specify
-            )
-            print(
-                f"Starting Temporal Worker on task queue '{self.config.task_queue}' with {len(activities)} activities."
-            )
-
-        await self._worker.run()