arcade-mcp/libs/arcade-serve/arcade_serve/mcp/server.py

import asyncio
import logging
import os
import uuid
from enum import Enum
from typing import Any, Callable, Union

from arcade_core.catalog import MaterializedTool, ToolCatalog
from arcade_core.executor import ToolExecutor
from arcade_core.schema import ToolAuthorizationContext, ToolContext
from arcadepy import ArcadeError, AsyncArcade
from arcadepy.types.auth_authorize_params import AuthRequirement, AuthRequirementOauth2
from arcadepy.types.shared import AuthorizationResponse

from arcade_serve.mcp.convert import convert_to_mcp_content, create_mcp_tool
from arcade_serve.mcp.logging import create_mcp_logging_middleware
from arcade_serve.mcp.message_processor import MCPMessageProcessor, create_message_processor
from arcade_serve.mcp.types import (
    CallToolRequest,
    CallToolResponse,
    CallToolResult,
    CancelRequest,
    Implementation,
    InitializeRequest,
    InitializeResponse,
    InitializeResult,
    JSONRPCError,
    JSONRPCResponse,
    ListPromptsRequest,
    ListPromptsResponse,
    ListResourcesRequest,
    ListResourcesResponse,
    ListToolsRequest,
    ListToolsResponse,
    ListToolsResult,
    PingRequest,
    PingResponse,
    ProgressNotification,
    ServerCapabilities,
    ShutdownRequest,
    ShutdownResponse,
    Tool,
)

logger = logging.getLogger("arcade.mcp")

MCP_PROTOCOL_VERSION = "2024-11-05"


class MessageMethod(str, Enum):
    """Enumeration of supported MCP message methods"""

    PING = "ping"
    INITIALIZE = "initialize"
    LIST_TOOLS = "tools/list"
    CALL_TOOL = "tools/call"
    PROGRESS = "progress"
    CANCEL = "$/cancelRequest"
    SHUTDOWN = "shutdown"
    LIST_RESOURCES = "resources/list"
    LIST_PROMPTS = "prompts/list"


class MCPServer:
    """
    Unified async MCP server that manages connections, middleware, and tool invocation.
    Handles protocol-level messages (ping, initialize, list_tools, call_tool, etc.).
    """

    def __init__(
        self,
        tool_catalog: Any,
        enable_logging: bool = True,
        **client_kwargs: dict[str, Any],
    ) -> None:
        """
        Initialize the MCP server.

        Args:
            tool_catalog: Catalog of available tools
            **client_kwargs: Additional arguments to pass to the AsyncArcade client
        """
        self.tool_catalog: ToolCatalog = tool_catalog
        self.message_processor: MCPMessageProcessor = create_message_processor()

        # Pop middleware_config from client_kwargs regardless of logging state,
        # as it's internal config not meant for AsyncArcade.
        middleware_config = client_kwargs.pop("middleware_config", {})

        if enable_logging:
            # Create and add the logging middleware if logging is enabled.
            # Note: enable_logging must be True for this middleware (and its stdio_mode behavior)
            # to be activated.
            self.message_processor.add_middleware(
                create_mcp_logging_middleware(**middleware_config)
            )

        self._shutdown: bool = False
        # Initialize AsyncArcade with the *remaining* client_kwargs
        self.arcade = AsyncArcade(**client_kwargs)  # type: ignore[arg-type]

        # Initialize handler dispatch table
        self._method_handlers: dict[str, Callable] = {
            MessageMethod.PING: self._handle_ping,
            MessageMethod.INITIALIZE: self._handle_initialize,
            MessageMethod.LIST_TOOLS: self._handle_list_tools,
            MessageMethod.CALL_TOOL: self._handle_call_tool,
            MessageMethod.PROGRESS: self._handle_progress,
            MessageMethod.CANCEL: self._handle_cancel,
            MessageMethod.SHUTDOWN: self._handle_shutdown,
            MessageMethod.LIST_RESOURCES: self._handle_list_resources,
            MessageMethod.LIST_PROMPTS: self._handle_list_prompts,
        }

    async def run_connection(
        self,
        read_stream: Any,
        write_stream: Any,
        init_options: Any,
    ) -> None:
        """
        Handle a single MCP connection (SSE or stdio).

        Args:
            read_stream: Async iterable yielding incoming messages.
            write_stream: Object with an async send(message) method.
            init_options: Initialization options for the connection.
        """
        # Generate a user ID if possible
        user_id = self._get_user_id(init_options)

        try:
            logger.info(f"Starting MCP connection for user {user_id}")

            async for message in read_stream:
                # Process the message
                response = await self.handle_message(message, user_id=user_id)

                # Skip sending responses for None (e.g., notifications)
                if response is None:
                    continue

                await self._send_response(write_stream, response)

        except asyncio.CancelledError:
            logger.info("Connection cancelled")
        except Exception:
            logger.exception("Error in connection")

    def _get_user_id(self, init_options: Any) -> str:
        """
        Get the user ID for a connection.

        Args:
            init_options: Initialization options for the connection

        Returns:
            A user ID string
        """
        try:
            from arcade_core.config import config

            # Prefer config.user.email if available
            if config.user and config.user.email:
                return config.user.email
        except ValueError:
            logger.debug("No logged in user for MCP Server")

        fallback = str(uuid.uuid4())
        if os.environ.get("ARCADE_USER_ID", None):
            return os.environ.get("ARCADE_USER_ID", fallback)
        elif isinstance(init_options, dict):
            user_id = init_options.get("user_id")
            if user_id:
                return str(user_id)
        # Fallback to random UUID
        return str(fallback)

    async def _send_response(self, write_stream: Any, response: Any) -> None:
        """
        Send a response to the client.

        Args:
            write_stream: Stream to write the response to
            response: Response object to send
        """
        # Ensure the response is properly serialized to JSON
        if hasattr(response, "model_dump_json"):
            # It's a Pydantic model, serialize it
            json_response = response.model_dump_json()
            # Ensure it ends with a newline for JSON-RPC-over-stdio
            if not json_response.endswith("\n"):
                json_response += "\n"
            logger.debug(f"Sending response: {json_response[:200]}...")
            await write_stream.send(json_response)
        elif isinstance(response, dict):
            # It's a dict, convert to JSON
            import json

            json_response = json.dumps(response)
            # Ensure it ends with a newline for JSON-RPC-over-stdio
            if not json_response.endswith("\n"):
                json_response += "\n"
            logger.debug(f"Sending response: {json_response[:200]}...")
            await write_stream.send(json_response)
        else:
            # It's already a string or something else
            response_str = str(response)
            # Ensure it ends with a newline for JSON-RPC-over-stdio
            if not response_str.endswith("\n"):
                response_str += "\n"
            logger.debug(f"Sending raw response type: {type(response)}")
            await write_stream.send(response_str)

    async def handle_message(self, message: Any, user_id: str | None = None) -> Any:
        """
        Handle an incoming MCP message. Processes it through middleware and dispatches
        to the appropriate handler based on the message method.

        Args:
            message: The raw incoming message
            user_id: Optional user ID for authentication

        Returns:
            A properly formatted response message
        """
        # Pre-process message through middleware
        processed = await self.message_processor.process_request(message)

        # Handle special case for JSON string initialize requests
        if isinstance(processed, str):
            try:
                import json

                parsed = json.loads(processed)
                if (
                    isinstance(parsed, dict)
                    and parsed.get("method") == MessageMethod.INITIALIZE
                    and "id" in parsed
                ):
                    # This is an initialize request
                    init_response = await self._handle_initialize(InitializeRequest(**parsed))
                    return init_response
            except Exception:
                logger.exception("Error processing JSON string")
                # Not parseable JSON, continue with normal processing
                pass

        # Check if it's a notification
        if hasattr(processed, "method"):
            method = getattr(processed, "method", None)

            # Handle notifications (methods starting with "notifications/")
            if method and method.startswith("notifications/"):
                await self._handle_notification(method, processed)
                return None

            # Handle regular methods using the dispatch table
            if method in self._method_handlers:
                # If it's a call_tool request, we need to pass the user_id
                if method == MessageMethod.CALL_TOOL:
                    return await self._method_handlers[method](processed, user_id=user_id)
                # For other methods, just pass the processed message
                return await self._method_handlers[method](processed)

            # Unknown method
            return JSONRPCError(
                id=getattr(processed, "id", None),
                error={
                    "code": -32601,
                    "message": f"Method not found: {method}",
                },
            )

        # If it's not a method request, just pass it through
        return processed

    async def _handle_notification(self, method: str, message: Any) -> None:
        """
        Handle notification messages.

        Args:
            method: The notification method
            message: The notification message
        """
        if method == "notifications/cancelled":
            logger.info(f"Request cancelled: {getattr(message, 'params', {})}")
        else:
            logger.debug(f"Received notification: {method}")

    async def _handle_ping(self, message: PingRequest) -> PingResponse:
        """
        Handle a ping request and return a pong response.

        Args:
            message: The ping request

        Returns:
            A properly formatted pong response
        """
        return PingResponse(id=message.id)

    async def _handle_initialize(self, message: InitializeRequest) -> InitializeResponse:
        """
        Handle an initialize request and return a proper initialize response.

        Args:
            message: The initialize request

        Returns:
            A properly formatted initialize response
        """
        # Create the result data
        result = InitializeResult(
            protocolVersion=MCP_PROTOCOL_VERSION,
            capabilities=ServerCapabilities(),
            serverInfo=Implementation(name="Arcade MCP Worker", version="0.1.0"),
            instructions="Arcade MCP Worker initialized.",
        )

        # Construct proper response with result field
        response = InitializeResponse(id=message.id, result=result)

        logger.debug(f"Initialize response: {response.model_dump_json()}")
        return response

    async def _handle_list_tools(
        self, message: ListToolsRequest
    ) -> Union[ListToolsResponse, JSONRPCError]:
        """
        Handle a tools/list request and return a list of available tools.

        Args:
            message: The tools/list request

        Returns:
            A properly formatted tools/list response or error
        """
        try:
            # Get all tools from the catalog
            tools = []
            tool_conversion_errors = []

            for tool in self.tool_catalog:
                try:
                    mcp_tool = create_mcp_tool(tool)
                    if mcp_tool:
                        tools.append(mcp_tool)
                except Exception:
                    tool_name = getattr(tool, "name", str(tool))
                    logger.exception(f"Error converting tool: {tool_name}")
                    tool_conversion_errors.append(tool_name)

            # Log summary if we had errors
            if tool_conversion_errors:
                logger.warning(
                    f"Failed to convert {len(tool_conversion_errors)} tools: {tool_conversion_errors}"
                )

            # Create tool objects with exception handling for each one
            tool_objects = []
            for t in tools:
                try:
                    # Make input schema optional if missing
                    tool_dict = dict(t)
                    if "inputSchema" not in tool_dict:
                        tool_dict["inputSchema"] = {"type": "object", "properties": {}}

                    tool_objects.append(Tool(**tool_dict))
                except Exception:
                    logger.exception(f"Error creating Tool object for {t.get('name', 'unknown')}")

            # Return successful response with the tools we were able to convert
            result = ListToolsResult(tools=tool_objects)
            response = ListToolsResponse(id=message.id, result=result)

        except Exception:
            logger.exception("Error listing tools")
            return JSONRPCError(
                id=message.id,
                error={
                    "code": -32603,
                    "message": "Internal error listing tools",
                },
            )
        return response

    async def _handle_call_tool(
        self, message: CallToolRequest, user_id: str | None = None
    ) -> CallToolResponse:
        """
        Handle a tools/call request to execute a tool.

        Args:
            message: The tools/call request
            user_id: Optional user ID for authentication

        Returns:
            A properly formatted tools/call response
        """
        tool_name: str = message.params["name"]
        # Extract input from the correct field
        input_params: dict[str, Any] = message.params.get("input", {})
        if not input_params:
            input_params = message.params.get("arguments", {})

        logger.info(f"Handling tool call for {tool_name}")

        try:
            tool = self.tool_catalog.get_tool_by_name(tool_name, separator="_")
            tool_context = ToolContext()

            # Set up context with secrets
            if tool.definition.requirements and tool.definition.requirements.secrets:
                self._setup_tool_secrets(tool, tool_context)

            # Handle authorization if needed
            requirement = self._get_auth_requirement(tool)
            if requirement:
                auth_result = await self._check_authorization(requirement, user_id=user_id)
                if auth_result.status != "completed":
                    return CallToolResponse(
                        id=message.id,
                        result=CallToolResult(content=[{"type": "text", "text": auth_result.url}]),
                    )
                else:
                    tool_context.authorization = ToolAuthorizationContext(
                        token=auth_result.context.token if auth_result.context else None,
                        user_info={"user_id": user_id} if user_id else {},
                    )

            # Execute the tool
            logger.debug(f"Executing tool {tool_name} with input: {input_params}")
            result = await ToolExecutor.run(
                func=tool.tool,
                definition=tool.definition,
                input_model=tool.input_model,
                output_model=tool.output_model,
                context=tool_context,
                **input_params,
            )
            logger.debug(f"Tool result: {result}")
            if result.value:
                return CallToolResponse(
                    id=message.id,
                    result=CallToolResult(content=convert_to_mcp_content(result.value)),
                )
            else:
                error = result.error or "Error calling tool"
                logger.error(f"Tool {tool_name} returned error: {error}")
                return CallToolResponse(
                    id=message.id,
                    result=CallToolResult(
                        content=[{"type": "text", "text": convert_to_mcp_content(error)}]
                    ),
                )
        except Exception as e:
            logger.exception(f"Error calling tool {tool_name}")
            error = f"Error calling tool {tool_name}: {e!s}"
            return CallToolResponse(
                id=message.id,
                result=CallToolResult(
                    content=[{"type": "text", "text": convert_to_mcp_content(error)}]
                ),
            )

    def _setup_tool_secrets(self, tool: Any, tool_context: ToolContext) -> None:
        """
        Set up tool secrets in the tool context.

        Args:
            tool: The tool to set up secrets for
            tool_context: The tool context to update
        """
        for secret in tool.definition.requirements.secrets:
            value = os.environ.get(secret.key)
            if value is not None:
                tool_context.set_secret(secret.key, value)

    async def _handle_progress(self, message: ProgressNotification) -> JSONRPCResponse:
        """
        Handle a progress notification.

        Args:
            message: The progress notification

        Returns:
            A response acknowledging the notification
        """
        return JSONRPCResponse(id=getattr(message, "id", None), result={"ok": True})

    async def _handle_cancel(self, message: CancelRequest) -> JSONRPCResponse:
        """
        Handle a cancel request.

        Args:
            message: The cancel request

        Returns:
            A response acknowledging the cancellation
        """
        return JSONRPCResponse(id=getattr(message, "id", None), result={"ok": True})

    async def _handle_shutdown(self, message: ShutdownRequest) -> ShutdownResponse:
        """
        Handle a shutdown request.

        Args:
            message: The shutdown request

        Returns:
            A response acknowledging the shutdown request
        """
        # Schedule a task to shutdown the server after sending the response
        proc = asyncio.create_task(self.shutdown())
        proc.add_done_callback(lambda _: logger.info("MCP server shutdown complete"))
        return ShutdownResponse(id=message.id, result={"ok": True})

    async def _handle_list_resources(self, message: ListResourcesRequest) -> ListResourcesResponse:
        """
        Handle a resources/list request.

        Args:
            message: The resources/list request

        Returns:
            A properly formatted resources/list response
        """
        return ListResourcesResponse(id=message.id, result={"resources": []})

    async def _handle_list_prompts(self, message: ListPromptsRequest) -> ListPromptsResponse:
        """
        Handle a prompts/list request.

        Args:
            message: The prompts/list request

        Returns:
            A properly formatted prompts/list response
        """
        return ListPromptsResponse(id=message.id, result={"prompts": []})

    def _get_auth_requirement(self, tool: MaterializedTool) -> AuthRequirement | None:
        """
        Get the authentication requirement for a tool.

        Args:
            tool: The tool to get the requirement for

        Returns:
            An authentication requirement or None if not required
        """
        req = tool.definition.requirements.authorization
        if not req:
            return None
        if not req.provider_id and not req.provider_type:
            return None
        if hasattr(req, "oauth2") and req.oauth2:
            return AuthRequirement(
                provider_id=str(req.provider_id),
                provider_type=str(req.provider_type),
                oauth2=AuthRequirementOauth2(scopes=req.oauth2.scopes or []),
            )
        return AuthRequirement(
            provider_id=str(req.provider_id),
            provider_type=str(req.provider_type),
        )

    async def _check_authorization(
        self, auth_requirement: AuthRequirement, user_id: str | None = None
    ) -> AuthorizationResponse:
        """
        Check if a tool is authorized for a user.

        Args:
            tool: The tool to check authorization for
            user_id: The user ID to check authorization for

        Returns:
            An authorization response

        Raises:
            RuntimeError: If the tool has no authorization requirement
            Exception: If authorization fails
        """
        try:
            response = await self.arcade.auth.authorize(
                auth_requirement=auth_requirement,
                user_id=user_id or "anonymous",
            )
            logger.debug(f"Authorization response: {response}")

        except ArcadeError:
            logger.exception("Error authorizing tool")
            raise
        return response

    async def shutdown(self) -> None:
        """Shutdown the server."""
        self._shutdown = True
        logger.info("MCP server shutdown complete")