Skip to content

datarobot_genai.core.mcp.config

config

MCPConfig

Bases: DataRobotAppFrameworkBaseSettings

Configuration for MCP server connection.

Derived values are exposed as properties rather than stored, avoiding Pydantic field validation/serialization concerns for internal helpers.

Source code in datarobot_genai/core/mcp/config.py
class MCPConfig(DataRobotAppFrameworkBaseSettings):
    """Configuration for MCP server connection.

    Derived values are exposed as properties rather than stored, avoiding
    Pydantic field validation/serialization concerns for internal helpers.
    """

    external_mcp_url: str | None = None
    external_mcp_headers: str | None = None
    external_mcp_transport: Literal["sse", "streamable-http"] = "streamable-http"
    mcp_deployment_id: str | None = None
    datarobot_endpoint: str | None = None
    datarobot_api_token: str | None = None
    authorization_context: dict[str, Any] | None = None
    forwarded_headers: dict[str, str] | None = None
    mcp_server_port: int | None = None

    _auth_context_handler: AuthContextHeaderHandler | None = None
    _server_config: dict[str, Any] | None = None

    @field_validator("external_mcp_headers", mode="before")
    @classmethod
    def validate_external_mcp_headers(cls, value: str | None) -> str | None:
        if value is None:
            return None

        candidate = value.strip()

        try:
            json.loads(candidate)
        except json.JSONDecodeError:
            msg = "external_mcp_headers must be valid JSON"
            logger.warning(msg)
            return None

        return candidate

    @field_validator("mcp_deployment_id", mode="before")
    @classmethod
    def validate_mcp_deployment_id(cls, value: str | None) -> str | None:
        if value is None:
            return None

        candidate = value.strip()

        if not re.fullmatch(r"[0-9a-fA-F]{24}", candidate):
            msg = "mcp_deployment_id must be a valid 24-character hex ID"
            logger.warning(msg)
            return None

        return candidate

    @field_validator("mcp_server_port", mode="after")
    @classmethod
    def validate_mcp_server_port(cls, value: int | None) -> int | None:
        # Pydantic already coerces/rejects the type; only the range is ours.
        if value is not None and not 1 <= value <= 65535:
            logger.warning("mcp_server_port must be between 1 and 65535; ignoring")
            return None
        return value

    def _authorization_bearer_header(self) -> dict[str, str]:
        """Return Authorization header with Bearer token or empty dict."""
        if not self.datarobot_api_token:
            return {}
        auth = (
            self.datarobot_api_token
            if self.datarobot_api_token.startswith("Bearer ")
            else f"Bearer {self.datarobot_api_token}"
        )
        return {"Authorization": auth}

    @property
    def auth_context_handler(self) -> AuthContextHeaderHandler:
        if self._auth_context_handler is None:
            self._auth_context_handler = AuthContextHeaderHandler()
        return self._auth_context_handler

    @property
    def server_config(self) -> dict[str, Any] | None:
        if self._server_config is None:
            self._server_config = self._build_server_config()
        return self._server_config

    def _config_kind(self) -> Literal["deployment", "external", "local"] | None:
        """Single source of truth for which MCP server config is active.

        Precedence: deployment > external > local.
        """
        if self.mcp_deployment_id:
            return "deployment"
        if self.external_mcp_url:
            return "external"
        if self.mcp_server_port:
            return "local"
        return None

    @property
    def is_local_server(self) -> bool:
        """True when the MCP server is a local process we start (resolved via
        mcp_server_port), as opposed to a DataRobot deployment or external URL.
        """
        return self._config_kind() == "local"

    def _authorization_context_header(self) -> dict[str, str]:
        """Return X-DataRobot-Authorization-Context header or empty dict."""
        try:
            return self.auth_context_handler.get_header(self.authorization_context)
        except (LookupError, RuntimeError):
            # Authorization context not available (e.g., in tests)
            return {}

    def _build_authenticated_headers(self) -> dict[str, str]:
        """Build headers for authenticated requests.

        Returns
        -------
            Dictionary containing forwarded headers (if available) and authentication headers.
        """
        headers: dict[str, str] = {}
        if self.forwarded_headers:
            headers.update(self.forwarded_headers)
        headers.update(self._authorization_bearer_header())
        headers.update(self._authorization_context_header())
        return headers

    def _build_server_config(self) -> dict[str, Any] | None:
        """
        Get MCP server configuration.

        Returns
        -------
            Server configuration dict with url, transport, and optional headers,
            or None if not configured.
        """
        kind = self._config_kind()

        if kind == "deployment":
            # DataRobot deployment ID - requires authentication
            if self.datarobot_endpoint is None:
                raise ValueError(
                    "When using a DataRobot hosted MCP deployment, datarobot_endpoint must be set."
                )
            if self.datarobot_api_token is None:
                raise ValueError(
                    "When using a DataRobot hosted MCP deployment, datarobot_api_token must be set."
                )

            base_url = self.datarobot_endpoint.rstrip("/")
            if not base_url.endswith("/api/v2"):
                base_url = f"{base_url}/api/v2"

            url = f"{base_url}/deployments/{self.mcp_deployment_id}/directAccess/mcp"
            headers = self._build_authenticated_headers()

            logger.info(f"Using DataRobot hosted MCP deployment: {url}")

            return {
                "url": url,
                "transport": "streamable-http",
                "headers": headers,
            }

        if kind == "external":
            # External MCP URL - no authentication needed
            headers = {}

            # Merge external headers if provided
            if self.external_mcp_headers:
                external_headers = json.loads(self.external_mcp_headers)
                headers.update(external_headers)

            logger.info(f"Using external MCP URL: {self.external_mcp_url}")

            return {
                # cast: kind == "external" guarantees external_mcp_url is set
                "url": cast(str, self.external_mcp_url).rstrip("/"),
                "transport": self.external_mcp_transport,
                "headers": headers,
            }

        if kind == "local":
            url = f"http://localhost:{self.mcp_server_port}"
            headers = self._build_authenticated_headers()
            logger.info(f"Using localhost MCP server: {url}")
            return {
                "url": f"{url}/mcp",
                "transport": "streamable-http",
                "headers": headers,
            }

        return None

is_local_server property

is_local_server: bool

True when the MCP server is a local process we start (resolved via mcp_server_port), as opposed to a DataRobot deployment or external URL.