classic-models-agent-assistant/src/assistant_agent/a2a_server.py

"""The A2A facade: how other agents reach this one.

Notes on a2a-sdk 1.1, none of which are guessable and all of which were
established against the running library:

- `a2a.types` are **protobuf** messages, not pydantic models.
- Use `DefaultRequestHandlerV2`. `DefaultRequestHandler` speaks A2A 0.3 and
  rejects 1.0 traffic.
- The JSON-RPC methods are `SendMessage`, `GetTask`, … The 0.3 spelling
  (`message/send`) is available through `enable_v0_3_compat=True`, kept on here
  so a client written against the older spelling still works.
- **A client must send `A2A-Version: 1.0`.** With the header absent the server
  assumes 0.3 and answers `VERSION_NOT_SUPPORTED` — a confusing failure, since
  nothing in the request is actually wrong.
- An executor must enqueue a `Task` before any status update when
  `context.current_task` is None, or the SDK answers
  "Agent should enqueue Task before TaskStatusUpdateEvent".
"""

from __future__ import annotations

import contextlib
import logging
from collections.abc import AsyncIterator

import httpx
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandlerV2
from a2a.server.routes.agent_card_routes import create_agent_card_routes
from a2a.server.routes.jsonrpc_routes import create_jsonrpc_routes
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    AgentInterface,
    AgentSkill,
    Part,
    Task,
    TaskState,
    TaskStatus,
)
from starlette.applications import Starlette

from .config import Settings
from .orchestrator import Assistant

logger = logging.getLogger(__name__)

# The bearer the caller presents, which this agent exchanges for MCP tokens.
# It is the END USER's token, not the calling agent's: the whole scoping model
# rests on `sub` staying the human throughout the chain while `azp` names the
# agent doing the asking.
SUBJECT_TOKEN_HEADER = "authorization"


def build_agent_card(settings: Settings) -> AgentCard:
    """One skill, and it is deliberately about routing rather than data.

    This agent holds nothing. What it offers is that a caller can ask one
    question and have the right specialists consulted — which is a capability,
    not a data surface. Naming it `answer_about_the_business` rather than
    `query` keeps the distinction visible.
    """
    return AgentCard(
        name="assistant",
        description=(
            "The entry point for the Classic Models estate. Understands a request, "
            "consults the specialists that can settle it, and composes one answer. "
            "Holds no data and no domain scope of its own."
        ),
        version="0.1.0",
        default_input_modes=["text/plain"],
        default_output_modes=["text/plain"],
        capabilities=AgentCapabilities(streaming=False),
        supported_interfaces=[
            AgentInterface(
                url=settings.public_url,
                protocol_binding="JSONRPC",
                protocol_version="1.0",
            )
        ],
        skills=[
            AgentSkill(
                id="answer_about_the_business",
                name="Answer a question about the business",
                description=(
                    "Route a question to the specialists that can settle it and compose "
                    "their answers. Every figure comes from a specialist, within what the "
                    "asking user is allowed to see."
                ),
                tags=["orchestration", "routing"],
                examples=[
                    "How many products are in the Classic Cars line, and which is dearest?",
                    "Can we ship order 10248?",
                ],
            )
        ],
    )


class AssistantExecutor(AgentExecutor):
    def __init__(self, agent: Assistant) -> None:
        self._agent = agent

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        updater = TaskUpdater(event_queue, context.task_id, context.context_id)

        # The SDK requires the Task to exist before any status update.
        if context.current_task is None:
            await event_queue.enqueue_event(
                Task(
                    id=context.task_id,
                    context_id=context.context_id,
                    status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
                )
            )
        await updater.start_work()

        subject_token = _subject_token(context)
        if subject_token is None:
            # Refused rather than answered from the model's own knowledge. This
            # agent holds no data at all, so an answer produced without a
            # specialist is necessarily invented — and would be indistinguishable
            # from one that came from the estate.
            await updater.failed(
                message=updater.new_agent_message(
                    parts=[Part(text="No bearer token on the request; cannot consult any specialist.")]
                )
            )
            return

        try:
            answer = await self._agent.run(
                context.get_user_input(), subject_token, context.context_id
            )
        except Exception as exc:  # noqa: BLE001 — surfaced to the caller, logged here
            logger.exception("turn failed")
            await updater.failed(
                message=updater.new_agent_message(parts=[Part(text=f"Turn failed: {exc}")])
            )
            return

        # Always `completed`, never `input-required`. A specialist holding a
        # write parks its own task and its message is relayed here verbatim —
        # the user answers that specialist through this agent's next turn, and
        # the orchestrator never owns the decision.
        await updater.complete(message=updater.new_agent_message(parts=[Part(text=answer)]))

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        await TaskUpdater(event_queue, context.task_id, context.context_id).cancel()


def _subject_token(context: RequestContext) -> str | None:
    headers = {}
    call = context.call_context
    if call is not None:
        headers = {k.lower(): v for k, v in (call.state.get("headers") or {}).items()}
    raw = headers.get(SUBJECT_TOKEN_HEADER)
    if not raw:
        return None
    return raw[7:] if raw.lower().startswith("bearer ") else raw


def build_app(settings: Settings) -> Starlette:
    http_client = httpx.AsyncClient(verify=settings.verify, timeout=settings.http_timeout)
    card = build_agent_card(settings)
    handler = DefaultRequestHandlerV2(
        agent_executor=AssistantExecutor(Assistant(settings, http_client)),
        task_store=InMemoryTaskStore(),
        agent_card=card,
    )

    # lifespan rather than on_shutdown: Starlette 1.x removed the latter.
    @contextlib.asynccontextmanager
    async def lifespan(_app: Starlette) -> AsyncIterator[None]:
        yield
        await http_client.aclose()

    return Starlette(
        routes=create_agent_card_routes(card)
        + create_jsonrpc_routes(handler, "/", enable_v0_3_compat=True),
        lifespan=lifespan,
    )