mcp-celery 0.1.0__py3-none-any.whl

mcp_celery/__init__.py

@@ -0,0 +1,80 @@
+"""mcp-celery — Expose Celery tasks as async MCP tools."""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+try:
+    __version__ = _version("mcp-celery")
+except PackageNotFoundError:
+    __version__ = "0.1.0.dev0"
+
+from typing import TYPE_CHECKING
+
+# These modules have no external dependencies — safe to import eagerly
+from mcp_celery.errors import BackendError, McpCeleryError, TaskNotFound, TaskNotReady
+from mcp_celery.lifecycle import map_celery_state
+from mcp_celery.registry import ToolRegistry
+from mcp_celery.schema import AsyncToolDef, TaskInfo, TaskStatus, format_response
+
+if TYPE_CHECKING:
+    from mcp_celery.backend.base import AbstractBackend
+    from mcp_celery.backend.celery import CeleryBackend
+    from mcp_celery.exposure.base import AbstractExposureStrategy, GeneratedTool
+    from mcp_celery.exposure.operation_resource import OperationResourceStrategy
+    from mcp_celery.exposure.polling import PollingExposureStrategy
+    from mcp_celery.server import AsyncToolServer
+    from mcp_celery.stores.celery_task_store import CeleryTaskStore
+
+
+def __getattr__(name: str):
+    """Lazy-load modules that depend on celery / mcp at runtime."""
+    _lazy_imports = {
+        "AbstractBackend": "mcp_celery.backend.base",
+        "CeleryBackend": "mcp_celery.backend.celery",
+        "AbstractExposureStrategy": "mcp_celery.exposure.base",
+        "GeneratedTool": "mcp_celery.exposure.base",
+        "OperationResourceStrategy": "mcp_celery.exposure.operation_resource",
+        "PollingExposureStrategy": "mcp_celery.exposure.polling",
+        "AsyncToolServer": "mcp_celery.server",
+        "CeleryTaskStore": "mcp_celery.stores.celery_task_store",
+    }
+    if name in _lazy_imports:
+        import importlib
+
+        module = importlib.import_module(_lazy_imports[name])
+        return getattr(module, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    # Version
+    "__version__",
+    # Main entry point
+    "AsyncToolServer",
+    # Schema / types
+    "AsyncToolDef",
+    "TaskInfo",
+    "TaskStatus",
+    "format_response",
+    # Backend
+    "AbstractBackend",
+    "CeleryBackend",
+    # Exposure strategies
+    "AbstractExposureStrategy",
+    "GeneratedTool",
+    "OperationResourceStrategy",
+    "PollingExposureStrategy",
+    # Stores
+    "CeleryTaskStore",
+    # Registry
+    "ToolRegistry",
+    # Lifecycle
+    "map_celery_state",
+    # Errors
+    "McpCeleryError",
+    "TaskNotReady",
+    "TaskNotFound",
+    "BackendError",
+]

mcp_celery/backend/__init__.py

@@ -0,0 +1,21 @@
+"""Execution backend layer."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from mcp_celery.backend.base import AbstractBackend
+
+if TYPE_CHECKING:
+    from mcp_celery.backend.celery import CeleryBackend
+
+
+def __getattr__(name: str):
+    if name == "CeleryBackend":
+        from mcp_celery.backend.celery import CeleryBackend
+
+        return CeleryBackend
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["AbstractBackend", "CeleryBackend"]

mcp_celery/backend/base.py

@@ -0,0 +1,52 @@
+"""Abstract execution backend interface."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from mcp_celery.schema import TaskInfo
+
+
+class AbstractBackend(ABC):
+    """Interface that every execution backend must implement.
+
+    The exposure layer depends only on this interface, never on a
+    concrete backend.  This allows swapping Celery for another task
+    runner (Dramatiq, ARQ, ...) without touching the rest of the
+    package.
+    """
+
+    @abstractmethod
+    def dispatch(
+        self,
+        task_ref: Any,
+        kwargs: dict[str, Any],
+        *,
+        result_ttl: int | None = None,
+    ) -> str:
+        """Start task execution.
+
+        Returns a token (task ID) that can be used to track the task.
+        """
+
+    @abstractmethod
+    def get_status(self, token: str) -> TaskInfo:
+        """Query the current state of a task."""
+
+    @abstractmethod
+    def get_result(self, token: str) -> TaskInfo:
+        """Retrieve the final result.
+
+        Should return a ``TaskInfo`` with status ``COMPLETED`` and a
+        populated ``result`` field, or ``FAILED`` with an ``error``
+        field.  If the task is still running, return the current status
+        without a result.
+        """
+
+    @abstractmethod
+    def cancel(self, token: str) -> TaskInfo:
+        """Request cancellation of a running task.
+
+        Cancellation is best-effort.  Returns the updated ``TaskInfo``.
+        """

mcp_celery/backend/celery.py

@@ -0,0 +1,68 @@
+"""Celery execution backend."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from celery import Celery
+from celery.result import AsyncResult
+
+from mcp_celery.backend.base import AbstractBackend
+from mcp_celery.lifecycle import map_celery_state
+from mcp_celery.schema import TaskInfo, TaskStatus
+
+
+class CeleryBackend(AbstractBackend):
+    """Dispatches and tracks tasks using a Celery application."""
+
+    def __init__(self, app: Celery) -> None:
+        self._app = app
+
+    def dispatch(
+        self,
+        task_ref: Any,
+        kwargs: dict[str, Any],
+        *,
+        result_ttl: int | None = None,
+    ) -> str:
+        result = task_ref.apply_async(kwargs=kwargs)
+        return result.id
+
+    def get_status(self, token: str) -> TaskInfo:
+        result = AsyncResult(token, app=self._app)
+        status = map_celery_state(result.state)
+
+        progress: float | None = None
+        metadata: dict[str, Any] = {}
+
+        if isinstance(result.info, dict) and status == TaskStatus.RUNNING:
+            progress = result.info.get("progress")
+            metadata = {k: v for k, v in result.info.items() if k != "progress"}
+
+        error: str | None = None
+        if status == TaskStatus.FAILED:
+            error = str(result.result)
+
+        return TaskInfo(
+            token=token,
+            status=status,
+            progress=progress,
+            metadata=metadata,
+            error=error,
+        )
+
+    def get_result(self, token: str) -> TaskInfo:
+        result = AsyncResult(token, app=self._app)
+        status = map_celery_state(result.state)
+
+        if status == TaskStatus.COMPLETED:
+            return TaskInfo(token=token, status=status, result=result.result)
+        if status == TaskStatus.FAILED:
+            return TaskInfo(token=token, status=status, error=str(result.result))
+
+        return TaskInfo(token=token, status=status)
+
+    def cancel(self, token: str) -> TaskInfo:
+        result = AsyncResult(token, app=self._app)
+        result.revoke(terminate=True)
+        return TaskInfo(token=token, status=TaskStatus.CANCELLED)

mcp_celery/errors.py

@@ -0,0 +1,32 @@
+"""Custom exceptions for mcp-celery."""
+
+from __future__ import annotations
+
+from mcp_celery.schema import TaskStatus
+
+
+class McpCeleryError(Exception):
+    """Base exception for all mcp-celery errors."""
+
+
+class TaskNotReady(McpCeleryError):
+    """Raised when a result is requested but the task has not completed."""
+
+    def __init__(self, token: str, current_status: TaskStatus):
+        self.token = token
+        self.current_status = current_status
+        super().__init__(
+            f"Task {token!r} not complete. Current status: {current_status.value}"
+        )
+
+
+class TaskNotFound(McpCeleryError):
+    """Raised when a token does not correspond to a known task."""
+
+    def __init__(self, token: str):
+        self.token = token
+        super().__init__(f"Task {token!r} not found")
+
+
+class BackendError(McpCeleryError):
+    """Raised when the execution backend encounters an unexpected error."""

mcp_celery/exposure/__init__.py

@@ -0,0 +1,17 @@
+from mcp_celery.exposure.base import AbstractExposureStrategy, GeneratedTool
+from mcp_celery.exposure.polling import PollingExposureStrategy
+
+
+def __getattr__(name: str):
+    if name == "OperationResourceStrategy":
+        from mcp_celery.exposure.operation_resource import OperationResourceStrategy
+        return OperationResourceStrategy
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "AbstractExposureStrategy",
+    "GeneratedTool",
+    "OperationResourceStrategy",
+    "PollingExposureStrategy",
+]

mcp_celery/exposure/base.py

@@ -0,0 +1,49 @@
+"""Abstract exposure strategy interface."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable
+
+from mcp_celery.backend.base import AbstractBackend
+from mcp_celery.schema import AsyncToolDef
+
+if TYPE_CHECKING:
+    from mcp.server.fastmcp import FastMCP
+
+
+@dataclass
+class GeneratedTool:
+    """An MCP tool definition ready to be registered on the server."""
+
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+    handler: Callable[..., Any]
+
+
+class AbstractExposureStrategy(ABC):
+    """Defines how an ``AsyncToolDef`` is exposed as one or more MCP tools.
+
+    Each strategy implementation decides:
+    - how many MCP tools to create per async tool
+    - what their schemas and handlers look like
+    - how the client interacts with them (polling, resources, etc.)
+    """
+
+    @abstractmethod
+    def generate_tools(
+        self,
+        tool_def: AsyncToolDef,
+        backend: AbstractBackend,
+    ) -> list[GeneratedTool]:
+        """Produce MCP tool definitions for the given async tool."""
+
+    def setup(self, mcp: FastMCP, backend: AbstractBackend) -> None:
+        """Optional server-level wiring (task infrastructure, resources, etc.).
+
+        Called once after all tools have been registered.  The default
+        implementation is a no-op.  Override in strategies that need
+        server-level setup (e.g. ``OperationResourceStrategy``).
+        """