A2A Protocol

# Terminal 1 — Agente A
langgraph dev --port 2024

# Terminal 2 — Agente B
langgraph dev --port 2025

from langgraph.graph import StateGraph, MessagesState
from langchain_anthropic import ChatAnthropic

# Instancia o LLM
llm = ChatAnthropic(model="claude-sonnet-4-20250514")

# Nó de processamento
def call_model(state: MessagesState):
    response = llm.invoke(state["messages"])
    return {"messages": [response]}

# Constrói o grafo
builder = StateGraph(MessagesState)
builder.add_node("agent", call_model)
builder.set_entry_point("agent")
builder.set_finish_point("agent")

graph = builder.compile()

import asyncio, aiohttp, uuid

async def send_message(session, url, message,
                         context_id=None, task_id=None):
    payload = {
        "jsonrpc": "2.0",
        "method": "message/send",
        "id": str(uuid.uuid4()),
        "params": {
            "message": {
                "role": "user",
                "parts": [{"type": "text", "text": message}],
                # Inclui contextId após primeira mensagem
                **({"contextId": context_id} if context_id else {}),
                **({"taskId": task_id}     if task_id    else {}),
            }
        }
    }
    async with session.post(url, json=payload) as resp:
        data = await resp.json()
        result  = data["result"]
        text    = result["artifacts"][0]["parts"][0]["text"]
        return text, result.get("contextId"), result.get("id")

async def run_conversation():
    agent_a = "http://localhost:2024/a2a/<AGENT_A_ID>"
    agent_b = "http://localhost:2025/a2a/<AGENT_B_ID>"
    message = "Olá, vamos colaborar numa tarefa!"
    context_id = task_id = None

    async with aiohttp.ClientSession() as session:
        for _ in range(3):
            message, context_id, task_id = await send_message(
                session, agent_a, message, context_id, task_id)
            message, context_id, task_id = await send_message(
                session, agent_b, message, context_id, task_id)

asyncio.run(run_conversation())

services:
  - name: agente-a
    url: http://langgraph-agent-a:2024
    routes:
      - name: a2a-agente-a
        paths: ["/agentes/a"]
        strip_path: true
    plugins:
      - name: jwt
        config:
          key_claim_name: iss
          claims_to_verify: [exp]
      - name: rate-limiting
        config:
          minute: 60
          policy: local

location /agentes/ {
    proxy_pass         http://langgraph_upstream;

    # Desabilita buffering — essencial para SSE
    proxy_buffering    off;
    proxy_cache        off;
    proxy_read_timeout 300s;

    # Headers para SSE
    proxy_set_header   Connection '';
    proxy_http_version 1.1;
    chunked_transfer_encoding on;
}

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
import httpx, jwt

app = FastAPI()

AGENTES = {
    "planejador": "http://agente-planejador:2024/a2a/assistente-abc",
    "executor":   "http://agente-executor:2025/a2a/assistente-xyz",
}

@app.post("/a2a/{agente}")
async def proxy_a2a(agente: str, request: Request):
    # 1. Autenticação JWT
    token = request.headers.get("Authorization", "").replace("Bearer ", "")
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Token inválido")

    # 2. Roteamento
    if agente not in AGENTES:
        raise HTTPException(status_code=404, detail="Agente não encontrado")

    # 3. Proxy com suporte a SSE
    body = await request.json()

    async def stream():
        async with httpx.AsyncClient() as client:
            async with client.stream("POST", AGENTES[agente], json=body) as r:
                async for chunk in r.aiter_bytes():
                    yield chunk

    return StreamingResponse(stream(), media_type="text/event-stream")

async def a2a_middleware(request, call_next):
    body = await request.json()

    # Extrai metadados do JSON-RPC
    method     = body.get("method")     # ex: "message/send"
    context_id = (
        body.get("params", {})
            .get("message", {})
            .get("contextId")
    )

    logger.info({
        "protocol":  "a2a",
        "method":    method,
        "context_id": context_id,
        "agent":    request.url.path,
        "timestamp": datetime.utcnow().isoformat(),
    })

    return await call_next(request)