class DRAgentFastApiFrontEndPluginWorker(FastApiFrontEndPluginWorker):
def __init__(self, *args: object, **kwargs: object) -> None:
super().__init__(*args, **kwargs)
self._a2a_worker: DRAgentA2AFrontEndPluginWorker | None = None
def get_step_adaptor(self) -> StepAdaptor:
return DRAgentNestedReasoningStepAdaptor(self.front_end_config.step_adaptor)
async def _create_session_manager(
self, builder: WorkflowBuilder, entry_function: str | None = None
) -> SessionManager:
"""Create and register a SessionManager."""
sm = await DRAgentAGUISessionManager.create(
config=self._config, shared_builder=builder, entry_function=entry_function
)
self._session_managers.append(sm)
return sm
@property
def _a2a_config(self) -> DRAgentA2AConfig | None:
"""Return the agent's A2A block, or ``None`` when it has none.
``getattr`` because ``build_app()`` runs for any worker, and one built on NAT's vanilla
``FastApiFrontEndConfig`` has no ``a2a`` field at all.
"""
return getattr(self.front_end_config, "a2a", None)
async def add_routes(self, app: FastAPI, builder: WorkflowBuilder) -> None:
# Cancel the agent run when a streaming client disconnects, so an abandoned run
# can't keep working (and exporting spans) alongside later requests.
patch_nat_streaming_disconnect_cancellation()
await super().add_routes(app, builder)
if a2a := self._a2a_config:
await self._add_a2a_routes(app, builder, a2a)
def _resolve_expected_audience(self) -> str | None:
"""Audience an inbound token must carry: the one callers obtain through the
cross-application-access flow this agent advertises.
``None`` when XAA is not configured or its audience is unset. Whether the check runs
at all is ``a2a.oauth_claim_validation``, not this.
"""
a2a = self._a2a_config
if a2a is None or a2a.cross_application_access is None:
return None
return a2a.cross_application_access.token_request.audience
def _add_audience_validation_middleware(self, app: FastAPI) -> None:
"""Install the claim check on the served app, if opted in. See ``claim_validation``.
One instance covers ``/a2a`` too: Starlette keeps the mount prefix in
``scope["path"]``, so the outer middleware sees A2A traffic.
Must run here, not from ``add_routes``: NAT calls that from inside the lifespan, by
which point Starlette has frozen the middleware stack and ``add_middleware`` raises.
"""
a2a = self._a2a_config
if a2a is None or not a2a.oauth_claim_validation:
logger.info("OAuth claim validation disabled (a2a.oauth_claim_validation is not true)")
return
expected_audience = self._resolve_expected_audience()
if not expected_audience:
# Opted in with nothing to enforce: fail loudly rather than run an agent that
# believes it is validating.
raise ValueError(
"a2a.oauth_claim_validation is true but no expected audience is configured. "
"Set a2a.cross_application_access.token_request.audience, or turn "
"oauth_claim_validation off."
)
app.add_middleware(
GeneralOAuthClaimValidationMiddleware, expected_audience=expected_audience
)
logger.info("OAuth claim validation enabled (aud)")
async def _add_a2a_routes(
self, app: FastAPI, builder: WorkflowBuilder, a2a: DRAgentA2AConfig
) -> None:
# A2AFrontEndPluginWorker reads config.general.front_end to get its front_end_config.
# Pass a full Config with the A2AFrontEndConfig substituted in, and inherit host/port
# from the FastAPI front end so the agent card URL matches where the app is mounted.
server_config = a2a.server.model_copy(
update={"host": self.front_end_config.host, "port": self.front_end_config.port}
)
nat_config = self._config.model_copy(
update={"general": self._config.general.model_copy(update={"front_end": server_config})}
)
self._a2a_worker = DRAgentA2AFrontEndPluginWorker(nat_config)
agent_card = await create_agent_card(
frontend_config=self._a2a_worker.front_end_config,
cross_app_access=a2a.cross_application_access,
skills=a2a.skills,
external=a2a.external,
mount_path=a2a.mount_path,
)
session_manager = await DRAgentAGUISessionManager.create(
config=self._config,
shared_builder=builder,
max_concurrency=self._a2a_worker.max_concurrency,
)
self._session_managers.append(session_manager)
agent_executor = _PerUserCompatibleAgentExecutor(session_manager)
a2a_server = self._a2a_worker.create_a2a_server(
agent_card,
agent_executor,
enable_unauthenticated_well_known_route=a2a.enable_unauthenticated_well_known_route,
redacted_agent_card_config=a2a.redacted_agent_card,
)
a2a_app = a2a_server.build()
# Before we add anything of our own, so what we inspect is exactly the set of
# foreign routes. (Our own fallback could not collide anyway — it lives under
# /.well-known/ and mount_path segments cannot start with a dot — but checking
# first means that argument isn't load-bearing.)
self._check_mount_path_conflicts(app, a2a.mount_path)
# Registered before the mount purely to keep this file's "specific routes before
# catch-alls" convention; the two don't actually overlap since mount_path is
# always a non-empty suffix and the fallback lives at the disjoint app root.
self._register_root_agent_card_fallback(app, a2a_server, a2a.mount_path)
app.mount(f"/{a2a.mount_path}", a2a_app)
logger.info(f"A2A endpoint URL: {agent_card.url}")
logger.info(f"A2A agent card URL: {agent_card.url}.well-known/agent-card.json")
def _check_mount_path_conflicts(self, app: FastAPI, mount_path: str) -> None:
"""Fail (or warn) when ``mount_path`` lands on ground another route already owns.
Checked here rather than in ``DRAgentA2AConfig`` because this is the one moment the
real answer is knowable: the A2A mount is the last thing registered on the app, so
``app.routes`` now holds every route DataRobot, NAT and the installed plugins
actually added — including endpoints composed from ``front_end.endpoints`` at
arbitrary paths. A hardcoded reserved-name list would both miss those and drift
from whatever NAT registers next.
An **exact** conflict is fatal. A mount is normally reachable without its trailing
slash because Starlette answers ``POST /{mount_path}`` with a 307 to
``/{mount_path}/``, preserving method and body. A route sitting on the mount point
pre-empts that redirect, so the collision does not merely add ambiguity — it
removes the mechanism that makes the slashless form work, and an A2A client using
it silently gets that other route's response instead of the JSON-RPC endpoint.
A **nested** conflict only warns: the foreign route keeps winning for its own path
(it was registered first) and A2A still works at its advertised URL, so the result
is confusing rather than broken — not worth refusing to boot over, especially since
a future NAT route could otherwise brick a previously working config.
"""
mount_point = f"/{mount_path}"
nested_prefix = f"{mount_point}/"
# Not every BaseRoute carries a ``path`` (Host and WebSocketRoute do not), so read
# it defensively and keep only the string ones.
paths = {
path for route in app.routes if isinstance(path := getattr(route, "path", None), str)
}
if mount_point in paths:
raise ValueError(
f"a2a.mount_path {mount_path!r} is already served by another route at "
f"{mount_point}. That route was registered first, so it wins there and "
f"also suppresses the redirect that would otherwise send {mount_point} "
f"to {nested_prefix} — an A2A client that omits the trailing slash would "
"silently get that route's response instead of the JSON-RPC endpoint. "
'Choose a mount path that is not already routed, e.g. "a2a" or "api/a2a".'
)
if nested := sorted(path for path in paths if path.startswith(nested_prefix)):
logger.warning(
"a2a.mount_path %r overlaps existing route(s): %s. Those keep working and "
"A2A is still reachable at %s, but paths under %s are now split between "
"two apps, which will be confusing to debug. Consider a mount path that "
"does not overlap.",
mount_path,
", ".join(nested),
nested_prefix,
nested_prefix,
)
def _register_root_agent_card_fallback(
self, app: FastAPI, a2a_server: DRAgentA2AStarletteApplication, mount_path: str
) -> None:
"""Serve the agent card at the app root as well, for discovery fallback.
Clients resolve an agent card by trying ``{url}/.well-known/agent-card.json`` and
falling back to the same path at the host root. That fallback only exists if the
root actually answers, which it does not by default — the card lives at
``/{mount_path}/.well-known/agent-card.json`` and nowhere else since A2A cannot be
mounted at the root itself (``DRAgentA2AConfig.mount_path`` rejects an empty value).
The route delegates to the *same* bound handler the mounted app uses, so the
unauthenticated-access policy and card redaction stay in one place instead of
being reimplemented (and drifting) here.
"""
async def agent_card_fallback(request: Request) -> Response:
return await a2a_server._handle_get_agent_card(request)
app.add_api_route(
path=AGENT_CARD_WELL_KNOWN_PATH,
endpoint=agent_card_fallback,
methods=_GET_AND_HEAD,
response_model=None,
description=(
"Agent card, served at the root as a discovery fallback for clients that "
f"do not know it is mounted under /{mount_path}/"
),
tags=["A2A"],
)
logger.info(f"Added root agent card fallback at {AGENT_CARD_WELL_KNOWN_PATH}")
def build_app(self) -> FastAPI:
"""Build the FastAPI app, wrapping the parent lifespan to clean up the A2A worker."""
app = super().build_app()
# Register DataRobot health routes (/, /ping, /ping/, /health, /health/).
# NAT 1.6 no longer calls self.add_health_route() so we register here.
self._register_health_routes(app)
self._register_agent_manifest_route(app)
self._add_audience_validation_middleware(app)
self._add_model_monitoring_header_middleware(app)
# app.router.lifespan_context is the lifespan set by the parent's build_app().
# We wrap it to ensure the A2A worker's httpx client is closed on shutdown.
# (app.add_event_handler("shutdown", ...) is silently ignored when a lifespan is set.)
parent_lifespan = app.router.lifespan_context
@asynccontextmanager
async def lifespan(lifespan_app: FastAPI) -> AsyncIterator[None]:
async with parent_lifespan(lifespan_app):
await warmup_registry_from_config(self._config)
async with registry_refresh_lifespan(self._config):
yield
if self._a2a_worker is not None:
await self._a2a_worker.cleanup()
logger.info("A2A worker resources cleaned up")
app.router.lifespan_context = lifespan
_instrument_fastapi_app(app)
setup_logging()
return app
def _chat_completion_paths(self) -> frozenset[str]:
"""Paths served by the OpenAI-compatible chat endpoints NAT registers.
Mirrors the route construction in ``nat.front_ends.fastapi.routes.chat.add_chat_routes``.
Built from config rather than hardcoded, since ``endpoints`` entries can add more of
these at arbitrary paths.
"""
disable_legacy_routes = self.front_end_config.disable_legacy_routes
paths: set[str] = set()
for endpoint in [self.front_end_config.workflow, *self.front_end_config.endpoints]:
if endpoint.openai_api_v1_path:
paths.add(endpoint.openai_api_v1_path)
if endpoint.openai_api_path and endpoint.openai_api_path != endpoint.openai_api_v1_path:
paths.add(endpoint.openai_api_path)
paths.add(f"{endpoint.openai_api_path}/stream")
if not disable_legacy_routes and endpoint.legacy_openai_api_path:
paths.add(endpoint.legacy_openai_api_path)
paths.add(f"{endpoint.legacy_openai_api_path}/stream")
return frozenset(paths)
def _add_model_monitoring_header_middleware(self, app: FastAPI) -> None:
"""Set X-DataRobot-Model-Monitoring on chat-completions responses.
Scoped to the OpenAI-compatible chat routes: predictions-gateway uses it to decide whether
to run its own chat-completions monitoring.
"""
chat_completion_paths = self._chat_completion_paths()
@app.middleware("http")
async def add_model_monitoring_header(
request: Request, call_next: RequestResponseEndpoint
) -> Response:
response = await call_next(request)
if _route_path(request) in chat_completion_paths:
response.headers[DATAROBOT_MODEL_MONITORING_HEADER] = "true"
return response
def _register_health_routes(self, app: FastAPI) -> None:
"""Register DataRobot health check endpoints."""
class HealthResponse(BaseModel):
status: str = Field(description="Health status of the server")
async def health_check() -> HealthResponse:
"""Health check endpoint for liveness/readiness probes."""
return HealthResponse(status="healthy")
for path in DATAROBOT_EXPECTED_HEALTH_ROUTES:
app.add_api_route(
path=path,
endpoint=health_check,
methods=_GET_AND_HEAD,
response_model=HealthResponse,
description="Health check endpoint for liveness/readiness probes",
tags=["Health"],
responses={
200: {
"description": "Server is healthy",
"content": {"application/json": {"example": {"status": "healthy"}}},
}
},
)
logger.info(f"Added health check endpoint at {path}")
def _register_agent_manifest_route(self, app: FastAPI) -> None:
"""Register the well-known Agent Manifest endpoint.
Built once from ``self._config`` here (not per-request) since the
manifest is a static reflection of the declared workflow.yaml
structure, not a live computation.
"""
manifest = build_agent_manifest(self._config)
async def agent_manifest() -> AgentManifest:
"""Return the declared components, tools, and root agent for this running agent."""
return manifest
app.add_api_route(
path="/.well-known/agent-manifest.json",
endpoint=agent_manifest,
methods=_GET_AND_HEAD,
response_model=AgentManifest,
description="Declared components, tools, and root agent for this running agent",
tags=["Agent Manifest"],
)
logger.info("Added Agent Manifest endpoint at /.well-known/agent-manifest.json")