datarobot_genai.dragent.frontends.a2a

a2a

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.

DRAgentA2AStarletteApplication

Bases: A2AStarletteApplication

A2A server that selects redacted vs extended cards on the public GET route.

Source code in datarobot_genai/dragent/frontends/a2a.py

class DRAgentA2AStarletteApplication(A2AStarletteApplication):
    """A2A server that selects redacted vs extended cards on the public GET route."""

    async def _handle_get_agent_card(self, request):  # type: ignore[no-untyped-def]
        headers = normalise_headers(dict(request.headers))
        token = _a2a_headers.set(headers)
        try:
            return await super()._handle_get_agent_card(request)
        finally:
            _a2a_headers.reset(token)

DRAgentA2AFrontEndPluginWorker

Bases: A2AFrontEndPluginWorker

A2A worker with identity-keyed public cards and an authenticated extended card.

Source code in datarobot_genai/dragent/frontends/a2a.py

class DRAgentA2AFrontEndPluginWorker(A2AFrontEndPluginWorker):
    """A2A worker with identity-keyed public cards and an authenticated extended card."""

    def create_a2a_server(
        self,
        agent_card: AgentCard,
        agent_executor: NATWorkflowAgentExecutor,
    ) -> DRAgentA2AStarletteApplication:
        """Create an A2A server with identity-keyed public and extended agent cards.

        The public ``GET /.well-known/agent-card.json`` route serves a redacted card
        to anonymous callers and the full card when gateway identity headers are
        present.  ``extended_agent_card`` is also wired for
        ``agent/getAuthenticatedExtendedCard`` clients.
        """
        base_server = super().create_a2a_server(agent_card, agent_executor)
        server = DRAgentA2AStarletteApplication(
            agent_card=base_server.agent_card,
            http_handler=base_server.handler.request_handler,
            extended_agent_card=agent_card,
            card_modifier=_public_card_modifier,
            extended_card_modifier=_extended_card_modifier,
        )
        logger.info("Created A2A server with identity-keyed public agent card")
        return server

create_a2a_server

create_a2a_server(agent_card: AgentCard, agent_executor: NATWorkflowAgentExecutor) -> DRAgentA2AStarletteApplication

Create an A2A server with identity-keyed public and extended agent cards.

The public GET /.well-known/agent-card.json route serves a redacted card to anonymous callers and the full card when gateway identity headers are present. extended_agent_card is also wired for agent/getAuthenticatedExtendedCard clients.

Source code in datarobot_genai/dragent/frontends/a2a.py

def create_a2a_server(
    self,
    agent_card: AgentCard,
    agent_executor: NATWorkflowAgentExecutor,
) -> DRAgentA2AStarletteApplication:
    """Create an A2A server with identity-keyed public and extended agent cards.

    The public ``GET /.well-known/agent-card.json`` route serves a redacted card
    to anonymous callers and the full card when gateway identity headers are
    present.  ``extended_agent_card`` is also wired for
    ``agent/getAuthenticatedExtendedCard`` clients.
    """
    base_server = super().create_a2a_server(agent_card, agent_executor)
    server = DRAgentA2AStarletteApplication(
        agent_card=base_server.agent_card,
        http_handler=base_server.handler.request_handler,
        extended_agent_card=agent_card,
        card_modifier=_public_card_modifier,
        extended_card_modifier=_extended_card_modifier,
    )
    logger.info("Created A2A server with identity-keyed public agent card")
    return server

get_a2a_endpoint_url

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)

resolve_oauth_endpoints 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

build_oauth_flow_from_server_auth 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

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

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

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

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},
    )

build_security_schemes 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}]

create_agent_card 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,
        supports_authenticated_extended_card=True,
    )

redact_agent_card

redact_agent_card(card: AgentCard) -> AgentCard

Return a public-safe view of an agent card.

Strips advertised skills and removes internal/external identity extensions while preserving auth and cross-application-access metadata needed for anonymous discovery.

Source code in datarobot_genai/dragent/frontends/a2a.py

def redact_agent_card(card: AgentCard) -> AgentCard:
    """Return a public-safe view of an agent card.

    Strips advertised skills and removes internal/external identity extensions
    while preserving auth and cross-application-access metadata needed for
    anonymous discovery.
    """
    extensions = card.capabilities.extensions
    filtered_extensions = None
    if extensions:
        filtered = [ext for ext in extensions if ext.uri not in _IDENTITY_EXTENSION_URIS]
        filtered_extensions = filtered or None

    return card.model_copy(
        update={
            "skills": [],
            "capabilities": card.capabilities.model_copy(
                update={"extensions": filtered_extensions}
            ),
        }
    )