class DRAgentFastApiFrontEndPluginWorker(FastApiFrontEndPluginWorker):
def __init__(self, *args: object, **kwargs: object) -> None:
super().__init__(*args, **kwargs)
self._a2a_worker: A2AFrontEndPluginWorker | 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
async def add_routes(self, app: FastAPI, builder: WorkflowBuilder) -> None:
await super().add_routes(app, builder)
if self.front_end_config.a2a:
await self._add_a2a_routes(app, builder, self.front_end_config.a2a)
async def _add_a2a_routes(
self, app: FastAPI, builder: WorkflowBuilder, a2a_config: A2AFrontEndConfig
) -> 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.
a2a_config = a2a_config.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": a2a_config})}
)
self._a2a_worker = A2AFrontEndPluginWorker(nat_config)
cross_app_access = (
self.front_end_config.a2a.cross_application_access
if self.front_end_config.a2a
else None
)
skills = self.front_end_config.a2a.skills if self.front_end_config.a2a else []
external = self.front_end_config.a2a.external if self.front_end_config.a2a else None
agent_card = await create_agent_card(
frontend_config=self._a2a_worker.front_end_config,
cross_app_access=cross_app_access,
skills=skills,
external=external,
)
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)
a2a_app = a2a_server.build()
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 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)
# 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):
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
setup_logging()
return app
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"],
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}")