a2a_server.py

"""A2A server support for ADCP handlers.

Bridges ADCPHandler to the a2a-sdk server framework so the same handler
can be served over both MCP and A2A transports.

    from adcp.server import ADCPHandler, serve
    serve(MyHandler(), name="my-agent", transport="a2a")

.. note::
    Function signatures here use ``ADCPHandler[Any]`` rather than a
    propagated ``TContext`` TypeVar. This module dispatches by tool
    name and never reads typed fields off the context, so ``Any`` is
    both correct and keeps the call sites tidy — downstream code that
    needs typed context (their own handler subclass) keeps the TypeVar
    all the way to dispatch via :class:`ADCPHandler`. See the matching
    note in :mod:`adcp.server.mcp_tools`.
"""

from __future__ import annotations

import json
import logging
import os
from typing import TYPE_CHECKING, Any
from uuid import uuid4

from a2a.server.agent_execution.agent_executor import AgentExecutor
from a2a.server.agent_execution.context import RequestContext
from a2a.server.events.event_queue import EventQueue
from a2a.server.request_handlers.default_request_handler import (
    DefaultRequestHandler,
)
from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    AgentSkill,
    Artifact,
    DataPart,
    Part,
    Task,
    TaskState,
    TaskStatus,
    TextPart,
)

from adcp.exceptions import ADCPError, ADCPTaskError
from adcp.server.base import ADCPHandler, ToolContext

if TYPE_CHECKING:
    from collections.abc import Sequence

    from a2a.server.tasks.push_notification_config_store import (
        PushNotificationConfigStore,
    )
    from a2a.server.tasks.task_store import TaskStore

    from adcp.server.serve import ContextFactory, SkillMiddleware
from adcp.server.helpers import STANDARD_ERROR_CODES
from adcp.server.mcp_tools import create_tool_caller, get_tools_for_handler
from adcp.server.test_controller import TestControllerStore, _handle_test_controller

logger = logging.getLogger(__name__)


class ADCPAgentExecutor(AgentExecutor):
    """Bridges ADCPHandler methods to the a2a-sdk AgentExecutor interface.

    Incoming A2A messages are parsed to extract the ADCP skill name and
    parameters, dispatched to the matching handler method, and the result
    is published back as A2A Task events.

    Expects the explicit skill invocation format used by A2AAdapter:
        DataPart(data={"skill": "get_products", "parameters": {...}})
    """

    def __init__(
        self,
        handler: ADCPHandler[Any],
        test_controller: TestControllerStore | None = None,
        *,
        context_factory: ContextFactory | None = None,
        middleware: Sequence[SkillMiddleware] | None = None,
    ) -> None:
        self._handler = handler
        self._context_factory = context_factory
        # Store as a tuple so the executor can't be mutated from underneath
        # at runtime (a flaky test or a handler reaching self._middleware
        # can't corrupt the dispatch chain). Tuple ordering = runtime
        # ordering; first entry wraps outermost (see ``SkillMiddleware``
        # docstring for the composition semantics).
        self._middleware: tuple[SkillMiddleware, ...] = tuple(middleware or ())
        self._tool_callers: dict[str, Any] = {}

        # Build tool callers for all tools this handler supports.
        # Skip comply_test_controller unless the seller passed a
        # TestControllerStore; otherwise we would advertise a skill
        # backed only by the handler's not-supported stub.
        tool_defs = get_tools_for_handler(handler)
        for tool_def in tool_defs:
            name = tool_def["name"]
            if name == "comply_test_controller" and test_controller is None:
                continue
            self._tool_callers[name] = create_tool_caller(handler, name)

        if test_controller is not None:
            self._register_test_controller(test_controller)

    @property
    def supported_skills(self) -> list[str]:
        """List of skill names this executor can handle."""
        return list(self._tool_callers.keys())

    def _register_test_controller(self, store: TestControllerStore) -> None:
        """Register comply_test_controller as a callable skill."""

        async def _call_test_controller(
            params: dict[str, Any], context: ToolContext | None = None
        ) -> Any:
            return await _handle_test_controller(store, params)

        self._tool_callers["comply_test_controller"] = _call_test_controller

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Execute an ADCP skill from an incoming A2A message."""
        skill_name, params = self._parse_request(context)

        if skill_name is None:
            await self._send_error(event_queue, context, "No skill specified in message")
            return

        if skill_name not in self._tool_callers:
            await self._send_error(event_queue, context, f"Unknown skill: {skill_name}")
            return

        tool_context = self._build_tool_context(skill_name, context)
        try:
            result = await self._dispatch_with_middleware(skill_name, params, tool_context)
            await self._send_result(event_queue, context, skill_name, result)
        except ADCPError as exc:
            # Application-layer AdCP error (IdempotencyConflictError etc.).
            # Emit a failed task with the adcp_error in a DataPart per
            # transport-errors.mdx §A2A Binding, plus a human-readable text
            # part. The JSON-RPC channel is reserved for transport-level
            # errors (auth rejected, rate-limited pre-dispatch).
            logger.info("AdCP application error for skill %s: %s", skill_name, exc)
            await self._send_adcp_error(event_queue, context, exc)
        except Exception:
            logger.exception("Error executing skill %s", skill_name)
            await self._send_error(event_queue, context, f"Skill execution failed: {skill_name}")

    async def _dispatch_with_middleware(
        self,
        skill_name: str,
        params: dict[str, Any],
        tool_context: ToolContext,
    ) -> Any:
        """Run the handler wrapped in the configured middleware chain.

        Middleware composes outermost-first: the first entry in
        ``self._middleware`` sees every call *before* the later entries
        and *before* the handler. This matches Starlette / ASGI
        conventions so sellers porting from those stacks aren't
        surprised. Composition is done via a small recursive dispatcher
        (no mutable indices, no lambdas closing over loop variables) —
        the chain reads the same whether you have zero or ten
        middlewares.

        Middleware exceptions propagate to the executor's normal error
        handling path in ``execute()``; this method does no try/except
        so short-circuiting, transform, and exception-observation all
        work the same way they do for the underlying handler.
        """
        if not self._middleware:
            return await self._tool_callers[skill_name](params, tool_context)

        async def _step(index: int) -> Any:
            if index >= len(self._middleware):
                return await self._tool_callers[skill_name](params, tool_context)
            middleware = self._middleware[index]

            async def call_next() -> Any:
                return await _step(index + 1)

            return await middleware(skill_name, params, tool_context, call_next)

        return await _step(0)

    def _build_tool_context(self, skill_name: str, request: RequestContext) -> ToolContext:
        """Build the :class:`ToolContext` handed to the skill dispatcher.

        When ``context_factory`` is configured, call it with a
        :class:`RequestMetadata` describing this A2A invocation; overlay the
        transport-derived ``caller_identity`` / ``request_id`` afterwards
        **only when the factory left them unset**, so factories that already
        know the principal (e.g. from a ContextVar the seller's auth layer
        populated) aren't clobbered.

        When no factory is configured, fall back to the A2A-only path that
        derives ``caller_identity`` from ``ServerCallContext.user`` —
        preserving behavior for sellers who haven't adopted
        ``context_factory=`` yet.
        """
        if self._context_factory is None:
            return _tool_context_from_request(request)

        from adcp.server.serve import RequestMetadata

        meta = RequestMetadata(
            tool_name=skill_name,
            transport="a2a",
            request_id=request.task_id,
        )
        ctx = self._context_factory(meta)
        if not isinstance(ctx, ToolContext):
            raise TypeError(
                f"context_factory for skill {skill_name!r} returned "
                f"{type(ctx).__name__}, not a ToolContext instance"
            )
        # Fill in transport-derived fields the factory didn't set. This
        # preserves the pre-factory A2A security invariant: if the seller
        # didn't explicitly populate caller_identity in their factory,
        # fall through to ServerCallContext.user (verified by the a2a-sdk
        # auth middleware) rather than silently sending None.
        if ctx.caller_identity is None:
            fallback = _tool_context_from_request(request)
            ctx.caller_identity = fallback.caller_identity
        if ctx.request_id is None:
            ctx.request_id = request.task_id
        return ctx

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        """ADCP operations are synchronous; cancellation sets state to canceled."""
        event = _make_task(
            context,
            state=TaskState.canceled,
            message="Task canceled",
        )
        await event_queue.enqueue_event(event)

    # ------------------------------------------------------------------
    # Message parsing
    # ------------------------------------------------------------------

    def _parse_request(self, context: RequestContext) -> tuple[str | None, dict[str, Any]]:
        """Extract skill name and parameters from the A2A message.

        Supports two formats:
        1. Explicit skill invocation via DataPart:
           DataPart(data={"skill": "get_products", "parameters": {...}})
        2. Natural language fallback via TextPart (best-effort parse)
        """
        msg = context.message
        if msg is None or not msg.parts:
            return None, {}

        # Try DataPart first (explicit skill invocation)
        for part in msg.parts:
            inner = part.root if hasattr(part, "root") else part
            if isinstance(inner, DataPart) and isinstance(inner.data, dict):
                skill = inner.data.get("skill")
                params = inner.data.get("parameters", {})
                if skill:
                    return str(skill), params if isinstance(params, dict) else {}

        # Fallback: try to parse TextPart as JSON
        for part in msg.parts:
            inner = part.root if hasattr(part, "root") else part
            if isinstance(inner, TextPart):
                parsed = self._parse_text_request(inner.text)
                if parsed[0] is not None:
                    return parsed

        return None, {}

    def _parse_text_request(self, text: str) -> tuple[str | None, dict[str, Any]]:
        """Best-effort parse of a text request for skill + params."""
        try:
            data = json.loads(text)
            if isinstance(data, dict) and "skill" in data:
                return str(data["skill"]), data.get("parameters", {})
        except (json.JSONDecodeError, TypeError):
            pass
        return None, {}

    # ------------------------------------------------------------------
    # Response helpers
    # ------------------------------------------------------------------

    async def _send_result(
        self,
        event_queue: EventQueue,
        context: RequestContext,
        skill_name: str,
        result: Any,
    ) -> None:
        """Publish a completed task with the skill result."""
        # Normalize result to a JSON-safe dict
        if hasattr(result, "model_dump"):
            data = result.model_dump(mode="json", exclude_none=True)
        elif not isinstance(result, dict):
            data = {"result": result}
        else:
            data = result

        task = _make_task(
            context,
            state=TaskState.completed,
            data=data,
            message=f"Completed {skill_name}",
        )
        await event_queue.enqueue_event(task)

    async def _send_error(
        self,
        event_queue: EventQueue,
        context: RequestContext,
        error_msg: str,
    ) -> None:
        """Publish a failed task."""
        task = _make_task(
            context,
            state=TaskState.failed,
            message=error_msg,
        )
        await event_queue.enqueue_event(task)

    async def _send_adcp_error(
        self,
        event_queue: EventQueue,
        context: RequestContext,
        exc: ADCPError,
    ) -> None:
        """Publish a failed task carrying an AdCP ``adcp_error`` payload.

        Follows transport-errors.mdx §A2A Binding: failed task with artifact
        containing a ``DataPart`` keyed under ``adcp_error`` plus a terse
        ``TextPart`` for human/LLM consumption.
        """
        # Derive the spec error code. ADCPTaskError carries a list of codes
        # (e.g. IdempotencyConflictError → IDEMPOTENCY_CONFLICT); fall back
        # to a generic INTERNAL_ERROR when the exception doesn't supply one.
        code = "INTERNAL_ERROR"
        if isinstance(exc, ADCPTaskError) and exc.error_codes:
            code = str(exc.error_codes[0])

        adcp_error: dict[str, Any] = {
            "code": code,
            "message": exc.message,
        }
        recovery = STANDARD_ERROR_CODES.get(code, {}).get("recovery")
        if recovery:
            adcp_error["recovery"] = recovery
        suggestion = getattr(exc, "suggestion", None)
        if suggestion:
            adcp_error["suggestion"] = suggestion

        task = _make_task(
            context,
            state=TaskState.failed,
            data={"adcp_error": adcp_error},
            message=exc.message,
        )
        await event_queue.enqueue_event(task)


# ------------------------------------------------------------------
# Request context helpers
# ------------------------------------------------------------------


def _tool_context_from_request(request: RequestContext) -> ToolContext:
    """Derive a :class:`ToolContext` from an A2A :class:`RequestContext`.

    Extracts the authenticated principal from ``request.call_context.user``
    when present. Unauthenticated / anonymous requests get a bare
    ``ToolContext`` — server middleware that requires a principal (e.g. the
    idempotency store's per-principal scoping) falls through to its
    no-principal default rather than collapsing everyone into a shared
    namespace.

    Security invariant: ``ServerCallContext`` is populated by the seller's
    server-side auth middleware from verified transport material (bearer
    token, mTLS cert, OAuth identity). A malicious client cannot flip
    ``is_authenticated`` or set ``user_name`` from the message payload.
    The ``is_authenticated and user_name`` gate below relies on this
    invariant — do not relax it.

    PII note: the ``user_name`` string becomes ``caller_identity``, which
    the idempotency middleware logs prefix-truncated at DEBUG. If your auth
    layer sets ``user_name`` to an email address, treat idempotency debug
    logs as containing PII. Prefer opaque principal IDs.
    """
    ctx = ToolContext(request_id=request.task_id)
    call_context = getattr(request, "call_context", None)
    user = getattr(call_context, "user", None)
    if user is not None:
        is_auth = getattr(user, "is_authenticated", False)
        user_name = getattr(user, "user_name", "") or ""
        if is_auth and user_name:
            ctx.caller_identity = user_name
    return ctx


# ------------------------------------------------------------------
# Task factory
# ------------------------------------------------------------------


def _make_task(
    context: RequestContext,
    *,
    state: TaskState,
    data: dict[str, Any] | None = None,
    message: str | None = None,
) -> Task:
    """Build an a2a Task event from context and result data."""
    parts: list[Part] = []
    if data is not None:
        parts.append(Part(root=DataPart(data=data)))
    if message:
        parts.append(Part(root=TextPart(text=message)))

    artifacts = []
    if parts:
        artifacts.append(
            Artifact(
                artifact_id=str(uuid4()),
                parts=parts,
            )
        )

    return Task(
        id=context.task_id or str(uuid4()),
        context_id=context.context_id or str(uuid4()),
        status=TaskStatus(state=state),
        artifacts=artifacts if artifacts else None,
    )


# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------


def _build_agent_card(
    handler: ADCPHandler[Any],
    *,
    name: str,
    port: int,
    description: str | None = None,
    version: str = "1.0.0",
    extra_skills: list[AgentSkill] | None = None,
) -> AgentCard:
    """Build an A2A AgentCard from an ADCPHandler's tool definitions.

    ``comply_test_controller`` is excluded from the card skills list unless
    the caller supplied it via ``extra_skills`` (which is how
    :func:`create_a2a_server` opts in when a ``TestControllerStore`` is
    wired). Extra skills are deduped by id so advertising the test
    controller never produces two entries.
    """
    tool_defs = get_tools_for_handler(handler)
    extra_ids = {s.id for s in extra_skills} if extra_skills else set()

    skills = [
        AgentSkill(
            id=td["name"],
            name=td["name"],
            description=td.get("description", td["name"]),
            tags=["adcp"],
        )
        for td in tool_defs
        if td["name"] != "comply_test_controller" and td["name"] not in extra_ids
    ]

    if extra_skills:
        skills.extend(extra_skills)

    return AgentCard(
        name=name,
        description=description or f"ADCP agent: {name}",
        url=f"http://localhost:{port}/",
        version=version,
        skills=skills,
        capabilities=AgentCapabilities(streaming=False),
        default_input_modes=["application/json"],
        default_output_modes=["application/json"],
    )


def create_a2a_server(
    handler: ADCPHandler[Any],
    *,
    name: str = "adcp-agent",
    port: int | None = None,
    description: str | None = None,
    version: str = "1.0.0",
    test_controller: TestControllerStore | None = None,
    context_factory: ContextFactory | None = None,
    task_store: TaskStore | None = None,
    push_config_store: PushNotificationConfigStore | None = None,
    middleware: Sequence[SkillMiddleware] | None = None,
) -> Any:
    """Create an A2A Starlette application from an ADCP handler.

    Args:
        handler: An ADCPHandler subclass instance.
        name: Agent name shown in the A2A agent card.
        port: Port number (used in the agent card URL).
        description: Agent description for the agent card.
        version: Agent version string.
        test_controller: Optional TestControllerStore for storyboard testing.
        context_factory: Optional callable invoked per skill call to build
            a :class:`ToolContext` from :class:`RequestMetadata`. Mirrors
            the MCP-side ``context_factory=`` on
            :func:`~adcp.server.create_mcp_server` so a single factory
            populates tenant/adapter fields on both transports. When
            unset, the executor falls back to deriving ``caller_identity``
            from ``ServerCallContext.user`` — preserving pre-factory
            behavior. See :data:`~adcp.server.ContextFactory` for the
            recommended contextvars pattern.
        task_store: Optional a2a-sdk :class:`~a2a.server.tasks.task_store.TaskStore`
            instance for persisting A2A task state. Defaults to
            :class:`~a2a.server.tasks.inmemory_task_store.InMemoryTaskStore`,
            which is single-process and non-durable — fine for demos and
            local development, but tasks vanish on restart and don't share
            across workers. Production agents pass a durable subclass
            (Postgres, Redis, etc.). See ``examples/a2a_db_tasks.py`` for
            a reference SQLite-backed implementation and
            ``docs/handler-authoring.md`` for the persistence caveats on
            the default store.
        push_config_store: Optional a2a-sdk
            :class:`~a2a.server.tasks.push_notification_config_store.PushNotificationConfigStore`
            instance for persisting push-notification configs that clients
            register via ``tasks/pushNotificationConfig/set``. **When
            unset, a2a-sdk surfaces push-notif endpoints as
            ``UnsupportedOperationError``** — clients cannot register
            subscriptions at all. Set this only when your agent is ready
            to accept push-notif subscriptions. See
            ``examples/a2a_db_tasks.py`` for a reference SQLite-backed
            implementation that pairs with the ``SqliteTaskStore`` there.

            Security note: unlike ``TaskStore``, a2a-sdk's
            ``PushNotificationConfigStore`` ABC does not pass a
            ``ServerCallContext`` to ``set_info`` / ``get_info`` /
            ``delete_info``. Scoping by principal has to happen out-of-band
            (via a ``ContextVar`` your auth middleware populates) or by
            composition with a tenant-scoped ``TaskStore`` — the reference
            impl shows the ContextVar pattern.
        middleware: Optional sequence of :data:`~adcp.server.SkillMiddleware`
            callables wrapping every A2A skill dispatch. Composes
            outermost-first (first entry sees the call before later
            entries and before the handler). Use for audit logging,
            activity-feed hooks, rate limiting, per-skill tracing. See
            :data:`~adcp.server.SkillMiddleware` for the signature,
            composition semantics, and the exception-capture pattern
            audit hooks need.

    Returns:
        A Starlette app ready to be run with uvicorn.
    """
    from a2a.server.apps.jsonrpc.starlette_app import A2AStarletteApplication

    resolved_port = port or int(os.environ.get("PORT", "3001"))

    executor = ADCPAgentExecutor(
        handler,
        test_controller=test_controller,
        context_factory=context_factory,
        middleware=middleware,
    )

    agent_card = _build_agent_card(
        handler,
        name=name,
        port=resolved_port,
        description=description,
        version=version,
        extra_skills=_test_controller_skills() if test_controller else None,
    )

    if task_store is None:
        task_store = InMemoryTaskStore()

    # DefaultRequestHandler stores push_config_store verbatim and treats
    # None as "push-notif endpoints unsupported" (UnsupportedOperationError
    # on tasks/pushNotificationConfig/*). Passing None is the correct
    # default; sellers opt in by wiring a store.
    request_handler = DefaultRequestHandler(
        agent_executor=executor,
        task_store=task_store,
        push_config_store=push_config_store,
    )

    a2a_app = A2AStarletteApplication(
        agent_card=agent_card,
        http_handler=request_handler,
    )

    return a2a_app.build()


def _test_controller_skills() -> list[AgentSkill]:
    """Build A2A skill definition for comply_test_controller."""
    return [
        AgentSkill(
            id="comply_test_controller",
            name="comply_test_controller",
            description="Compliance test controller. Sandbox only, not for production use.",
            tags=["adcp", "testing"],
        )
    ]