rest_adapter.py

import functools
import json
import logging

from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
from typing import TYPE_CHECKING, Any


if TYPE_CHECKING:
    from sse_starlette.sse import EventSourceResponse
    from starlette.requests import Request
    from starlette.responses import JSONResponse, Response

    from a2a.server.request_handlers.request_handler import RequestHandler
    from a2a.types.a2a_pb2 import AgentCard

    _package_starlette_installed = True
else:
    try:
        from sse_starlette.sse import EventSourceResponse
        from starlette.requests import Request
        from starlette.responses import JSONResponse, Response

        _package_starlette_installed = True
    except ImportError:
        EventSourceResponse = Any
        Request = Any
        JSONResponse = Any
        Response = Any

        _package_starlette_installed = False


from a2a.compat.v0_3 import conversions
from a2a.compat.v0_3.rest_handler import REST03Handler
from a2a.server.apps.rest.rest_adapter import RESTAdapterInterface
from a2a.server.context import ServerCallContext
from a2a.server.routes import CallContextBuilder, DefaultCallContextBuilder
from a2a.utils.error_handlers import (
    rest_error_handler,
    rest_stream_error_handler,
)
from a2a.utils.errors import (
    ExtendedAgentCardNotConfiguredError,
    InvalidRequestError,
)
from a2a.utils.helpers import maybe_await


logger = logging.getLogger(__name__)


class REST03Adapter(RESTAdapterInterface):
    """Adapter to make RequestHandler work with v0.3 RESTful API.

    Defines v0.3 REST request processors and their routes, as well as managing response generation including Server-Sent Events (SSE).
    """

    def __init__(  # noqa: PLR0913
        self,
        agent_card: 'AgentCard',
        http_handler: 'RequestHandler',
        extended_agent_card: 'AgentCard | None' = None,
        context_builder: 'CallContextBuilder | None' = None,
        card_modifier: 'Callable[[AgentCard], Awaitable[AgentCard] | AgentCard] | None' = None,
        extended_card_modifier: 'Callable[[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard] | None' = None,
    ):
        self.agent_card = agent_card
        self.extended_agent_card = extended_agent_card
        self.card_modifier = card_modifier
        self.extended_card_modifier = extended_card_modifier
        self.handler = REST03Handler(
            agent_card=agent_card, request_handler=http_handler
        )
        self._context_builder = context_builder or DefaultCallContextBuilder()

    @rest_error_handler
    async def _handle_request(
        self,
        method: 'Callable[[Request, ServerCallContext], Awaitable[Any]]',
        request: Request,
    ) -> Response:
        call_context = self._context_builder.build(request)
        response = await method(request, call_context)
        return JSONResponse(content=response)

    @rest_stream_error_handler
    async def _handle_streaming_request(
        self,
        method: 'Callable[[Request, ServerCallContext], AsyncIterable[Any]]',
        request: Request,
    ) -> EventSourceResponse:
        try:
            await request.body()
        except (ValueError, RuntimeError, OSError) as e:
            raise InvalidRequestError(
                message=f'Failed to pre-consume request body: {e}'
            ) from e

        call_context = self._context_builder.build(request)

        async def event_generator(
            stream: AsyncIterable[Any],
        ) -> AsyncIterator[str]:
            async for item in stream:
                yield json.dumps(item)

        return EventSourceResponse(
            event_generator(method(request, call_context))
        )

    async def handle_get_agent_card(
        self, request: Request, call_context: ServerCallContext | None = None
    ) -> dict[str, Any]:
        """Handles GET requests for the agent card endpoint."""
        card_to_serve = self.agent_card
        if self.card_modifier:
            card_to_serve = await maybe_await(self.card_modifier(card_to_serve))
        v03_card = conversions.to_compat_agent_card(card_to_serve)
        return v03_card.model_dump(mode='json', exclude_none=True)

    async def handle_authenticated_agent_card(
        self, request: Request, call_context: ServerCallContext | None = None
    ) -> dict[str, Any]:
        """Hook for per credential agent card response."""
        if not self.agent_card.capabilities.extended_agent_card:
            raise ExtendedAgentCardNotConfiguredError(
                message='Authenticated card not supported'
            )
        card_to_serve = self.extended_agent_card

        if not card_to_serve:
            card_to_serve = self.agent_card

        if self.extended_card_modifier:
            context = self._context_builder.build(request)
            card_to_serve = await maybe_await(
                self.extended_card_modifier(card_to_serve, context)
            )
        elif self.card_modifier:
            card_to_serve = await maybe_await(self.card_modifier(card_to_serve))

        v03_card = conversions.to_compat_agent_card(card_to_serve)
        return v03_card.model_dump(mode='json', exclude_none=True)

    def routes(self) -> dict[tuple[str, str], Callable[[Request], Any]]:
        """Constructs a dictionary of API routes and their corresponding handlers."""
        routes: dict[tuple[str, str], Callable[[Request], Any]] = {
            ('/v1/message:send', 'POST'): functools.partial(
                self._handle_request, self.handler.on_message_send
            ),
            ('/v1/message:stream', 'POST'): functools.partial(
                self._handle_streaming_request,
                self.handler.on_message_send_stream,
            ),
            ('/v1/tasks/{id}:cancel', 'POST'): functools.partial(
                self._handle_request, self.handler.on_cancel_task
            ),
            ('/v1/tasks/{id}:subscribe', 'GET'): functools.partial(
                self._handle_streaming_request,
                self.handler.on_subscribe_to_task,
            ),
            ('/v1/tasks/{id}:subscribe', 'POST'): functools.partial(
                self._handle_streaming_request,
                self.handler.on_subscribe_to_task,
            ),
            ('/v1/tasks/{id}', 'GET'): functools.partial(
                self._handle_request, self.handler.on_get_task
            ),
            (
                '/v1/tasks/{id}/pushNotificationConfigs/{push_id}',
                'GET',
            ): functools.partial(
                self._handle_request, self.handler.get_push_notification
            ),
            (
                '/v1/tasks/{id}/pushNotificationConfigs',
                'POST',
            ): functools.partial(
                self._handle_request, self.handler.set_push_notification
            ),
            (
                '/v1/tasks/{id}/pushNotificationConfigs',
                'GET',
            ): functools.partial(
                self._handle_request, self.handler.list_push_notifications
            ),
            ('/v1/tasks', 'GET'): functools.partial(
                self._handle_request, self.handler.list_tasks
            ),
        }
        if self.agent_card.capabilities.extended_agent_card:
            routes[('/v1/card', 'GET')] = functools.partial(
                self._handle_request, self.handle_authenticated_agent_card
            )

        return routes