arcade-mcp/libs/arcade-mcp-server/arcade_mcp_server/mcp_app.py

"""
MCPApp - A FastAPI-like interface for MCP servers.

Provides a clean, minimal API for building MCP servers with lazy initialization.
"""

from __future__ import annotations

import asyncio
import os
import re
import subprocess
import sys
from pathlib import Path
from types import ModuleType
from typing import Any, Callable, Literal, ParamSpec, TypeVar, cast

from arcade_core.catalog import MaterializedTool, ToolCatalog, ToolDefinitionError
from arcade_core.metadata import ToolMetadata
from arcade_core.subprocess_utils import (
    get_windows_no_window_creationflags,
    graceful_terminate_process,
)
from arcade_core.utils import snake_to_pascal_case
from arcade_tdk.auth import ToolAuthorization
from arcade_tdk.error_adapters import ErrorAdapter
from arcade_tdk.tool import tool as tool_decorator
from loguru import logger
from watchfiles import watch

from arcade_mcp_server._validation import normalize_version
from arcade_mcp_server.exceptions import ServerError
from arcade_mcp_server.logging_utils import intercept_standard_logging
from arcade_mcp_server.managers.resource import (
    _is_template_uri,
    make_file_handler,
    make_text_handler,
)
from arcade_mcp_server.resource_server.base import ResourceServerValidator
from arcade_mcp_server.server import MCPServer
from arcade_mcp_server.settings import MCPSettings, ServerSettings, find_env_file
from arcade_mcp_server.types import Annotations, Prompt, PromptMessage, Resource, ResourceTemplate
from arcade_mcp_server.usage import ServerTracker
from arcade_mcp_server.worker import create_arcade_mcp, serve_with_force_quit

P = ParamSpec("P")
T = TypeVar("T")

TransportType = Literal["http", "stdio"]


class MCPApp:
    """
    A FastAPI-like interface for building MCP servers.

    The app collects tools and configuration, then lazily creates the server
    and transport when run() is called.

    Example:
        ```python
        from arcade_mcp_server import MCPApp

        app = MCPApp(name="my_server", version="1.0.0")

        @app.tool
        def greet(name: str) -> str:
            return f"Hello, {name}!"

        # Runtime CRUD once you have a server bound to the app:
        # app.server = mcp_server
        # await app.tools.add(materialized_tool)
        # await app.prompts.add(prompt, handler)
        # await app.resources.add(resource)

        app.run(host="127.0.0.1", port=8000)
        ```
    """

    def __init__(
        self,
        name: str = "ArcadeMCP",
        version: str = "0.1.0",
        title: str | None = None,
        instructions: str | None = None,
        log_level: str = "INFO",
        transport: TransportType = "stdio",
        host: str = "127.0.0.1",
        port: int = 8000,
        reload: bool = False,
        auth: ResourceServerValidator | None = None,
        **kwargs: Any,
    ):
        """
        Initialize the MCP app.

        Args:
            name: Server name
            version: Server version
            title: Server title for display
            instructions: Server instructions
            log_level: Logging level (DEBUG, INFO, WARNING, ERROR)
            transport: Transport type ("stdio")
            host: Host for transport
            port: Port for transport
            reload: Enable auto-reload for development
            auth: Resource Server validator for front-door authentication
            **kwargs: Additional server configuration
        """
        self._name = self._validate_name(name)
        self._version = self._validate_version(version)
        self.title = title or name
        self.instructions = instructions
        self.log_level = log_level
        self.resource_server_validator = auth
        self.server_kwargs = kwargs
        self.transport = transport
        self.host = host
        self.port = port
        self.reload = reload

        # Tool collection (build-time)
        self._catalog = ToolCatalog()
        self._toolkit_name = name

        # Resource collection (build-time)
        self._initial_resources: list[
            tuple[Resource | ResourceTemplate, Callable[..., Any] | None]
        ] = []

        # Tool _meta extensions (build-time) — keyed by tool FQN
        self._tool_meta_extensions: dict[str, dict[str, Any]] = {}

        # Public handle to the MCPServer (set by caller for runtime ops)
        self.server: MCPServer | None = None

        server_settings_kwargs = {
            "name": self._name,
            "version": self._version,
            "title": self.title,
        }
        if self.instructions:
            server_settings_kwargs["instructions"] = self.instructions

        self._mcp_settings = MCPSettings(server=ServerSettings(**server_settings_kwargs))

        # Store the actual instructions that ended up in ServerSettings
        self.instructions = self._mcp_settings.server.instructions

        if not logger._core.handlers:  # type: ignore[attr-defined]
            self._setup_logging(transport == "stdio")

    def _validate_name(self, name: str) -> str:
        """
        Validate that the name follows the required pattern:
        - Alphanumeric characters and underscores only
        - Must end with alphanumeric character
        - Cannot start with underscore
        - Cannot have consecutive underscores

        Args:
            name: The name to validate

        Returns:
            The validated name

        Raises:
            TypeError: If the name is not a string
            ValueError: If the name doesn't follow the required pattern
        """
        if not isinstance(name, str):
            raise TypeError("MCPApp's name must be a string")

        if not name:
            raise ValueError("MCPApp's name cannot be empty")

        if not re.match(r"^[a-zA-Z0-9_]+$", name):
            raise ValueError(
                "MCPApp's name must contain only alphanumeric characters and underscores"
            )

        if name.startswith("_"):
            raise ValueError("MCPApp's name cannot start with an underscore")

        if "__" in name:
            raise ValueError("MCPApp's name cannot have consecutive underscores")

        if not re.match(r".*[a-zA-Z0-9]$", name):
            raise ValueError("MCPApp's name must end with an alphanumeric character")

        return name

    def _validate_version(self, version: str) -> str:
        """Validate and normalize version to canonical semver."""
        try:
            return normalize_version(version)
        except TypeError:
            raise TypeError("MCPApp's version must be a string")
        except ValueError as e:
            raise ValueError(
                f"MCPApp's version must be a valid semver string "
                f"(e.g., '1.0.0', '1.2.3-beta.1'), got '{e}'"
            )

    # Properties (exposed below initializer)
    @property
    def name(self) -> str:
        """Get the server name."""
        return self._name

    @name.setter
    def name(self, value: str) -> None:
        """Set the server name with validation."""
        self._name = self._validate_name(value)

    @property
    def version(self) -> str:
        """Get the server version."""
        return self._version

    @version.setter
    def version(self, value: str) -> None:
        """Set the server version with validation."""
        self._version = self._validate_version(value)

    @property
    def tools(self) -> _ToolsAPI:
        """Runtime and build-time tools API: add/update/remove/list."""
        return _ToolsAPI(self)

    @property
    def prompts(self) -> _PromptsAPI:
        """Runtime prompts API: add/remove/list."""
        return _PromptsAPI(self)

    @property
    def resources(self) -> _ResourcesAPI:
        """Runtime resources API: add/remove/list."""
        return _ResourcesAPI(self)

    def _setup_logging(self, stdio_mode: bool = False) -> None:
        logger.remove()

        # In stdio mode, use stderr (stdout is reserved for JSON-RPC)
        sink = sys.stderr if stdio_mode else sys.stdout

        if self.log_level == "DEBUG":
            format_str = "<level>{level: <8}</level> | <green>{time:HH:mm:ss}</green> | <cyan>{name}:{line}</cyan> | <level>{message}</level>"
        else:
            format_str = "<level>{level: <8}</level> | <green>{time:HH:mm:ss}</green> | <level>{message}</level>"
        logger.add(
            sink,
            format=format_str,
            level=self.log_level,
            colorize=(not stdio_mode),
            diagnose=(self.log_level == "DEBUG"),
        )

        # Intercept standard logging and route through Loguru
        intercept_standard_logging()

    def add_tool(
        self,
        func: Callable[P, T],
        desc: str | None = None,
        name: str | None = None,
        requires_auth: ToolAuthorization | None = None,
        requires_secrets: list[str] | None = None,
        requires_metadata: list[str] | None = None,
        adapters: list[ErrorAdapter] | None = None,
        metadata: ToolMetadata | None = None,
        meta: dict[str, Any] | None = None,
    ) -> Callable[P, T]:
        """Add a tool for build-time materialization (pre-server)."""
        if meta and "arcade" in meta:
            raise ToolDefinitionError(
                "The 'arcade' key in meta is reserved. "
                "Use the 'metadata' parameter (ToolMetadata) instead."
            )

        if not hasattr(func, "__tool_name__"):
            func = tool_decorator(
                func,
                desc=desc,
                name=name,
                requires_auth=requires_auth,
                requires_secrets=requires_secrets,
                requires_metadata=requires_metadata,
                adapters=adapters,
                metadata=metadata,
            )
        try:
            self._catalog.add_tool(
                func,
                self._toolkit_name,
                toolkit_version=self.version,
                toolkit_description=self.instructions,
            )
        except ToolDefinitionError as e:
            raise e.with_context(func.__name__) from e

        # Store _meta extensions for the tool
        if meta:
            tool_name = snake_to_pascal_case(getattr(func, "__tool_name__", func.__name__))
            fqn = None
            for mat_tool in self._catalog:
                if mat_tool.definition.name == tool_name:
                    fqn = str(mat_tool.definition.fully_qualified_name)
                    break
            if fqn:
                self._tool_meta_extensions[fqn] = meta

        logger.debug(f"Added tool: {func.__name__}")
        return func

    def add_tools_from_module(self, module: ModuleType) -> None:
        """Add all the tools in a module to the catalog."""
        self._catalog.add_module(
            module, self._toolkit_name, version=self.version, description=self.instructions
        )

    def add_resource(
        self,
        uri: str,
        *,
        name: str | None = None,
        title: str | None = None,
        description: str | None = None,
        mime_type: str | None = None,
        handler: Callable[..., Any] | None = None,
        annotations: Annotations | None = None,
        meta: dict[str, Any] | None = None,
    ) -> None:
        """Register a resource at build time (before server start).

        If the URI contains ``{`` it is treated as a URI template and a
        :class:`ResourceTemplate` is stored instead of a :class:`Resource`.
        """
        common_kwargs: dict[str, Any] = {
            "name": name or uri,
            "title": title,
            "description": description,
            "mimeType": mime_type,
            "annotations": annotations,
        }
        if meta is not None:
            common_kwargs["_meta"] = meta

        if _is_template_uri(uri):
            item: Resource | ResourceTemplate = ResourceTemplate(uriTemplate=uri, **common_kwargs)
        else:
            item = Resource(uri=uri, **common_kwargs)
        self._initial_resources.append((item, handler))
        logger.debug(f"Added resource: {uri}")

    def resource(
        self,
        uri: str,
        *,
        name: str | None = None,
        title: str | None = None,
        description: str | None = None,
        mime_type: str | None = None,
        annotations: Annotations | None = None,
        meta: dict[str, Any] | None = None,
    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
        """Decorator for registering a resource with a handler at build time."""

        def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
            self.add_resource(
                uri,
                name=name or func.__name__,
                title=title,
                description=description,
                mime_type=mime_type,
                handler=func,
                annotations=annotations,
                meta=meta,
            )
            return func

        return decorator

    def add_text_resource(
        self,
        uri: str,
        *,
        text: str,
        name: str | None = None,
        title: str | None = None,
        description: str | None = None,
        mime_type: str = "text/plain",
    ) -> None:
        """Register a static text resource at build time."""
        if _is_template_uri(uri):
            raise ValueError(
                f"Template URIs are not supported for static text resources: '{uri}'. "
                "Use add_resource() with a handler that accepts template parameters instead."
            )
        self.add_resource(
            uri,
            name=name,
            title=title,
            description=description,
            mime_type=mime_type,
            handler=make_text_handler(text),
        )

    def add_file_resource(
        self,
        uri: str,
        *,
        path: str | Path,
        name: str | None = None,
        title: str | None = None,
        description: str | None = None,
        mime_type: str | None = None,
    ) -> None:
        """Register a file-backed resource at build time."""
        if _is_template_uri(uri):
            raise ValueError(
                f"Template URIs are not supported for file resources: '{uri}'. "
                "Use add_resource() with a handler that accepts template parameters instead."
            )
        self.add_resource(
            uri,
            name=name,
            title=title,
            description=description,
            mime_type=mime_type,
            handler=make_file_handler(path),
        )

    def tool(
        self,
        func: Callable[P, T] | None = None,
        desc: str | None = None,
        name: str | None = None,
        requires_auth: ToolAuthorization | None = None,
        requires_secrets: list[str] | None = None,
        requires_metadata: list[str] | None = None,
        adapters: list[ErrorAdapter] | None = None,
        metadata: ToolMetadata | None = None,
        meta: dict[str, Any] | None = None,
    ) -> Callable[[Callable[P, T]], Callable[P, T]] | Callable[P, T]:
        """Decorator for adding tools with optional parameters."""

        def decorator(f: Callable[P, T]) -> Callable[P, T]:
            return self.add_tool(
                f,
                desc=desc,
                name=name,
                requires_auth=requires_auth,
                requires_secrets=requires_secrets,
                requires_metadata=requires_metadata,
                adapters=adapters,
                metadata=metadata,
                meta=meta,
            )

        if func is not None:
            return decorator(func)
        return decorator

    def run(
        self,
        host: str = "127.0.0.1",
        port: int = 8000,
        reload: bool = False,
        transport: TransportType = "stdio",
        **kwargs: Any,
    ) -> None:
        if len(self._catalog) == 0 and len(self._initial_resources) == 0:
            logger.error(
                "No tools or resources added. Use @app.tool, app.add_tool(), @app.resource, or app.add_resource()."
            )
            sys.exit(1)

        host, port, transport, reload = MCPApp._get_configuration_overrides(
            host, port, transport, reload
        )

        # Since the transport could have changed since __init__, we need to setup logging again
        self._setup_logging(transport == "stdio")

        if os.getenv("ARCADE_MCP_CHILD_PROCESS") == "1":
            # parent watcher has already been setup
            reload = False

        logger.info(
            f"Starting {self._name} v{self.version}"
            f" with {len(self._catalog)} tools and {len(self._initial_resources)} resources"
        )

        if transport in ["http", "streamable-http", "streamable"]:
            resource_server_auth_enabled = isinstance(
                self.resource_server_validator, ResourceServerValidator
            )
            if resource_server_auth_enabled:
                logger.info("Resource Server authentication is enabled. MCP routes are protected.")
            else:
                logger.warning(
                    "Resource Server authentication is disabled. MCP routes are not protected, so tools requiring auth or secrets will fail."
                )
            if (
                isinstance(self.resource_server_validator, ResourceServerValidator)
                and self.resource_server_validator.supports_oauth_discovery()
            ):
                metadata = self.resource_server_validator.get_resource_metadata()
                if metadata:
                    auth_servers = metadata.get("authorization_servers", [])
                    logger.info(f"Accepted authorization server(s): {', '.join(auth_servers)}")

            if reload:
                self._run_with_reload(host, port)
            else:
                self._create_and_run_server(host, port)
        elif transport == "stdio":
            from arcade_mcp_server.stdio_runner import run_stdio_server

            tracker = ServerTracker()
            tracker.track_server_start(
                transport="stdio",
                host=None,
                port=None,
                tool_count=len(self._catalog),
                resource_server_validator=self.resource_server_validator,
            )
            asyncio.run(
                run_stdio_server(
                    catalog=self._catalog,
                    settings=self._mcp_settings,
                    initial_resources=self._initial_resources,
                    tool_meta_extensions=self._tool_meta_extensions,
                    **self.server_kwargs,
                )
            )
        else:
            raise ServerError(f"Invalid transport: {transport}")

    def _run_with_reload(self, host: str, port: int) -> None:
        """
        Run with file watching for auto-reload.

        This method runs as the parent process that watches for file changes
        and spawns/restarts child processes to run the actual server.
        """
        env_file_path = find_env_file()

        def start_server_process() -> subprocess.Popen:
            """Start a child process running the server."""
            env = os.environ.copy()
            env["ARCADE_MCP_CHILD_PROCESS"] = "1"

            creationflags = get_windows_no_window_creationflags(new_process_group=True)

            return subprocess.Popen(
                [sys.executable, *sys.argv],
                env=env,
                creationflags=creationflags,
            )

        def shutdown_server_process(process: subprocess.Popen, reason: str = "reload") -> None:
            """Shutdown server process gracefully with fallback to force kill.

            On Windows, ``process.terminate()`` calls ``TerminateProcess`` which
            kills the child immediately — there is no graceful shutdown.  To
            allow the child to clean up we first try sending ``CTRL_BREAK_EVENT``
            (requires ``CREATE_NEW_PROCESS_GROUP``), which Python's default
            ``SIGINT`` handler will catch as ``KeyboardInterrupt``.  If that
            doesn't work we fall back to ``terminate()`` / ``kill()``.
            """
            logger.info(f"Shutting down server for {reason}...")

            graceful_terminate_process(process)

            try:
                process.wait(timeout=5)
                logger.info("Server shut down gracefully")
            except subprocess.TimeoutExpired:
                logger.warning(
                    "Server did not shut down within 5 seconds (likely due to active client connections). "
                    "Force killing server process..."
                )
                process.kill()
                process.wait()
                logger.info("Server force killed")

        logger.info("Starting file watcher for auto-reload")
        process = start_server_process()

        try:

            def watch_filter(change: Any, path: str) -> bool:
                # Watch Python files and the .env file (if one was found)
                return path.endswith(".py") or (
                    env_file_path is not None and Path(path) == env_file_path
                )

            # Watch current directory, plus the .env file if it's outside cwd
            paths_to_watch: list[str] = ["."]
            if env_file_path is not None:
                paths_to_watch.append(str(env_file_path))

            for changes in watch(*paths_to_watch, watch_filter=watch_filter):
                logger.info(f"Detected changes in {len(changes)} file(s), restarting server...")
                shutdown_server_process(process, reason="reload")
                process = start_server_process()
        except KeyboardInterrupt:
            logger.info("Received shutdown signal")
            shutdown_server_process(process, reason="shutdown")
            logger.info("File watcher stopped")

    def _create_and_run_server(self, host: str, port: int) -> None:
        """
        Create and run the server directly without reload.

        This is used when reload=False or when running as a child process.
        """
        debug = self.log_level == "DEBUG"
        log_level = "debug" if debug else "info"

        app = create_arcade_mcp(
            catalog=self._catalog,
            mcp_settings=self._mcp_settings,
            debug=debug,
            resource_server_validator=self.resource_server_validator,
            initial_resources=self._initial_resources,
            tool_meta_extensions=self._tool_meta_extensions,
            **self.server_kwargs,
        )

        tracker = ServerTracker()
        tracker.track_server_start(
            transport="http",
            host=host,
            port=port,
            tool_count=len(self._catalog),
            resource_server_validator=self.resource_server_validator,
        )

        asyncio.run(serve_with_force_quit(app=app, host=host, port=port, log_level=log_level))

    @staticmethod
    def _get_configuration_overrides(
        host: str, port: int, transport: TransportType, reload: bool
    ) -> tuple[str, int, TransportType, bool]:
        """Get configuration overrides from environment variables."""
        if envvar_transport := os.getenv("ARCADE_SERVER_TRANSPORT"):
            transport = cast(TransportType, envvar_transport)
            logger.debug(
                f"Using '{transport}' as transport from ARCADE_SERVER_TRANSPORT environment variable"
            )

        # host and port are only relevant for HTTP Streamable transport
        if transport in ["http", "streamable-http", "streamable"]:
            if envvar_host := os.getenv("ARCADE_SERVER_HOST"):
                host = envvar_host
                logger.debug(f"Using '{host}' as host from ARCADE_SERVER_HOST environment variable")

            if envvar_port := os.getenv("ARCADE_SERVER_PORT"):
                try:
                    port = int(envvar_port)
                except ValueError:
                    logger.warning(
                        f"Invalid port: '{envvar_port}' from ARCADE_SERVER_PORT environment variable. Using default port {port}"
                    )
                else:
                    logger.debug(
                        f"Using '{port}' as port from ARCADE_SERVER_PORT environment variable"
                    )

            if envvar_reload := os.getenv("ARCADE_SERVER_RELOAD"):
                if envvar_reload.lower() not in ["0", "1"]:
                    logger.warning(
                        f"Invalid reload: '{envvar_reload}' from ARCADE_SERVER_RELOAD environment variable. Using default reload {reload}"
                    )
                else:
                    reload = bool(int(envvar_reload))
                    logger.debug(
                        f"Using '{reload}' as reload from ARCADE_SERVER_RELOAD environment variable"
                    )

        return host, port, transport, reload


class _ToolsAPI:
    """Unified tools API for MCPApp (build-time and runtime)."""

    def __init__(self, app: MCPApp) -> None:
        self._app = app

    async def add(self, tool: MaterializedTool) -> None:
        """Add or update a tool at runtime if server is bound; otherwise queue via app.add_tool decorator."""
        if self._app.server is None:
            raise ServerError("No server bound to app. Set app.server to use runtime tools API.")
        await self._app.server.tools.add_tool(tool)

    async def update(self, tool: MaterializedTool) -> None:
        if self._app.server is None:
            raise ServerError("No server bound to app. Set app.server to use runtime tools API.")
        await self._app.server.tools.update_tool(tool)

    async def remove(self, name: str) -> MaterializedTool:
        if self._app.server is None:
            raise ServerError("No server bound to app. Set app.server to use runtime tools API.")
        return await self._app.server.tools.remove_tool(name)

    async def list(self) -> list[Any]:
        if self._app.server is None:
            raise ServerError("No server bound to app. Set app.server to use runtime tools API.")
        return await self._app.server.tools.list_tools()


class _PromptsAPI:
    """Unified prompts API for MCPApp (runtime)."""

    def __init__(self, app: MCPApp) -> None:
        self._app = app

    async def add(
        self, prompt: Prompt, handler: Callable[[dict[str, str]], list[PromptMessage]] | None = None
    ) -> None:
        if self._app.server is None:
            raise ServerError("No server bound to app. Set app.server to use runtime prompts API.")
        await self._app.server.prompts.add_prompt(prompt, handler)

    async def remove(self, name: str) -> Prompt:
        if self._app.server is None:
            raise ServerError("No server bound to app. Set app.server to use runtime prompts API.")
        return await self._app.server.prompts.remove_prompt(name)

    async def list(self) -> list[Prompt]:
        if self._app.server is None:
            raise ServerError("No server bound to app. Set app.server to use runtime prompts API.")
        return await self._app.server.prompts.list_prompts()


class _ResourcesAPI:
    """Unified resources API for MCPApp (runtime)."""

    def __init__(self, app: MCPApp) -> None:
        self._app = app

    async def add(self, resource: Resource, handler: Callable[[str], Any] | None = None) -> None:
        if self._app.server is None:
            raise ServerError(
                "No server bound to app. Set app.server to use runtime resources API."
            )
        await self._app.server.resources.add_resource(resource, handler)

    async def remove(self, uri: str) -> Resource:
        if self._app.server is None:
            raise ServerError(
                "No server bound to app. Set app.server to use runtime resources API."
            )
        return await self._app.server.resources.remove_resource(uri)

    async def list(self) -> list[Resource]:
        if self._app.server is None:
            raise ServerError(
                "No server bound to app. Set app.server to use runtime resources API."
            )
        return await self._app.server.resources.list_resources()