Skip to content

datarobot_genai.dragent.frontends.claim_validation

claim_validation

Rejects an inbound IdP token whose aud claim does not name this agent.

The DataRobot API Gateway already validated that the token is well-formed, signed and trusted, then forwarded it as-is. This answers what is left: was it issued for us?

Not an authentication check. A request with no IdP token passes through -- an agent may be called with a DataRobot API token instead (see a2a.py's bearerAuth scheme), and whether a caller is authenticated is the gateway's business. Signature, issuer and expiry are not re-verified either; claims are decoded unverified and only read.

Covers every route with no exemptions, /a2a and agent-card discovery included: NAT copies inbound headers into the workflow context on every route, and the cross-application-access provider reads the token from there regardless of which route it arrived on.

Agent-card discovery needs no exemption because auth there is optional, which the pass-through above already models: an unauthenticated request reaches _handle_get_agent_card, which applies enable_unauthenticated_well_known_route and serves a redacted card or a generic 404. A request that does present a token gets the full check first -- a token naming another agent is rejected rather than earning a card.

Installed by fastapi.DRAgentFastApiFrontEndPluginWorker.build_app.

GeneralOAuthClaimValidationMiddleware

Bases: BaseHTTPMiddleware

Audience check on every inbound request. See the module docstring for scope.

Sibling of drmcp.core.middleware.GeneralOAuthClaimValidationMiddleware, duplicated so dragent does not depend on the MCP subpackages.

HTTP-only is sufficient: DRAgentFastApiFrontEndConfig overrides NAT's workflow default without a websocket_path, so no WebSocket route is registered. If a future scope-validation phase needs to hand claims to the endpoint, switch to pure ASGI -- BaseHTTPMiddleware runs call_next in a separate task, so ContextVars set here do not reach it.

Source code in datarobot_genai/dragent/frontends/claim_validation.py 
 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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
class GeneralOAuthClaimValidationMiddleware(BaseHTTPMiddleware):
    """Audience check on every inbound request.  See the module docstring for scope.

    Sibling of ``drmcp.core.middleware.GeneralOAuthClaimValidationMiddleware``, duplicated so
    ``dragent`` does not depend on the MCP subpackages.

    HTTP-only is sufficient: ``DRAgentFastApiFrontEndConfig`` overrides NAT's ``workflow``
    default without a ``websocket_path``, so no WebSocket route is registered.  If a future
    scope-validation phase needs to hand claims to the endpoint, switch to pure ASGI --
    ``BaseHTTPMiddleware`` runs ``call_next`` in a separate task, so ``ContextVar``s set here
    do not reach it.
    """

    def __init__(self, app: ASGIApp, expected_audience: str) -> None:
        super().__init__(app)
        self._expected_audience = expected_audience

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        rejection = self._reject(request)
        if rejection is not None:
            return rejection
        return await call_next(request)

    def _reject(self, request: Request) -> JSONResponse | None:
        """Response to send instead of calling the app, or ``None`` to allow."""
        token = find_idp_token(request.headers)
        carrier = OAUTH_ACCESS_TOKEN_HEADER if OAUTH_ACCESS_TOKEN_HEADER in request.headers else "-"

        if token is None:
            return None  # No claim validation When using standard datarobot api tokens

        try:
            claims = decode_jwt_claims_unverified(token)
        except ValueError as ex:
            message = f"Malformed authorization token: {ex}"
            logger.warning(
                "Inbound token rejected: reason=undecodable_token carrier=%s detail=%s",
                carrier,
                ex,  # the parser naming the defect; it never echoes the token
            )
            return _error(HTTPStatus.UNPROCESSABLE_ENTITY, message)

        # `audience` is a list, so `in` is exact equality per entry. If _audience_values ever
        # returned the raw claim, a str would make this a substring match and let "-aaa"
        # satisfy "aaa" -- see TestExactAudienceMatching.
        audience = _audience_values(claims)
        if self._expected_audience not in audience:
            message = "Authorization audience claim validation failed"  # no token/claim values
            logger.warning(
                "Inbound token rejected: reason=audience_mismatch expected_audience=%s "
                "token_aud=%s carrier=%s iss=%s has_sub=%s has_scp=%s has_ext=%s",
                self._expected_audience,
                audience,
                carrier,
                claims.get("iss"),
                bool(claims.get("sub")),
                bool(claims.get("scp") or claims.get("scope")),
                bool(claims.get("ext")),
            )
            return _error(HTTPStatus.UNAUTHORIZED, message)

        logger.debug("OAuth audience claim validation succeeded")
        return None