`datarobot_genai.dragent.plugins.okta_a2a_auth`

okta_a2a_auth

Cross-Application Access (XAA) auth provider for A2A agent communication.

Registered as _type: okta_cross_app_access in workflow YAML.

Overview

Implements the Identity Assertion Authorization Grant (ID-JAG) flow for secure cross-application access between AI agents. The flow involves two authorization servers:

Org Authorization Server — issues and validates user identity; handles the initial token exchange (Step 1).
Custom Authorization Server — protects the target resource/API; issues the final scoped access token with custom audience, scopes, claims, and policies required by that resource (Step 2).

Discovery phase ~~~~~~~~~~~~~~~ Forwards the incoming x-datarobot-external-access-token header directly to the agent card endpoint as a Bearer token.

Call phase ~~~~~~~~~~ Executes a two-step XAA token exchange:

Step 1 — RFC 8693 Token Exchange (Org AS): Exchange the caller's access token for an ID-JAG token. The ID-JAG is a short-lived, single-use identity assertion that encodes who (human user) is acting through which AI agent (workload principal).

Step 2 — RFC 7523 JWT Bearer Grant (Custom AS): Exchange the ID-JAG for a scoped access token at the custom authorization server. The resulting token contains the sub (human user), cid (AI agent workload principal), and requested scopes.

Both steps authenticate the client using private_key_jwt — a signed JWT client assertion built from the agent's RSA private key (IDP_AGENT_PRIVATE_KEY_JWK) and principal ID (IDP_AGENT_ID).

Two implementations are available, selected via XAA_TOKEN_EXCHANGE_IMPL:

okta_sdk (default) — delegates to okta-client-python's CrossAppAccessFlow which handles both steps internally.
http — direct HTTP calls via httpx + PyJWT, giving full control over request parameters.

Credentials (principal_id, private_jwk) are loaded automatically from environment variables / Runtime Parameters / .env / file_secrets via :class:_OktaSettings. They do not need to be specified in workflow.yaml — the minimal YAML is::

authentication:
  okta_auth:
    _type: okta_cross_app_access

function_groups:
  remote_agent:
    _type: authenticated_a2a_client
    url: "https://..."
    auth_provider: okta_auth

Environment variables ~~~~~~~~~~~~~~~~~~~~~ * IDP_AGENT_ID — Okta AI agent workload principal ID * IDP_AGENT_PRIVATE_KEY_JWK — base64-encoded or raw-JSON RSA private JWK * XAA_TOKEN_EXCHANGE_IMPL — okta_sdk (default) or http

OAuth2CrossApplicationAccessAuthProviderConfig

Bases: AuthProviderBaseConfig

Configuration for the XAA auth provider.

principal_id and private_jwk are auto-loaded from env vars / Runtime Parameters / .env / file_secrets — no workflow.yaml entries needed for credentials.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

class OAuth2CrossApplicationAccessAuthProviderConfig(
    AuthProviderBaseConfig,
    name="okta_cross_app_access",
):  # type: ignore[call-arg]
    """Configuration for the XAA auth provider.

    ``principal_id`` and ``private_jwk`` are auto-loaded from env vars /
    Runtime Parameters / ``.env`` / ``file_secrets`` — no ``workflow.yaml``
    entries needed for credentials.
    """

    okta_token_header: str = Field(
        default="x-datarobot-external-access-token",
        description=(
            "Incoming header carrying the caller's Okta access token. "
            "Forwarded as Bearer for discovery; used as subject token in Step 1. "
            "Matched case-insensitively."
        ),
    )
    fallback_token_headers: list[str] = Field(
        default=["authorization"],
        description=(
            "Fallback headers to try (in order) when ``okta_token_header`` is absent. "
            "Useful for local development without an API gateway that remaps "
            "``Authorization`` → ``x-datarobot-external-access-token``. "
            "If the value starts with 'Bearer ', the prefix is stripped automatically."
        ),
    )
    principal_id: str | None = Field(
        default_factory=_get_default_principal_id,
        description=(
            "Okta AI agent principal ID (env: ``IDP_AGENT_ID``). "
            "Used as ``iss``/``sub`` in the JWT client assertion."
        ),
    )
    private_jwk: OptionalSecretStr = Field(
        default_factory=_get_default_private_jwk,
        description=(
            "Base64-encoded or raw-JSON RSA private JWK (env: ``IDP_AGENT_PRIVATE_KEY_JWK``)."
        ),
    )

XAATokenExchangeImpl

Bases: StrEnum

Allowed values for XAA_TOKEN_EXCHANGE_IMPL.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

class XAATokenExchangeImpl(StrEnum):
    """Allowed values for ``XAA_TOKEN_EXCHANGE_IMPL``."""

    HTTP = "http"
    OKTA_SDK = "okta_sdk"

XAATokenExchange

Bases: Protocol

Protocol for XAA token exchange implementations.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

class XAATokenExchange(Protocol):
    """Protocol for XAA token exchange implementations."""

    config: OAuth2CrossApplicationAccessAuthProviderConfig

    async def exchange_token(self, params: _CrossAppFlowParams, subject_token: str) -> str:
        """Execute the full two-step XAA flow and return the final access token."""
        ...

exchange_token `async`

exchange_token(params: _CrossAppFlowParams, subject_token: str) -> str

Execute the full two-step XAA flow and return the final access token.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

async def exchange_token(self, params: _CrossAppFlowParams, subject_token: str) -> str:
    """Execute the full two-step XAA flow and return the final access token."""
    ...

OktaTokenExchange

Bases: XAATokenExchange

Token exchange via the okta-client-python SDK.

Delegates both XAA steps to CrossAppAccessFlow which internally:

Calls flow.start(token, audience, scope) — builds a private_key_jwt client assertion, posts RFC 8693 token exchange to the org AS, and returns the ID-JAG token.
Calls flow.resume() — builds a second client assertion for the custom AS, posts RFC 7523 JWT Bearer grant, and returns the final scoped access token.

The SDK handles assertion signing, endpoint discovery, and request formatting. This implementation maps our _CrossAppFlowParams to the SDK's configuration:

OAuth2ClientConfiguration(issuer=trusted_issuer) — org AS domain
JWTBearerClaims(audience=org_token_url) — aud for client assertion
CrossAppAccessTarget(issuer=exchange_audience) — custom AS issuer
flow.start(audience=exchange_audience, scope=id_jag_scopes) — ID-JAG params

Requires pip install okta-client-python.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

class OktaTokenExchange(XAATokenExchange):
    """Token exchange via the ``okta-client-python`` SDK.

    Delegates both XAA steps to ``CrossAppAccessFlow`` which internally:

    1. Calls ``flow.start(token, audience, scope)`` — builds a ``private_key_jwt``
       client assertion, posts RFC 8693 token exchange to the org AS, and returns
       the ID-JAG token.
    2. Calls ``flow.resume()`` — builds a second client assertion for the custom AS,
       posts RFC 7523 JWT Bearer grant, and returns the final scoped access token.

    The SDK handles assertion signing, endpoint discovery, and request formatting.
    This implementation maps our ``_CrossAppFlowParams`` to the SDK's configuration:

    * ``OAuth2ClientConfiguration(issuer=trusted_issuer)`` — org AS domain
    * ``JWTBearerClaims(audience=org_token_url)`` — ``aud`` for client assertion
    * ``CrossAppAccessTarget(issuer=exchange_audience)`` — custom AS issuer
    * ``flow.start(audience=exchange_audience, scope=id_jag_scopes)`` — ID-JAG params

    Requires ``pip install okta-client-python``.
    """

    def __init__(self, config: OAuth2CrossApplicationAccessAuthProviderConfig) -> None:
        self.config = config

    async def exchange_token(self, params: _CrossAppFlowParams, subject_token: str) -> str:
        if not _HAS_OKTA_SDK:
            raise RuntimeError(
                "okta-client-python is not installed. "
                "Install it (`pip install okta-client-python`) or "
                "set XAA_TOKEN_EXCHANGE_IMPL=http to use the HTTP implementation."
            )

        if not self.config.principal_id:
            raise ValueError("principal_id is required for the XAA flow")
        if not self.config.private_jwk:
            raise ValueError("private_jwk is required for the XAA flow")

        private_jwk = _parse_private_jwk(self.config.private_jwk.get_secret_value())
        org_token_url = params.trusted_issuer.rstrip("/") + "/oauth2/v1/token"

        key_provider = LocalKeyProvider(
            key=private_jwk,
            algorithm=private_jwk.get("alg", "RS256"),
            key_id=private_jwk.get("kid"),
        )
        claims = JWTBearerClaims(
            issuer=self.config.principal_id,
            subject=self.config.principal_id,
            audience=org_token_url,
            expires_in=300,
        )
        additional_parameters = {"resource": params.target_audience}
        client = OAuth2Client(
            configuration=OAuth2ClientConfiguration(
                issuer=params.trusted_issuer,
                additional_parameters=additional_parameters,
                client_authorization=ClientAssertionAuthorization(
                    assertion_claims=claims,
                    key_provider=key_provider,
                ),
            )
        )
        target = CrossAppAccessTarget(issuer=params.exchange_audience)
        flow = CrossAppAccessFlow(client=client, target=target)

        logger.info(
            "OktaTokenExchange: starting XAA flow (trusted_issuer=%s, exchange_audience=%s)",
            params.trusted_issuer,
            params.exchange_audience,
        )

        result = await flow.start(
            token=subject_token,
            token_type="access_token",
            audience=params.exchange_audience,
            scope=params.id_jag_scopes,
        )
        if result.resume_assertion_claims is not None:
            logger.info("OktaTokenExchange: SDK requires resume step")
        token = await flow.resume()

        logger.info("OktaTokenExchange: XAA flow completed successfully")
        return token.access_token

ApiTokenExchange

Bases: XAATokenExchange

Token exchange via direct HTTP calls (httpx + PyJWT).

Implements both XAA steps explicitly, giving full control over request parameters (unlike the SDK which derives some from AS policy):

Step 1 — RFC 8693 Token Exchange at {trusted_issuer}/oauth2/v1/token: Sends subject_token (caller's access token), audience (custom AS issuer), resource (target API), and scope to obtain an ID-JAG. The ID-JAG is a short-lived, single-use JWT with typ: oauth-id-jag+jwt that encodes the user identity and agent principal.

Step 2 — RFC 7523 JWT Bearer Grant at token_url (custom AS): Sends the ID-JAG as assertion to obtain the final scoped access token. The token carries sub (human user) and cid (agent workload principal).

Both steps use a private_key_jwt client assertion signed with the agent's RSA private key.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

class ApiTokenExchange(XAATokenExchange):
    """Token exchange via direct HTTP calls (``httpx`` + ``PyJWT``).

    Implements both XAA steps explicitly, giving full control over request
    parameters (unlike the SDK which derives some from AS policy):

    **Step 1 — RFC 8693 Token Exchange** at ``{trusted_issuer}/oauth2/v1/token``:
      Sends ``subject_token`` (caller's access token), ``audience`` (custom AS
      issuer), ``resource`` (target API), and ``scope`` to obtain an ID-JAG.
      The ID-JAG is a short-lived, single-use JWT with ``typ: oauth-id-jag+jwt``
      that encodes the user identity and agent principal.

    **Step 2 — RFC 7523 JWT Bearer Grant** at ``token_url`` (custom AS):
      Sends the ID-JAG as ``assertion`` to obtain the final scoped access token.
      The token carries ``sub`` (human user) and ``cid`` (agent workload principal).

    Both steps use a ``private_key_jwt`` client assertion signed with the
    agent's RSA private key.
    """

    def __init__(self, config: OAuth2CrossApplicationAccessAuthProviderConfig) -> None:
        self.config = config

    async def exchange_token(self, params: _CrossAppFlowParams, subject_token: str) -> str:
        if not self.config.principal_id:
            raise ValueError("principal_id is required for the XAA flow")
        if not self.config.private_jwk:
            raise ValueError("private_jwk is required for the XAA flow")

        private_jwk = _parse_private_jwk(self.config.private_jwk.get_secret_value())

        org_as_token_url = params.trusted_issuer.rstrip("/") + "/oauth2/v1/token"
        custom_as_token_url = params.token_url

        # Build both client assertions up front — pure crypto, no I/O
        assertion1 = _make_client_assertion(
            principal_id=self.config.principal_id,
            token_url=org_as_token_url,
            private_jwk=private_jwk,
        )
        assertion2 = _make_client_assertion(
            principal_id=self.config.principal_id,
            token_url=custom_as_token_url,
            private_jwk=private_jwk,
        )

        async with httpx.AsyncClient() as http:
            # Step 1: RFC 8693 token exchange → ID-JAG
            logger.info(
                "ApiTokenExchange Step 1: token exchange → ID-JAG (org_as=%s, audience=%s)",
                org_as_token_url,
                params.exchange_audience,
            )
            resp1 = await http.post(
                org_as_token_url,
                data={
                    "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
                    "subject_token": subject_token,
                    "subject_token_type": _SUBJECT_TOKEN_TYPE,
                    "requested_token_type": _REQUESTED_TOKEN_TYPE,
                    "audience": params.exchange_audience,
                    "resource": params.target_audience,
                    "scope": " ".join(params.id_jag_scopes),
                    "client_assertion_type": _CLIENT_ASSERTION_TYPE,
                    "client_assertion": assertion1,
                },
                headers={"Content-Type": "application/x-www-form-urlencoded"},
                timeout=_TOKEN_EXCHANGE_TIMEOUT,
            )
            _raise_token_error(resp1, step=1, url=org_as_token_url)
            id_jag = resp1.json()["access_token"]

            # Step 2: RFC 7523 JWT Bearer grant → scoped agent token
            logger.info(
                "ApiTokenExchange Step 2: JWT Bearer grant → scoped token "
                "(custom_as=%s, target=%s)",
                custom_as_token_url,
                params.target_audience,
            )
            resp2 = await http.post(
                custom_as_token_url,
                data={
                    "grant_type": _JWT_BEARER_GRANT_TYPE,
                    "assertion": id_jag,
                    "client_assertion_type": _CLIENT_ASSERTION_TYPE,
                    "client_assertion": assertion2,
                },
                headers={"Content-Type": "application/x-www-form-urlencoded"},
                timeout=_TOKEN_EXCHANGE_TIMEOUT,
            )
            _raise_token_error(resp2, step=2, url=custom_as_token_url)
            return resp2.json()["access_token"]

OAuth2CrossApplicationAccessOAuth2AuthProvider

Bases: A2ADiscoveryAuthMixin, AuthProviderBase[OAuth2CrossApplicationAccessAuthProviderConfig]

Auth provider for XAA A2A calls.

Discovery — forwards the incoming bearer token as Authorization: Bearer.
Call — delegates to :func:get_token_exchange which returns either :class:OktaTokenExchange (SDK) or :class:ApiTokenExchange (HTTP) based on XAA_TOKEN_EXCHANGE_IMPL.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

class OAuth2CrossApplicationAccessOAuth2AuthProvider(
    A2ADiscoveryAuthMixin,
    AuthProviderBase[OAuth2CrossApplicationAccessAuthProviderConfig],
):
    """Auth provider for XAA A2A calls.

    * **Discovery** — forwards the incoming bearer token as ``Authorization: Bearer``.
    * **Call** — delegates to :func:`get_token_exchange` which returns either
      :class:`OktaTokenExchange` (SDK) or :class:`ApiTokenExchange` (HTTP)
      based on ``XAA_TOKEN_EXCHANGE_IMPL``.
    """

    def __init__(self, config: OAuth2CrossApplicationAccessAuthProviderConfig) -> None:
        super().__init__(config)
        self._flow_params: _CrossAppFlowParams | None = None
        self._forward_inbound_http_headers: bool = False

    def set_cross_app_flow_params(self, cross_app_flow_params: _CrossAppFlowParams) -> None:
        self._flow_params = cross_app_flow_params

    def set_forward_inbound_http_headers(self, enabled: bool) -> None:
        self._forward_inbound_http_headers = enabled

    async def authenticate_for_discovery(self, user_id: str | None = None) -> dict[str, str]:
        token = self._extract_token()
        logger.info("Forwarding token from '%s' for discovery", self.config.okta_token_header)
        return {"Authorization": f"Bearer {token}"}

    def set_agent_card(self, card: AgentCard) -> None:
        """Parse the agent card and store :class:`_CrossAppFlowParams`.

        Called by ``_AuthenticatedA2ABaseClient`` after discovery, before
        ``authenticate()``.
        """
        self._flow_params = _parse_cross_app_params(card)
        logger.info(
            "Agent card parsed: trusted_issuer=%s, exchange_audience=%s, "
            "target_audience=%s, auth_method=%s, scopes=%s",
            self._flow_params.trusted_issuer,
            self._flow_params.exchange_audience,
            self._flow_params.target_audience,
            self._flow_params.token_endpoint_auth_method,
            self._flow_params.id_jag_scopes,
        )

    def get_non_forwardable_header_keys(self) -> set[str]:
        excluded_header_keys = {self.config.okta_token_header.lower()}
        excluded_header_keys.update(header.lower() for header in self.config.fallback_token_headers)
        return excluded_header_keys

    def get_forwardable_headers_from_inbound_request(self) -> list[HeaderCred]:
        headers: dict[str, str] = Context.get().metadata.headers or {}

        return [
            HeaderCred(name=header_key, value=SecretStr(header_value))
            for header_key, header_value in headers.items()
            if header_key not in self.get_non_forwardable_header_keys() and header_value is not None
        ]

    async def get_exchanged_token(self) -> BearerTokenCred:
        if self._flow_params is None:
            raise RuntimeError(
                "authenticate() called before set_agent_card(). "
                "Ensure the provider is used with authenticated_a2a_client."
            )

        subject_token = self._extract_token()
        impl = get_token_exchange(self.config)
        exchanged_token = await impl.exchange_token(self._flow_params, subject_token)
        return BearerTokenCred(token=exchanged_token)

    async def authenticate(self, user_id: str | None = None, **kwargs: Any) -> AuthResult | None:
        """Obtain a scoped agent token via the XAA flow.

        Delegates to the configured :class:`XAATokenExchange` implementation.
        Requires :meth:`set_agent_card` to have been called first.
        """
        auth_request_credentials: list[HeaderCred | BearerTokenCred] = []

        if self._forward_inbound_http_headers:
            forwardable_header_creds = self.get_forwardable_headers_from_inbound_request()
            auth_request_credentials.extend(forwardable_header_creds)
        bearer_token_cred = await self.get_exchanged_token()
        auth_request_credentials.append(bearer_token_cred)

        return AuthResult(credentials=auth_request_credentials)

    def _extract_token(self) -> str:
        """Extract the access token from NAT request context headers.

        Tries ``okta_token_header`` first, then each ``fallback_token_headers``
        entry (stripping ``Bearer `` prefix if present).
        """
        headers: dict[str, str] = Context.get().metadata.headers or {}

        token = headers.get(self.config.okta_token_header.lower())
        if token:
            return token

        for fallback in self.config.fallback_token_headers:
            value = headers.get(fallback.lower())
            if value:
                if value.lower().startswith("bearer "):
                    value = value[len("bearer ") :]
                logger.debug("Using fallback header '%s'", fallback)
                return value

        raise RuntimeError(
            f"Header '{self.config.okta_token_header}' not found in request context "
            f"(also tried fallbacks: {self.config.fallback_token_headers}). "
            "The access token must be forwarded with every agent call."
        )

set_agent_card

set_agent_card(card: AgentCard) -> None

Parse the agent card and store :class:_CrossAppFlowParams.

Called by _AuthenticatedA2ABaseClient after discovery, before authenticate().

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

def set_agent_card(self, card: AgentCard) -> None:
    """Parse the agent card and store :class:`_CrossAppFlowParams`.

    Called by ``_AuthenticatedA2ABaseClient`` after discovery, before
    ``authenticate()``.
    """
    self._flow_params = _parse_cross_app_params(card)
    logger.info(
        "Agent card parsed: trusted_issuer=%s, exchange_audience=%s, "
        "target_audience=%s, auth_method=%s, scopes=%s",
        self._flow_params.trusted_issuer,
        self._flow_params.exchange_audience,
        self._flow_params.target_audience,
        self._flow_params.token_endpoint_auth_method,
        self._flow_params.id_jag_scopes,
    )

authenticate `async`

authenticate(user_id: str | None = None, **kwargs: Any) -> AuthResult | None

Obtain a scoped agent token via the XAA flow.

Delegates to the configured :class:XAATokenExchange implementation. Requires :meth:set_agent_card to have been called first.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

async def authenticate(self, user_id: str | None = None, **kwargs: Any) -> AuthResult | None:
    """Obtain a scoped agent token via the XAA flow.

    Delegates to the configured :class:`XAATokenExchange` implementation.
    Requires :meth:`set_agent_card` to have been called first.
    """
    auth_request_credentials: list[HeaderCred | BearerTokenCred] = []

    if self._forward_inbound_http_headers:
        forwardable_header_creds = self.get_forwardable_headers_from_inbound_request()
        auth_request_credentials.extend(forwardable_header_creds)
    bearer_token_cred = await self.get_exchanged_token()
    auth_request_credentials.append(bearer_token_cred)

    return AuthResult(credentials=auth_request_credentials)

get_token_exchange

get_token_exchange(config: OAuth2CrossApplicationAccessAuthProviderConfig) -> XAATokenExchange

Return the configured :class:XAATokenExchange implementation.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

def get_token_exchange(
    config: OAuth2CrossApplicationAccessAuthProviderConfig,
) -> XAATokenExchange:
    """Return the configured :class:`XAATokenExchange` implementation."""
    if _get_token_exchange_impl() is XAATokenExchangeImpl.HTTP:
        return ApiTokenExchange(config)
    return OktaTokenExchange(config)

okta_cross_app_access_auth_provider `async`

okta_cross_app_access_auth_provider(config: OAuth2CrossApplicationAccessAuthProviderConfig, builder: Builder) -> AsyncGenerator[OAuth2CrossApplicationAccessOAuth2AuthProvider, None]

NAT auth provider factory for Cross-Application Access.

Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py

@register_auth_provider(config_type=OAuth2CrossApplicationAccessAuthProviderConfig)
async def okta_cross_app_access_auth_provider(
    config: OAuth2CrossApplicationAccessAuthProviderConfig, builder: Builder
) -> AsyncGenerator[OAuth2CrossApplicationAccessOAuth2AuthProvider, None]:
    """NAT auth provider factory for Cross-Application Access."""
    yield OAuth2CrossApplicationAccessOAuth2AuthProvider(config=config)

datarobot_genai.dragent.plugins.okta_a2a_auth

okta_a2a_auth

Overview

OAuth2CrossApplicationAccessAuthProviderConfig

XAATokenExchangeImpl

XAATokenExchange

exchange_token async

OktaTokenExchange

ApiTokenExchange

OAuth2CrossApplicationAccessOAuth2AuthProvider

set_agent_card

authenticate async

get_token_exchange

okta_cross_app_access_auth_provider async

`datarobot_genai.dragent.plugins.okta_a2a_auth`

exchange_token `async`

authenticate `async`

okta_cross_app_access_auth_provider `async`