task_context.py

"""ServerTaskContext - Server-integrated task context with elicitation and sampling.

This wraps the pure TaskContext and adds server-specific functionality:
- Elicitation (task.elicit())
- Sampling (task.create_message())
- Status notifications
"""

from typing import Any

import anyio

from mcp.server.experimental.task_result_handler import TaskResultHandler
from mcp.server.session import ServerSession
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
from mcp.shared.exceptions import MCPError
from mcp.shared.experimental.tasks.capabilities import (
    require_task_augmented_elicitation,
    require_task_augmented_sampling,
)
from mcp.shared.experimental.tasks.context import TaskContext
from mcp.shared.experimental.tasks.message_queue import QueuedMessage, TaskMessageQueue
from mcp.shared.experimental.tasks.resolver import Resolver
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.types import (
    INVALID_REQUEST,
    TASK_STATUS_INPUT_REQUIRED,
    TASK_STATUS_WORKING,
    ClientCapabilities,
    CreateMessageResult,
    CreateTaskResult,
    ElicitationCapability,
    ElicitRequestedSchema,
    ElicitResult,
    IncludeContext,
    ModelPreferences,
    RequestId,
    Result,
    SamplingCapability,
    SamplingMessage,
    Task,
    TaskMetadata,
    TaskStatusNotification,
    TaskStatusNotificationParams,
    Tool,
    ToolChoice,
)


class ServerTaskContext:
    """Server-integrated task context with elicitation and sampling.

    This wraps a pure TaskContext and adds server-specific functionality:
    - elicit() for sending elicitation requests to the client
    - create_message() for sampling requests
    - Status notifications via the session

    Example:
        async def my_task_work(task: ServerTaskContext) -> CallToolResult:
            await task.update_status("Starting...")

            result = await task.elicit(
                message="Continue?",
                requested_schema={"type": "object", "properties": {"ok": {"type": "boolean"}}}
            )

            if result.content.get("ok"):
                return CallToolResult(content=[TextContent(text="Done!")])
            else:
                return CallToolResult(content=[TextContent(text="Cancelled")])
    """

    def __init__(
        self,
        *,
        task: Task,
        store: TaskStore,
        session: ServerSession,
        queue: TaskMessageQueue,
        handler: TaskResultHandler | None = None,
    ):
        """Create a ServerTaskContext.

        Args:
            task: The Task object
            store: The task store
            session: The server session
            queue: The message queue for elicitation/sampling
            handler: The result handler for response routing (required for elicit/create_message)
        """
        self._ctx = TaskContext(task=task, store=store)
        self._session = session
        self._queue = queue
        self._handler = handler
        self._store = store

    # Delegate pure properties to inner context

    @property
    def task_id(self) -> str:
        """The task identifier."""
        return self._ctx.task_id

    @property
    def task(self) -> Task:
        """The current task state."""
        return self._ctx.task

    @property
    def is_cancelled(self) -> bool:
        """Whether cancellation has been requested."""
        return self._ctx.is_cancelled

    def request_cancellation(self) -> None:
        """Request cancellation of this task."""
        self._ctx.request_cancellation()

    # Enhanced methods with notifications

    async def update_status(self, message: str, *, notify: bool = True) -> None:
        """Update the task's status message.

        Args:
            message: The new status message
            notify: Whether to send a notification to the client
        """
        await self._ctx.update_status(message)
        if notify:
            await self._send_notification()

    async def complete(self, result: Result, *, notify: bool = True) -> None:
        """Mark the task as completed with the given result.

        Args:
            result: The task result
            notify: Whether to send a notification to the client
        """
        await self._ctx.complete(result)
        if notify:
            await self._send_notification()

    async def fail(self, error: str, *, notify: bool = True) -> None:
        """Mark the task as failed with an error message.

        Args:
            error: The error message
            notify: Whether to send a notification to the client
        """
        await self._ctx.fail(error)
        if notify:
            await self._send_notification()

    async def _send_notification(self) -> None:
        """Send a task status notification to the client."""
        task = self._ctx.task
        await self._session.send_notification(
            TaskStatusNotification(
                params=TaskStatusNotificationParams(
                    task_id=task.task_id,
                    status=task.status,
                    status_message=task.status_message,
                    created_at=task.created_at,
                    last_updated_at=task.last_updated_at,
                    ttl=task.ttl,
                    poll_interval=task.poll_interval,
                )
            )
        )

    # Server-specific methods: elicitation and sampling

    def _check_elicitation_capability(self) -> None:
        """Check if the client supports elicitation."""
        if not self._session.check_client_capability(ClientCapabilities(elicitation=ElicitationCapability())):
            raise MCPError(code=INVALID_REQUEST, message="Client does not support elicitation capability")

    def _check_sampling_capability(self) -> None:
        """Check if the client supports sampling."""
        if not self._session.check_client_capability(ClientCapabilities(sampling=SamplingCapability())):
            raise MCPError(code=INVALID_REQUEST, message="Client does not support sampling capability")

    async def elicit(
        self,
        message: str,
        requested_schema: ElicitRequestedSchema,
    ) -> ElicitResult:
        """Send an elicitation request via the task message queue.

        This method:
        1. Checks client capability
        2. Updates task status to "input_required"
        3. Queues the elicitation request
        4. Waits for the response (delivered via tasks/result round-trip)
        5. Updates task status back to "working"
        6. Returns the result

        Args:
            message: The message to present to the user
            requested_schema: Schema defining the expected response structure

        Returns:
            The client's response

        Raises:
            MCPError: If client doesn't support elicitation capability
        """
        self._check_elicitation_capability()

        if self._handler is None:
            raise RuntimeError("handler is required for elicit(). Pass handler= to ServerTaskContext.")

        # Update status to input_required
        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)

        # Build the request using session's helper
        request = self._session._build_elicit_form_request(  # pyright: ignore[reportPrivateUsage]
            message=message,
            requested_schema=requested_schema,
            related_task_id=self.task_id,
        )
        request_id: RequestId = request.id

        resolver: Resolver[dict[str, Any]] = Resolver()
        self._handler._pending_requests[request_id] = resolver  # pyright: ignore[reportPrivateUsage]

        queued = QueuedMessage(
            type="request",
            message=request,
            resolver=resolver,
            original_request_id=request_id,
        )
        await self._queue.enqueue(self.task_id, queued)

        try:
            # Wait for response (routed back via TaskResultHandler)
            response_data = await resolver.wait()
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            return ElicitResult.model_validate(response_data)
        except anyio.get_cancelled_exc_class():
            # This path is tested in test_elicit_restores_status_on_cancellation
            # which verifies status is restored to "working" after cancellation.
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            raise

    async def elicit_url(
        self,
        message: str,
        url: str,
        elicitation_id: str,
    ) -> ElicitResult:
        """Send a URL mode elicitation request via the task message queue.

        This directs the user to an external URL for out-of-band interactions
        like OAuth flows, credential collection, or payment processing.

        This method:
        1. Checks client capability
        2. Updates task status to "input_required"
        3. Queues the elicitation request
        4. Waits for the response (delivered via tasks/result round-trip)
        5. Updates task status back to "working"
        6. Returns the result

        Args:
            message: Human-readable explanation of why the interaction is needed
            url: The URL the user should navigate to
            elicitation_id: Unique identifier for tracking this elicitation

        Returns:
            The client's response indicating acceptance, decline, or cancellation

        Raises:
            MCPError: If client doesn't support elicitation capability
            RuntimeError: If handler is not configured
        """
        self._check_elicitation_capability()

        if self._handler is None:
            raise RuntimeError("handler is required for elicit_url(). Pass handler= to ServerTaskContext.")

        # Update status to input_required
        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)

        # Build the request using session's helper
        request = self._session._build_elicit_url_request(  # pyright: ignore[reportPrivateUsage]
            message=message,
            url=url,
            elicitation_id=elicitation_id,
            related_task_id=self.task_id,
        )
        request_id: RequestId = request.id

        resolver: Resolver[dict[str, Any]] = Resolver()
        self._handler._pending_requests[request_id] = resolver  # pyright: ignore[reportPrivateUsage]

        queued = QueuedMessage(
            type="request",
            message=request,
            resolver=resolver,
            original_request_id=request_id,
        )
        await self._queue.enqueue(self.task_id, queued)

        try:
            # Wait for response (routed back via TaskResultHandler)
            response_data = await resolver.wait()
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            return ElicitResult.model_validate(response_data)
        except anyio.get_cancelled_exc_class():  # pragma: no cover
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            raise

    async def create_message(
        self,
        messages: list[SamplingMessage],
        *,
        max_tokens: int,
        system_prompt: str | None = None,
        include_context: IncludeContext | None = None,
        temperature: float | None = None,
        stop_sequences: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        model_preferences: ModelPreferences | None = None,
        tools: list[Tool] | None = None,
        tool_choice: ToolChoice | None = None,
    ) -> CreateMessageResult:
        """Send a sampling request via the task message queue.

        This method:
        1. Checks client capability
        2. Updates task status to "input_required"
        3. Queues the sampling request
        4. Waits for the response (delivered via tasks/result round-trip)
        5. Updates task status back to "working"
        6. Returns the result

        Args:
            messages: The conversation messages for sampling
            max_tokens: Maximum tokens in the response
            system_prompt: Optional system prompt
            include_context: Context inclusion strategy
            temperature: Sampling temperature
            stop_sequences: Stop sequences
            metadata: Additional metadata
            model_preferences: Model selection preferences
            tools: Optional list of tools the LLM can use during sampling
            tool_choice: Optional control over tool usage behavior

        Returns:
            The sampling result from the client

        Raises:
            MCPError: If client doesn't support sampling capability or tools
            ValueError: If tool_use or tool_result message structure is invalid
        """
        self._check_sampling_capability()
        client_caps = self._session.client_params.capabilities if self._session.client_params else None
        validate_sampling_tools(client_caps, tools, tool_choice)
        validate_tool_use_result_messages(messages)

        if self._handler is None:
            raise RuntimeError("handler is required for create_message(). Pass handler= to ServerTaskContext.")

        # Update status to input_required
        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)

        # Build the request using session's helper
        request = self._session._build_create_message_request(  # pyright: ignore[reportPrivateUsage]
            messages=messages,
            max_tokens=max_tokens,
            system_prompt=system_prompt,
            include_context=include_context,
            temperature=temperature,
            stop_sequences=stop_sequences,
            metadata=metadata,
            model_preferences=model_preferences,
            tools=tools,
            tool_choice=tool_choice,
            related_task_id=self.task_id,
        )
        request_id: RequestId = request.id

        resolver: Resolver[dict[str, Any]] = Resolver()
        self._handler._pending_requests[request_id] = resolver  # pyright: ignore[reportPrivateUsage]

        queued = QueuedMessage(
            type="request",
            message=request,
            resolver=resolver,
            original_request_id=request_id,
        )
        await self._queue.enqueue(self.task_id, queued)

        try:
            # Wait for response (routed back via TaskResultHandler)
            response_data = await resolver.wait()
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            return CreateMessageResult.model_validate(response_data)
        except anyio.get_cancelled_exc_class():
            # This path is tested in test_create_message_restores_status_on_cancellation
            # which verifies status is restored to "working" after cancellation.
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            raise

    async def elicit_as_task(
        self,
        message: str,
        requested_schema: ElicitRequestedSchema,
        *,
        ttl: int = 60000,
    ) -> ElicitResult:
        """Send a task-augmented elicitation via the queue, then poll client.

        This is for use inside a task-augmented tool call when you want the client
        to handle the elicitation as its own task. The elicitation request is queued
        and delivered when the client calls tasks/result. After the client responds
        with CreateTaskResult, we poll the client's task until complete.

        Args:
            message: The message to present to the user
            requested_schema: Schema defining the expected response structure
            ttl: Task time-to-live in milliseconds for the client's task

        Returns:
            The client's elicitation response

        Raises:
            MCPError: If client doesn't support task-augmented elicitation
            RuntimeError: If handler is not configured
        """
        client_caps = self._session.client_params.capabilities if self._session.client_params else None
        require_task_augmented_elicitation(client_caps)

        if self._handler is None:
            raise RuntimeError("handler is required for elicit_as_task()")

        # Update status to input_required
        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)

        request = self._session._build_elicit_form_request(  # pyright: ignore[reportPrivateUsage]
            message=message,
            requested_schema=requested_schema,
            related_task_id=self.task_id,
            task=TaskMetadata(ttl=ttl),
        )
        request_id: RequestId = request.id

        resolver: Resolver[dict[str, Any]] = Resolver()
        self._handler._pending_requests[request_id] = resolver  # pyright: ignore[reportPrivateUsage]

        queued = QueuedMessage(
            type="request",
            message=request,
            resolver=resolver,
            original_request_id=request_id,
        )
        await self._queue.enqueue(self.task_id, queued)

        try:
            # Wait for initial response (CreateTaskResult from client)
            response_data = await resolver.wait()
            create_result = CreateTaskResult.model_validate(response_data)
            client_task_id = create_result.task.task_id

            # Poll the client's task using session.experimental
            async for _ in self._session.experimental.poll_task(client_task_id):
                pass

            # Get final result from client
            result = await self._session.experimental.get_task_result(
                client_task_id,
                ElicitResult,
            )

            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            return result

        except anyio.get_cancelled_exc_class():  # pragma: no cover
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            raise

    async def create_message_as_task(
        self,
        messages: list[SamplingMessage],
        *,
        max_tokens: int,
        ttl: int = 60000,
        system_prompt: str | None = None,
        include_context: IncludeContext | None = None,
        temperature: float | None = None,
        stop_sequences: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        model_preferences: ModelPreferences | None = None,
        tools: list[Tool] | None = None,
        tool_choice: ToolChoice | None = None,
    ) -> CreateMessageResult:
        """Send a task-augmented sampling request via the queue, then poll client.

        This is for use inside a task-augmented tool call when you want the client
        to handle the sampling as its own task. The request is queued and delivered
        when the client calls tasks/result. After the client responds with
        CreateTaskResult, we poll the client's task until complete.

        Args:
            messages: The conversation messages for sampling
            max_tokens: Maximum tokens in the response
            ttl: Task time-to-live in milliseconds for the client's task
            system_prompt: Optional system prompt
            include_context: Context inclusion strategy
            temperature: Sampling temperature
            stop_sequences: Stop sequences
            metadata: Additional metadata
            model_preferences: Model selection preferences
            tools: Optional list of tools the LLM can use during sampling
            tool_choice: Optional control over tool usage behavior

        Returns:
            The sampling result from the client

        Raises:
            MCPError: If client doesn't support task-augmented sampling or tools
            ValueError: If tool_use or tool_result message structure is invalid
            RuntimeError: If handler is not configured
        """
        client_caps = self._session.client_params.capabilities if self._session.client_params else None
        require_task_augmented_sampling(client_caps)
        validate_sampling_tools(client_caps, tools, tool_choice)
        validate_tool_use_result_messages(messages)

        if self._handler is None:
            raise RuntimeError("handler is required for create_message_as_task()")

        # Update status to input_required
        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)

        # Build request WITH task field for task-augmented sampling
        request = self._session._build_create_message_request(  # pyright: ignore[reportPrivateUsage]
            messages=messages,
            max_tokens=max_tokens,
            system_prompt=system_prompt,
            include_context=include_context,
            temperature=temperature,
            stop_sequences=stop_sequences,
            metadata=metadata,
            model_preferences=model_preferences,
            tools=tools,
            tool_choice=tool_choice,
            related_task_id=self.task_id,
            task=TaskMetadata(ttl=ttl),
        )
        request_id: RequestId = request.id

        resolver: Resolver[dict[str, Any]] = Resolver()
        self._handler._pending_requests[request_id] = resolver  # pyright: ignore[reportPrivateUsage]

        queued = QueuedMessage(
            type="request",
            message=request,
            resolver=resolver,
            original_request_id=request_id,
        )
        await self._queue.enqueue(self.task_id, queued)

        try:
            # Wait for initial response (CreateTaskResult from client)
            response_data = await resolver.wait()
            create_result = CreateTaskResult.model_validate(response_data)
            client_task_id = create_result.task.task_id

            # Poll the client's task using session.experimental
            async for _ in self._session.experimental.poll_task(client_task_id):
                pass

            # Get final result from client
            result = await self._session.experimental.get_task_result(
                client_task_id,
                CreateMessageResult,
            )

            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            return result

        except anyio.get_cancelled_exc_class():  # pragma: no cover
            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
            raise