simbiot-harness/venv/lib/python3.12/site-packages/strands/tools/mcp/mcp_tasks.py

"""Types and configuration for MCP task-augmented tool execution.

This surface is experimental and subject to change. The finalized SEP-2663
models require mcp 2.x (``MCPCreateTaskResult``, ``MCPGetTaskResult``, and the
other task result types). On the runtime pin ``mcp<2.0.0`` they cannot
round-trip server JSON; the corresponding client methods raise ``RuntimeError``.
"""

from collections.abc import Mapping
from datetime import datetime, timedelta
from typing import Annotated, Any, Literal

from mcp.types import (
    CallToolResult,
    CreateMessageRequest,
    CreateMessageResult,
    CreateMessageResultWithTools,
    ElicitRequest,
    ElicitResult,
    ListRootsRequest,
    ListRootsResult,
    Request,
    RequestParams,
    Result,
)
from pydantic import Field, field_validator, model_validator
from typing_extensions import Self, TypedDict

_TASKS_EXTENSION = "io.modelcontextprotocol/tasks"
_TASKS_PROTOCOL_VERSION = "2026-07-28"

MCPTaskStatus = Literal["working", "input_required", "completed", "failed", "cancelled"]
MCPCallToolResult = CallToolResult
MCPInputRequest = CreateMessageRequest | ListRootsRequest | ElicitRequest
MCPInputResponse = CreateMessageResult | CreateMessageResultWithTools | ListRootsResult | ElicitResult
MCPInputRequests = dict[str, MCPInputRequest]
MCPInputResponses = dict[str, MCPInputResponse]
_NonNegativeInteger = Annotated[int, Field(ge=0, strict=True)]


def _parse_task_timestamp(value: str) -> datetime:
    """Parse an ISO 8601 task timestamp, accepting both ``Z`` and numeric UTC offsets."""
    return datetime.fromisoformat(value.replace("Z", "+00:00"))


class TasksConfig(TypedDict, total=False):
    """Configuration for MCP task-augmented tool execution.

    Experimental: this configuration and the task lifecycle it enables are
    subject to change as MCP Tasks evolve.

    On MCP 2.x, enabling this configuration advertises the SEP-2663 Tasks
    extension and automatically completes task handles returned by tools.
    On MCP 1.x, the legacy 2025-11-25 task flow remains supported.

    Attributes:
        poll_timeout: Overall timeout for task completion. Defaults to 5 minutes.
        request_timeout: Timeout for each task lifecycle request. Defaults to 1 minute.
        poll_interval: Polling delay when the server omits ``pollIntervalMs``.
            Defaults to 1 second.
        ttl: Legacy 2025-11-25 task time-to-live. Defaults to 1 minute.
    """

    poll_timeout: timedelta
    request_timeout: timedelta
    poll_interval: timedelta
    ttl: timedelta


class MCPTaskError(Result):
    """JSON-RPC error stored by a failed MCP task."""

    code: int
    message: str
    data: Any = None


class MCPTask(Result):
    """Operational state shared by every SEP-2663 task result."""

    task_id: Annotated[str, Field(min_length=1)]
    status: MCPTaskStatus
    status_message: str | None = None
    created_at: str
    last_updated_at: str
    # SEP-2663 requires ttlMs on every task object; null means the server did not
    # advertise a TTL. Omitting the field is a validation error.
    ttl_ms: _NonNegativeInteger | None
    poll_interval_ms: _NonNegativeInteger | None = None

    @field_validator("created_at", "last_updated_at")
    @classmethod
    def validate_timestamp(cls, value: str) -> str:
        """Validate an ISO 8601 timestamp with a UTC offset."""
        parsed = _parse_task_timestamp(value)
        if parsed.tzinfo is None:
            raise ValueError("MCP task timestamps require a UTC offset")
        return value

    @model_validator(mode="after")
    def validate_chronology(self) -> Self:
        """Validate that the task update does not predate its creation."""
        created_at = _parse_task_timestamp(self.created_at)
        last_updated_at = _parse_task_timestamp(self.last_updated_at)
        if last_updated_at < created_at:
            raise ValueError("MCP task lastUpdatedAt must not precede createdAt")
        return self


class MCPCreateTaskResult(MCPTask):
    """Task handle returned instead of an immediate tool result."""

    result_type: Literal["task"]


class MCPGetTaskResult(MCPTask):
    """Status-specific task state returned by ``tasks/get``."""

    result_type: Literal["complete"]
    input_requests: MCPInputRequests | None = None
    result: CallToolResult | None = None
    error: MCPTaskError | None = None

    @model_validator(mode="before")
    @classmethod
    def validate_status_fields(cls, value: Any) -> Any:
        """Reject status-specific fields on every other task status."""
        if not isinstance(value, Mapping):
            return value
        status = value.get("status")
        allowed_by_status = {
            "input_required": "inputRequests",
            "completed": "result",
            "failed": "error",
        }
        allowed = allowed_by_status.get(status) if isinstance(status, str) else None
        for field in ("inputRequests", "result", "error"):
            if field != allowed and field in value:
                raise ValueError(f'MCP task status "{status}" contains an unexpected status-specific payload')
        return value

    @model_validator(mode="after")
    def validate_status_payload(self) -> Self:
        """Validate the payload associated with the task's current status."""
        payloads = {
            "input_required": self.input_requests,
            "completed": self.result,
            "failed": self.error,
        }
        if self.status in payloads and payloads[self.status] is None:
            raise ValueError(f'MCP task status "{self.status}" requires its status-specific payload')
        return self


class _MCPTaskAcknowledgement(Result):
    """Validated empty acknowledgement for a task lifecycle operation."""

    result_type: Literal["complete"]

    @model_validator(mode="before")
    @classmethod
    def validate_empty_acknowledgement(cls, value: Any) -> Any:
        """Reject task state in an operation's empty acknowledgement."""
        task_fields = {
            "taskId",
            "status",
            "statusMessage",
            "createdAt",
            "lastUpdatedAt",
            "ttlMs",
            "pollIntervalMs",
            "inputRequests",
            "result",
            "error",
        }
        if isinstance(value, Mapping) and task_fields.intersection(value):
            raise ValueError("MCP task acknowledgement must not contain task state")
        return value


class MCPUpdateTaskResult(_MCPTaskAcknowledgement):
    """Acknowledgement returned by ``tasks/update``."""


class MCPCancelTaskResult(_MCPTaskAcknowledgement):
    """Acknowledgement returned by ``tasks/cancel``."""


class _GetTaskRequestParams(RequestParams):
    task_id: str


class _GetTaskRequest(Request[_GetTaskRequestParams, Literal["tasks/get"]]):
    method: Literal["tasks/get"] = "tasks/get"
    params: _GetTaskRequestParams


_GetTaskRequest.name_param = "taskId"  # type: ignore[attr-defined]


class _UpdateTaskRequestParams(RequestParams):
    task_id: str
    input_responses: MCPInputResponses


class _UpdateTaskRequest(Request[_UpdateTaskRequestParams, Literal["tasks/update"]]):
    method: Literal["tasks/update"] = "tasks/update"
    params: _UpdateTaskRequestParams


_UpdateTaskRequest.name_param = "taskId"  # type: ignore[attr-defined]


class _CancelTaskRequestParams(RequestParams):
    task_id: str


class _CancelTaskRequest(Request[_CancelTaskRequestParams, Literal["tasks/cancel"]]):
    method: Literal["tasks/cancel"] = "tasks/cancel"
    params: _CancelTaskRequestParams


_CancelTaskRequest.name_param = "taskId"  # type: ignore[attr-defined]


DEFAULT_TASK_TTL = timedelta(minutes=1)
DEFAULT_TASK_POLL_TIMEOUT = timedelta(minutes=5)
DEFAULT_TASK_REQUEST_TIMEOUT = timedelta(minutes=1)
DEFAULT_TASK_POLL_INTERVAL = timedelta(seconds=1)
DEFAULT_TASK_CONFIG = TasksConfig(
    ttl=DEFAULT_TASK_TTL,
    poll_timeout=DEFAULT_TASK_POLL_TIMEOUT,
    request_timeout=DEFAULT_TASK_REQUEST_TIMEOUT,
    poll_interval=DEFAULT_TASK_POLL_INTERVAL,
)


__all__ = [
    "DEFAULT_TASK_CONFIG",
    "DEFAULT_TASK_POLL_INTERVAL",
    "DEFAULT_TASK_POLL_TIMEOUT",
    "DEFAULT_TASK_REQUEST_TIMEOUT",
    "DEFAULT_TASK_TTL",
    "MCPCallToolResult",
    "MCPCancelTaskResult",
    "MCPCreateTaskResult",
    "MCPGetTaskResult",
    "MCPInputRequest",
    "MCPInputRequests",
    "MCPInputResponse",
    "MCPInputResponses",
    "MCPTask",
    "MCPTaskError",
    "MCPTaskStatus",
    "MCPUpdateTaskResult",
    "TasksConfig",
]