core/a2a_layer.py

"""A2A protocol layer (plan §7.2), built on the official a2a-sdk (0.3.x, pydantic).

- Each agent publishes an AgentCard at /.well-known/agent-card.json.
- Messaging is JSON-RPC 2.0 `message/send`; payloads travel as DataPart.
- `serve_agent` runs an agent as a Starlette app on its own port (background thread).
- `send_payload` is the client: resolves the peer's card, sends, returns the reply dict.
"""
from __future__ import annotations

import asyncio
import threading
import uuid
from collections.abc import Callable

import httpx
import uvicorn
from a2a.client import A2ACardResolver, ClientConfig, ClientFactory
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.apps import A2AStarletteApplication
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    AgentSkill,
    DataPart,
    Message,
    Part,
    Role,
)


def make_card(name: str, description: str, port: int, skills: list[dict],
              streaming: bool = False) -> AgentCard:
    return AgentCard(
        protocol_version="0.3.0",
        name=name,
        description=description,
        version="1.0.0",
        url=f"http://127.0.0.1:{port}",
        capabilities=AgentCapabilities(streaming=streaming),
        default_input_modes=["application/json"],
        default_output_modes=["application/json"],
        skills=[AgentSkill(id=s["id"], name=s.get("name", s["id"]),
                           description=s["description"], tags=s.get("tags", []))
                for s in skills],
    )


def extract_payload(message: Message) -> dict | None:
    """First DataPart payload in a message, else None."""
    for part in message.parts or []:
        inner = part.root if isinstance(part, Part) else part
        if isinstance(inner, DataPart):
            return inner.data
    return None


def payload_message(payload: dict, role: Role = Role.agent) -> Message:
    return Message(
        message_id=uuid.uuid4().hex,
        role=role,
        parts=[Part(root=DataPart(data=payload))],
    )


class PayloadExecutor(AgentExecutor):
    """Adapts a plain `handler(payload: dict) -> dict` into an A2A AgentExecutor."""

    def __init__(self, handler: Callable[[dict], dict]):
        self.handler = handler

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        payload = extract_payload(context.message) if context.message else None
        try:
            if asyncio.iscoroutinefunction(self.handler):
                reply = await self.handler(payload or {})
            else:
                reply = await asyncio.to_thread(self.handler, payload or {})
        except Exception as e:  # noqa: BLE001 — surface agent errors as structured replies
            reply = {"status": "error", "detail": str(e)}
        await event_queue.enqueue_event(payload_message(reply))

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        await event_queue.enqueue_event(payload_message({"status": "cancelled"}))


class AgentServer:
    """Runs one agent as an A2A Starlette app on a background uvicorn thread."""

    def __init__(self, card: AgentCard, handler: Callable[[dict], dict], port: int):
        self.card = card
        self.port = port
        handler_ = DefaultRequestHandler(
            agent_executor=PayloadExecutor(handler),
            task_store=InMemoryTaskStore(),
        )
        self.app = A2AStarletteApplication(agent_card=card, http_handler=handler_).build()
        self._server: uvicorn.Server | None = None
        self._thread: threading.Thread | None = None

    def start(self, wait_ready_s: float = 10.0) -> None:
        config = uvicorn.Config(self.app, host="127.0.0.1", port=self.port,
                                log_level="warning")
        self._server = uvicorn.Server(config)
        self._thread = threading.Thread(target=self._server.run, daemon=True)
        self._thread.start()
        # block until the card endpoint answers
        deadline = wait_ready_s
        import time
        while deadline > 0:
            try:
                r = httpx.get(f"http://127.0.0.1:{self.port}/.well-known/agent-card.json",
                              timeout=1.0)
                if r.status_code == 200:
                    return
            except httpx.HTTPError:
                pass
            time.sleep(0.2)
            deadline -= 0.2
        raise RuntimeError(f"A2A agent on port {self.port} did not become ready")

    def stop(self) -> None:
        if self._server:
            self._server.should_exit = True
        if self._thread:
            self._thread.join(timeout=5)


async def send_payload(base_url: str, payload: dict, timeout_s: float = 120.0) -> dict:
    """A2A client: resolve the peer's AgentCard, send one message/send, return reply."""
    async with httpx.AsyncClient(timeout=timeout_s) as http:
        card = await A2ACardResolver(http, base_url).get_agent_card()
        client = ClientFactory(ClientConfig(httpx_client=http, streaming=False)).create(card)
        message = payload_message(payload, role=Role.user)
        reply_payload: dict | None = None
        async for event in client.send_message(message):
            if isinstance(event, Message):
                reply_payload = extract_payload(event)
            elif isinstance(event, tuple):  # (Task, UpdateEvent)
                task = event[0]
                if task and task.artifacts:
                    for artifact in task.artifacts:
                        for part in artifact.parts:
                            inner = part.root if isinstance(part, Part) else part
                            if isinstance(inner, DataPart):
                                reply_payload = inner.data
                if task and task.status and task.status.message:
                    got = extract_payload(task.status.message)
                    if got is not None:
                        reply_payload = got
        if reply_payload is None:
            raise RuntimeError(f"no DataPart reply from {base_url}")
        return reply_payload


def export_card(card: AgentCard, path) -> None:
    """Write the AgentCard JSON (spec camelCase) to protocols/a2a_cards/."""
    import json
    from pathlib import Path
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(
        json.dumps(card.model_dump(by_alias=True, exclude_none=True), indent=2),
        encoding="utf-8",
    )