datarobot_genai.drmcputils.auth

auth

set_request_headers

set_request_headers(headers: dict[str, str]) -> None

Inject request headers from any runtime adapter (FastMCP, Starlette, LangChain, tests).

Source code in datarobot_genai/drmcputils/auth.py

def set_request_headers(headers: dict[str, str]) -> None:
    """Inject request headers from any runtime adapter (FastMCP, Starlette, LangChain, tests)."""
    _request_headers_ctx.set({str(k).lower(): v for k, v in headers.items() if isinstance(v, str)})

get_request_headers

get_request_headers() -> dict[str, str]

Return headers injected for the current request context.

Source code in datarobot_genai/drmcputils/auth.py

def get_request_headers() -> dict[str, str]:
    """Return headers injected for the current request context."""
    return _request_headers_ctx.get() or {}

set_auth_context

set_auth_context(auth_context: AuthCtx | dict[str, Any] | None) -> None

Inject authorization context from a runtime adapter.

Source code in datarobot_genai/drmcputils/auth.py

def set_auth_context(auth_context: AuthCtx | dict[str, Any] | None) -> None:
    """Inject authorization context from a runtime adapter."""
    _auth_context_ctx.set(auth_context)

get_auth_context

get_auth_context() -> AuthCtx | None

Return authorization context injected for the current request.

Source code in datarobot_genai/drmcputils/auth.py

def get_auth_context() -> AuthCtx | None:
    """Return authorization context injected for the current request."""
    auth_ctx = _auth_context_ctx.get()
    if auth_ctx is None:
        return None
    if isinstance(auth_ctx, AuthCtx):
        return auth_ctx
    if isinstance(auth_ctx, dict):
        try:
            return AuthCtx(**auth_ctx)
        except Exception:
            return None
    return None

must_get_auth_context async

must_get_auth_context() -> AuthCtx

Retrieve the AuthCtx from the current request context or raise error.

Source code in datarobot_genai/drmcputils/auth.py

async def must_get_auth_context() -> AuthCtx:
    """Retrieve the AuthCtx from the current request context or raise error."""
    auth_ctx = get_auth_context()
    if not auth_ctx:
        raise RuntimeError("Could not retrieve authorization context.")
    return auth_ctx

extract_auth_context_from_headers

extract_auth_context_from_headers(headers: dict[str, str]) -> AuthCtx | None

Parse x-datarobot-authorization-context into an :class:AuthCtx.

Source code in datarobot_genai/drmcputils/auth.py

def extract_auth_context_from_headers(headers: dict[str, str]) -> AuthCtx | None:
    """Parse ``x-datarobot-authorization-context`` into an :class:`AuthCtx`."""
    auth_handler = AuthContextHeaderHandler()
    return auth_handler.get_context(headers)

get_access_token async

get_access_token(provider_type: str | None = None) -> str

Retrieve access token from the DataRobot OAuth Provider Service.

Source code in datarobot_genai/drmcputils/auth.py

async def get_access_token(provider_type: str | None = None) -> str:
    """Retrieve access token from the DataRobot OAuth Provider Service."""
    auth_ctx = await must_get_auth_context()
    logger.debug("Retrieved authorization context")

    oauth_token_provider = AsyncOAuthTokenProvider(auth_ctx)
    oauth_access_token = await oauth_token_provider.get_token(
        auth_type=ToolAuth.OBO,
        provider_type=provider_type,
    )
    return oauth_access_token

oauth_access_token_header_name

oauth_access_token_header_name(header_segment: str) -> str

Build x-datarobot-<segment>-access-token from header_segment (normalized).

Source code in datarobot_genai/drmcputils/auth.py

def oauth_access_token_header_name(header_segment: str) -> str:
    """Build ``x-datarobot-<segment>-access-token`` from *header_segment* (normalized)."""
    segment = header_segment.strip().lower().replace("_", "-")
    return f"x-datarobot-{segment}-access-token"

resolve_datarobot_token

resolve_datarobot_token() -> str | None

Resolve the DataRobot API token according to auth_resolution_strategy.

Source code in datarobot_genai/drmcputils/auth.py

def resolve_datarobot_token() -> str | None:
    """Resolve the DataRobot API token according to ``auth_resolution_strategy``."""
    creds = get_credentials()
    strategy = creds.auth_resolution_strategy

    if strategy == AuthResolutionStrategy.HTTP:
        return _extract_token_from_headers_with_fallback(_safe_request_headers())

    return creds.datarobot.datarobot_api_token or None

resolve_secret

resolve_secret(header_name: str, config_value: str) -> str | None

Resolve a secret from request headers and/or config per auth_resolution_strategy.

Source code in datarobot_genai/drmcputils/auth.py

def resolve_secret(header_name: str, config_value: str) -> str | None:
    """Resolve a secret from request headers and/or config per ``auth_resolution_strategy``."""
    creds = get_credentials()
    strategy = creds.auth_resolution_strategy

    use_headers = strategy != AuthResolutionStrategy.CONFIG
    header_value = (
        _extract_header_value(_safe_request_headers(), header_name) if use_headers else None
    )

    use_config = strategy != AuthResolutionStrategy.HTTP
    config_secret = (config_value or None) if use_config else None

    return _resolve_by_strategy(
        strategy=strategy,
        header_value=header_value,
        config_value=config_secret,
    )

resolve_oauth_access_token_from_headers

resolve_oauth_access_token_from_headers(header_segment: str) -> str | None

Read x-datarobot-<segment>-access-token for header_segment (raw or Bearer).

Source code in datarobot_genai/drmcputils/auth.py

def resolve_oauth_access_token_from_headers(header_segment: str) -> str | None:
    """Read ``x-datarobot-<segment>-access-token`` for *header_segment* (raw or Bearer)."""
    raw = _extract_header_value(
        _safe_request_headers(), oauth_access_token_header_name(header_segment)
    )
    return _parse_optional_bearer_token(raw) if raw else None

get_oauth_access_token_with_header_fallback async

get_oauth_access_token_with_header_fallback(provider_type: str, *, display_name: str, access_token_header_segment: str) -> str | ToolError

Resolve an OAuth access token via OBO or request header (http strategy only).

Source code in datarobot_genai/drmcputils/auth.py

async def get_oauth_access_token_with_header_fallback(
    provider_type: str,
    *,
    display_name: str,
    access_token_header_segment: str,
) -> str | ToolError:
    """Resolve an OAuth access token via OBO or request header (``http`` strategy only)."""
    creds = get_credentials()
    strategy = creds.auth_resolution_strategy
    pt = provider_type.strip()

    if strategy == AuthResolutionStrategy.CONFIG:
        header = oauth_access_token_header_name(access_token_header_segment)
        return ToolError(
            f"{display_name} does not support auth_resolution_strategy=config. "
            f"Use auth_resolution_strategy=http with OAuth OBO or pass an access token "
            f"via the {header} header.",
            kind=ToolErrorKind.AUTHENTICATION,
        )

    access_token, oauth_exc = await _try_get_oauth_access_token(pt)
    if access_token:
        return access_token

    if header_token := resolve_oauth_access_token_from_headers(access_token_header_segment):
        return header_token

    header = oauth_access_token_header_name(access_token_header_segment)
    return _oauth_access_token_failure_error(
        display_name=display_name,
        provider_type=pt,
        header=header,
        oauth_exc=oauth_exc,
    )