A2A na AWS
API Gateway, SigV4 & Bearer JWT

# ─────────────────────────────────────────────
# CONTA B — Specialist Agent (exemplo genérico)
# ─────────────────────────────────────────────

# ── API Gateway REST API ──────────────────────
resource "aws_api_gateway_rest_api" "agent_api" {
  name        = "specialist-agent-b"
  description = "A2A endpoint for Specialist Agent B"

  endpoint_configuration {
    types = ["REGIONAL"]
  }
}

resource "aws_api_gateway_resource" "invoke" {
  rest_api_id = aws_api_gateway_rest_api.agent_api.id
  parent_id   = aws_api_gateway_rest_api.agent_api.root_resource_id
  path_part   = "invoke"
}

resource "aws_api_gateway_method" "post_invoke" {
  rest_api_id   = aws_api_gateway_rest_api.agent_api.id
  resource_id   = aws_api_gateway_resource.invoke.id
  http_method   = "POST"
  authorization = "AWS_IAM"   # SigV4 — sem API Key exposta
}

resource "aws_api_gateway_integration" "lambda_integration" {
  rest_api_id             = aws_api_gateway_rest_api.agent_api.id
  resource_id             = aws_api_gateway_resource.invoke.id
  http_method             = aws_api_gateway_method.post_invoke.http_method
  integration_http_method = "POST"
  type                    = "AWS_PROXY"
  uri                     = aws_lambda_function.agent.invoke_arn
}

resource "aws_api_gateway_deployment" "prod" {
  depends_on  = [aws_api_gateway_integration.lambda_integration]
  rest_api_id = aws_api_gateway_rest_api.agent_api.id
  stage_name  = "prod"
}

# ── Resource Policy: só deixa a Conta A chamar ──
resource "aws_api_gateway_rest_api_policy" "allow_supervisor_account" {
  rest_api_id = aws_api_gateway_rest_api.agent_api.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Effect    = "Allow"
        Principal = {
          AWS = "arn:aws:iam::ACCOUNT_A_ID:root"
        }
        Action   = "execute-api:Invoke"
        Resource = "${aws_api_gateway_rest_api.agent_api.execution_arn}/*"
      }
    ]
  })
}

# ── Lambda do Agent B ────────────────────────
resource "aws_lambda_function" "agent" {
  function_name = "specialist-agent-b"
  role          = aws_iam_role.lambda_exec.arn
  runtime       = "python3.12"
  handler       = "handler.lambda_handler"
  filename      = "agent_b.zip"

  environment {
    variables = {
      AGENT_NAME          = "specialist-b"
      OPENAI_API_KEY      = var.openai_api_key
      OPENSEARCH_ENDPOINT = var.opensearch_endpoint
    }
  }
}

# ── IAM Role da Lambda ───────────────────────
resource "aws_iam_role" "lambda_exec" {
  name = "agent-b-lambda-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect    = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
      Action    = "sts:AssumeRole"
    }]
  })
}

# ─────────────────────────────────────────────
# CONTA A — Cross-Account Role para o Supervisor
# ─────────────────────────────────────────────
resource "aws_iam_role" "supervisor_invoker" {
  name = "a2a-supervisor-invoker"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect    = "Allow"
      Principal = {
        AWS = "arn:aws:iam::ACCOUNT_A_ID:role/supervisor-task-role"
      }
      Action = "sts:AssumeRole"
    }]
  })
}

resource "aws_iam_role_policy" "invoke_agents" {
  name = "invoke-remote-agents"
  role = aws_iam_role.supervisor_invoker.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Effect   = "Allow"
        Action   = "execute-api:Invoke"
        Resource = [
          "arn:aws:execute-api:us-east-1:ACCOUNT_B_ID:*/prod/POST/invoke",
          "arn:aws:execute-api:us-east-1:ACCOUNT_C_ID:*/prod/POST/invoke",
        ]
      }
    ]
  })
}

# ──────────────────────────────────────────────────────────
# Lado do Supervisor (Conta A) — consumindo os agents via A2A
# ──────────────────────────────────────────────────────────
import json, boto3, requests
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
from langchain_core.tools import tool
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from typing import TypedDict, Annotated
import operator

# ── Helper: chama API Gateway com SigV4 ──────────────────
def call_agent_endpoint(
    endpoint_url: str,
    payload: dict,
    region: str = "us-east-1",
    role_arn: str | None = None,
) -> dict:
    """
    Faz POST autenticado via SigV4.
    Se role_arn for passado, assume a cross-account role primeiro.
    """
    session = boto3.Session()

    if role_arn:
        sts = session.client("sts")
        assumed = sts.assume_role(
            RoleArn=role_arn,
            RoleSessionName="supervisor-a2a-call",
        )["Credentials"]
        creds = Credentials(
            access_key=assumed["AccessKeyId"],
            secret_key=assumed["SecretAccessKey"],
            token=assumed["SessionToken"],
        )
    else:
        creds = session.get_credentials().get_frozen_credentials()
        creds = Credentials(creds.access_key, creds.secret_key, creds.token)

    body    = json.dumps(payload)
    request = AWSRequest(
        method="POST",
        url=endpoint_url,
        data=body,
        headers={"Content-Type": "application/json"},
    )
    SigV4Auth(creds, "execute-api", region).add_auth(request)

    resp = requests.post(
        endpoint_url,
        data=body,
        headers=dict(request.headers),
        timeout=30,
    )
    resp.raise_for_status()
    return resp.json()


# ── Config dos agentes remotos ────────────────────────────
REMOTE_AGENTS = {
    "agent_b": {
        "url":      "https://API_ID_B.execute-api.us-east-1.amazonaws.com/prod/invoke",
        "role_arn": "arn:aws:iam::ACCOUNT_B_ID:role/a2a-supervisor-invoker",
    },
    "agent_c": {
        "url":      "https://API_ID_C.execute-api.us-east-1.amazonaws.com/prod/invoke",
        "role_arn": "arn:aws:iam::ACCOUNT_C_ID:role/a2a-supervisor-invoker",
    },
}


# ── LangChain Tools wrapping os agents remotos ───────────
@tool
def invoke_agent_b(task: str, context: dict = {}) -> str:
    """Invoca o Specialist Agent B para tarefas de domínio X."""
    result = call_agent_endpoint(
        endpoint_url=REMOTE_AGENTS["agent_b"]["url"],
        payload={"task": task, "context": context},
        role_arn=REMOTE_AGENTS["agent_b"]["role_arn"],
    )
    return result.get("output", str(result))


@tool
def invoke_agent_c(task: str, context: dict = {}) -> str:
    """Invoca o Specialist Agent C para tarefas de domínio Y."""
    result = call_agent_endpoint(
        endpoint_url=REMOTE_AGENTS["agent_c"]["url"],
        payload={"task": task, "context": context},
        role_arn=REMOTE_AGENTS["agent_c"]["role_arn"],
    )
    return result.get("output", str(result))


# ── LangGraph Supervisor ──────────────────────────────────
class SupervisorState(TypedDict):
    messages: Annotated[list, operator.add]

from langchain_openai import ChatOpenAI

tools     = [invoke_agent_b, invoke_agent_c]
llm       = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools)
tool_node = ToolNode(tools)

def supervisor_node(state: SupervisorState):
    return {"messages": [llm.invoke(state["messages"])]}

def should_continue(state: SupervisorState):
    last = state["messages"][-1]
    return "tools" if getattr(last, "tool_calls", None) else END

graph = StateGraph(SupervisorState)
graph.add_node("supervisor", supervisor_node)
graph.add_node("tools", tool_node)
graph.set_entry_point("supervisor")
graph.add_conditional_edges("supervisor", should_continue)
graph.add_edge("tools", "supervisor")
app = graph.compile()


# ── Lambda handler do Agent B (Conta B) ──────────────────
import os
from langchain_core.messages import HumanMessage, SystemMessage

def lambda_handler(event, context):
    try:
        body   = json.loads(event.get("body", "{}"))
        task   = body["task"]
        ctx    = body.get("context", {})

        llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0,
            openai_api_key=os.environ["OPENAI_API_KEY"],
        )
        messages = [
            SystemMessage(content="Você é o Specialist Agent B, especializado em domínio X."),
            HumanMessage(content=f"Task: {task}\nContext: {json.dumps(ctx)}"),
        ]
        result = llm.invoke(messages)

        return {
            "statusCode": 200,
            "headers": {"Content-Type": "application/json"},
            "body": json.dumps({
                "agent":    "agent_b",
                "output":   result.content,
                "metadata": {"tokens": result.usage_metadata},
            }),
        }
    except Exception as e:
        return {"statusCode": 500, "body": json.dumps({"error": str(e)})}

{
  "task": "Analise os dados de vendas do Q1",
  "context": {
    "session_id": "sess_abc123",
    "redis_keys": ["wm:sess_abc123:sales_data"],
    "user_id": "user_42"
  },
  "metadata": {
    "trace_id": "trace_xyz",
    "caller": "supervisor",
    "timestamp": "2026-06-04T10:00:00Z"
  }
}

{
  "agent": "sales-analyst",
  "output": "Análise completa: ...",
  "artifacts": {
    "stored_key": "wm:sess_abc123:sales_analysis"
  },
  "metadata": {
    "model": "gpt-4o-mini",
    "latency_ms": 1240
  }
}

# ──────────────────────────────────────────────────────────────────
# CONTA B — API Gateway HTTP API com JWT Authorizer nativo
# ──────────────────────────────────────────────────────────────────
# NOTA: HTTP API (v2) — não REST API (v1)
# JWT Authorizer é nativo no HTTP API, sem Lambda Authorizer extra

variable "idp_issuer" {
  description = "Issuer URL do IdP. Ex: https://your-tenant.okta.com/oauth2/default"
  type        = string
}

variable "idp_audience" {
  description = "Audience esperado no JWT. Ex: api://agent-b"
  type        = string
}

# ── HTTP API ────────────────────────────────────────────────────────
resource "aws_apigatewayv2_api" "agent_api" {
  name          = "specialist-agent-b"
  protocol_type = "HTTP"

  cors_configuration {
    allow_origins = ["*"]
    allow_methods = ["POST", "OPTIONS"]
    allow_headers = ["Authorization", "Content-Type"]
  }
}

# ── JWT Authorizer (valida contra JWKS do IdP automaticamente) ─────
resource "aws_apigatewayv2_authorizer" "jwt_auth" {
  api_id           = aws_apigatewayv2_api.agent_api.id
  authorizer_type  = "JWT"
  identity_sources = ["$request.header.Authorization"]
  name             = "idp-jwt-authorizer"

  jwt_configuration {
    issuer   = var.idp_issuer
    # API GW busca JWKS em {issuer}/.well-known/jwks.json automaticamente
    audience = [var.idp_audience]
  }
}

# ── Rota POST /invoke protegida ────────────────────────────────────
resource "aws_apigatewayv2_route" "invoke" {
  api_id             = aws_apigatewayv2_api.agent_api.id
  route_key          = "POST /invoke"
  authorization_type = "JWT"
  authorizer_id      = aws_apigatewayv2_authorizer.jwt_auth.id
  target             = "integrations/${aws_apigatewayv2_integration.lambda.id}"
}

# ── Integração Lambda ──────────────────────────────────────────────
resource "aws_apigatewayv2_integration" "lambda" {
  api_id                 = aws_apigatewayv2_api.agent_api.id
  integration_type       = "AWS_PROXY"
  integration_uri        = aws_lambda_function.agent.invoke_arn
  payload_format_version = "2.0"
}

# ── Stage (auto-deploy) ────────────────────────────────────────────
resource "aws_apigatewayv2_stage" "prod" {
  api_id      = aws_apigatewayv2_api.agent_api.id
  name        = "prod"
  auto_deploy = true

  access_log_settings {
    destination_arn = aws_cloudwatch_log_group.api_logs.arn
    format = jsonencode({
      requestId      = "$context.requestId"
      sourceIp       = "$context.identity.sourceIp"
      httpMethod     = "$context.httpMethod"
      routeKey       = "$context.routeKey"
      status         = "$context.status"
      authorizeError = "$context.authorizer.error"
      jwtClaims      = "$context.authorizer.claims"
    })
  }
}

resource "aws_cloudwatch_log_group" "api_logs" {
  name              = "/aws/apigw/specialist-agent-b"
  retention_in_days = 7
}

# ── Lambda permission ──────────────────────────────────────────────
resource "aws_lambda_permission" "apigw" {
  statement_id  = "AllowAPIGatewayInvoke"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.agent.function_name
  principal     = "apigateway.amazonaws.com"
  source_arn    = "${aws_apigatewayv2_api.agent_api.execution_arn}/*/*"
}

# ── Lambda do Agent ────────────────────────────────────────────────
resource "aws_lambda_function" "agent" {
  function_name = "specialist-agent-b"
  role          = aws_iam_role.lambda_exec.arn
  runtime       = "python3.12"
  handler       = "handler.lambda_handler"
  filename      = "agent_b.zip"
  timeout       = 60

  environment {
    variables = {
      AGENT_NAME        = "specialist-b"
      OPENAI_API_KEY    = data.aws_secretsmanager_secret_version.openai.secret_string
      EXPECTED_AUDIENCE = var.idp_audience
      EXPECTED_ISSUER   = var.idp_issuer
    }
  }
}

# ── Secrets Manager — client_id/secret do Supervisor (Conta A) ────
resource "aws_secretsmanager_secret" "idp_creds" {
  provider = aws.account_a
  name     = "a2a/idp/agent-b"
}

resource "aws_secretsmanager_secret_version" "idp_creds" {
  provider  = aws.account_a
  secret_id = aws_secretsmanager_secret.idp_creds.id
  secret_string = jsonencode({
    client_id     = "<client_id_do_IdP>"
    client_secret = "<client_secret_do_IdP>"
    token_url     = "${var.idp_issuer}/v1/token"
    audience      = var.idp_audience
  })
}

output "agent_b_endpoint" {
  value = "${aws_apigatewayv2_stage.prod.invoke_url}/invoke"
}

# ──────────────────────────────────────────────────────────────────
# Supervisor (Conta A) — Bearer token com cache + LangGraph tools
# ──────────────────────────────────────────────────────────────────
import time, json, threading
import boto3, requests
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from typing import TypedDict, Annotated
import operator


# ── Token Manager — busca e cacheia JWT do IdP ────────────────────
class IdPTokenManager:
    """
    Faz client_credentials flow contra o IdP.
    Cacheia o token em memória; renova automaticamente antes de expirar.
    Em produção: usar Redis (ElastiCache) para cache compartilhado.
    """

    def __init__(self, secret_name: str, region: str = "us-east-1"):
        self._lock   = threading.Lock()
        self._token  = None
        self._expiry = 0.0
        self._creds  = self._load_secret(secret_name, region)

    def _load_secret(self, name: str, region: str) -> dict:
        client = boto3.client("secretsmanager", region_name=region)
        raw = client.get_secret_value(SecretId=name)["SecretString"]
        return json.loads(raw)

    def get_token(self) -> str:
        with self._lock:
            # Renova 60s antes de expirar
            if time.time() < self._expiry - 60 and self._token:
                return self._token

            resp = requests.post(
                self._creds["token_url"],
                data={
                    "grant_type":    "client_credentials",
                    "client_id":     self._creds["client_id"],
                    "client_secret": self._creds["client_secret"],
                    "audience":      self._creds["audience"],  # Auth0
                    # Okta usa "scope" em vez de "audience":
                    # "scope": "agent:invoke",
                },
                timeout=10,
            )
            resp.raise_for_status()
            data = resp.json()

            self._token  = data["access_token"]
            self._expiry = time.time() + data.get("expires_in", 3600)
            return self._token


# ── Uma instância por agent (cada um tem seu client_id/audience) ──
_token_managers: dict[str, IdPTokenManager] = {}

def get_token_manager(agent_id: str) -> IdPTokenManager:
    if agent_id not in _token_managers:
        _token_managers[agent_id] = IdPTokenManager(
            secret_name=f"a2a/idp/{agent_id}"
        )
    return _token_managers[agent_id]


# ── Helper de chamada HTTP ─────────────────────────────────────────
def call_agent(agent_id: str, endpoint_url: str, payload: dict) -> dict:
    token = get_token_manager(agent_id).get_token()

    resp = requests.post(
        endpoint_url,
        json=payload,
        headers={
            "Authorization":  f"Bearer {token}",
            "Content-Type":   "application/json",
            "X-Agent-Caller": "supervisor",
        },
        timeout=30,
    )
    resp.raise_for_status()
    return resp.json()


# ── Config ─────────────────────────────────────────────────────────
AGENTS = {
    "agent_b": "https://API_ID_B.execute-api.us-east-1.amazonaws.com/prod/invoke",
    "agent_c": "https://API_ID_C.execute-api.us-east-1.amazonaws.com/prod/invoke",
}


# ── LangChain Tools ────────────────────────────────────────────────
@tool
def invoke_agent_b(task: str, context: dict = {}) -> str:
    """Invoca o Specialist Agent B para tarefas de domínio X."""
    result = call_agent("agent-b", AGENTS["agent_b"], {
        "task": task, "context": context,
    })
    return result.get("output", str(result))


@tool
def invoke_agent_c(task: str, context: dict = {}) -> str:
    """Invoca o Specialist Agent C para tarefas de domínio Y."""
    result = call_agent("agent-c", AGENTS["agent_c"], {
        "task": task, "context": context,
    })
    return result.get("output", str(result))


# ── LangGraph Supervisor ──────────────────────────────────────────
class SupervisorState(TypedDict):
    messages: Annotated[list, operator.add]

tools     = [invoke_agent_b, invoke_agent_c]
llm       = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools)
tool_node = ToolNode(tools)

def supervisor_node(state: SupervisorState):
    return {"messages": [llm.invoke(state["messages"])]}

def should_continue(state: SupervisorState):
    last = state["messages"][-1]
    return "tools" if getattr(last, "tool_calls", None) else END

graph = StateGraph(SupervisorState)
graph.add_node("supervisor", supervisor_node)
graph.add_node("tools", tool_node)
graph.set_entry_point("supervisor")
graph.add_conditional_edges("supervisor", should_continue)
graph.add_edge("tools", "supervisor")
app = graph.compile()

# ──────────────────────────────────────────────────────────────────
# Agent B — Lambda handler (Conta B)
# API GW já validou o JWT antes de chegar aqui.
# Claims chegam em event["requestContext"]["authorizer"]["jwt"]["claims"]
# ──────────────────────────────────────────────────────────────────
import json, os
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage


def lambda_handler(event, context):
    # ── Claims do JWT (já validados pelo API GW) ──────────────────
    auth_ctx = event.get("requestContext", {}).get("authorizer", {}).get("jwt", {})
    claims   = auth_ctx.get("claims", {})

    caller   = claims.get("sub", "unknown")   # client_id do Supervisor
    audience = claims.get("aud", "")
    scope    = claims.get("scope", "")

    # Verificação extra de audience (defense in depth)
    expected_aud = os.environ["EXPECTED_AUDIENCE"]
    aud_list = audience if isinstance(audience, list) else [audience]
    if expected_aud not in aud_list:
        return _error(403, f"Invalid audience: {audience}")

    # ── Payload ───────────────────────────────────────────────────
    try:
        body = json.loads(event.get("body", "{}"))
        task = body["task"]
        ctx  = body.get("context", {})
    except (KeyError, json.JSONDecodeError) as e:
        return _error(400, f"Bad request: {e}")

    # ── Agent logic ───────────────────────────────────────────────
    try:
        llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0,
            openai_api_key=os.environ["OPENAI_API_KEY"],
        )

        messages = [
            SystemMessage(content=(
                "Você é o Specialist Agent B, especializado em domínio X. "
                f"Caller: {caller}"
            )),
            HumanMessage(content=f"Task: {task}\nContext: {json.dumps(ctx)}"),
        ]

        result = llm.invoke(messages)

        return {
            "statusCode": 200,
            "headers": {"Content-Type": "application/json"},
            "body": json.dumps({
                "agent":  "agent_b",
                "output": result.content,
                "caller": caller,   # auditoria
            }),
        }

    except Exception as e:
        return _error(500, str(e))


def _error(status: int, msg: str) -> dict:
    return {
        "statusCode": status,
        "body": json.dumps({"error": msg}),
    }


# ──────────────────────────────────────────────────────────────────
# Estrutura do event (HTTP API payload format 2.0)
# ──────────────────────────────────────────────────────────────────
#
# {
#   "routeKey": "POST /invoke",
#   "body": '{"task": "...", "context": {...}}',
#   "requestContext": {
#     "authorizer": {
#       "jwt": {
#         "claims": {
#           "sub":   "supervisor-client-id",
#           "iss":   "https://your-tenant.okta.com/oauth2/default",
#           "aud":   "api://agent-b",
#           "exp":   1749000000,
#           "scope": "agent:invoke"
#         },
#         "scopes": ["agent:invoke"]
#       }
#     }
#   }
# }