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

mcp_agent/executor/temporal.py

@@ -0,0 +1,405 @@
+"""
+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 (
+    Any,
+    AsyncIterator,
+    Callable,
+    Coroutine,
+    Dict,
+    List,
+    Optional,
+    TYPE_CHECKING,
+)
+
+from pydantic import ConfigDict
+from temporalio import activity, workflow, exceptions
+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):
+            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):
+                # 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):
+        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):
+        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,
+    ):
+        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):
+        """
+        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()

mcp_agent/executor/workflow.py

@@ -0,0 +1,197 @@
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    TypeVar,
+    Union,
+)
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from mcp_agent.executor.executor import Executor
+
+T = TypeVar("T")
+
+
+class WorkflowState(BaseModel):
+    """
+    Simple container for persistent workflow state.
+    This can hold fields that should persist across tasks.
+    """
+
+    status: str = "initialized"
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    updated_at: float | None = None
+    error: Dict[str, Any] | None = None
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+    def record_error(self, error: Exception) -> None:
+        self.error = {
+            "type": type(error).__name__,
+            "message": str(error),
+            "timestamp": datetime.utcnow().timestamp(),
+        }
+
+
+class WorkflowResult(BaseModel, Generic[T]):
+    value: Union[T, None] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    start_time: float | None = None
+    end_time: float | None = None
+
+
+class Workflow(ABC, Generic[T]):
+    """
+    Base class for user-defined workflows.
+    Handles execution and state management.
+    Some key notes:
+        - To enable the executor engine to recognize and orchestrate the workflow,
+            - the class MUST be decorated with @workflow.
+            - the main entrypoint method MUST be decorated with @workflow_run.
+            - any task methods MUST be decorated with @workflow_task.
+
+        - Persistent state: Provides a simple `state` object for storing data across tasks.
+    """
+
+    def __init__(
+        self,
+        executor: Executor,
+        name: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+        **kwargs: Any,
+    ):
+        self.executor = executor
+        self.name = name or self.__class__.__name__
+        self.init_kwargs = kwargs
+        # TODO: handle logging
+        # self._logger = logging.getLogger(self.name)
+
+        # A simple workflow state object
+        # If under Temporal, storing it as a field on this class
+        # means it can be replayed automatically
+        self.state = WorkflowState(name=name, metadata=metadata or {})
+
+    @abstractmethod
+    async def run(self, *args: Any, **kwargs: Any) -> "WorkflowResult[T]":
+        """
+        Main workflow implementation. Must be overridden by subclasses.
+        """
+
+    async def update_state(self, **kwargs):
+        """Syntactic sugar to update workflow state."""
+        for key, value in kwargs.items():
+            self.state[key] = value
+            setattr(self.state, key, value)
+
+        self.state.updated_at = datetime.utcnow().timestamp()
+
+    async def wait_for_input(self, description: str = "Provide input") -> str:
+        """
+        Convenience method for human input. Uses `human_input` signal
+        so we can unify local (console input) and Temporal signals.
+        """
+        return await self.executor.wait_for_signal(
+            "human_input", description=description
+        )
+
+
+# ############################
+# # Example: DocumentWorkflow
+# ############################
+
+
+# @workflow_defn  # <-- This becomes @temporal_workflow.defn if in Temporal mode, else no-op
+# class DocumentWorkflow(Workflow[List[Dict[str, Any]]]):
+#     """
+#     Example workflow with persistent state.
+#     If run locally, `self.state` is ephemeral.
+#     If run in Temporal mode, `self.state` is replayed automatically.
+#     """
+
+#     @workflow_task(
+#         schedule_to_close_timeout=timedelta(minutes=10),
+#         retry_policy={"initial_interval": 1, "max_attempts": 3},
+#     )
+#     async def process_document(self, doc_id: str) -> Dict[str, Any]:
+#         """Activity that simulates document processing."""
+#         await asyncio.sleep(1)
+#         # Optionally mutate workflow state
+#         self.state.metadata.setdefault("processed_docs", []).append(doc_id)
+#         return {
+#             "doc_id": doc_id,
+#             "status": "processed",
+#             "timestamp": datetime.utcnow().isoformat(),
+#         }
+
+#     @workflow_run  # <-- This becomes @temporal_workflow.run(...) if Temporal is used
+#     async def _run_impl(
+#         self, documents: List[str], batch_size: int = 2
+#     ) -> List[Dict[str, Any]]:
+#         """Main workflow logic, which becomes the official 'run' in Temporal mode."""
+#         self._logger.info("Workflow starting, state=%s", self.state)
+#         self.state.update_status("running")
+
+#         all_results = []
+#         for i in range(0, len(documents), batch_size):
+#             batch = documents[i : i + batch_size]
+#             tasks = [self.process_document(doc) for doc in batch]
+#             results = await self.executor.execute(*tasks)
+
+#             for res in results:
+#                 if isinstance(res.value, Exception):
+#                     self._logger.error(
+#                         f"Error processing document: {res.metadata.get('error')}"
+#                     )
+#                 else:
+#                     all_results.append(res.value)
+
+#         self.state.update_status("completed")
+#         return all_results
+
+
+# ########################
+# # 12. Example Local Usage
+# ########################
+
+
+# async def run_example_local():
+#     from . import AsyncIOExecutor, DocumentWorkflow  # if in a package
+
+#     executor = AsyncIOExecutor()
+#     wf = DocumentWorkflow(executor)
+
+#     documents = ["doc1", "doc2", "doc3", "doc4"]
+#     result = await wf.run(documents, batch_size=2)
+
+#     print("Local results:", result.value)
+#     print("Local workflow final state:", wf.state)
+#     # Notice `wf.state.metadata['processed_docs']` has the processed doc IDs.
+
+
+# ########################
+# # Example Temporal Usage
+# ########################
+
+
+# async def run_example_temporal():
+#     from . import TemporalExecutor, DocumentWorkflow  # if in a package
+
+#     # 1) Create a TemporalExecutor (client side)
+#     executor = TemporalExecutor(task_queue="my_task_queue")
+#     await executor.ensure_client()
+
+#     # 2) Start a worker in the same process (or do so in a separate process)
+#     asyncio.create_task(executor.start_worker())
+#     await asyncio.sleep(2)  # Wait for worker to be up
+
+#     # 3) Now we can run the workflow by normal means if we like,
+#     #    or rely on the Worker picking it up. Typically, you'd do:
+#     #    handle = await executor._client.start_workflow(...)
+#     #    but let's keep it simple and show conceptually
+#     #    that 'DocumentWorkflow' is now recognized as a real Temporal workflow
+#     print(
+#         "Temporal environment is running. Use the Worker logs or CLI to start 'DocumentWorkflow'."
+#     )