Skip to content

datarobot_genai.drmcpbase.middleware

middleware

FastMCP middleware utilities shared across DataRobot MCP servers.

These middleware classes are runtime-specific (FastMCP). Callers inject header and auth-context hooks so each server can wire its own auth resolution (for example drmcputils.auth).

RequestHeadersMiddleware

Bases: Middleware

Read FastMCP request headers and forward them to an injector callback.

Source code in datarobot_genai/drmcpbase/middleware.py 
42
43
44
45
46
47
48
49
50
51
class RequestHeadersMiddleware(Middleware):
    """Read FastMCP request headers and forward them to an injector callback."""

    def __init__(self, inject_headers: HeaderInjector) -> None:
        self._inject_headers = inject_headers

    async def on_call_tool(self, context: Any, call_next: Any) -> Any:
        headers = _safe_read_http_headers()
        self._inject_headers(headers)
        return await call_next(context)

OAuthMiddleWare

Bases: Middleware

Inject HTTP headers and authorization context for FastMCP tool calls.

Parameters

inject_headers: Called with request headers (for example drmcputils.auth.set_request_headers). extract_auth_context: Parses authorization context from headers; return None when unavailable. set_auth_context: Stores parsed authorization context for the current request. auth_context_state_key: Key used when attaching auth context to FastMCP context state. attach_to_fastmcp_state: When True, also persist auth context on context.fastmcp_context state.

Source code in datarobot_genai/drmcpbase/middleware.py 
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
class OAuthMiddleWare(Middleware):
    """Inject HTTP headers and authorization context for FastMCP tool calls.

    Parameters
    ----------
    inject_headers:
        Called with request headers (for example ``drmcputils.auth.set_request_headers``).
    extract_auth_context:
        Parses authorization context from headers; return ``None`` when unavailable.
    set_auth_context:
        Stores parsed authorization context for the current request.
    auth_context_state_key:
        Key used when attaching auth context to FastMCP context state.
    attach_to_fastmcp_state:
        When True, also persist auth context on ``context.fastmcp_context`` state.
    """

    def __init__(
        self,
        *,
        inject_headers: HeaderInjector,
        extract_auth_context: AuthContextExtractor,
        set_auth_context: AuthContextSetter,
        auth_context_state_key: str = "authorization_context",
        attach_to_fastmcp_state: bool = True,
    ) -> None:
        self._inject_headers = inject_headers
        self._extract_auth_context = extract_auth_context
        self._set_auth_context = set_auth_context
        self._auth_context_state_key = auth_context_state_key
        self._attach_to_fastmcp_state = attach_to_fastmcp_state

    async def on_call_tool(self, context: Any, call_next: Any) -> Any:
        headers = _safe_read_http_headers()
        self._inject_headers(headers)

        auth_context = self._safe_extract_auth_context(headers)
        if not auth_context:
            logger.debug("No valid authorization context extracted from request headers.")

        self._set_auth_context(auth_context)

        if self._attach_to_fastmcp_state and context.fastmcp_context is not None:
            await context.fastmcp_context.set_state(self._auth_context_state_key, auth_context)
            logger.debug("Authorization context attached to FastMCP state.")

        return await call_next(context)

    def _safe_extract_auth_context(self, headers: dict[str, str]) -> Any | None:
        try:
            return self._extract_auth_context(headers)
        except (ValueError, KeyError, TypeError) as exc:
            logger.warning("Failed to extract auth context from headers: %s", exc, exc_info=True)
            return None
        except Exception as exc:
            logger.error("Unexpected error extracting auth context: %s", exc, exc_info=True)
            return None

read_http_headers 

read_http_headers() -> dict[str, str]

Return current FastMCP HTTP headers, including authorization headers.

Source code in datarobot_genai/drmcpbase/middleware.py 
36
37
38
39
def read_http_headers() -> dict[str, str]:
    """Return current FastMCP HTTP headers, including authorization headers."""
    # fastmcp 3.x strips authorization headers by default; include them for auth
    return get_http_headers(include_all=True)

register_oauth_middleware 

register_oauth_middleware(mcp: Any, middleware: Middleware) -> None

Register a middleware instance with a FastMCP server.

Source code in datarobot_genai/drmcpbase/middleware.py 
121
122
123
124
def register_oauth_middleware(mcp: Any, middleware: Middleware) -> None:
    """Register a middleware instance with a FastMCP server."""
    mcp.add_middleware(middleware)
    logger.info("OAuth middleware registered successfully")