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 gates public agent-card access on per-agent developer opt-in.

Unauthenticated access also depends on platform-level opt-in per cluster to route anonymous traffic to the agent; this class enforces only the agent-side policy once a request reaches the process.

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

class DRAgentA2AStarletteApplication(A2AStarletteApplication):
    """A2A server that gates public agent-card access on per-agent developer opt-in.

    Unauthenticated access also depends on platform-level opt-in per cluster to
    route anonymous traffic to the agent; this class enforces only the agent-side
    policy once a request reaches the process.
    """

    def __init__(
        self,
        agent_card: AgentCard,
        http_handler: RequestHandler,
        extended_agent_card: AgentCard | None = None,
        context_builder: CallContextBuilder | None = None,
        card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard] | None = None,
        extended_card_modifier: Callable[
            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
        ]
        | None = None,
        max_content_length: int | None = 10 * 1024 * 1024,
        *,
        enable_unauthenticated_well_known_route: bool = False,
    ) -> None:
        self._enable_unauthenticated_well_known_route = enable_unauthenticated_well_known_route
        super().__init__(
            agent_card=agent_card,
            http_handler=http_handler,
            extended_agent_card=extended_agent_card,
            context_builder=context_builder,
            card_modifier=card_modifier,
            extended_card_modifier=extended_card_modifier,
            max_content_length=max_content_length,
        )

    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:
            if (
                resolve_identity_from_headers(headers, on_invalid_auth_context="none") is None
                and not self._enable_unauthenticated_well_known_route
            ):
                # The reason belongs in the log, not the response: the caller is anonymous and
                # telling it why it was refused is what leaks the agent's existence.
                logger.info(
                    "Serving 404 for unauthenticated agent-card request: "
                    "enable_unauthenticated_well_known_route is disabled for this agent "
                    "(set it to true in workflow.yaml to allow anonymous access; also "
                    "requires platform-level opt-in per cluster)"
                )
                return JSONResponse(status_code=404, content=AGENT_CARD_NOT_FOUND_BODY)
            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,
        *,
        enable_unauthenticated_well_known_route: bool = False,
        redacted_agent_card_config: DRAgentA2ARedactedAgentCardConfig | None = None,
    ) -> DRAgentA2AStarletteApplication:
        """Create an A2A server with identity-keyed public and extended agent cards.

        Unauthenticated ``GET /.well-known/agent-card.json`` access requires opt-in
        at two levels: platform administrators must enable unauthenticated routing
        per cluster, and ``enable_unauthenticated_well_known_route`` must be set
        in the agent's ``workflow.yaml``. This method enforces the agent-side
        flag only.

        When the agent flag is disabled (default), unauthenticated callers receive
        the generic 404 -- indistinguishable from a nonexistent agent, so the
        refusal is not an existence oracle. When enabled, they receive a redacted
        card, shaped by ``redacted_agent_card_config`` (e.g. its ``skills`` list is
        empty unless ``redacted_agent_card_config.enable_skills`` opts in).
        Authenticated callers always receive the full card. ``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=_make_public_card_modifier(redacted_agent_card_config),
            extended_card_modifier=_extended_card_modifier,
            enable_unauthenticated_well_known_route=enable_unauthenticated_well_known_route,
        )
        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, *, enable_unauthenticated_well_known_route: bool = False, redacted_agent_card_config: DRAgentA2ARedactedAgentCardConfig | None = None) -> DRAgentA2AStarletteApplication

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

Unauthenticated GET /.well-known/agent-card.json access requires opt-in at two levels: platform administrators must enable unauthenticated routing per cluster, and enable_unauthenticated_well_known_route must be set in the agent's workflow.yaml. This method enforces the agent-side flag only.

When the agent flag is disabled (default), unauthenticated callers receive the generic 404 -- indistinguishable from a nonexistent agent, so the refusal is not an existence oracle. When enabled, they receive a redacted card, shaped by redacted_agent_card_config (e.g. its skills list is empty unless redacted_agent_card_config.enable_skills opts in). Authenticated callers always receive the full card. 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,
    *,
    enable_unauthenticated_well_known_route: bool = False,
    redacted_agent_card_config: DRAgentA2ARedactedAgentCardConfig | None = None,
) -> DRAgentA2AStarletteApplication:
    """Create an A2A server with identity-keyed public and extended agent cards.

    Unauthenticated ``GET /.well-known/agent-card.json`` access requires opt-in
    at two levels: platform administrators must enable unauthenticated routing
    per cluster, and ``enable_unauthenticated_well_known_route`` must be set
    in the agent's ``workflow.yaml``. This method enforces the agent-side
    flag only.

    When the agent flag is disabled (default), unauthenticated callers receive
    the generic 404 -- indistinguishable from a nonexistent agent, so the
    refusal is not an existence oracle. When enabled, they receive a redacted
    card, shaped by ``redacted_agent_card_config`` (e.g. its ``skills`` list is
    empty unless ``redacted_agent_card_config.enable_skills`` opts in).
    Authenticated callers always receive the full card. ``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=_make_public_card_modifier(redacted_agent_card_config),
        extended_card_modifier=_extended_card_modifier,
        enable_unauthenticated_well_known_route=enable_unauthenticated_well_known_route,
    )
    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, mount_path: str = A2A_MOUNT_PATH) -> str

Construct the A2A endpoint URL for the running server.

Three tiers, most specific first:

1. Behind an Envoy API gateway (DR_WORKLOAD_EXTERNAL_URL_HOST and DR_WORKLOAD_EXTERNAL_URL_PREFIX are both injected), the gateway's own route is the only externally reachable one, so it wins in every hosting mode.
2. In a DataRobot deployment (MLOPS_DEPLOYMENT_ID is set) or workload (WORKLOAD_ID), composes the platform URL from DATAROBOT_PUBLIC_API_ENDPOINT / DATAROBOT_ENDPOINT.
3. Otherwise falls back to the local http://{host}:{port}/a2a/ URL.

mount_path is the suffix A2A is actually mounted under inside the container. Every tier forwards the full prefixed path through to the app, so the suffix has to appear in the advertised URL or clients would address a route that does not exist. DRAgentA2AConfig never passes an empty mount_path (it rejects mounting A2A at the application root), but this function stays a generic URL composer rather than re-asserting that policy: an empty value here just yields the tier's own base with a single trailing slash.

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

def get_a2a_endpoint_url(host: str, port: int, mount_path: str = A2A_MOUNT_PATH) -> str:
    """Construct the A2A endpoint URL for the running server.

    Three tiers, most specific first:

    1. Behind an Envoy API gateway (``DR_WORKLOAD_EXTERNAL_URL_HOST`` and
       ``DR_WORKLOAD_EXTERNAL_URL_PREFIX`` are both injected), the gateway's own
       route is the only externally reachable one, so it wins in every hosting
       mode.
    2. In a DataRobot deployment (``MLOPS_DEPLOYMENT_ID`` is set) or workload
       (``WORKLOAD_ID``), composes the platform URL from
       ``DATAROBOT_PUBLIC_API_ENDPOINT`` / ``DATAROBOT_ENDPOINT``.
    3. Otherwise falls back to the local ``http://{host}:{port}/a2a/`` URL.

    ``mount_path`` is the suffix A2A is actually mounted under inside the container.
    Every tier forwards the full prefixed path through to the app, so the suffix has
    to appear in the advertised URL or clients would address a route that does not
    exist. ``DRAgentA2AConfig`` never passes an empty ``mount_path`` (it rejects
    mounting A2A at the application root), but this function stays a generic URL
    composer rather than re-asserting that policy: an empty value here just yields
    the tier's own base with a single trailing slash.
    """
    if external_base := resolve_external_workload_base():
        return join_mount_path(external_base, mount_path)

    deployment_id = get_deployment_id()
    workload_id = get_workload_id()

    if not (deployment_id or workload_id):
        return join_mount_path(f"http://{host}:{port}", 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, mount_path)
    assert workload_id is not None  # non-None guaranteed by the early-return above
    return build_workload_a2a_url(datarobot_endpoint, workload_id, mount_path)

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_default_bearer_security_schemes

build_default_bearer_security_schemes() -> tuple[dict[str, SecurityScheme], list[dict[str, list[str]]]]

Return the default HTTP Bearer scheme for agents without explicit OAuth config.

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

def build_default_bearer_security_schemes() -> tuple[
    dict[str, SecurityScheme], list[dict[str, list[str]]]
]:
    """Return the default HTTP Bearer scheme for agents without explicit OAuth config."""
    security_schemes = {
        BEARER_SECURITY_SCHEME_NAME: SecurityScheme(
            root=HTTPAuthSecurityScheme(
                type="http",
                scheme="bearer",
                description=BEARER_SECURITY_DESCRIPTION,
            )
        )
    }
    return security_schemes, [{BEARER_SECURITY_SCHEME_NAME: []}]

build_security_schemes async

build_security_schemes(frontend_config: A2AFrontEndConfig, cross_app_access: CrossApplicationAccessConfig | None) -> tuple[dict[str, SecurityScheme], list[dict[str, list[str]]]]

Assemble A2A security schemes, merging up to two auth sources.

• server_auth → authorization_code flow.
• cross_app_access → client_credentials flow.
• Neither configured → HTTP Bearer scheme for DataRobot API tokens.

Always returns a populated securitySchemes map.

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], list[dict[str, list[str]]]]:
    """Assemble A2A security schemes, merging up to two auth sources.

    * ``server_auth`` → authorization_code flow.
    * ``cross_app_access`` → client_credentials flow.
    * Neither configured → HTTP Bearer scheme for DataRobot API tokens.

    Always returns a populated ``securitySchemes`` map.
    """
    server_auth = frontend_config.server_auth

    if not server_auth and not cross_app_access:
        return build_default_bearer_security_schemes()

    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, mount_path: str = A2A_MOUNT_PATH) -> 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.

mount_path must match where the A2A app is actually mounted, so the card's url points at the live RPC endpoint rather than the default /a2a/.

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,
    mount_path: str = A2A_MOUNT_PATH,
) -> 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``.

    ``mount_path`` must match where the A2A app is actually mounted, so the card's
    ``url`` points at the live RPC endpoint rather than the default ``/a2a/``.
    """
    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, mount_path)

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

redact_agent_card

redact_agent_card(card: AgentCard, *, include_skills: bool = False) -> AgentCard

Return a public-safe view of an agent card.

Strips advertised skills, unless include_skills opts back in, 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, *, include_skills: bool = False) -> AgentCard:
    """Return a public-safe view of an agent card.

    Strips advertised skills, unless ``include_skills`` opts back in, 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": card.skills if include_skills else [],
            "capabilities": card.capabilities.model_copy(
                update={"extensions": filtered_extensions}
            ),
        }
    )