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.
    """

A2ACredentialServiceWithDisabledCache

Bases: A2ACredentialService

A2ACredentialService with disabled cache.

The caching implementation of NAT A2ACredentialService is problematic. If the token is cached, the cached token is not keyed on specific claims. The side effect is that once the first token is cached, all following requests with different token claims will reuse the first cached token. This is not an expected caching behavior. This class override _authenticate method by getting rid of using its own caching.

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

class A2ACredentialServiceWithDisabledCache(A2ACredentialService):
    """A2ACredentialService with disabled cache.

    The caching implementation of NAT A2ACredentialService is problematic. If the token is cached,
    the cached token is not keyed on specific claims. The side effect is that once the first token
    is cached, all following requests with different token claims will reuse the first cached token.
    This is not an expected caching behavior. This class override _authenticate method by getting
    rid of using its own caching.
    """

    async def _authenticate(self, user_id: str | None) -> AuthResult | None:
        """
        Authenticate and get credentials from NAT auth provider. This implementation doesn't rely
        on A2ACredentialService._cached_auth_result.

        Handles token expiration by triggering re-authentication if needed.
        Uses a lock to prevent concurrent authentication requests and race conditions.

        Args:
            user_id: User identifier for authentication

        Returns
        -------
            AuthResult with credentials or None on failure
        """
        try:
            # Acquire lock to serialize authentication attempts
            async with self._auth_lock:
                # Call NAT auth provider (provider is responsible for token refresh/validity)
                # Caching happens here on the level of auth provider
                # (e.g., OAuth2CrossApplicationAccessOAuth2AuthProvider)
                auth_result = await self._auth_provider.authenticate(user_id=user_id)

                # Warn if provider returned expired credentials (provider bug)
                if auth_result and auth_result.is_expired():
                    logger.warning(
                        "Auth provider returned already-expired credentials. "
                        "This may indicate a bug in the auth provider's token refresh logic."
                    )

                return auth_result

        except Exception as e:
            logger.error("Authentication failed: %s", _sanitize_a2a_error(e))
            logger.debug("Authentication exception detail: %s: %s", type(e).__name__, e)
            return None

AgentCardRegistryLookup

Bases: BaseModel

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

Exactly one of deployment_id, external_id or workload_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. With Okta cross-application access, request-time lookups use the caller's Okta token instead.

Example YAML::

registry:
  deployment_id: "64a1b2c3d4e5f6a7b8c9d0e1"

An agent served by the Workload API runtime is keyed by its workload ID instead::

registry:
  workload_id: "64a1b2c3d4e5f6a7b8c9d0e2"

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``, ``external_id`` or ``workload_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.
    With Okta cross-application access, request-time lookups use the caller's
    Okta token instead.

    Example YAML::

        registry:
          deployment_id: "64a1b2c3d4e5f6a7b8c9d0e1"

    An agent served by the Workload API runtime is keyed by its workload ID
    instead::

        registry:
          workload_id: "64a1b2c3d4e5f6a7b8c9d0e2"
    """

    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.",
    )
    workload_id: str | None = Field(
        default=None,
        description="DataRobot workload ID to look up in the central agent card registry.",
    )

    @model_validator(mode="after")
    def _exactly_one_identifier(self) -> "AgentCardRegistryLookup":
        identifiers = {
            "deployment_id": self.deployment_id,
            "external_id": self.external_id,
            "workload_id": self.workload_id,
        }
        provided = sorted(name for name, value in identifiers.items() if value)
        if len(provided) > 1:
            raise ValueError(
                f"Specify exactly one of 'deployment_id', 'external_id' or 'workload_id' "
                f"inside 'registry', not {len(provided)} ({', '.join(provided)}). "
                "Each identifies the agent card differently in the central registry."
            )
        if not provided:
            raise ValueError(
                "The 'registry' block requires exactly one of 'deployment_id', 'external_id' "
                "or 'workload_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 one of deployment_id, external_id or workload_id. The card is fetched from the DataRobot central agent card registry using DATAROBOT_API_TOKEN (the caller's Okta token with cross-application access), 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 one of ``deployment_id``,
    ``external_id`` or ``workload_id``.  The card is fetched from the DataRobot
    central agent card registry using ``DATAROBOT_API_TOKEN`` (the caller's Okta
    token with cross-application access), 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]
                workload_id=self.registry.workload_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)
                if isinstance(auth_provider, AgentCardAware):
                    # get_auth_provider returns one instance per configured name, shared
                    # by every function group and every user.  A provider that keeps
                    # agent-card state would otherwise serve whichever group connected
                    # last, so this group takes its own copy — config and credentials
                    # stay shared.
                    auth_provider = copy.copy(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()
                # Okta XAA looks the card up as the caller first.
                auth_headers = (
                    await _discovery_headers_for_registry(auth_provider)
                    if isinstance(auth_provider, A2ADiscoveryAuthMixin)
                    else None
                )
                pre_resolved_card = await registry.get(
                    deployment_id=config.registry.deployment_id,
                    external_id=config.registry.external_id,
                    workload_id=config.registry.workload_id,
                    auth_headers=auth_headers,
                )
            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, "
                "workload_id=%s), RPC base URL derived from card: %s",
                config.registry.deployment_id,
                config.registry.external_id,
                config.registry.workload_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