fastmcp 2.13.3__py3-none-any.whl → 2.14.1__py3-none-any.whl

fastmcp/client/tasks.py

@@ -0,0 +1,551 @@
+"""SEP-1686 client Task classes."""
+
+from __future__ import annotations
+
+import abc
+import asyncio
+import inspect
+import time
+import weakref
+from collections.abc import Awaitable, Callable
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+import mcp.types
+from mcp.types import GetTaskResult, TaskStatusNotification
+
+from fastmcp.client.messages import Message, MessageHandler
+from fastmcp.utilities.logging import get_logger
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from fastmcp.client.client import CallToolResult, Client
+
+
+class TaskNotificationHandler(MessageHandler):
+    """MessageHandler that routes task status notifications to Task objects."""
+
+    def __init__(self, client: Client):
+        super().__init__()
+        self._client_ref: weakref.ref[Client] = weakref.ref(client)
+
+    async def dispatch(self, message: Message) -> None:
+        """Dispatch messages, including task status notifications."""
+        if isinstance(message, mcp.types.ServerNotification):
+            if isinstance(message.root, TaskStatusNotification):
+                client = self._client_ref()
+                if client:
+                    client._handle_task_status_notification(message.root)
+
+        await super().dispatch(message)
+
+
+TaskResultT = TypeVar("TaskResultT")
+
+
+class Task(abc.ABC, Generic[TaskResultT]):
+    """
+    Abstract base class for MCP background tasks (SEP-1686).
+
+    Provides a uniform API whether the server accepts background execution
+    or executes synchronously (graceful degradation per SEP-1686).
+
+    Subclasses:
+        - ToolTask: For tool calls (result type: CallToolResult)
+        - PromptTask: For prompts (future, result type: GetPromptResult)
+        - ResourceTask: For resources (future, result type: ReadResourceResult)
+    """
+
+    def __init__(
+        self,
+        client: Client,
+        task_id: str,
+        immediate_result: TaskResultT | None = None,
+    ):
+        """
+        Create a Task wrapper.
+
+        Args:
+            client: The FastMCP client
+            task_id: The task identifier
+            immediate_result: If server executed synchronously, the immediate result
+        """
+        self._client = client
+        self._task_id = task_id
+        self._immediate_result = immediate_result
+        self._is_immediate = immediate_result is not None
+
+        # Notification-based optimization (SEP-1686 notifications/tasks/status)
+        self._status_cache: GetTaskResult | None = None
+        self._status_event: asyncio.Event | None = None  # Lazy init
+        self._status_callbacks: list[
+            Callable[[GetTaskResult], None | Awaitable[None]]
+        ] = []
+        self._cached_result: TaskResultT | None = None
+
+    def _check_client_connected(self) -> None:
+        """Validate that client context is still active.
+
+        Raises:
+            RuntimeError: If accessed outside client context (unless immediate)
+        """
+        if self._is_immediate:
+            return  # Already resolved, no client needed
+
+        try:
+            _ = self._client.session
+        except RuntimeError as e:
+            raise RuntimeError(
+                "Cannot access task results outside client context. "
+                "Task futures must be used within 'async with client:' block."
+            ) from e
+
+    @property
+    def task_id(self) -> str:
+        """Get the task ID."""
+        return self._task_id
+
+    @property
+    def returned_immediately(self) -> bool:
+        """Check if server executed the task immediately.
+
+        Returns:
+            True if server executed synchronously (graceful degradation or no task support)
+            False if server accepted background execution
+        """
+        return self._is_immediate
+
+    def _handle_status_notification(self, status: GetTaskResult) -> None:
+        """Process incoming notifications/tasks/status (internal).
+
+        Called by Client when a notification is received for this task.
+        Updates cache, triggers events, and invokes user callbacks.
+
+        Args:
+            status: Task status from notification
+        """
+        # Update cache for next status() call
+        self._status_cache = status
+
+        # Wake up any wait() calls
+        if self._status_event is not None:
+            self._status_event.set()
+
+        # Invoke user callbacks
+        for callback in self._status_callbacks:
+            try:
+                result = callback(status)
+                if inspect.isawaitable(result):
+                    # Fire and forget async callbacks
+                    asyncio.create_task(result)  # type: ignore[arg-type] # noqa: RUF006
+            except Exception as e:
+                logger.warning(f"Task callback error: {e}", exc_info=True)
+
+    def on_status_change(
+        self,
+        callback: Callable[[GetTaskResult], None | Awaitable[None]],
+    ) -> None:
+        """Register callback for status change notifications.
+
+        The callback will be invoked when a notifications/tasks/status is received
+        for this task (optional server feature per SEP-1686 lines 436-444).
+
+        Supports both sync and async callbacks (auto-detected).
+
+        Args:
+            callback: Function to call with GetTaskResult when status changes.
+                     Can return None (sync) or Awaitable[None] (async).
+
+        Example:
+            >>> task = await client.call_tool("slow_operation", {}, task=True)
+            >>>
+            >>> def on_update(status: GetTaskResult):
+            ...     print(f"Task {status.taskId} is now {status.status}")
+            >>>
+            >>> task.on_status_change(on_update)
+            >>> result = await task  # Callback fires when status changes
+        """
+        self._status_callbacks.append(callback)
+
+    async def status(self) -> GetTaskResult:
+        """Get current task status.
+
+        If server executed immediately, returns synthetic completed status.
+        Otherwise queries the server for current status.
+        """
+        self._check_client_connected()
+
+        if self._is_immediate:
+            # Return synthetic completed status
+            now = datetime.now(timezone.utc)
+            return GetTaskResult(
+                taskId=self._task_id,
+                status="completed",
+                createdAt=now,
+                lastUpdatedAt=now,
+                ttl=None,
+                pollInterval=1000,
+            )
+
+        # Return cached status if available (from notification)
+        if self._status_cache is not None:
+            cached = self._status_cache
+            # Don't clear cache - keep it for next call
+            return cached
+
+        # Query server and cache the result
+        self._status_cache = await self._client.get_task_status(self._task_id)
+        return self._status_cache
+
+    @abc.abstractmethod
+    async def result(self) -> TaskResultT:
+        """Wait for and return the task result.
+
+        Must be implemented by subclasses to return the appropriate result type.
+        """
+        ...
+
+    async def wait(
+        self, *, state: str | None = None, timeout: float = 300.0
+    ) -> GetTaskResult:
+        """Wait for task to reach a specific state or complete.
+
+        Uses event-based waiting when notifications are available (fast),
+        with fallback to polling (reliable). Optimally wakes up immediately
+        on status changes when server sends notifications/tasks/status.
+
+        Args:
+            state: Desired state ('submitted', 'working', 'completed', 'failed').
+                   If None, waits for any terminal state (completed/failed)
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            GetTaskResult: Final task status
+
+        Raises:
+            TimeoutError: If desired state not reached within timeout
+        """
+        self._check_client_connected()
+
+        if self._is_immediate:
+            # Already done
+            return await self.status()
+
+        # Initialize event for notification wake-ups
+        if self._status_event is None:
+            self._status_event = asyncio.Event()
+
+        start = time.time()
+        terminal_states = {"completed", "failed", "cancelled"}
+        poll_interval = 0.5  # Fallback polling interval (500ms)
+
+        while True:
+            # Check cached status first (updated by notifications)
+            if self._status_cache:
+                current = self._status_cache.status
+                if state is None:
+                    if current in terminal_states:
+                        return self._status_cache
+                elif current == state:
+                    return self._status_cache
+
+            # Check timeout
+            elapsed = time.time() - start
+            if elapsed >= timeout:
+                raise TimeoutError(
+                    f"Task {self._task_id} did not reach {state or 'terminal state'} within {timeout}s"
+                )
+
+            remaining = timeout - elapsed
+
+            # Wait for notification event OR poll timeout
+            try:
+                await asyncio.wait_for(
+                    self._status_event.wait(), timeout=min(poll_interval, remaining)
+                )
+                self._status_event.clear()
+            except asyncio.TimeoutError:
+                # Fallback: poll server (notification didn't arrive in time)
+                self._status_cache = await self._client.get_task_status(self._task_id)
+
+    async def cancel(self) -> None:
+        """Cancel this task, transitioning it to cancelled state.
+
+        Sends a tasks/cancel protocol request. The server will attempt to halt
+        execution and move the task to cancelled state.
+
+        Note: If server executed immediately (graceful degradation), this is a no-op
+        as there's no server-side task to cancel.
+        """
+        if self._is_immediate:
+            # No server-side task to cancel
+            return
+        self._check_client_connected()
+        await self._client.cancel_task(self._task_id)
+        # Invalidate cache to force fresh status fetch
+        self._status_cache = None
+
+    def __await__(self):
+        """Allow 'await task' to get result."""
+        return self.result().__await__()
+
+
+class ToolTask(Task["CallToolResult"]):
+    """
+    Represents a tool call that may execute in background or immediately.
+
+    Provides a uniform API whether the server accepts background execution
+    or executes synchronously (graceful degradation per SEP-1686).
+
+    Usage:
+        task = await client.call_tool_as_task("analyze", args)
+
+        # Check status
+        status = await task.status()
+
+        # Wait for completion
+        await task.wait()
+
+        # Get result (waits if needed)
+        result = await task.result()  # Returns CallToolResult
+
+        # Or just await the task directly
+        result = await task
+    """
+
+    def __init__(
+        self,
+        client: Client,
+        task_id: str,
+        tool_name: str,
+        immediate_result: CallToolResult | None = None,
+    ):
+        """
+        Create a ToolTask wrapper.
+
+        Args:
+            client: The FastMCP client
+            task_id: The task identifier
+            tool_name: Name of the tool being executed
+            immediate_result: If server executed synchronously, the immediate result
+        """
+        super().__init__(client, task_id, immediate_result)
+        self._tool_name = tool_name
+
+    async def result(self) -> CallToolResult:
+        """Wait for and return the tool result.
+
+        If server executed immediately, returns the immediate result.
+        Otherwise waits for background task to complete and retrieves result.
+
+        Returns:
+            CallToolResult: The parsed tool result (same as call_tool returns)
+        """
+        # Check cache first
+        if self._cached_result is not None:
+            return self._cached_result
+
+        if self._is_immediate:
+            assert self._immediate_result is not None  # Type narrowing
+            result = self._immediate_result
+        else:
+            # Check client connected
+            self._check_client_connected()
+
+            # Wait for completion using event-based wait (respects notifications)
+            await self.wait()
+
+            # Get the raw result (dict or CallToolResult)
+            raw_result = await self._client.get_task_result(self._task_id)
+
+            # Convert to CallToolResult if needed and parse
+            if isinstance(raw_result, dict):
+                # Raw dict from get_task_result - parse as CallToolResult
+                mcp_result = mcp.types.CallToolResult.model_validate(raw_result)
+                result = await self._client._parse_call_tool_result(
+                    self._tool_name, mcp_result, raise_on_error=True
+                )
+            elif isinstance(raw_result, mcp.types.CallToolResult):
+                # Already a CallToolResult from MCP protocol - parse it
+                result = await self._client._parse_call_tool_result(
+                    self._tool_name, raw_result, raise_on_error=True
+                )
+            else:
+                # Legacy ToolResult format - convert to MCP type
+                if hasattr(raw_result, "content") and hasattr(
+                    raw_result, "structured_content"
+                ):
+                    mcp_result = mcp.types.CallToolResult(
+                        content=raw_result.content,
+                        structuredContent=raw_result.structured_content,  # type: ignore[arg-type]
+                        _meta=raw_result.meta,
+                    )
+                    result = await self._client._parse_call_tool_result(
+                        self._tool_name, mcp_result, raise_on_error=True
+                    )
+                else:
+                    # Unknown type - just return it
+                    result = raw_result  # type: ignore[assignment]
+
+        # Cache before returning
+        self._cached_result = result
+        return result
+
+
+class PromptTask(Task[mcp.types.GetPromptResult]):
+    """
+    Represents a prompt call that may execute in background or immediately.
+
+    Provides a uniform API whether the server accepts background execution
+    or executes synchronously (graceful degradation per SEP-1686).
+
+    Usage:
+        task = await client.get_prompt_as_task("analyze", args)
+        result = await task  # Returns GetPromptResult
+    """
+
+    def __init__(
+        self,
+        client: Client,
+        task_id: str,
+        prompt_name: str,
+        immediate_result: mcp.types.GetPromptResult | None = None,
+    ):
+        """
+        Create a PromptTask wrapper.
+
+        Args:
+            client: The FastMCP client
+            task_id: The task identifier
+            prompt_name: Name of the prompt being executed
+            immediate_result: If server executed synchronously, the immediate result
+        """
+        super().__init__(client, task_id, immediate_result)
+        self._prompt_name = prompt_name
+
+    async def result(self) -> mcp.types.GetPromptResult:
+        """Wait for and return the prompt result.
+
+        If server executed immediately, returns the immediate result.
+        Otherwise waits for background task to complete and retrieves result.
+
+        Returns:
+            GetPromptResult: The prompt result with messages and description
+        """
+        # Check cache first
+        if self._cached_result is not None:
+            return self._cached_result
+
+        if self._is_immediate:
+            assert self._immediate_result is not None
+            result = self._immediate_result
+        else:
+            # Check client connected
+            self._check_client_connected()
+
+            # Wait for completion using event-based wait (respects notifications)
+            await self.wait()
+
+            # Get the raw MCP result
+            mcp_result = await self._client.get_task_result(self._task_id)
+
+            # Parse as GetPromptResult
+            result = mcp.types.GetPromptResult.model_validate(mcp_result)
+
+        # Cache before returning
+        self._cached_result = result
+        return result
+
+
+class ResourceTask(
+    Task[list[mcp.types.TextResourceContents | mcp.types.BlobResourceContents]]
+):
+    """
+    Represents a resource read that may execute in background or immediately.
+
+    Provides a uniform API whether the server accepts background execution
+    or executes synchronously (graceful degradation per SEP-1686).
+
+    Usage:
+        task = await client.read_resource_as_task("file://data.txt")
+        contents = await task  # Returns list[ReadResourceContents]
+    """
+
+    def __init__(
+        self,
+        client: Client,
+        task_id: str,
+        uri: str,
+        immediate_result: list[
+            mcp.types.TextResourceContents | mcp.types.BlobResourceContents
+        ]
+        | None = None,
+    ):
+        """
+        Create a ResourceTask wrapper.
+
+        Args:
+            client: The FastMCP client
+            task_id: The task identifier
+            uri: URI of the resource being read
+            immediate_result: If server executed synchronously, the immediate result
+        """
+        super().__init__(client, task_id, immediate_result)
+        self._uri = uri
+
+    async def result(
+        self,
+    ) -> list[mcp.types.TextResourceContents | mcp.types.BlobResourceContents]:
+        """Wait for and return the resource contents.
+
+        If server executed immediately, returns the immediate result.
+        Otherwise waits for background task to complete and retrieves result.
+
+        Returns:
+            list[ReadResourceContents]: The resource contents
+        """
+        # Check cache first
+        if self._cached_result is not None:
+            return self._cached_result
+
+        if self._is_immediate:
+            assert self._immediate_result is not None
+            result = self._immediate_result
+        else:
+            # Check client connected
+            self._check_client_connected()
+
+            # Wait for completion using event-based wait (respects notifications)
+            await self.wait()
+
+            # Get the raw MCP result
+            mcp_result = await self._client.get_task_result(self._task_id)
+
+            # Parse as ReadResourceResult or extract contents
+            if isinstance(mcp_result, mcp.types.ReadResourceResult):
+                # Already parsed by TasksResponse - extract contents
+                result = list(mcp_result.contents)
+            elif isinstance(mcp_result, dict) and "contents" in mcp_result:
+                # Dict format - parse each content item
+                parsed_contents = []
+                for item in mcp_result["contents"]:
+                    if isinstance(item, dict):
+                        if "blob" in item:
+                            parsed_contents.append(
+                                mcp.types.BlobResourceContents.model_validate(item)
+                            )
+                        else:
+                            parsed_contents.append(
+                                mcp.types.TextResourceContents.model_validate(item)
+                            )
+                    else:
+                        parsed_contents.append(item)
+                result = parsed_contents
+            else:
+                # Fallback - might be the list directly
+                result = mcp_result if isinstance(mcp_result, list) else [mcp_result]
+
+        # Cache before returning
+        self._cached_result = result
+        return result