# Copyright (c) Microsoft. All rights reserved. """FastAPI endpoint creation for AG-UI agents.""" from __future__ import annotations import asyncio import copy import logging import uuid from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence from contextlib import suppress from inspect import isawaitable from typing import Any, TypeAlias, cast from ag_ui.core import RunErrorEvent from ag_ui.encoder import EventEncoder from agent_framework import CheckpointStorage, SupportsAgentRun, Workflow from fastapi import FastAPI, HTTPException, Request from fastapi.params import Depends from fastapi.responses import JSONResponse, Response, StreamingResponse from ._agent import AgentFrameworkAgent from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY from ._run_common import _is_snapshot_hydration_request from ._session_ownership import workflow_agent_tool_error, workflow_agent_tools from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshotStore, SnapshotScopeResolver, ) from ._types import AGUIRequest from ._workflow import AgentFrameworkWorkflow logger = logging.getLogger(__name__) AuthorizationScopeResolver: TypeAlias = Callable[..., str | Awaitable[str]] """FastAPI dependency callable that returns the authenticated AG-UI authorization scope.""" _DETACHED_READER_START_TIMEOUT_SECONDS = 1.0 _DETACHED_STREAM_QUEUE_SIZE = 16 _KEEPALIVE_COMMENT = "keepalive" _AUTHORIZATION_SCOPE_STATE_KEY = "_agent_framework_ag_ui_authorization_scope" _SINGLE_USER_SCOPE_PREFIX = "ag-ui:single-user:" def _get_snapshot_store( protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow, ) -> AGUIThreadSnapshotStore | None: if isinstance(protocol_runner, AgentFrameworkAgent): return protocol_runner.config.snapshot_store return protocol_runner.snapshot_store def _set_snapshot_store( protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow, snapshot_store: AGUIThreadSnapshotStore, ) -> None: if isinstance(protocol_runner, AgentFrameworkAgent): protocol_runner.config.snapshot_store = snapshot_store return protocol_runner.snapshot_store = snapshot_store def _configure_snapshot_persistence( protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow, *, snapshot_store: AGUIThreadSnapshotStore | None, ) -> None: existing_snapshot_store = _get_snapshot_store(protocol_runner) if snapshot_store is not None: if existing_snapshot_store is not None and existing_snapshot_store is not snapshot_store: raise ValueError("snapshot_store is already configured on the AG-UI runner.") if existing_snapshot_store is None: _set_snapshot_store(protocol_runner, snapshot_store) def _validate_keepalive_seconds(keepalive_seconds: float | None) -> None: if keepalive_seconds is not None and not keepalive_seconds > 0: raise ValueError("keepalive_seconds must be positive or None.") def _validate_detached_run_options(max_detached_runs: int, detached_run_timeout_seconds: float) -> None: if max_detached_runs < 1: raise ValueError("max_detached_runs must be greater than 0.") if detached_run_timeout_seconds <= 0: raise ValueError("detached_run_timeout_seconds must be positive.") def _validate_authorization_scope_configuration( *, authorization_scope_resolver: AuthorizationScopeResolver | None, snapshot_scope_resolver: SnapshotScopeResolver | None, single_user: bool, ) -> None: if not isinstance(single_user, bool): raise TypeError("single_user must be a bool.") configured_modes = sum( ( authorization_scope_resolver is not None, snapshot_scope_resolver is not None, single_user, ) ) if configured_modes == 0: raise ValueError( "AG-UI endpoint authorization must be explicit. " "Provide authorization_scope_resolver for authenticated applications, " "reuse snapshot_scope_resolver for compatibility, or declare single_user=True." ) if configured_modes > 1: raise ValueError( "Configure exactly one AG-UI authorization mode: authorization_scope_resolver, " "snapshot_scope_resolver, or single_user=True." ) def _single_user_scope() -> str: """Return one opaque scope that stays stable for an endpoint registration.""" return f"{_SINGLE_USER_SCOPE_PREFIX}{uuid.uuid4().hex}" def _reject_workflow_agent_tools(protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow) -> None: """Refuse a runner that exposes a ``WorkflowAgent`` as a tool, so that the mistake shows at startup.""" # A runner subclass that skips the base constructor has nothing to inspect, and a workflow factory builds its # workflow for each request, so a run checks what the factory built. if isinstance(protocol_runner, AgentFrameworkAgent): subject = getattr(protocol_runner, "agent", None) else: subject = getattr(protocol_runner, "workflow", None) tool_names = workflow_agent_tools(subject) if tool_names: runner_name = getattr(protocol_runner, "name", None) or type(protocol_runner).__name__ raise ValueError(workflow_agent_tool_error(runner_name, tool_names)) def add_agent_framework_fastapi_endpoint( app: FastAPI, agent: SupportsAgentRun | AgentFrameworkAgent | Workflow | AgentFrameworkWorkflow, path: str = "/", state_schema: Any | None = None, predict_state_config: dict[str, dict[str, str]] | None = None, allow_origins: list[str] | None = None, default_state: dict[str, Any] | None = None, tags: list[str] | None = None, dependencies: Sequence[Depends] | None = None, snapshot_store: AGUIThreadSnapshotStore | None = None, snapshot_scope_resolver: SnapshotScopeResolver | None = None, checkpoint_storage: CheckpointStorage | None = None, keepalive_seconds: float | None = 15, a2ui_config: dict[str, Any] | None = None, detached_runs: bool = False, max_detached_runs: int = 32, detached_run_timeout_seconds: float = 3600, authorization_scope_resolver: AuthorizationScopeResolver | None = None, single_user: bool = False, ) -> None: """Add an AG-UI endpoint to a FastAPI app. Args: app: The FastAPI application agent: The agent to expose (can be raw SupportsAgentRun or wrapped). The endpoint raises ``ValueError`` for an agent, or a workflow, that exposes a ``WorkflowAgent`` as a tool (``WorkflowAgent.as_tool()``), and a run ends with a ``WORKFLOW_AGENT_TOOL_UNSUPPORTED`` ``RUN_ERROR`` when a workflow factory builds one. path: The endpoint path state_schema: Optional state schema for shared state management; accepts dict or Pydantic model/class predict_state_config: Optional predictive state update configuration. Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}} allow_origins: CORS origins (not yet implemented) default_state: Optional initial state to seed when the client does not provide state keys tags: OpenAPI tags for endpoint categorization (defaults to ["AG-UI"]) dependencies: Optional FastAPI dependencies for authentication/authorization. These dependencies run before the endpoint handler. Use this to add authentication checks, rate limiting, or other middleware-like behavior. Their return values do not establish the AG-UI authorization scope; use authorization_scope_resolver for that. Example: `dependencies=[Depends(verify_api_key)]` snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence is opt-in and requires either an authorization resolver, the compatible Snapshot Scope resolver, or explicit single-user mode. snapshot_scope_resolver: Compatible request-body resolver for the application-defined Snapshot Scope. Prefer authorization_scope_resolver for new FastAPI applications because it can consume authenticated dependency results. This resolver also scopes internal Agent Sessions, in-memory workflow_factory instances, and approval state. It must return a non-empty string derived from authorized request context; invalid results fail the request before accessing state or invoking the runner. The shared-workflow rule described for authorization_scope_resolver applies to this resolver as well. checkpoint_storage: Optional workflow checkpoint storage, applied when the endpoint exposes a workflow. When provided, each run creates a checkpoint at the end of every superstep, and a run may resume from a persisted checkpoint by supplying its id in the AG-UI forwarded props (``forwarded_props: {"checkpoint_id": ...}``). keepalive_seconds: Endpoint SSE keepalive interval in seconds. Defaults to 15. Positive values emit fixed SSE comments while the stream is open. None disables keepalive and preserves the non-keepalive response path. Keepalive comments are transport traffic and do not change AG-UI events. a2ui_config: Optional backend A2UI config used when the runtime auto-injects the surface-generation tool (``forwardedProps.injectA2UITool``). Keys: ``inject_a2ui_tool`` (required backend opt-in), ``default_catalog_id``, ``catalog``, ``guidelines``, ``recovery``, ``default_surface_id``. detached_runs: Whether agent/workflow execution continues after the SSE client disconnects. Defaults to False. When enabled, the endpoint owns a bounded background producer and completes runner persistence even if the HTTP reader is cancelled. Request-scoped disposable resources may be released after disconnect, so detached work must use values resolved before streaming rather than retaining request-owned clients or sessions. This option does not provide resumable event replay. max_detached_runs: Maximum number of detached producers retained by this endpoint registration. Defaults to 32. Additional requests receive HTTP 503 until capacity is released. detached_run_timeout_seconds: Maximum time a producer may remain active after its SSE reader disconnects or never starts. Defaults to 3600 seconds. Expired producers are cancelled and release endpoint capacity. authorization_scope_resolver: Optional FastAPI dependency callable that returns a non-empty authorization scope derived from authenticated request context. The callable may declare normal FastAPI dependencies, headers, security schemes, or the current Request. Its result scopes snapshots, internal Agent Sessions, workflow instances and checkpoints, detached-run admission, and server-side approval state. An ``AgentExecutor`` owns one ``AgentSession`` (history, pending approvals, standing approval rules), so a shared ``Workflow`` or ``WorkflowAgent`` that holds one serves only the first scope that runs it. A run in another scope ends with a ``WORKFLOW_SCOPE_CONFLICT`` ``RUN_ERROR`` before it reaches the model. Pass ``AgentFrameworkWorkflow(workflow_factory=...)`` with a factory that builds a new workflow, including new ``AgentExecutor`` and ``AgentSession`` instances, for every call to serve several scopes. single_user: Explicitly declare that every caller of this endpoint belongs to one authorization principal. The endpoint receives one opaque scope that remains stable for that registration and separates equal route paths registered by other applications. Use an explicit resolver when scope continuity must survive process restarts or when storage is shared beyond one registration. A single-user application can return its own stable application scope from authorization_scope_resolver. Never enable this mode for callers with distinct authorization. A shared workflow keeps one conversation state for every thread of that principal; use ``workflow_factory`` when threads must stay independent. """ _validate_keepalive_seconds(keepalive_seconds) _validate_detached_run_options(max_detached_runs, detached_run_timeout_seconds) protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow if isinstance(agent, AgentFrameworkWorkflow): protocol_runner = agent elif isinstance(agent, AgentFrameworkAgent): protocol_runner = agent elif isinstance(agent, Workflow): protocol_runner = AgentFrameworkWorkflow(workflow=agent) elif isinstance(agent, SupportsAgentRun): protocol_runner = AgentFrameworkAgent( agent=agent, state_schema=state_schema, predict_state_config=predict_state_config, snapshot_store=snapshot_store, a2ui_config=a2ui_config, ) else: raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.") _reject_workflow_agent_tools(protocol_runner) _validate_authorization_scope_configuration( authorization_scope_resolver=authorization_scope_resolver, snapshot_scope_resolver=snapshot_scope_resolver, single_user=single_user, ) if checkpoint_storage is not None: if not isinstance(protocol_runner, AgentFrameworkWorkflow): raise ValueError("checkpoint_storage is only supported when the endpoint exposes a workflow.") # A pre-wrapped runner without storage adopts the endpoint's; a runner that # already carries a different storage is a configuration conflict. if ( protocol_runner.checkpoint_storage is not None and protocol_runner.checkpoint_storage is not checkpoint_storage ): raise ValueError("checkpoint_storage is already configured on the AG-UI workflow runner.") protocol_runner.checkpoint_storage = checkpoint_storage _configure_snapshot_persistence( protocol_runner, snapshot_store=snapshot_store, ) background_tasks: set[asyncio.Task[Any]] = set() producer_tasks: set[asyncio.Task[None]] = set() active_runs: dict[tuple[str | None, str], asyncio.Task[None]] = {} route_dependencies = list(dependencies or ()) single_user_scope = _single_user_scope() if single_user else None if authorization_scope_resolver is not None: async def capture_authorization_scope( request: Request, resolved_scope: Any = Depends(authorization_scope_resolver), ) -> None: if isawaitable(resolved_scope): resolved_scope = await resolved_scope if not isinstance(resolved_scope, str) or not resolved_scope: logger.error("authorization_scope_resolver must return a non-empty string.") raise HTTPException(status_code=500, detail="An internal error has occurred.") setattr(request.state, _AUTHORIZATION_SCOPE_STATE_KEY, resolved_scope) route_dependencies.append(Depends(capture_authorization_scope)) def retain_background_task(task: asyncio.Task[Any]) -> None: background_tasks.add(task) def task_done(completed: asyncio.Task[Any]) -> None: background_tasks.discard(completed) if completed.cancelled(): return exception = completed.exception() if exception is not None: logger.error( "[%s] Detached stream task failed", path, exc_info=(type(exception), exception, exception.__traceback__), ) task.add_done_callback(task_done) @app.post(path, tags=tags or ["AG-UI"], dependencies=route_dependencies, response_model=None) # type: ignore[arg-type] async def agent_endpoint(request_body: AGUIRequest, request: Request) -> Response: """Handle AG-UI agent requests. Note: Function is accessed via FastAPI's decorator registration, despite appearing unused to static analysis. """ try: input_data = request_body.model_dump(exclude_none=True) snapshot_persistence_active = _get_snapshot_store(protocol_runner) is not None snapshot_scope: str | None = None if authorization_scope_resolver is not None: snapshot_scope = cast(str, getattr(request.state, _AUTHORIZATION_SCOPE_STATE_KEY)) elif snapshot_scope_resolver is not None: resolved_scope = snapshot_scope_resolver(request_body) if isawaitable(resolved_scope): resolved_scope = await resolved_scope if not isinstance(resolved_scope, str) or not resolved_scope: raise ValueError("snapshot_scope_resolver must return a non-empty string.") snapshot_scope = resolved_scope elif single_user: snapshot_scope = single_user_scope if snapshot_scope is not None: input_data[_APPROVAL_SCOPE_INPUT_KEY] = snapshot_scope input_data[_SNAPSHOT_SCOPE_INPUT_KEY] = snapshot_scope if default_state: if snapshot_persistence_active: # Defer default application to the runner so defaults only fill keys # missing from both the stored snapshot state and the request state. input_data[_DEFAULT_STATE_INPUT_KEY] = copy.deepcopy(default_state) else: state = input_data.setdefault("state", {}) for key, value in default_state.items(): if key not in state: state[key] = copy.deepcopy(value) logger.debug( f"[{path}] Received request - Run ID: {input_data.get('run_id', 'no-run-id')}, " f"Thread ID: {input_data.get('thread_id', 'no-thread-id')}, " f"Messages: {len(input_data.get('messages', []))}" ) logger.info(f"Received request at {path}: {input_data.get('run_id', 'no-run-id')}") keepalive_enabled = keepalive_seconds is not None snapshot_hydration_request = _is_snapshot_hydration_request( input_data, snapshot_enabled=snapshot_persistence_active, supports_checkpoint_resume=isinstance(protocol_runner, AgentFrameworkWorkflow), ) active_run_key: tuple[str | None, str] | None = None if detached_runs and request_body.thread_id is not None and not snapshot_hydration_request: active_run_key = (snapshot_scope, request_body.thread_id) active_task = active_runs.get(active_run_key) if active_task is not None and not active_task.done(): return JSONResponse( status_code=409, content={"detail": "An AG-UI run is already active for this scoped thread."}, ) active_runs.pop(active_run_key, None) if detached_runs and not snapshot_hydration_request: for completed_task in tuple(producer_tasks): if completed_task.done(): producer_tasks.discard(completed_task) if len(producer_tasks) >= max_detached_runs: return JSONResponse( status_code=503, content={"detail": "AG-UI detached run capacity is exhausted."}, ) def prepare_frame(encoded: str) -> str | bytes: if keepalive_enabled: return encoded.encode("utf-8") return encoded async def event_generator() -> AsyncGenerator[str | bytes]: encoder = EventEncoder() event_count = 0 try: async for event in protocol_runner.run(input_data): event_count += 1 event_type_name = getattr(event, "type", type(event).__name__) if "TOOL_CALL" in str(event_type_name) or "RUN" in str(event_type_name): logger.info("[%s] Event %d: %s", path, event_count, event_type_name) try: encoded = encoder.encode(event) except Exception as encode_error: logger.exception("[%s] Failed to encode event %s", path, event_type_name) run_error = RunErrorEvent( message="An internal error has occurred while streaming events.", code=type(encode_error).__name__, ) try: yield prepare_frame(encoder.encode(run_error)) except Exception: logger.exception("[%s] Failed to encode RUN_ERROR event", path) return logger.debug( f"[{path}] Encoded as: {encoded[:200]}..." if len(encoded) > 200 else f"[{path}] Encoded as: {encoded}" ) yield prepare_frame(encoded) logger.info(f"[{path}] Completed streaming {event_count} events") except Exception as stream_error: logger.exception("[%s] Streaming failed", path) run_error = RunErrorEvent( message="An internal error has occurred while streaming events.", code=type(stream_error).__name__, ) try: yield prepare_frame(encoder.encode(run_error)) except Exception: logger.exception("[%s] Failed to encode RUN_ERROR event", path) async def drain_detached_stream(queue: asyncio.Queue[str | bytes | None]) -> None: while await queue.get() is not None: pass stream: AsyncGenerator[str | bytes] if detached_runs and not snapshot_hydration_request: queue: asyncio.Queue[str | bytes | None] = asyncio.Queue(maxsize=_DETACHED_STREAM_QUEUE_SIZE) reader_started = asyncio.Event() reader_abandoned = asyncio.Event() async def produce_events() -> None: try: try: await asyncio.wait_for( reader_started.wait(), timeout=_DETACHED_READER_START_TIMEOUT_SECONDS, ) except TimeoutError: reader_abandoned.set() async for frame in event_generator(): if reader_abandoned.is_set(): continue await queue.put(frame) except asyncio.CancelledError: reader_abandoned.set() raise finally: current_task = asyncio.current_task() if active_run_key is not None and active_runs.get(active_run_key) is current_task: active_runs.pop(active_run_key, None) if reader_started.is_set(): await queue.put(None) producer_task = asyncio.create_task( produce_events(), name=f"ag-ui-run-{input_data.get('run_id', 'generated')}", ) if active_run_key is not None: active_runs[active_run_key] = producer_task producer_tasks.add(producer_task) producer_task.add_done_callback(producer_tasks.discard) retain_background_task(producer_task) async def expire_abandoned_run() -> None: if producer_task.done(): return if not reader_abandoned.is_set(): abandoned_wait = asyncio.create_task(reader_abandoned.wait()) done, _ = await asyncio.wait( {producer_task, abandoned_wait}, return_when=asyncio.FIRST_COMPLETED, ) if producer_task in done: abandoned_wait.cancel() with suppress(asyncio.CancelledError): await abandoned_wait return if producer_task.done(): return try: await asyncio.wait_for( asyncio.shield(producer_task), timeout=detached_run_timeout_seconds, ) except TimeoutError: logger.error( "[%s] Detached run exceeded %.1f seconds after reader disconnect; cancelling", path, detached_run_timeout_seconds, ) reader_abandoned.set() producer_task.cancel() with suppress(asyncio.CancelledError): await producer_task expiry_task = asyncio.create_task( expire_abandoned_run(), name=f"ag-ui-expiry-{input_data.get('run_id', 'generated')}", ) retain_background_task(expiry_task) async def detached_event_generator() -> AsyncGenerator[str | bytes]: completed = False try: reader_started.set() if reader_abandoned.is_set(): return while True: item = await queue.get() if item is None: completed = True return yield item finally: if not completed and not producer_task.done(): reader_abandoned.set() drain_task = asyncio.create_task( drain_detached_stream(queue), name=f"ag-ui-drain-{input_data.get('run_id', 'generated')}", ) retain_background_task(drain_task) stream = detached_event_generator() else: stream = event_generator() headers = { "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", } if keepalive_seconds is not None: from sse_starlette.event import ServerSentEvent from sse_starlette.sse import EventSourceResponse return EventSourceResponse( stream, ping=cast(int, keepalive_seconds), ping_message_factory=lambda: ServerSentEvent(comment=_KEEPALIVE_COMMENT), headers=headers, media_type="text/event-stream", ) return StreamingResponse( stream, media_type="text/event-stream", headers=headers, ) except Exception as e: logger.error(f"Error in agent endpoint: {e}", exc_info=True) raise HTTPException(status_code=500, detail="An internal error has occurred.") from e