_agui_runner.py

# Tencent is pleased to support the open source community by making tRPC-Agent-Python available.
#
# Copyright (C) 2026 Tencent. All rights reserved.
#
# tRPC-Agent-Python is licensed under Apache-2.0.
"""Runner manager: owns the AG-UI manager for the FastAPI server with cancel support.
"""
from contextlib import asynccontextmanager
from typing import Any

from ag_ui.core import RunAgentInput
from fastapi import FastAPI
from pydantic import BaseModel

from trpc_agent_sdk.agents import BaseAgent
from trpc_agent_sdk.log import logger
from trpc_agent_sdk.server.ag_ui import AgUiAgent
from trpc_agent_sdk.server.ag_ui import AgUiManager
from trpc_agent_sdk.server.ag_ui import AgUiService
from trpc_agent_sdk.version import __version__


class HealthResponse(BaseModel):
    """Response body for GET /health."""

    status: str = "ok"
    app_name: str
    version: str = __version__


class AguiRunner:
    """AG-UI runner: owns the AG-UI manager for the FastAPI server."""

    def __init__(
        self,
        app_name: str,
    ) -> None:
        self._app_name = app_name
        self._agui_manager = AgUiManager()
        self._app = self._create_app()

    @property
    def app(self) -> FastAPI:
        """Get the FastAPI app for the AG-UI runner."""
        return self._app

    def register_service(self, service_name: str, service: AgUiService) -> None:
        """Register an AG-UI service."""
        self._agui_manager.register_service(service_name, service)

    def run(self, host: str, port: int, **kwargs: Any) -> None:
        """Run the AG-UI runner."""
        self._app.get("/health", response_model=HealthResponse, tags=["meta"])(self.health)
        self._agui_manager.set_app(self._app)
        self._agui_manager.run(host, port, **kwargs)

    @asynccontextmanager
    async def _lifespan(self, app: FastAPI):  # noqa: ARG001
        """Startup / shutdown hook: close the runner on exit."""
        logger.info("TRPC AG-UI Server (with cancel) starting up.")
        yield
        logger.info("TRPC AG-UI Server (with cancel) shutting down.")
        await self._agui_manager.close()

    def _create_app(self) -> FastAPI:
        app = FastAPI(
            title="TRPC AG-UI Server (Cancel Demo)",
            description="HTTP API for TRPC AG-UI Server with Cancel support",
            version=__version__,
            lifespan=self._lifespan,
        )
        return app

    async def health(self) -> HealthResponse:
        """Liveness check - always returns 200 while the server is up."""
        return HealthResponse(app_name=self._app_name)


def _create_agui_agent(name: str, root_agent: BaseAgent, **kwargs) -> AgUiAgent:
    """Create AgUiAgent with cancel support.

    Args:
        name: Name of the agent
        root_agent: Root agent instance
    Returns:
        AgUiAgent instance
    """
    agui_agent = AgUiAgent(
        trpc_agent=root_agent,
        app_name=name,
        cancel_wait_timeout=3.0,
        **kwargs,
    )
    return agui_agent


def create_agui_runner(app_name: str, service_name: str, uri: str, **kwargs: Any) -> AguiRunner:
    """Create AgUiService and add agent to it.

    Args:
        app_name: Name of the app
        service_name: Name of the service
        uri: URI of the agent
        kwargs: Additional keyword arguments to pass to the AgUiAgent constructor
    Returns:
        AguiRunner instance
    """
    ag_ui_runner: AguiRunner = AguiRunner(app_name)
    agui_service = AgUiService(service_name, app=ag_ui_runner.app)
    agui_agent = _create_agui_agent(app_name, **kwargs)
    agui_service.add_agent(uri, agui_agent)
    ag_ui_runner.register_service(service_name, agui_service)
    return ag_ui_runner