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
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 | None = Field(
default=None,
deprecated=True,
description=(
f"Deprecated. The inbound token header is fixed at "
f"``{OAUTH_ACCESS_TOKEN_HEADER}`` so that audience validation and this provider "
f"cannot read different headers. Accepted only when set to that value; removal "
f"in a future release."
),
)
fallback_token_headers: list[str] | None = Field(
default=None,
deprecated=True,
description=(
f"Deprecated. ``{OAUTH_ACCESS_TOKEN_FALLBACK_HEADER}`` is no longer read as an "
f"IdP-token carrier. Accepted only when set to that value; removal in a future "
f"release."
),
)
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``)."
),
)
cache_max_size: int = Field(
default=100,
description=(
"Max size of cache recording access token retrieved through XAA protocol. Each cache "
"record will be kept no longer than ttl_cache_in_second config."
),
)
cache_max_ttl_in_second: int = Field(
default=3600,
description=(
"TTL (in second) of cache recording access token retrieved through XAA protocol."
),
)
@model_validator(mode="after")
def _retire_token_header_overrides(self) -> "OAuth2CrossApplicationAccessAuthProviderConfig":
"""Accept a header override only where it still describes what happens.
Audience validation and this provider must read the same headers, so the set is now
fixed (``inbound_token.TOKEN_HEADERS``). A value equal to the old default is harmless
and kept working; anything else changed meaning, and ignoring it silently would leave
the config lying -- ``fallback_token_headers: []`` used to *disable* the fallback, so
accepting it would quietly switch it back on.
Warnings go to the log, not ``warnings.warn``: emitted from inside a validator, a
``DeprecationWarning`` is attributed to pydantic's module and filtered out in
production. ``Field(deprecated=True)`` above is for the JSON schema only -- it fires on
attribute access, and nothing reads these fields.
"""
for name, supported in (
("okta_token_header", OAUTH_ACCESS_TOKEN_HEADER),
("fallback_token_headers", [OAUTH_ACCESS_TOKEN_FALLBACK_HEADER]),
):
# __dict__, not getattr: Field(deprecated=True) installs a data descriptor whose
# __get__ warns, and a null means "no value" - model_dump() emits the None
# defaults explicitly, so a dump/reload round trip must not trip this.
value = self.__dict__.get(name)
if name not in self.model_fields_set or value is None:
continue
if value == supported:
logger.warning(
"%s is deprecated and no longer has any effect; it will be removed in a "
"future release. Delete it from workflow.yaml.",
name,
)
continue
raise ValueError(
f"{name} is no longer configurable and the value you set is not what the "
f"agent does. The access token is read only from "
f"'{OAUTH_ACCESS_TOKEN_HEADER}' -- inbound audience validation inspects "
f"exactly that, so reading any other header would exchange a token it never "
f"checked. Delete {name} from workflow.yaml and forward the standard header."
)
return self
|
Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py
| class XAATokenExchangeCacheManager:
def __init__(self, config: OAuth2CrossApplicationAccessAuthProviderConfig):
self.config = config
@staticmethod
def get_jwt_json_without_validating(jwt_str: str) -> dict[str, Any]:
return jwt.decode(jwt_str, options={"verify_signature": False})
@staticmethod
def get_hash(string_to_hash: str, bit_to_keep: int | None = None) -> str:
hash_value = hashlib.sha256(string_to_hash.encode()).hexdigest()
return hash_value if bit_to_keep is None else hash_value[:bit_to_keep]
@staticmethod
def get_ttl_from_xaa_access_token(xaa_access_token: str) -> float:
claims = XAATokenExchangeCacheManager.get_jwt_json_without_validating(xaa_access_token)
ttl = claims["exp"] - time.time()
return ttl
def get_ttl_or_fallback(self, xaa_access_token: str) -> float:
configured_ceiling = float(self.config.cache_max_ttl_in_second)
try:
ttl = self.get_ttl_from_xaa_access_token(xaa_access_token)
except (jwt.DecodeError, KeyError, TypeError):
logger.warning(
"XAA cache: exchanged token has no readable 'exp' claim; "
"falling back to configured TTL (%.0fs)",
configured_ceiling,
)
return configured_ceiling
return min(ttl, configured_ceiling)
def get_expires_at_from_xaa_access_token(self, xaa_access_token: str) -> datetime:
return datetime.now(UTC) + timedelta(seconds=self.get_ttl_or_fallback(xaa_access_token))
def ttu_handler(self, _: str, xaa_access_token: str, now: float) -> float:
return now + self.get_ttl_or_fallback(xaa_access_token)
@staticmethod
def get_stable_subject_token_string(subject_token: str) -> str:
"""Return stable subject JWT without information such as
- jti: JWT ID (unique identifier for this token)
- iat: JWT issue time
- exp: JWT expiry time
- etc.
"""
claims = XAATokenExchangeCacheManager.get_jwt_json_without_validating(subject_token)
stable = {
"sub": claims.get("sub"),
"iss": claims.get("iss"),
"cid": claims.get("cid"),
"aud": claims.get("aud"),
"scp": sorted(claims.get("scp", [])),
"act": claims.get("act"),
"sub_profile": claims.get("sub_profile"),
}
return json.dumps(stable, sort_keys=True, separators=(",", ":"))
@staticmethod
def get_subject_token_digest(subject_token: str) -> str:
stable_subject_token = XAATokenExchangeCacheManager.get_stable_subject_token_string(
subject_token
)
return XAATokenExchangeCacheManager.get_hash(stable_subject_token, 16)
@staticmethod
def get_nat_cross_app_params_digest(cross_app_params: _CrossAppFlowParams) -> str:
cross_app_params_str = cross_app_params.get_json_string_with_sorted_scope_claims()
return XAATokenExchangeCacheManager.get_hash(cross_app_params_str, 16)
@staticmethod
def get_safe_cache_key(
cross_app_params: _CrossAppFlowParams,
subject_token: str,
) -> str:
subject_token_digest = XAATokenExchangeCacheManager.get_subject_token_digest(subject_token)
params_digest = XAATokenExchangeCacheManager.get_nat_cross_app_params_digest(
cross_app_params
)
return XAATokenExchangeCacheManager.get_hash(subject_token_digest + params_digest)
def get_tlru_cache(self) -> TLRUCache:
return TLRUCache(
self.config.cache_max_size,
ttu=self.ttu_handler,
)
|
staticmethod
get_stable_subject_token_string(subject_token: str) -> str
Return stable subject JWT without information such as
- jti: JWT ID (unique identifier for this token)
- iat: JWT issue time
- exp: JWT expiry time
- etc.
Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py
| @staticmethod
def get_stable_subject_token_string(subject_token: str) -> str:
"""Return stable subject JWT without information such as
- jti: JWT ID (unique identifier for this token)
- iat: JWT issue time
- exp: JWT expiry time
- etc.
"""
claims = XAATokenExchangeCacheManager.get_jwt_json_without_validating(subject_token)
stable = {
"sub": claims.get("sub"),
"iss": claims.get("iss"),
"cid": claims.get("cid"),
"aud": claims.get("aud"),
"scp": sorted(claims.get("scp", [])),
"act": claims.get("act"),
"sub_profile": claims.get("sub_profile"),
}
return json.dumps(stable, sort_keys=True, separators=(",", ":"))
|
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"
|
Bases: ABC
Base class for XAA token exchange implementations.
Owns the caching logic shared by every implementation; subclasses only
implement the actual two-step exchange via :meth:_exchange_token_impl.
Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py
| class XAATokenExchange(abc.ABC):
"""Base class for XAA token exchange implementations.
Owns the caching logic shared by every implementation; subclasses only
implement the actual two-step exchange via :meth:`_exchange_token_impl`.
"""
def __init__(
self, config: OAuth2CrossApplicationAccessAuthProviderConfig, cache: TLRUCache
) -> None:
self.config = config
self.cache = cache
async def exchange_token(self, params: _CrossAppFlowParams, subject_token: str) -> str:
key = XAATokenExchangeCacheManager.get_safe_cache_key(params, subject_token)
try:
cached_access_token = self.cache[key]
logger.info("Return cached XAA access token.")
return cached_access_token
except KeyError:
logger.info("No cached XAA access token found.")
pass
token = await self._exchange_token_impl(params, subject_token)
self.cache[key] = token
logger.info("Add new cached XAA access token.")
return token
@abc.abstractmethod
async def _exchange_token_impl(self, params: _CrossAppFlowParams, subject_token: str) -> str:
"""Execute the full two-step XAA flow and return the final access token."""
...
|
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 domainJWTBearerClaims(audience=org_token_url) — aud for client assertionCrossAppAccessTarget(issuer=exchange_audience) — custom AS issuerflow.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``.
"""
@staticmethod
def get_oauth2_client_additional_parameters(
cross_app_flow_params: _CrossAppFlowParams,
) -> dict[str, str]:
additional_parameters = {}
if cross_app_flow_params.target_audience:
additional_parameters.update({"resource": cross_app_flow_params.target_audience})
return additional_parameters
async def _exchange_token_impl(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,
)
client = OAuth2Client(
configuration=OAuth2ClientConfiguration(
issuer=params.trusted_issuer,
additional_parameters=self.get_oauth2_client_additional_parameters(params),
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
|
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.
"""
@staticmethod
def get_xaa_token_exchange_request_payload(
cross_app_flow_params: _CrossAppFlowParams,
subject_token: str,
client_assertion: str,
) -> dict[str, str]:
payload = {
"grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
"subject_token": subject_token,
"subject_token_type": _SUBJECT_TOKEN_TYPE,
"requested_token_type": _REQUESTED_TOKEN_TYPE,
"audience": cross_app_flow_params.exchange_audience,
"scope": " ".join(cross_app_flow_params.id_jag_scopes),
"client_assertion_type": _CLIENT_ASSERTION_TYPE,
"client_assertion": client_assertion,
}
if cross_app_flow_params.target_audience:
payload["resource"] = cross_app_flow_params.target_audience
return payload
async def _exchange_token_impl(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=self.get_xaa_token_exchange_request_payload(params, subject_token, 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"]
|
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_x_datarobot_http_headers: bool = False
self._cache_manager = XAATokenExchangeCacheManager(self.config)
self._token_exchange_cache = self._cache_manager.get_tlru_cache()
self.token_exchange = get_token_exchange(self.config, self._token_exchange_cache)
def set_cross_app_flow_params(self, cross_app_flow_params: _CrossAppFlowParams) -> None:
self._flow_params = cross_app_flow_params
def set_forward_inbound_x_datarobot_http_headers(self, enabled: bool) -> None:
self._forward_inbound_x_datarobot_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", OAUTH_ACCESS_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]:
"""Token carriers are never forwarded; downstream agents get an exchanged token."""
return set(TOKEN_HEADERS)
def get_forwardable_x_datarobot_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.find("x-datarobot-") == 0
and 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()
exchanged_token = await self.token_exchange.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_x_datarobot_http_headers:
forwardable_header_creds = (
self.get_forwardable_x_datarobot_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)
token_expires_at = self._cache_manager.get_expires_at_from_xaa_access_token(
bearer_token_cred.token.get_secret_value()
)
return AuthResult(credentials=auth_request_credentials, token_expires_at=token_expires_at)
def _extract_token(self) -> str:
"""Extract the caller's IdP access token from NAT request context headers.
Uses ``inbound_token.find_idp_token``, the same function audience validation uses, so
this cannot exchange a token from a header the validator never inspected.
"""
headers: dict[str, str] = Context.get().metadata.headers or {}
token = find_idp_token(headers)
if token is not None:
return token
# Read for the message only, never for extraction: an operator who still has the
# deprecated override in workflow.yaml would otherwise see their own header name
# there and conclude the library is broken.
override = (
" A deprecated okta_token_header override is configured and ignored."
if "okta_token_header" in self.config.model_fields_set
else ""
)
raise RuntimeError(
f"No IdP access token in request context (looked in "
f"'{OAUTH_ACCESS_TOKEN_HEADER}').{override} The access token must be forwarded "
f"with every agent call."
)
|
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,
)
|
get_non_forwardable_header_keys() -> set[str]
Token carriers are never forwarded; downstream agents get an exchanged token.
Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py
| def get_non_forwardable_header_keys(self) -> set[str]:
"""Token carriers are never forwarded; downstream agents get an exchanged token."""
return set(TOKEN_HEADERS)
|
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_x_datarobot_http_headers:
forwardable_header_creds = (
self.get_forwardable_x_datarobot_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)
token_expires_at = self._cache_manager.get_expires_at_from_xaa_access_token(
bearer_token_cred.token.get_secret_value()
)
return AuthResult(credentials=auth_request_credentials, token_expires_at=token_expires_at)
|
get_token_exchange(config: OAuth2CrossApplicationAccessAuthProviderConfig, token_exchange_cache: TLRUCache) -> XAATokenExchange
Return the configured :class:XAATokenExchange implementation.
Source code in datarobot_genai/dragent/plugins/okta_a2a_auth.py
| def get_token_exchange(
config: OAuth2CrossApplicationAccessAuthProviderConfig,
token_exchange_cache: TLRUCache,
) -> XAATokenExchange:
"""Return the configured :class:`XAATokenExchange` implementation."""
if _get_token_exchange_impl() is XAATokenExchangeImpl.HTTP:
return ApiTokenExchange(config, token_exchange_cache)
return OktaTokenExchange(config, token_exchange_cache)
|
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)
|