Token Propagation
Architecture

from __future__ import annotations
from typing import Annotated, Optional
from datetime import datetime
from enum import Enum
import operator
from pydantic import BaseModel, Field
from langchain_core.messages import BaseMessage


# ── UserContext ───────────────────────────────────────
# Montado uma única vez a partir do JWT.
# Viaja imutável pelo grafo inteiro.

class UserContext(BaseModel):
    # Identidade
    user_id: str
    email:   Optional[str] = None
    name:    Optional[str] = None

    # Claims de autorização (vindas do EntraID)
    groups: list[str] = Field(default_factory=list)
    roles:  list[str] = Field(default_factory=list)
    scopes: list[str] = Field(default_factory=list)

    # Metadados do tenant/app
    tenant_id: Optional[str] = None
    client_id: Optional[str] = None  # azp — app que emitiu o token

    # Rastreabilidade — propagados em todo state e payload
    session_id: str
    thread_id:  str

    # Janela de validade (apenas para logging)
    token_issued_at:  Optional[datetime] = None
    token_expires_at: Optional[datetime] = None

    @classmethod
    def from_jwt_claims(
        cls,
        claims: dict,
        session_id: str,
        thread_id: str,
    ) -> "UserContext":
        """
        Allowlist explícita — só extrai o que está mapeado aqui.
        Qualquer claim fora dessa lista é descartada silenciosamente.
        """
        return cls(
            user_id   = claims["sub"],
            email     = claims.get("preferred_username"),
            name      = claims.get("name"),
            groups    = claims.get("groups", []),
            roles     = claims.get("roles", []),
            scopes    = claims.get("scp", "").split() if claims.get("scp") else [],
            tenant_id = claims.get("tid"),
            client_id = claims.get("azp"),
            session_id = session_id,
            thread_id  = thread_id,
            token_issued_at  = datetime.fromtimestamp(claims["iat"]) if "iat" in claims else None,
            token_expires_at = datetime.fromtimestamp(claims["exp"]) if "exp" in claims else None,
        )

    def to_propagatable_payload(self) -> dict:
        """
        Segunda camada de allowlist: só o que o agente precisa.
        Raw token NUNCA entra aqui.
        """
        return {
            "user_id":    self.user_id,
            "groups":     self.groups,
            "roles":      self.roles,
            "session_id": self.session_id,
            "thread_id":  self.thread_id,
        }


# ── AccessDenied ──────────────────────────────────────

class AccessDeniedReason(str, Enum):
    AGENT_RESTRICTED       = "AGENT_RESTRICTED"
    CLIENT_DATA_RESTRICTED = "CLIENT_DATA_RESTRICTED"
    TOOL_RESTRICTED        = "TOOL_RESTRICTED"


class AccessDeniedContext(BaseModel):
    reason:     AccessDeniedReason
    agent_id:   str
    tool_name:  Optional[str] = None
    client_id:  Optional[str] = None
    detail:     str  # mensagem técnica para log, nunca exposta ao usuário


# ── AgentState ────────────────────────────────────────
# Annotated[list, operator.add] → LangGraph concatena
# em vez de substituir (acumulação de mensagens)

class AgentState(TypedDict):
    user_context:        UserContext           # imutável após inject
    session_id:          str                    # duplicado para acesso rápido
    thread_id:           str
    messages:            Annotated[list[BaseMessage], operator.add]
    intent:              Optional[str]
    active_agent:        Optional[str]
    access_denied:       Optional[AccessDeniedContext]
    specialist_response: Optional[dict]

import time, logging, requests, boto3
from enum import Enum
from typing import Optional
from supervisor.graph.state import UserContext

logger = logging.getLogger(__name__)


class AccessScope(str, Enum):
    AGENT       = "ACCESS"
    CLIENT_DATA = "CLIENT_DATA"
    TOOL        = "TOOL"


class ACLCache:
    """
    Estratégia de resolução (em ordem):
    1. DynamoDB (cache local)
    2. MISS/expirado → API de parametrização → popula DynamoDB
    3. API falhou   → nega acesso (fail-closed)
    """

    def __init__(self, table_name: str, acl_api_url: str, ttl_seconds: int = 300):
        self.table       = boto3.resource("dynamodb").Table(table_name)
        self.acl_api_url = acl_api_url
        self.ttl_seconds = ttl_seconds

    def check_access(
        self,
        agent_id: str,
        scope: AccessScope,
        user_context: UserContext,
        tool_name: Optional[str] = None,
    ) -> bool:
        sk   = f"TOOL#{tool_name}" if scope == AccessScope.TOOL else scope.value
        pk   = f"AGENT#{agent_id}"
        rule = self._get_from_dynamo(pk, sk)

        if rule is None:
            logger.info("acl_cache_miss pk=%s sk=%s user=%s", pk, sk, user_context.user_id)
            rule = self._fetch_from_api_and_cache(pk, sk, agent_id, scope, tool_name)

        if rule is None:
            logger.error("acl_resolution_failed pk=%s sk=%s — denying", pk, sk)
            return False  # fail-closed

        allowed_groups = set(rule.get("allowed_groups", []))
        allowed_roles  = set(rule.get("allowed_roles", []))
        user_groups    = set(user_context.groups)
        user_roles     = set(user_context.roles)

        return bool(user_groups & allowed_groups or user_roles & allowed_roles)

    def _get_from_dynamo(self, pk: str, sk: str) -> Optional[dict]:
        try:
            resp = self.table.get_item(Key={"PK": pk, "SK": sk})
            item = resp.get("Item")
            if not item:
                return None
            # Verifica TTL manualmente (DynamoDB pode demorar para limpar)
            if item.get("ttl", 0) < time.time():
                return None
            return item
        except Exception as e:
            logger.error("dynamo_read_error pk=%s sk=%s error=%s", pk, sk, e)
            return None

    def _fetch_from_api_and_cache(
        self, pk, sk, agent_id, scope, tool_name
    ) -> Optional[dict]:
        try:
            params = {"agent_id": agent_id, "scope": scope.value}
            if tool_name:
                params["tool"] = tool_name

            resp = requests.get(
                f"{self.acl_api_url}/rules",
                params=params, timeout=3
            )
            resp.raise_for_status()
            rule = resp.json()

            self.table.put_item(Item={
                "PK": pk, "SK": sk,
                **rule,
                "ttl": int(time.time()) + self.ttl_seconds,
            })
            logger.info("acl_cache_populated pk=%s sk=%s", pk, sk)
            return rule
        except Exception as e:
            logger.error("acl_api_error agent=%s scope=%s error=%s", agent_id, scope, e)
            return None

import time, logging, requests
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class CachedToken:
    access_token: str
    expires_at:   float  # unix timestamp

    def is_valid(self, buffer_seconds: int = 60) -> bool:
        """Considera inválido 60s antes de expirar."""
        return time.time() < (self.expires_at - buffer_seconds)


class AgentTokenProvider:
    """
    Gerencia tokens client_credentials para o Supervisor.
    Cache in-memory por scope — renovação automática ao expirar.

    Este token representa o SUPERVISOR, não o usuário.
    O contexto do usuário viaja separadamente no payload.
    """

    def __init__(self, tenant_id: str, client_id: str, client_secret: str):
        self.token_url     = f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
        self.client_id     = client_id
        self.client_secret = client_secret
        self._cache: dict[str, CachedToken] = {}

    def get_token(self, scope: str) -> str:
        cached = self._cache.get(scope)
        if cached and cached.is_valid():
            logger.debug("agent_token_cache_hit scope=%s", scope)
            return cached.access_token
        return self._fetch_and_cache(scope)

    def _fetch_and_cache(self, scope: str) -> str:
        logger.info("agent_token_refreshing scope=%s", scope)
        try:
            resp = requests.post(self.token_url, data={
                "grant_type":    "client_credentials",
                "client_id":     self.client_id,
                "client_secret": self.client_secret,
                "scope":         scope,
            }, timeout=5)
            resp.raise_for_status()
            data  = resp.json()
            token = CachedToken(
                access_token = data["access_token"],
                expires_at   = time.time() + data["expires_in"],
            )
            self._cache[scope] = token
            logger.info("agent_token_acquired scope=%s expires_in=%ds", scope, data["expires_in"])
            return token.access_token
        except Exception as e:
            logger.error("agent_token_fetch_error scope=%s error=%s", scope, e)
            raise

import logging, requests
from langchain_core.messages import AIMessage
from supervisor.graph.state import (
    AgentState, UserContext,
    AccessDeniedContext, AccessDeniedReason,
)
from supervisor.acl.cache import ACLCache, AccessScope
from supervisor.auth.tokens import AgentTokenProvider

logger = logging.getLogger(__name__)

# Dependências singleton — injetadas na inicialização
acl_cache:      ACLCache            = None
token_provider: AgentTokenProvider  = None
llm                                  = None
AGENT_ENDPOINTS: dict[str, str]      = {}
AGENT_SCOPES:    dict[str, str]      = {}


# ── Nó 1: parse_user_context ──────────────────────────
# Valida e loga o início do request. Não modifica o state.

def parse_user_context(state: AgentState) -> dict:
    u = state["user_context"]
    logger.info(
        "request_received user_id=%s session=%s thread=%s groups=%s roles=%s",
        u.user_id, u.session_id, u.thread_id, u.groups, u.roles,
    )
    return {}


# ── Nó 2: resolve_intent ──────────────────────────────
# LLM lê a mensagem e retorna o identificador do agente.
# Separado do check de acesso para logar "tentou acessar X".

def resolve_intent(state: AgentState) -> dict:
    u    = state["user_context"]
    msg  = state["messages"][-1].content
    resp = llm.invoke([
        {"role": "system", "content": """
Analise a mensagem e retorne APENAS o identificador do agente:
- financial: questões financeiras, relatórios, pagamentos
- hr: recursos humanos, folha, benefícios
- crm: dados de clientes, accounts, oportunidades
- general: sem agente específico
Retorne só o identificador."""},
        {"role": "user", "content": msg},
    ])
    intent = resp.content.strip().lower()
    logger.info("intent_resolved user=%s session=%s intent=%s", u.user_id, u.session_id, intent)
    return {"intent": intent}


# ── Nó 3: check_agent_access ──────────────────────────
# Primeiro portão de segurança. Popula access_denied ou active_agent.

def check_agent_access(state: AgentState) -> dict:
    u, intent = state["user_context"], state["intent"]

    if not intent or intent == "general":
        return {"active_agent": "general"}

    if not acl_cache.check_access(intent, AccessScope.AGENT, u):
        denied = AccessDeniedContext(
            reason   = AccessDeniedReason.AGENT_RESTRICTED,
            agent_id = intent,
            detail   = f"groups={u.groups} roles={u.roles}",
        )
        _log_access_denied(denied, u)
        return {"access_denied": denied}

    return {"active_agent": intent}


# ── Nó 4: invoke_specialist ───────────────────────────
# Back-to-back: obtém token client_credentials, chama o agente.
# UserContext viaja no payload, não no token.

def invoke_specialist(state: AgentState) -> dict:
    u, agent_id = state["user_context"], state["active_agent"]

    if agent_id == "general":
        return {"specialist_response": {"content": "Posso ajudar com informações gerais."}}

    try:
        token = token_provider.get_token(AGENT_SCOPES[agent_id])
        logger.info("invoking_specialist agent=%s user=%s session=%s thread=%s",
                    agent_id, u.user_id, u.session_id, u.thread_id)

        resp = requests.post(
            AGENT_ENDPOINTS[agent_id],
            headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
            json={
                "user_context": u.to_propagatable_payload(),
                "messages": [{"role": m.type, "content": m.content} for m in state["messages"]],
            },
            timeout=30,
        )
        resp.raise_for_status()
        result = resp.json()

        # Captura erros estruturados de ACL vindos do agente
        if result.get("error_code") in {"CLIENT_DATA_RESTRICTED", "TOOL_RESTRICTED"}:
            denied = AccessDeniedContext(
                reason    = AccessDeniedReason(result["error_code"]),
                agent_id  = agent_id,
                tool_name = result.get("tool_name"),
                client_id = result.get("client_id"),
                detail    = result.get("detail", ""),
            )
            _log_access_denied(denied, u)
            return {"access_denied": denied, "specialist_response": result}

        logger.info("specialist_ok agent=%s user=%s session=%s", agent_id, u.user_id, u.session_id)
        return {"specialist_response": result}

    except requests.RequestException as e:
        logger.error("specialist_error agent=%s user=%s error=%s", agent_id, u.user_id, e)
        return {"specialist_response": {"content": "Não foi possível processar sua solicitação."}}


# ── Nó 5: process_response ────────────────────────────
# Nó terminal. Normaliza qualquer resultado em AIMessage.
# Detalhes técnicos ficam nos logs — nunca expostos ao usuário.

def process_response(state: AgentState) -> dict:
    denied = state.get("access_denied")
    if denied:
        content = _build_denied_message(denied)
    else:
        content = state.get("specialist_response", {}).get("content", "Sem resposta disponível.")
    return {"messages": [AIMessage(content=content)]}


# ── Helpers ───────────────────────────────────────────

def _log_access_denied(denied: AccessDeniedContext, u: UserContext):
    logger.warning(
        "access_denied reason=%s agent=%s tool=%s client=%s user=%s session=%s thread=%s detail=%s",
        denied.reason.value, denied.agent_id, denied.tool_name,
        denied.client_id, u.user_id, u.session_id, u.thread_id, denied.detail,
    )

def _build_denied_message(denied: AccessDeniedContext) -> str:
    return {
        AccessDeniedReason.AGENT_RESTRICTED: (
            "Você não possui acesso ao módulo solicitado. "
            "Entre em contato com o administrador para solicitar permissão."
        ),
        AccessDeniedReason.CLIENT_DATA_RESTRICTED: (
            "Você não possui acesso aos dados deste cliente. "
            "Apenas o account responsável pode consultar essas informações."
        ),
        AccessDeniedReason.TOOL_RESTRICTED: (
            "Você não possui permissão para executar esta operação. "
            "Sua conta permite apenas consultas neste módulo."
        ),
    }.get(denied.reason, "Acesso negado.")

from supervisor.graph.state import AgentState


def after_access_check(state: AgentState) -> str:
    """
    Após check_agent_access:
    - access_denied preenchido → encerra em process_response
    - aprovado                 → vai para invoke_specialist
    """
    if state.get("access_denied"):
        return "process_response"
    return "invoke_specialist"


def after_invoke(state: AgentState) -> str:
    # Sempre normaliza — erros do agente já foram capturados no nó anterior
    return "process_response"

from langgraph.graph import StateGraph, START, END
from supervisor.graph.state import AgentState
from supervisor.graph.nodes import (
    parse_user_context, resolve_intent,
    check_agent_access, invoke_specialist, process_response,
)
from supervisor.graph.edges import after_access_check


def build_supervisor_graph():
    builder = StateGraph(AgentState)

    # Registra nós
    builder.add_node("parse_user_context", parse_user_context)
    builder.add_node("resolve_intent",     resolve_intent)
    builder.add_node("check_agent_access", check_agent_access)
    builder.add_node("invoke_specialist",  invoke_specialist)
    builder.add_node("process_response",   process_response)

    # Fluxo principal
    builder.add_edge(START,                "parse_user_context")
    builder.add_edge("parse_user_context", "resolve_intent")
    builder.add_edge("resolve_intent",     "check_agent_access")

    # Edge condicional: acesso negado ou aprovado
    builder.add_conditional_edges(
        "check_agent_access",
        after_access_check,
        {
            "invoke_specialist": "invoke_specialist",
            "process_response":  "process_response",
        },
    )

    builder.add_edge("invoke_specialist", "process_response")
    builder.add_edge("process_response",  END)

    return builder.compile()


# Singleton — compilado uma vez, reutilizado em todos os requests
supervisor_graph = build_supervisor_graph()

import json, logging
import jwt  # PyJWT
from langchain_core.messages import HumanMessage
from supervisor.graph.builder import supervisor_graph
from supervisor.graph.state import UserContext

logger = logging.getLogger(__name__)


def handler(event: dict, context) -> dict:
    """
    AWS Lambda / ECS handler para requests do API Gateway.
    JWT já validado pelo JWT Authorizer — apenas decodificamos claims.
    """
    try:
        # Extrai claims do JWT (assinatura já validada pelo API GW)
        raw_token = event["headers"].get("Authorization", "").removeprefix("Bearer ")
        claims    = jwt.decode(raw_token, options={"verify_signature": False})

        body       = json.loads(event.get("body", "{}"))
        session_id = body.get("session_id") or event["headers"].get("X-Session-Id")
        thread_id  = body.get("thread_id")  or event["headers"].get("X-Thread-Id")

        if not session_id or not thread_id:
            return _error(400, "session_id e thread_id são obrigatórios")

        # Monta UserContext via allowlist — ponto único de extração do JWT
        user_context = UserContext.from_jwt_claims(claims, session_id, thread_id)

        # Invoca o grafo com state inicial completo
        result = supervisor_graph.invoke({
            "user_context":        user_context,
            "session_id":          session_id,
            "thread_id":           thread_id,
            "messages":            [HumanMessage(content=body["message"])],
            "intent":              None,
            "active_agent":        None,
            "access_denied":       None,
            "specialist_response": None,
        })

        final = result["messages"][-1].content
        return {
            "statusCode": 200,
            "body": json.dumps({
                "message":    final,
                "session_id": session_id,
                "thread_id":  thread_id,
            }),
        }

    except KeyError as e:
        logger.error("missing_field error=%s", e)
        return _error(400, f"Campo obrigatório ausente: {e}")
    except Exception as e:
        logger.error("unhandled_error error=%s", e, exc_info=True)
        return _error(500, "Erro interno")


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