datarobot_genai.drmcp.core.middleware

middleware

Wire drmcpbase FastMCP middleware to drtools auth resolution.

BaseAuthZMiddleware

Bases: BaseHTTPMiddleware, ABC

Applies the mcp_enable_oauth_claim_validation gate before any AuthZ work runs.

Subclasses implement run_authz and in order not to repeat the gate check, the dispatch method is overridden to call the run_authz method.

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

class BaseAuthZMiddleware(BaseHTTPMiddleware, ABC):
    """Applies the ``mcp_enable_oauth_claim_validation`` gate before any AuthZ work runs.

    Subclasses implement ``run_authz`` and in order not to repeat the gate check, the
    dispatch method is overridden to call the ``run_authz`` method.
    """

    async def dispatch(
        self,
        request: Request,
        call_next: Any,
    ) -> Response:
        if not should_run_claim_validation(request):
            return await call_next(request)
        return await self.run_authz(request, call_next)

    @abstractmethod
    async def run_authz(
        self,
        request: Request,
        call_next: Any,
    ) -> Response: ...

OAuthJWTTokenHandlerMiddleware

Bases: BaseAuthZMiddleware

ASGI middleware parsing OAuth JWT token and setting it in request scope. The parsed JWT token can be reused in the downstream validation (e.g., audience, tool scope).

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

class OAuthJWTTokenHandlerMiddleware(BaseAuthZMiddleware):
    """ASGI middleware parsing OAuth JWT token and setting it in request scope.
    The parsed JWT token can be reused in the downstream validation (e.g., audience, tool scope).
    """

    HTTP_HEADER_TO_VALIDATE = "x-datarobot-external-access-token"

    @staticmethod
    def update_scope_with_authenticated_user(
        scope: Scope,
        access_token: AccessToken,
    ) -> None:
        scope["user"] = AuthenticatedUser(access_token)

    @staticmethod
    def update_scope_with_auth_credentials(
        scope: Scope,
        access_token: AccessToken,
    ) -> None:
        scope["auth"] = AuthCredentials(access_token.scopes)

    async def run_authz(
        self,
        request: Request,
        call_next: Any,
    ) -> Response:
        try:
            access_token = JWTTokenHandler.parse_to_access_token(
                self.HTTP_HEADER_TO_VALIDATE,
                request.headers,
            )
        except (
            NoJWTBearerTokenKeyInRequestHeaderError,
            JWTDecodeError,
        ) as ex:
            error_message = str(ex)
            logger.info(error_message)
            return build_http_response_from_auth_error(
                status_code=HTTPStatus.UNAUTHORIZED,
                auth_error_response=AuthErrorResponse(
                    resource_metadata=build_well_known_protected_resource_url(request),
                    error_code=ErrorCodeInAuthErrorResponse.INVALID_TOKEN,
                    error_description=error_message,
                ),
            )

        scope = request.scope
        self.update_scope_with_auth_credentials(scope, access_token)
        self.update_scope_with_authenticated_user(scope, access_token)

        return await call_next(request)

GeneralOAuthClaimValidationMiddleware

Bases: BaseAuthZMiddleware

ASGI middleware validating the audience claim in a JWT token. The naming general in this context means claim validations in this middleware is not MCP protocol specific check to be only triggerd by a specific MCP protocol (e.g., call a tool). This middleware is expected to be triggered after OAuthJWTTokenHandlerMiddleware which parses JWT token and sets it the request scope which is to be validated in this middleware.

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

class GeneralOAuthClaimValidationMiddleware(BaseAuthZMiddleware):
    """ASGI middleware validating the audience claim in a JWT token. The naming general in this
    context means claim validations in this middleware is not MCP protocol specific check to be only
    triggerd by a specific MCP protocol (e.g., call a tool).
    This middleware is expected to be triggered after OAuthJWTTokenHandlerMiddleware which
    parses JWT token and sets it the request scope which is to be validated in this middleware.
    """

    @staticmethod
    def get_expected_audience_claim() -> str | None:
        mcp_server_config = get_config()
        return mcp_server_config.mcp_xaa_token_audience

    async def run_authz(
        self,
        request: Request,
        call_next: Any,
    ) -> Response:
        user = get_user_from_request_scope(request)
        if not user:
            return await call_next(request)

        try:
            claims_validator = JWTTokenClaimsValidator(user)
            claims_validator.validate_audience_claim(self.get_expected_audience_claim())
        except AudienceClaimValidationError as ex:
            error_message = str(ex)
            logger.info(error_message)
            return build_http_response_from_auth_error(
                status_code=HTTPStatus.FORBIDDEN,
                auth_error_response=AuthErrorResponse(
                    resource_metadata=build_well_known_protected_resource_url(request),
                    error_code=ErrorCodeInAuthErrorResponse.INVALID_TOKEN,
                    error_description=error_message,
                ),
            )

        return await call_next(request)

OAuthMCPToolCallScopeValidationMiddleware

Bases: BaseAuthZMiddleware

ASGI middleware validating the scope claim in a JWT token. It is now enforced only on tools/call requests. This middleware is expected to be triggered after OAuthJWTTokenHandlerMiddleware which parses JWT token and sets it the request scope which is to be validated in this middleware.

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

class OAuthMCPToolCallScopeValidationMiddleware(BaseAuthZMiddleware):
    """ASGI middleware validating the scope claim in a JWT token. It is now enforced only on
    ``tools/call`` requests.
    This middleware is expected to be triggered after OAuthJWTTokenHandlerMiddleware which
    parses JWT token and sets it the request scope which is to be validated in this middleware.
    """

    @staticmethod
    async def get_mcp_tool_name_in_request(request: Request) -> str | None:
        if request.method != HTTPMethod.POST:
            return None
        request_body = await request.body()
        try:
            payload = json.loads(request_body)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return None

        if not isinstance(payload, dict) or payload.get("method") != "tools/call":
            return None
        params = payload.get("params")
        name = params.get("name") if isinstance(params, dict) else None
        return name if isinstance(name, str) else None

    async def run_authz(
        self,
        request: Request,
        call_next: Any,
    ) -> Response:
        mcp_tool_name = await self.get_mcp_tool_name_in_request(request)
        if not mcp_tool_name:
            return await call_next(request)

        declared_scopes = await declared_scopes_for_one_tool(
            request.app.state.fastmcp_server,
            mcp_tool_name,
        )
        if not declared_scopes:
            return await call_next(request)

        user = get_user_from_request_scope(request)
        if not user:
            return await call_next(request)

        try:
            claims_validator = JWTTokenClaimsValidator(user)
            claims_validator.validate_mcp_tool_scope_claims(declared_scopes)
        except MCPToolScopeClaimValidationError as ex:
            error_message = str(ex)
            logger.info(error_message)
            return build_http_response_from_auth_error(
                status_code=HTTPStatus.FORBIDDEN,
                auth_error_response=AuthErrorResponse(
                    resource_metadata=build_well_known_protected_resource_url(request),
                    error_code=ErrorCodeInAuthErrorResponse.INSUFFICIENT_SCOPE,
                    error_description=error_message,
                    scopes=sorted(declared_scopes),
                ),
            )

        return await call_next(request)

is_path_exempt_from_oauth_validation

is_path_exempt_from_oauth_validation(path: str) -> bool

Please refer datarobot_genai/drmcp/core/routes.py for the route definition.

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

def is_path_exempt_from_oauth_validation(path: str) -> bool:
    """Please refer datarobot_genai/drmcp/core/routes.py for the route definition."""
    path = _normalize_path(path)
    health_path = _normalize_path(prefix_mount_path("/"))
    well_known_prefix = _normalize_path(prefix_mount_path("/.well-known"))
    return path == health_path or path.startswith(well_known_prefix + "/")

should_run_claim_validation

should_run_claim_validation(request: Request) -> bool

Whether the AuthZ validators run for this request.

mcp_enable_oauth_claim_validation gates all of them; exempt paths never run one.

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

def should_run_claim_validation(request: Request) -> bool:
    """Whether the AuthZ validators run for this request.

    ``mcp_enable_oauth_claim_validation`` gates all of them; exempt paths never run one.
    """
    return (
        get_config().mcp_enable_oauth_claim_validation
        and not is_path_exempt_from_oauth_validation(request.url.path)
    )

create_oauth_middleware

create_oauth_middleware(extract_auth_context: AuthContextExtractor | None = None) -> OAuthMiddleWare

Build OAuth middleware wired to drtools request/auth context injection.

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

def create_oauth_middleware(
    extract_auth_context: AuthContextExtractor | None = None,
) -> OAuthMiddleWare:
    """Build OAuth middleware wired to drtools request/auth context injection."""
    return OAuthMiddleWare(
        inject_headers=set_request_headers,
        extract_auth_context=extract_auth_context or extract_auth_context_from_headers,
        set_auth_context=set_auth_context,
        auth_context_state_key=AUTH_CTX_KEY,
    )

initialize_oauth_middleware

initialize_oauth_middleware(mcp: Any) -> None

Register OAuth middleware with the template MCP server.

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

def initialize_oauth_middleware(mcp: Any) -> None:
    """Register OAuth middleware with the template MCP server."""
    register_oauth_middleware(mcp, create_oauth_middleware())