datarobot_genai.dragent.plugins.auth_a2a_client

auth_a2a_client

AgentCardAware

Bases: Protocol

Protocol for auth providers that need the agent card before authenticate().

:class:_AuthenticatedA2ABaseClient checks for this via isinstance and calls :meth:set_agent_card automatically after card resolution.

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

@runtime_checkable
class AgentCardAware(Protocol):
    """Protocol for auth providers that need the agent card before ``authenticate()``.

    :class:`_AuthenticatedA2ABaseClient` checks for this via ``isinstance``
    and calls :meth:`set_agent_card` automatically after card resolution.
    """

    def set_agent_card(self, card: AgentCard) -> None:
        """Receive the resolved agent card before ``authenticate()`` is invoked."""
        ...

set_agent_card

set_agent_card(card: AgentCard) -> None

Receive the resolved agent card before authenticate() is invoked.

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

def set_agent_card(self, card: AgentCard) -> None:
    """Receive the resolved agent card before ``authenticate()`` is invoked."""
    ...

A2ADiscoveryAuthMixin

Bases: ABC

Mixin for auth providers that need different credentials for agent card discovery.

Without this mixin, _AuthenticatedA2ABaseClient calls authenticate() for both discovery and call phases. Implement authenticate_for_discovery() to supply separate headers for the agent card GET.

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

class A2ADiscoveryAuthMixin(abc.ABC):
    """Mixin for auth providers that need different credentials for agent card discovery.

    Without this mixin, ``_AuthenticatedA2ABaseClient`` calls ``authenticate()`` for
    both discovery and call phases.  Implement ``authenticate_for_discovery()`` to
    supply separate headers for the agent card GET.
    """

    @abc.abstractmethod
    async def authenticate_for_discovery(self, user_id: str | None = None) -> dict[str, str]:
        """Return HTTP headers for the agent card GET request.

        Returns
        -------
        dict[str, str]
            Header name → value pairs to attach to the agent card HTTP request.
        """

authenticate_for_discovery abstractmethod async

authenticate_for_discovery(user_id: str | None = None) -> dict[str, str]

Return HTTP headers for the agent card GET request.

Returns

dict[str, str] Header name → value pairs to attach to the agent card HTTP request.

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

@abc.abstractmethod
async def authenticate_for_discovery(self, user_id: str | None = None) -> dict[str, str]:
    """Return HTTP headers for the agent card GET request.

    Returns
    -------
    dict[str, str]
        Header name → value pairs to attach to the agent card HTTP request.
    """

AgentCardRegistryLookup

Bases: BaseModel

Identifies an agent card in the central DataRobot agent card registry.

Exactly one of deployment_id or external_id must be specified. The registry is queried using standard DataRobot API-token authentication (DATAROBOT_API_TOKEN), which avoids the chicken-and-egg problem where the agent's own card endpoint requires per-agent AuthN/AuthZ.

Example YAML::

registry:
  deployment_id: "64a1b2c3d4e5f6a7b8c9d0e1"

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

class AgentCardRegistryLookup(BaseModel):
    """Identifies an agent card in the central DataRobot agent card registry.

    Exactly one of ``deployment_id`` or ``external_id`` must be specified.
    The registry is queried using standard DataRobot API-token authentication
    (``DATAROBOT_API_TOKEN``), which avoids the chicken-and-egg problem where the
    agent's own card endpoint requires per-agent AuthN/AuthZ.

    Example YAML::

        registry:
          deployment_id: "64a1b2c3d4e5f6a7b8c9d0e1"
    """

    deployment_id: str | None = Field(
        default=None,
        description="DataRobot deployment ID to look up in the central agent card registry.",
    )
    external_id: str | None = Field(
        default=None,
        description="External agent identifier to look up in the central agent card registry.",
    )

    @model_validator(mode="after")
    def _exactly_one_identifier(self) -> "AgentCardRegistryLookup":
        if self.deployment_id and self.external_id:
            raise ValueError(
                "Specify exactly one of 'deployment_id' or 'external_id' inside 'registry', "
                "not both. Each identifies the agent card differently in the central registry."
            )
        if not self.deployment_id and not self.external_id:
            raise ValueError(
                "The 'registry' block requires exactly one of 'deployment_id' or 'external_id' "
                "to identify the agent card in the central registry."
            )
        return self

AuthenticatedA2AClientConfig

Bases: A2AClientConfig

A2A client config with separate discovery and call-phase auth.

Supports two modes for agent card resolution:

Direct fetch (existing behaviour) — set url to the agent's base URL and the card is fetched from {url}/.well-known/agent-card.json::

function_groups:
  my_agent:
    _type: authenticated_a2a_client
    url: "http://agent.example.com:8080"
    auth_provider: my_auth

Central registry lookup — set registry with either deployment_id or external_id. The card is fetched from the DataRobot central agent card registry using DATAROBOT_API_TOKEN, and the agent's RPC URL is derived from the card's advertised url::

function_groups:
  my_agent:
    _type: authenticated_a2a_client
    registry:
      deployment_id: "64a1b2c3d4e5f6a7b8c9d0e1"
    auth_provider: okta_auth

The two modes are mutually exclusive.

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

class AuthenticatedA2AClientConfig(A2AClientConfig, name="authenticated_a2a_client"):  # type: ignore[call-arg]
    """A2A client config with separate discovery and call-phase auth.

    Supports two modes for agent card resolution:

    **Direct fetch** (existing behaviour) — set ``url`` to the agent's base URL
    and the card is fetched from ``{url}/.well-known/agent-card.json``::

        function_groups:
          my_agent:
            _type: authenticated_a2a_client
            url: "http://agent.example.com:8080"
            auth_provider: my_auth

    **Central registry lookup** — set ``registry`` with either ``deployment_id``
    or ``external_id``.  The card is fetched from the DataRobot central agent card
    registry using ``DATAROBOT_API_TOKEN``, and the agent's RPC URL is derived
    from the card's advertised ``url``::

        function_groups:
          my_agent:
            _type: authenticated_a2a_client
            registry:
              deployment_id: "64a1b2c3d4e5f6a7b8c9d0e1"
            auth_provider: okta_auth

    The two modes are mutually exclusive.
    """

    url: Any = Field(  # type: ignore[assignment]
        default=None,
        description="Base URL of the A2A agent for direct agent card fetch. "
        "Mutually exclusive with 'registry'.",
    )

    registry: AgentCardRegistryLookup | None = Field(
        default=None,
        description="Central DataRobot agent card registry lookup. Mutually exclusive with 'url'.",
    )

    @model_validator(mode="after")
    def _url_xor_registry(self) -> "AuthenticatedA2AClientConfig":
        has_url = self.url is not None
        has_registry = self.registry is not None
        if has_url and has_registry:
            raise ValueError(
                "Provide either 'url' for direct agent card fetch or 'registry' for "
                "central registry lookup, not both. "
                "When 'registry' is set, the agent's RPC URL is derived from the "
                "agent card's advertised URL."
            )
        if not has_url and not has_registry:
            raise ValueError(
                "Either 'url' or 'registry' must be provided. "
                "Use 'url' to fetch the agent card directly from the agent, "
                "or 'registry' to look it up in the central DataRobot agent card registry."
            )
        # Eager registration for batch prefetch (N+1 optimisation).
        # register() here is sync, no I/O — just queues the ID.
        # First async get() in __aenter__ flushes all pending IDs
        # in ≤2 HTTP calls; subsequent gets hit the warm cache.
        # See AgentCardRegistry docstring for full details.
        if has_registry:
            registry = get_default_registry_sync()
            registry.register(
                deployment_id=self.registry.deployment_id,  # type: ignore[union-attr]
                external_id=self.registry.external_id,  # type: ignore[union-attr]
            )
        return self

AuthenticatedA2AClientFunctionGroup

Bases: A2AClientFunctionGroup

Uses :class:_AuthenticatedA2ABaseClient so both A2A phases are authenticated.

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

class AuthenticatedA2AClientFunctionGroup(A2AClientFunctionGroup):
    """Uses :class:`_AuthenticatedA2ABaseClient` so both A2A phases are authenticated."""

    def add_function(self, name: str, fn: Any, **kwargs: Any) -> None:  # type: ignore[override]
        """Intercept function registration to wrap *fn* with error handling.

        In degraded mode (registry lookup failed), replaces *fn* entirely
        with one that raises the stored error — so ``_wrap_a2a_function``
        catches it and returns an actionable message to the LLM.
        This avoids coupling to NAT's client method signatures or protocols.
        """
        registry_error = getattr(self, "_registry_error", None)
        if registry_error is not None:
            err = registry_error

            @functools.wraps(fn)
            async def _raise_registry_error(*args: Any, **kwargs: Any) -> Any:
                raise err

            fn = _raise_registry_error
        super().add_function(name, _wrap_a2a_function(fn), **kwargs)

    async def __aenter__(self) -> "AuthenticatedA2AClientFunctionGroup":
        config: AuthenticatedA2AClientConfig = self._config  # type: ignore[assignment]

        user_id = Context.get().user_id
        if not user_id:
            raise RuntimeError("User ID not found in context")

        auth_provider: AuthProviderBase | None = None
        if config.auth_provider:
            try:
                auth_provider = await self._builder.get_auth_provider(config.auth_provider)
                logger.info(
                    "Resolved authentication provider '%s' for A2A client",
                    config.auth_provider,
                )
            except Exception as e:
                logger.error(
                    "Failed to resolve auth provider '%s': %s",
                    config.auth_provider,
                    e,
                )
                raise RuntimeError(f"Failed to resolve auth provider: {e}") from e

        # -------------------------------------------------------------------
        # Resolve agent card: registry lookup vs. direct fetch
        # -------------------------------------------------------------------
        pre_resolved_card: AgentCard | None = None

        if config.registry:
            # Catch AgentCardRegistryError so the function group initialises
            # in a degraded state instead of crashing (generic JSON-RPC -32603).
            try:
                registry = await get_default_registry()
                pre_resolved_card = await registry.get(
                    deployment_id=config.registry.deployment_id,
                    external_id=config.registry.external_id,
                )
            except AgentCardRegistryError as exc:
                error_msg = str(exc)
                logger.error(
                    "Agent card registry lookup failed: %s",
                    error_msg,
                )
                # Store the error so add_function replaces every registered
                # function with one that raises it (caught by _wrap_a2a_function).
                self._registry_error = exc
                placeholder_card = _make_placeholder_agent_card(error_msg)
                self._client = _FailedRegistryClient(placeholder_card)  # type: ignore[assignment]
                self._register_functions()
                return self

            base_url = str(pre_resolved_card.url)
            logger.info(
                "Agent card resolved via central registry (deployment_id=%s, external_id=%s), "
                "RPC base URL derived from card: %s",
                config.registry.deployment_id,
                config.registry.external_id,
                base_url,
            )
        else:
            base_url = str(config.url)

        self._client = _AuthenticatedA2ABaseClient(
            base_url=base_url,
            agent_card_path=config.agent_card_path,
            task_timeout=config.task_timeout,
            streaming=config.streaming,
            auth_provider=auth_provider,
        )

        # Inject pre-resolved card so _AuthenticatedA2ABaseClient skips discovery
        if pre_resolved_card:
            self._client._agent_card = pre_resolved_card

        await self._client.__aenter__()

        logger.info(
            "Connected to A2A agent at %s (auth_provider: %s, user_id: %s, registry: %s)",
            base_url,
            config.auth_provider or "none",
            user_id,
            bool(config.registry),
        )

        self._register_functions()
        return self

add_function

add_function(name: str, fn: Any, **kwargs: Any) -> None

Intercept function registration to wrap fn with error handling.

In degraded mode (registry lookup failed), replaces fn entirely with one that raises the stored error — so _wrap_a2a_function catches it and returns an actionable message to the LLM. This avoids coupling to NAT's client method signatures or protocols.

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

def add_function(self, name: str, fn: Any, **kwargs: Any) -> None:  # type: ignore[override]
    """Intercept function registration to wrap *fn* with error handling.

    In degraded mode (registry lookup failed), replaces *fn* entirely
    with one that raises the stored error — so ``_wrap_a2a_function``
    catches it and returns an actionable message to the LLM.
    This avoids coupling to NAT's client method signatures or protocols.
    """
    registry_error = getattr(self, "_registry_error", None)
    if registry_error is not None:
        err = registry_error

        @functools.wraps(fn)
        async def _raise_registry_error(*args: Any, **kwargs: Any) -> Any:
            raise err

        fn = _raise_registry_error
    super().add_function(name, _wrap_a2a_function(fn), **kwargs)

authenticated_a2a_client async

authenticated_a2a_client(config: AuthenticatedA2AClientConfig, _builder: Builder) -> AsyncGenerator[Any, None]

NAT factory for :class:AuthenticatedA2AClientFunctionGroup.

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

@register_per_user_function_group(config_type=AuthenticatedA2AClientConfig)  # type: ignore[untyped-decorator]
async def authenticated_a2a_client(
    config: AuthenticatedA2AClientConfig, _builder: Builder
) -> AsyncGenerator[Any, None]:
    """NAT factory for :class:`AuthenticatedA2AClientFunctionGroup`."""
    async with AuthenticatedA2AClientFunctionGroup(config, _builder) as group:
        yield group