Multi-Tenant A2A Routing¶

"""Snippet: LitestarA2A with A2A 1.0 HTTP+JSON REST routes, tenant header resolution, and context_builder."""

# start-example
from typing import Any
from uuid import uuid4

from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.context import ServerCallContext
from a2a.server.events import EventQueue
from a2a.types import AgentCapabilities, AgentCard, AgentInterface, AgentSkill, Message, Part, Role
from a2a.utils.errors import UnsupportedOperationError
from litestar import Litestar, Request, get
from litestar.exceptions import PermissionDeniedException

from litestar_mcp.a2a import A2AConfig, LitestarA2A

ALLOWED_TENANTS = frozenset({"acme", "globex"})


@get(
    "/tickets/{ticket_id:str}",
    opt={
        "a2a_skill": {
            "id": "ticket_lookup",
            "name": "Ticket Lookup",
            "description": "Retrieve tenant support ticket status by identifier.",
            "tags": ["support", "tickets"],
        }
    },
)
async def get_ticket(ticket_id: str) -> dict[str, str]:
    """Retrieve a tenant support ticket by identifier."""
    return {"ticket_id": ticket_id, "status": "open"}


class TenantSupportExecutor(AgentExecutor):
    """Execute A2A requests within the validated tenant context."""

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Enqueue a tenant-scoped agent response message."""
        tenant = context.call_context.tenant if context.call_context is not None else "default"
        reply_text = f"[{tenant}] Processed: {context.get_user_input()}"
        await event_queue.enqueue_event(
            Message(
                message_id=str(uuid4()),
                role=Role.ROLE_AGENT,
                parts=[Part(text=reply_text)],
            )
        )

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Report cancellation as unsupported for stateless message responses."""
        raise UnsupportedOperationError(message="Cancellation is not supported for stateless messages")


async def authorize_tenant(request: Request[Any, Any, Any], context: ServerCallContext) -> ServerCallContext:
    """Validate the resolved A2A tenant from URL path or X-A2A-Tenant header."""
    del request
    if context.tenant not in ALLOWED_TENANTS:
        raise PermissionDeniedException(detail="Authorized tenant is required")
    context.state["tenant"] = context.tenant
    return context


card = AgentCard(
    name="Multi-Tenant Support Agent",
    description="Handles tenant-isolated support inquiries over JSON-RPC 2.0 and HTTP+JSON REST",
    version="1.0.0",
    default_input_modes=["text/plain"],
    default_output_modes=["text/plain"],
    skills=[
        AgentSkill(
            id="ticket_lookup",
            name="Ticket Lookup",
            description="Retrieve tenant support ticket status by identifier.",
            tags=["support", "tickets"],
        )
    ],
    capabilities=AgentCapabilities(streaming=True),
    supported_interfaces=[
        AgentInterface(
            url="https://agents.example.com/a2a",
            protocol_binding="JSONRPC",
            protocol_version="1.0",
        ),
        AgentInterface(
            url="https://agents.example.com/a2a",
            protocol_binding="HTTP+JSON",
            protocol_version="1.0",
        ),
    ],
)

config = A2AConfig(
    path="/a2a",
    enable_http_json=True,
    tenant_header="X-A2A-Tenant",
    route_opt={"requires_tenant": True},
    agent_card_route_opt={"a2a_skill": "public_card"},
    context_builder=authorize_tenant,
)

app = Litestar(
    route_handlers=[get_ticket],
    plugins=[LitestarA2A(card, TenantSupportExecutor(), config)],
)
# end-example


def build() -> Litestar:
    """Return the configured Litestar application for snippet smoke testing."""
    return app

@dataclass(frozen=True)
class Principal:
    id: str
    workspaces: frozenset[str] = frozenset({"alpha", "beta"})

    async def authorize_workspace(self, workspace: str) -> str:
        """Replace with an awaited membership lookup in the existing application."""
        if workspace not in self.workspaces:
            raise PermissionDeniedException(detail="Workspace membership required")
        return workspace


class DemoAuthentication(AbstractAuthenticationMiddleware):
    async def authenticate_request(self, connection: "ASGIConnection[Any, Any, Any, Any]") -> AuthenticationResult:
        credentials = {"Bearer demo-alice": Principal("alice"), "Bearer demo-bob": Principal("bob")}
        principal = credentials.get(connection.headers.get("authorization", ""))
        if principal is None:
            raise NotAuthorizedException(
                detail="Valid demonstration credentials required", headers={"WWW-Authenticate": "Bearer"}
            )
        return AuthenticationResult(user=principal, auth=None)


async def authorize_context(request: "Request[Principal, Any, Any]", context: ServerCallContext) -> ServerCallContext:
    context.tenant = await request.user.authorize_workspace(context.tenant)
    if EXTENSION_URI in context.requested_extensions:
        context.state["a2a_activated_extensions"] = {EXTENSION_URI}
    return context


def resolve_owner(context: ServerCallContext) -> str:
    """Share the same collision-safe owner key between both SDK stores."""
    if not context.user.is_authenticated or not context.user.user_name or not context.tenant:
        raise NotAuthorizedException(
            detail="An authenticated workspace owner is required", headers={"WWW-Authenticate": "Bearer"}
        )
    return json.dumps([context.tenant, context.user.user_name], separators=(",", ":"))

def create_app(
    service_factory: Callable[[ServerCallContext], AbstractAsyncContextManager[ApplicationService]] = service_scope,
    *,
    task_store: TaskStore | None = None,
    push_store: PushNotificationConfigStore | None = None,
    push_sender: PushNotificationSender | None = None,
) -> Litestar:
    card = AgentCard(
        name="Workspace agent",
        description="Authenticated application services through the official A2A SDK",
        version="1.0.0",
        security_schemes={
            "bearerAuth": SecurityScheme(http_auth_security_scheme=HTTPAuthSecurityScheme(scheme="Bearer"))
        },
        security_requirements=[SecurityRequirement(schemes={"bearerAuth": StringList()})],
        default_input_modes=["text/plain"],
        default_output_modes=["text/plain"],
        skills=[
            AgentSkill(id="answer", name="Answer", description="Answer within an authorized workspace", tags=["demo"])
        ],
        capabilities=AgentCapabilities(
            streaming=True,
            push_notifications=True,
            extended_agent_card=True,
            extensions=[
                AgentExtension(uri=EXTENSION_URI, description="Prefix the answer with the authorized workspace")
            ],
        ),
        supported_interfaces=[
            AgentInterface(url="http://localhost:8000/a2a", protocol_binding="JSONRPC", protocol_version="1.0")
        ],
    )
    if push_store is None:
        push_store = InMemoryPushNotificationConfigStore(owner_resolver=resolve_owner)
    if task_store is None:
        task_store = InMemoryTaskStore(owner_resolver=resolve_owner)
    if push_sender is None:
        push_sender = RecordingPushSender(push_store)
    handler = WorkspaceTaskHandler(
        agent_executor=ApplicationExecutor(service_factory),
        task_store=task_store,
        agent_card=card,
        push_config_store=push_store,
        push_sender=push_sender,
        extended_agent_card=card,
    )
    return Litestar(
        middleware=[DefineMiddleware(DemoAuthentication)],
        plugins=[LitestarA2A(card, handler, A2AConfig(context_builder=authorize_context))],
    )


app = create_app()