A2A server helpers for DataRobot-hosted agents.
This module owns the A2A protocol layer: agent card construction, OAuth2
security scheme assembly, Cross-Application Access capability extensions,
and endpoint URL resolution. The FastAPI framework glue lives in
:mod:~datarobot_genai.dragent.frontends.fastapi.
get_a2a_endpoint_url(host: str, port: int) -> str
Construct the A2A endpoint URL for the running server.
In a DataRobot deployment (MLOPS_DEPLOYMENT_ID is set), uses the
deployment's direct-access URL built from DATAROBOT_PUBLIC_API_ENDPOINT
/ DATAROBOT_ENDPOINT. Otherwise falls back to the local
http://{host}:{port}/a2a/ URL.
Source code in datarobot_genai/dragent/frontends/a2a.py
| def get_a2a_endpoint_url(host: str, port: int) -> str:
"""Construct the A2A endpoint URL for the running server.
In a DataRobot deployment (``MLOPS_DEPLOYMENT_ID`` is set), uses the
deployment's direct-access URL built from ``DATAROBOT_PUBLIC_API_ENDPOINT``
/ ``DATAROBOT_ENDPOINT``. Otherwise falls back to the local
``http://{host}:{port}/a2a/`` URL.
"""
deployment_id = get_deployment_id()
workload_id = get_workload_id()
if not (deployment_id or workload_id):
return f"http://{host}:{port}/{A2A_MOUNT_PATH}/"
datarobot_endpoint = resolve_datarobot_endpoint(require=True)
assert datarobot_endpoint is not None # guaranteed by require=True
if deployment_id:
return build_deployment_a2a_url(datarobot_endpoint, deployment_id)
assert workload_id is not None # non-None guaranteed by the early-return above
return build_workload_a2a_url(datarobot_endpoint, workload_id)
|
async
resolve_oauth_endpoints(server_auth_config: OAuth2ResourceServerConfig) -> tuple[str, str]
Resolve (authorization_url, token_url) from an OAuth2ResourceServerConfig.
Uses OIDC discovery when discovery_url is set, otherwise derives from issuer_url.
Source code in datarobot_genai/dragent/frontends/a2a.py
| async def resolve_oauth_endpoints(
server_auth_config: OAuth2ResourceServerConfig,
) -> tuple[str, str]:
"""Resolve ``(authorization_url, token_url)`` from an OAuth2ResourceServerConfig.
Uses OIDC discovery when ``discovery_url`` is set, otherwise derives from ``issuer_url``.
"""
if server_auth_config.discovery_url:
try:
async with httpx.AsyncClient() as client:
response = await client.get(server_auth_config.discovery_url, timeout=5.0)
response.raise_for_status()
metadata = response.json()
auth_url = metadata.get("authorization_endpoint")
token_url = metadata.get("token_endpoint")
if auth_url and token_url:
logger.info(
"Resolved OAuth endpoints via discovery: %s",
server_auth_config.discovery_url,
)
return auth_url, token_url
except Exception as e:
logger.warning("Failed to discover OAuth endpoints: %s", e)
issuer = server_auth_config.issuer_url.rstrip("/")
auth_url = f"{issuer}/oauth/authorize"
token_url = f"{issuer}/oauth/token"
logger.info("Using derived OAuth endpoints from issuer: %s", issuer)
return auth_url, token_url
|
async
build_oauth_flow_from_server_auth(server_auth: OAuth2ResourceServerConfig) -> tuple[AuthorizationCodeOAuthFlow, list[str]]
Build the authorization_code OAuth2 flow and scopes from a NAT server_auth config.
Source code in datarobot_genai/dragent/frontends/a2a.py
| async def build_oauth_flow_from_server_auth(
server_auth: OAuth2ResourceServerConfig,
) -> tuple[AuthorizationCodeOAuthFlow, list[str]]:
"""Build the authorization_code OAuth2 flow and scopes from a NAT server_auth config."""
auth_url, token_url = await resolve_oauth_endpoints(server_auth)
flow = AuthorizationCodeOAuthFlow(
authorization_url=auth_url,
token_url=token_url,
scopes={scope: f"Permission: {scope}" for scope in server_auth.scopes},
)
return flow, list(server_auth.scopes)
|
build_oauth_flow_from_cross_app_access(config: CrossApplicationAccessConfig) -> tuple[ClientCredentialsOAuthFlow, list[str]]
Build the client_credentials flow and scopes from a CrossApplicationAccessConfig.
Extracts the OpenAPI-standard fields (token_url, scopes) only.
Cross-Application Access extension parameters are handled separately by
:func:build_cross_app_capability_extension and MUST NOT appear here.
Source code in datarobot_genai/dragent/frontends/a2a.py
| def build_oauth_flow_from_cross_app_access(
config: CrossApplicationAccessConfig,
) -> tuple[ClientCredentialsOAuthFlow, list[str]]:
"""Build the client_credentials flow and scopes from a CrossApplicationAccessConfig.
Extracts the OpenAPI-standard fields (``token_url``, ``scopes``) only.
Cross-Application Access extension parameters are handled separately by
:func:`build_cross_app_capability_extension` and MUST NOT appear here.
"""
flow = ClientCredentialsOAuthFlow(
token_url=config.token_request.token_url,
scopes={scope: f"Permission: {scope}" for scope in config.token_request.scopes},
)
return flow, list(config.token_request.scopes)
|
build_cross_app_capability_extension(config: CrossApplicationAccessConfig) -> list[AgentExtension]
Build the Cross-Application Access extension entry for capabilities.extensions.
Only extension-bound fields go in params; token_url and scopes
are intentionally omitted — they belong to OpenAPI securitySchemes.
Source code in datarobot_genai/dragent/frontends/a2a.py
| def build_cross_app_capability_extension(
config: CrossApplicationAccessConfig,
) -> list[AgentExtension]:
"""Build the Cross-Application Access extension entry for ``capabilities.extensions``.
Only extension-bound fields go in ``params``; ``token_url`` and ``scopes``
are intentionally omitted — they belong to OpenAPI ``securitySchemes``.
"""
params: dict = {
"ref": {
"scheme": CROSS_APP_SECURITY_SCHEME_REF,
"flow": CROSS_APP_SECURITY_SCHEME_FLOW_REF,
},
"tokenEndpointAuthMethod": config.token_endpoint_auth_method,
"tokenExchange": {
"grantType": TOKEN_EXCHANGE_GRANT_TYPE_URI,
"requestedTokenType": TOKEN_EXCHANGE_REQUESTED_TOKEN_TYPE,
"trustedIssuer": config.token_exchange.trusted_issuer,
"audience": config.token_exchange.audience,
},
"tokenRequest": {
"grantType": JWT_BEARER_GRANT_TYPE_URI,
"audience": config.token_request.audience,
},
}
return [
AgentExtension(
uri=JWT_BEARER_GRANT_TYPE_URI,
description=CROSS_APP_EXTENSION_DESCRIPTION,
params=params,
)
]
|
build_internal_identity_extension() -> AgentExtension | None
Build the internal identity extension for the current runtime, or None in local dev.
In a deployment container (MLOPS_DEPLOYMENT_ID) the params carry
deployment_id; in a workload container (WORKLOAD_ID) they carry
workload_id. Returns None when neither identity is present.
Source code in datarobot_genai/dragent/frontends/a2a.py
| def build_internal_identity_extension() -> AgentExtension | None:
"""Build the internal identity extension for the current runtime, or None in local dev.
In a deployment container (``MLOPS_DEPLOYMENT_ID``) the params carry
``deployment_id``; in a workload container (``WORKLOAD_ID``) they carry
``workload_id``. Returns *None* when neither identity is present.
"""
if dep_id := get_deployment_id():
params = {"deployment_id": dep_id}
elif wl_id := get_workload_id():
params = {"workload_id": wl_id}
else:
return None
return AgentExtension(
uri=INTERNAL_IDENTITY_URI,
description=INTERNAL_IDENTITY_DESCRIPTION,
required=True,
params=params,
)
|
build_external_identity_extension(external_id: str) -> AgentExtension
Build the external identity extension for catalog discovery.
Source code in datarobot_genai/dragent/frontends/a2a.py
| def build_external_identity_extension(external_id: str) -> AgentExtension:
"""Build the external identity extension for catalog discovery."""
return AgentExtension(
uri=EXTERNAL_IDENTITY_URI,
description=EXTERNAL_IDENTITY_DESCRIPTION,
required=False,
params={"id": external_id},
)
|
async
build_security_schemes(frontend_config: A2AFrontEndConfig, cross_app_access: CrossApplicationAccessConfig | None) -> tuple[dict[str, SecurityScheme] | None, list[dict[str, list[str]]] | None]
Assemble A2A security schemes, merging up to two auth sources.
server_auth → authorization_code flow.cross_app_access → client_credentials flow.
Returns (security_schemes, security_requirements), both None
when neither source is configured.
Source code in datarobot_genai/dragent/frontends/a2a.py
| async def build_security_schemes(
frontend_config: A2AFrontEndConfig,
cross_app_access: CrossApplicationAccessConfig | None,
) -> tuple[
dict[str, SecurityScheme] | None,
list[dict[str, list[str]]] | None,
]:
"""Assemble A2A security schemes, merging up to two auth sources.
* ``server_auth`` → authorization_code flow.
* ``cross_app_access`` → client_credentials flow.
Returns ``(security_schemes, security_requirements)``, both ``None``
when neither source is configured.
"""
server_auth = frontend_config.server_auth
if not server_auth and not cross_app_access:
return None, None
auth_code_flow, server_auth_scopes = (
await build_oauth_flow_from_server_auth(server_auth) if server_auth else (None, [])
)
client_creds_flow, cross_app_scopes = (
build_oauth_flow_from_cross_app_access(cross_app_access) if cross_app_access else (None, [])
)
all_scopes = list(dict.fromkeys(server_auth_scopes + cross_app_scopes))
security_schemes = {
"oauth2": SecurityScheme(
root=OAuth2SecurityScheme(
type="oauth2",
description=OAUTH2_SECURITY_DESCRIPTION_WITH_TOKEN_EXCHANGE,
flows=OAuthFlows(
authorization_code=auth_code_flow,
client_credentials=client_creds_flow,
),
)
)
}
return security_schemes, [{"oauth2": all_scopes}]
|
async
create_agent_card(frontend_config: A2AFrontEndConfig, cross_app_access: CrossApplicationAccessConfig | None, skills: list[AgentSkill], external: DRAgentA2AExternalConfig | None = None) -> AgentCard
Build an :class:~a2a.types.AgentCard for a DataRobot-hosted A2A agent.
When skills is empty, a single default skill is generated from
frontend_config.name / frontend_config.description.
Source code in datarobot_genai/dragent/frontends/a2a.py
| async def create_agent_card(
frontend_config: A2AFrontEndConfig,
cross_app_access: CrossApplicationAccessConfig | None,
skills: list[AgentSkill],
external: DRAgentA2AExternalConfig | None = None,
) -> AgentCard:
"""Build an :class:`~a2a.types.AgentCard` for a DataRobot-hosted A2A agent.
When ``skills`` is empty, a single default skill is generated from
``frontend_config.name`` / ``frontend_config.description``.
"""
security_schemes, security = await build_security_schemes(frontend_config, cross_app_access)
extensions = _collect_extensions(cross_app_access, external)
resolved_skills = skills or [
AgentSkill(
id="call",
name=frontend_config.name,
description=frontend_config.description,
tags=[],
examples=[],
)
]
url = _resolve_url(frontend_config, external)
return AgentCard(
name=frontend_config.name,
description=frontend_config.description,
url=url,
version=frontend_config.version,
default_input_modes=frontend_config.default_input_modes,
default_output_modes=frontend_config.default_output_modes,
capabilities=AgentCapabilities(
streaming=frontend_config.capabilities.streaming,
push_notifications=frontend_config.capabilities.push_notifications,
extensions=extensions,
),
skills=resolved_skills,
security_schemes=security_schemes or None,
security=security or None,
)
|