datarobot_genai.drmcp.core.dr_mcp_server

dr_mcp_server

DataRobotMCPServer

DataRobot MCP server implementation using FastMCP framework.

This server can be extended by providing custom configuration, credentials, and lifecycle handlers.

Source code in datarobot_genai/drmcp/core/dr_mcp_server.py

class DataRobotMCPServer:
    """
    DataRobot MCP server implementation using FastMCP framework.

    This server can be extended by providing custom configuration, credentials,
    and lifecycle handlers.
    """

    def __init__(
        self,
        mcp: FastMCP,
        transport: str = "streamable-http",
        config_factory: Callable[[], Any] | None = None,
        credentials_factory: Callable[[], Any] | None = None,
        lifecycle: BaseServerLifecycle | None = None,
        additional_module_paths: list[tuple[str, str]] | None = None,
        load_native_mcp_tools: bool = False,
    ):
        """
        Initialize the server.

        Args:
            mcp: FastMCP instance
            transport: Transport type ("streamable-http" or "stdio")
            config_factory: Optional factory function for user config
            credentials_factory: Optional factory function for user credentials
            lifecycle: Optional lifecycle handler (defaults to BaseServerLifecycle())
            additional_module_paths: Optional list of (directory, package_prefix) tuples for
                loading additional modules
            load_native_mcp_tools: If True, load tools from datarobot_genai.drtools (default False)
        """
        # Initialize config and logging
        self._config = get_config()
        MCPLogging(self._config.app_log_level)
        self._logger = logging.getLogger(self.__class__.__name__)
        self._logger.info(f"Config initialized: {self._config}")

        # Initialize credentials
        self._credentials = get_credentials()
        self._logger.info("Credentials initialized")

        self._user_config = config_factory() if config_factory else None
        self._logger.info(f"User config initialized: {self._user_config}")
        self._user_credentials = credentials_factory() if credentials_factory else None
        self._logger.info("User credentials initialized")

        # Initialize lifecycle
        self._lifecycle = lifecycle if lifecycle else BaseServerLifecycle()
        self._logger.info("Lifecycle initialized")

        self._mcp = mcp
        self._mcp_transport = transport

        # Configure MCP server capabilities
        self._configure_mcp_capabilities()

        # Initialize telemetry
        initialize_telemetry(mcp)

        # Initialize OAuth middleware
        initialize_oauth_middleware(mcp)

        register_mcp_catalog_transform(mcp)

        # Load native MCP tools modules (only when load_native_mcp_tools is True)
        base_dir = Path(__file__).parent.parent
        if load_native_mcp_tools:
            self._logger.info("Loading native MCP tools")
            for tool_type, tool_config in TOOL_CONFIGS.items():
                if is_tool_enabled(tool_type, self._config):
                    self._logger.debug(f"Loading {tool_type} tools from {tool_config['directory']}")
                    _import_modules_from_dir(
                        os.path.join(base_dir.parent, "drtools", tool_config["directory"]),
                        tool_config["package_prefix"],
                    )
                else:
                    self._logger.debug(f"Skipping {tool_type} tools - not enabled")

            # Panel resources live in drmcpbase (shared with global-mcp), so they
            # are registered explicitly onto this instance rather than discovered
            # via the drtools tool registry. Gated on the same panels enablement.
            if is_tool_enabled(ToolType.PANELS, self._config):
                self._logger.debug("Registering panel resources")
                register_panel_resources(mcp)

        # Load additional recipe user modules if provided
        if additional_module_paths:
            for directory, package_prefix in additional_module_paths:
                self._logger.info(f"Loading additional modules from {directory}")
                _import_modules_from_dir(directory, package_prefix)

        # Register HTTP routes if using streamable-http transport
        if transport == "streamable-http":
            register_routes(self._mcp)

    def _configure_mcp_capabilities(self) -> None:
        """Configure MCP capabilities that FastMCP doesn't expose directly.

        See: https://github.com/modelcontextprotocol/python-sdk/issues/1126
        """
        server = self._mcp._mcp_server

        # Declare prompts_changed capability  (capabilities.prompts.listChanged: true)
        server.notification_options.prompts_changed = True

        # Declare experimental capabilities ( experimental.dynamic_prompts: true)
        server.experimental_capabilities = {"dynamic_prompts": {"enabled": True}}

        # Patch to include experimental_capabilities (FastMCP doesn't expose this)
        original = server.create_initialization_options

        def patched(
            notification_options: Any = None,
            experimental_capabilities: dict[str, dict[str, Any]] | None = None,
            **kwargs: Any,
        ) -> Any:
            if experimental_capabilities is None:
                experimental_capabilities = getattr(server, "experimental_capabilities", None)
            return original(
                notification_options=notification_options,
                experimental_capabilities=experimental_capabilities,
                **kwargs,
            )

        server.create_initialization_options = patched

    def run(self, show_banner: bool = False) -> None:
        """Run the DataRobot MCP server synchronously."""
        try:
            # Validate configuration
            if not self._credentials.has_datarobot_credentials():
                self._logger.warning(
                    "DataRobot credentials not configured; skipping dynamic tool/prompt "
                    "registration and MCP lineage sync. The server will start with only "
                    "statically registered items."
                )
            else:
                if self._config.mcp_server_register_dynamic_tools_on_startup:
                    self._logger.info("Registering dynamic tools from deployments...")
                    asyncio.run(register_tools_of_datarobot_deployments())

                if self._config.mcp_server_register_dynamic_prompts_on_startup:
                    self._logger.info("Registering dynamic prompts from prompt management...")
                    asyncio.run(register_prompts_from_datarobot_prompt_management())

                try:
                    linear_manager = LineageManager(self._mcp)
                    asyncio.run(linear_manager.sync_mcp_tools())
                    asyncio.run(linear_manager.sync_mcp_prompts())
                    self._logger.info("Sync with MCP lineage")
                except LRSEnvVarIsNotSetError as error:
                    error_message = f"MCP item metadata is not sync. {str(error)}"
                    self._logger.warning(error_message)

            # Execute pre-server start actions
            asyncio.run(self._lifecycle.pre_server_start(self._mcp))

            # List registered tools, prompts, and resources before starting server
            tools = asyncio.run(self._mcp.list_tools())
            prompts = asyncio.run(self._mcp.list_prompts())
            resources = asyncio.run(self._mcp.list_resources())

            tools_count = len(tools)
            prompts_count = len(prompts)
            resources_count = len(resources)

            self._logger.info(f"Registered tools: {tools_count}")
            for tool in tools:
                self._logger.info(f" > {tool.name}")
            self._logger.info(f"Registered prompts: {prompts_count}")
            for prompt in prompts:
                self._logger.info(f" > {prompt.name}")
            self._logger.info(f"Registered resources: {resources_count}")
            for resource in resources:
                self._logger.info(f" > {resource.name}")

            # Create event loop for async operations
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

            async def run_server(show_banner: bool = show_banner) -> None:
                # Start server in background based on transport type

                if show_banner:
                    log_server_custom_banner(
                        self._mcp,
                        self._mcp_transport,
                        port=self._config.mcp_server_port,
                        tools_count=tools_count,
                        prompts_count=prompts_count,
                        resources_count=resources_count,
                    )

                if self._mcp_transport == "stdio":
                    server_task = asyncio.create_task(self._mcp.run_stdio_async(show_banner=False))
                elif self._mcp_transport == "streamable-http":
                    server_task = asyncio.create_task(
                        self._mcp.run_http_async(
                            transport="http",
                            middleware=[
                                # Request headers in context for REST routes only (skips MCP path).
                                Middleware(RequestHeadersMiddleware),
                                Middleware(OtelASGIMiddleware),
                            ],
                            show_banner=False,
                            port=self._config.mcp_server_port,
                            log_level=self._config.mcp_server_log_level,
                            host=self._config.mcp_server_host,
                            stateless_http=True,
                            path=prefix_mount_path(MCP_PATH_ENDPOINT),
                        )
                    )
                else:
                    raise ValueError(f"Unsupported transport: {self._mcp_transport}")

                # Give the server a moment to initialize
                await asyncio.sleep(1)

                # Execute post-server start actions
                await self._lifecycle.post_server_start(self._mcp)

                # Wait for server to complete
                await server_task

            # Start the server
            self._logger.info("Starting MCP server...")
            try:
                loop.run_until_complete(run_server(show_banner=show_banner))
            except KeyboardInterrupt:
                self._logger.info("Server interrupted by user")
            finally:
                # Execute pre-shutdown actions
                self._logger.info("Shutting down server...")
                loop.run_until_complete(self._lifecycle.pre_server_shutdown(self._mcp))
                loop.close()

        except Exception as e:
            self._logger.error(f"Server error: {e}")
            raise

    async def get_tools(self) -> dict[str, Tool]:
        return {t.name: t for t in await self._mcp.list_tools()}

    async def get_prompts(self) -> dict[str, Prompt]:
        return {p.name: p for p in await self._mcp.list_prompts()}

    async def get_resources(self) -> dict[str, Resource]:
        return {r.name: r for r in await self._mcp.list_resources()}

__init__

__init__(mcp: FastMCP, transport: str = 'streamable-http', config_factory: Callable[[], Any] | None = None, credentials_factory: Callable[[], Any] | None = None, lifecycle: BaseServerLifecycle | None = None, additional_module_paths: list[tuple[str, str]] | None = None, load_native_mcp_tools: bool = False)

Initialize the server.

Parameters:

Name	Type	Description	Default
mcp	FastMCP	FastMCP instance	required
transport	str	Transport type ("streamable-http" or "stdio")	'streamable-http'
config_factory	Callable[[], Any] | None	Optional factory function for user config	None
credentials_factory	Callable[[], Any] | None	Optional factory function for user credentials	None
lifecycle	BaseServerLifecycle | None	Optional lifecycle handler (defaults to BaseServerLifecycle())	None
additional_module_paths	list[tuple[str, str]] | None	Optional list of (directory, package_prefix) tuples for loading additional modules	None
load_native_mcp_tools	bool	If True, load tools from datarobot_genai.drtools (default False)	False

Source code in datarobot_genai/drmcp/core/dr_mcp_server.py

def __init__(
    self,
    mcp: FastMCP,
    transport: str = "streamable-http",
    config_factory: Callable[[], Any] | None = None,
    credentials_factory: Callable[[], Any] | None = None,
    lifecycle: BaseServerLifecycle | None = None,
    additional_module_paths: list[tuple[str, str]] | None = None,
    load_native_mcp_tools: bool = False,
):
    """
    Initialize the server.

    Args:
        mcp: FastMCP instance
        transport: Transport type ("streamable-http" or "stdio")
        config_factory: Optional factory function for user config
        credentials_factory: Optional factory function for user credentials
        lifecycle: Optional lifecycle handler (defaults to BaseServerLifecycle())
        additional_module_paths: Optional list of (directory, package_prefix) tuples for
            loading additional modules
        load_native_mcp_tools: If True, load tools from datarobot_genai.drtools (default False)
    """
    # Initialize config and logging
    self._config = get_config()
    MCPLogging(self._config.app_log_level)
    self._logger = logging.getLogger(self.__class__.__name__)
    self._logger.info(f"Config initialized: {self._config}")

    # Initialize credentials
    self._credentials = get_credentials()
    self._logger.info("Credentials initialized")

    self._user_config = config_factory() if config_factory else None
    self._logger.info(f"User config initialized: {self._user_config}")
    self._user_credentials = credentials_factory() if credentials_factory else None
    self._logger.info("User credentials initialized")

    # Initialize lifecycle
    self._lifecycle = lifecycle if lifecycle else BaseServerLifecycle()
    self._logger.info("Lifecycle initialized")

    self._mcp = mcp
    self._mcp_transport = transport

    # Configure MCP server capabilities
    self._configure_mcp_capabilities()

    # Initialize telemetry
    initialize_telemetry(mcp)

    # Initialize OAuth middleware
    initialize_oauth_middleware(mcp)

    register_mcp_catalog_transform(mcp)

    # Load native MCP tools modules (only when load_native_mcp_tools is True)
    base_dir = Path(__file__).parent.parent
    if load_native_mcp_tools:
        self._logger.info("Loading native MCP tools")
        for tool_type, tool_config in TOOL_CONFIGS.items():
            if is_tool_enabled(tool_type, self._config):
                self._logger.debug(f"Loading {tool_type} tools from {tool_config['directory']}")
                _import_modules_from_dir(
                    os.path.join(base_dir.parent, "drtools", tool_config["directory"]),
                    tool_config["package_prefix"],
                )
            else:
                self._logger.debug(f"Skipping {tool_type} tools - not enabled")

        # Panel resources live in drmcpbase (shared with global-mcp), so they
        # are registered explicitly onto this instance rather than discovered
        # via the drtools tool registry. Gated on the same panels enablement.
        if is_tool_enabled(ToolType.PANELS, self._config):
            self._logger.debug("Registering panel resources")
            register_panel_resources(mcp)

    # Load additional recipe user modules if provided
    if additional_module_paths:
        for directory, package_prefix in additional_module_paths:
            self._logger.info(f"Loading additional modules from {directory}")
            _import_modules_from_dir(directory, package_prefix)

    # Register HTTP routes if using streamable-http transport
    if transport == "streamable-http":
        register_routes(self._mcp)

run

run(show_banner: bool = False) -> None

Run the DataRobot MCP server synchronously.

Source code in datarobot_genai/drmcp/core/dr_mcp_server.py

def run(self, show_banner: bool = False) -> None:
    """Run the DataRobot MCP server synchronously."""
    try:
        # Validate configuration
        if not self._credentials.has_datarobot_credentials():
            self._logger.warning(
                "DataRobot credentials not configured; skipping dynamic tool/prompt "
                "registration and MCP lineage sync. The server will start with only "
                "statically registered items."
            )
        else:
            if self._config.mcp_server_register_dynamic_tools_on_startup:
                self._logger.info("Registering dynamic tools from deployments...")
                asyncio.run(register_tools_of_datarobot_deployments())

            if self._config.mcp_server_register_dynamic_prompts_on_startup:
                self._logger.info("Registering dynamic prompts from prompt management...")
                asyncio.run(register_prompts_from_datarobot_prompt_management())

            try:
                linear_manager = LineageManager(self._mcp)
                asyncio.run(linear_manager.sync_mcp_tools())
                asyncio.run(linear_manager.sync_mcp_prompts())
                self._logger.info("Sync with MCP lineage")
            except LRSEnvVarIsNotSetError as error:
                error_message = f"MCP item metadata is not sync. {str(error)}"
                self._logger.warning(error_message)

        # Execute pre-server start actions
        asyncio.run(self._lifecycle.pre_server_start(self._mcp))

        # List registered tools, prompts, and resources before starting server
        tools = asyncio.run(self._mcp.list_tools())
        prompts = asyncio.run(self._mcp.list_prompts())
        resources = asyncio.run(self._mcp.list_resources())

        tools_count = len(tools)
        prompts_count = len(prompts)
        resources_count = len(resources)

        self._logger.info(f"Registered tools: {tools_count}")
        for tool in tools:
            self._logger.info(f" > {tool.name}")
        self._logger.info(f"Registered prompts: {prompts_count}")
        for prompt in prompts:
            self._logger.info(f" > {prompt.name}")
        self._logger.info(f"Registered resources: {resources_count}")
        for resource in resources:
            self._logger.info(f" > {resource.name}")

        # Create event loop for async operations
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        async def run_server(show_banner: bool = show_banner) -> None:
            # Start server in background based on transport type

            if show_banner:
                log_server_custom_banner(
                    self._mcp,
                    self._mcp_transport,
                    port=self._config.mcp_server_port,
                    tools_count=tools_count,
                    prompts_count=prompts_count,
                    resources_count=resources_count,
                )

            if self._mcp_transport == "stdio":
                server_task = asyncio.create_task(self._mcp.run_stdio_async(show_banner=False))
            elif self._mcp_transport == "streamable-http":
                server_task = asyncio.create_task(
                    self._mcp.run_http_async(
                        transport="http",
                        middleware=[
                            # Request headers in context for REST routes only (skips MCP path).
                            Middleware(RequestHeadersMiddleware),
                            Middleware(OtelASGIMiddleware),
                        ],
                        show_banner=False,
                        port=self._config.mcp_server_port,
                        log_level=self._config.mcp_server_log_level,
                        host=self._config.mcp_server_host,
                        stateless_http=True,
                        path=prefix_mount_path(MCP_PATH_ENDPOINT),
                    )
                )
            else:
                raise ValueError(f"Unsupported transport: {self._mcp_transport}")

            # Give the server a moment to initialize
            await asyncio.sleep(1)

            # Execute post-server start actions
            await self._lifecycle.post_server_start(self._mcp)

            # Wait for server to complete
            await server_task

        # Start the server
        self._logger.info("Starting MCP server...")
        try:
            loop.run_until_complete(run_server(show_banner=show_banner))
        except KeyboardInterrupt:
            self._logger.info("Server interrupted by user")
        finally:
            # Execute pre-shutdown actions
            self._logger.info("Shutting down server...")
            loop.run_until_complete(self._lifecycle.pre_server_shutdown(self._mcp))
            loop.close()

    except Exception as e:
        self._logger.error(f"Server error: {e}")
        raise

create_mcp_server

create_mcp_server(config_factory: Callable[[], Any] | None = None, credentials_factory: Callable[[], Any] | None = None, lifecycle: BaseServerLifecycle | None = None, additional_module_paths: list[tuple[str, str]] | None = None, transport: str = 'streamable-http', load_native_mcp_tools: bool = False) -> DataRobotMCPServer

Create a DataRobot MCP server.

Parameters:

Name	Type	Description	Default
config_factory	Callable[[], Any] | None	Optional factory function for user config	None
credentials_factory	Callable[[], Any] | None	Optional factory function for user credentials	None
lifecycle	BaseServerLifecycle | None	Optional lifecycle handler	None
additional_module_paths	list[tuple[str, str]] | None	Optional list of (directory, package_prefix) tuples	None
transport	str	Transport type ("streamable-http" or "stdio")	'streamable-http'
load_native_mcp_tools	bool	If True, load tools from datarobot_genai.drtools (default False)	False

Returns

Configured DataRobotMCPServer instance

Example

# Basic usage with defaults
server = create_mcp_server()
server.run()

# With custom configuration
from myapp.config import get_my_config
from myapp.lifecycle import MyLifecycle

server = create_mcp_server(
    config_factory=get_my_config,
    lifecycle=MyLifecycle(),
    additional_module_paths=[
        ("/path/to/my/tools", "myapp.tools"),
        ("/path/to/my/prompts", "myapp.prompts"),
    ]
)
server.run()

Source code in datarobot_genai/drmcp/core/dr_mcp_server.py

def create_mcp_server(
    config_factory: Callable[[], Any] | None = None,
    credentials_factory: Callable[[], Any] | None = None,
    lifecycle: BaseServerLifecycle | None = None,
    additional_module_paths: list[tuple[str, str]] | None = None,
    transport: str = "streamable-http",
    load_native_mcp_tools: bool = False,
) -> DataRobotMCPServer:
    """
    Create a DataRobot MCP server.

    Args:
        config_factory: Optional factory function for user config
        credentials_factory: Optional factory function for user credentials
        lifecycle: Optional lifecycle handler
        additional_module_paths: Optional list of (directory, package_prefix) tuples
        transport: Transport type ("streamable-http" or "stdio")
        load_native_mcp_tools: If True, load tools from datarobot_genai.drtools (default False)

    Returns
    -------
        Configured DataRobotMCPServer instance

    Example:
        ```python
        # Basic usage with defaults
        server = create_mcp_server()
        server.run()

        # With custom configuration
        from myapp.config import get_my_config
        from myapp.lifecycle import MyLifecycle

        server = create_mcp_server(
            config_factory=get_my_config,
            lifecycle=MyLifecycle(),
            additional_module_paths=[
                ("/path/to/my/tools", "myapp.tools"),
                ("/path/to/my/prompts", "myapp.prompts"),
            ]
        )
        server.run()
        ```
    """
    # Use the global mcp instance that tools are registered with

    return DataRobotMCPServer(
        mcp=mcp,
        transport=transport,
        config_factory=config_factory,
        credentials_factory=credentials_factory,
        lifecycle=lifecycle,
        additional_module_paths=additional_module_paths,
        load_native_mcp_tools=load_native_mcp_tools,
    )