API Reference¶

@dataclass(frozen=True)
class AgentClaims:
    """Immutable representation of the agentic claims carried in the
    `X-Agent-Authorization` JWT.

    Required claims (the issuer must populate these):
      - `sub`: agent identifier (e.g. "purchaser-agent-instance-42")
      - `task_type`: free-form ("interactive_purchase", "scheduled_audit")
      - `user_id`: the human in whose name the agent acts
      - `scopes`: list of scope strings the agent is allowed to invoke
                  (e.g. "treasury:write:budget:opex")

    Optional but recommended:
      - `session_id`, `conversation_id`: the conversational context
      - `task_id`: per-task identifier (for spend-limit accounting)
      - `org_id`: the customer's tenant
      - `role`: an agent role (separate from the service role)
      - `spend_limit`: decimal string — total budget for this task
      - `parent_chain`: tuple of mcp-ids that delegated this token down
                       (innermost first: ["io.acme.head-agent",
                       "io.acme.purchaser"])

    `raw_claims` keeps the original dict so the consumer can read
    custom fields without a round-trip through this dataclass."""

    sub: str = ""
    task_type: str = ""
    user_id: str = ""
    session_id: str = ""
    conversation_id: str = ""
    task_id: str = ""
    org_id: str = ""
    role: str = ""
    spend_limit: str = ""
    scopes: tuple[str, ...] = ()
    parent_chain: tuple[str, ...] = ()
    raw_claims: dict[str, Any] | None = None

    @classmethod
    def from_claims_dict(cls, raw: dict[str, Any]) -> "AgentClaims":
        """Build an `AgentClaims` from a JWT payload dict."""
        return cls(
            sub=str(raw.get("sub", "")),
            task_type=str(raw.get("task_type", "")),
            user_id=str(raw.get("user_id", "")),
            session_id=str(raw.get("session_id", "")),
            conversation_id=str(raw.get("conversation_id", "")),
            task_id=str(raw.get("task_id", "")),
            org_id=str(raw.get("org_id", "")),
            role=str(raw.get("role", "")),
            spend_limit=str(raw.get("spend_limit", "")),
            scopes=tuple(raw.get("scopes") or []),
            parent_chain=tuple(raw.get("parent_chain") or []),
            raw_claims=dict(raw),
        )

class AgentContext:
    """Singleton-style accessor for the per-request agent identity.

    Mirrors `SecurityContext` API. All methods are static; storage uses
    a `ContextVar` so each async task / thread has its own copy."""

    @staticmethod
    def get() -> Optional[AgentClaims]:
        return _agent_var.get()

    @staticmethod
    def is_present() -> bool:
        return _agent_var.get() is not None

    @staticmethod
    def get_scopes() -> tuple[str, ...]:
        agent = _agent_var.get()
        return agent.scopes if agent else ()

    @staticmethod
    def has_scope(scope: str) -> bool:
        return scope in AgentContext.get_scopes()

    @staticmethod
    def set(agent: AgentClaims) -> None:
        _agent_var.set(agent)

    @staticmethod
    def clear() -> None:
        _agent_var.set(None)

@configured(target="self", prefix="auth_client", mapping="tree")
@dataclass
class AuthClientSettings:
    """Type-safe settings for the auth client, loaded from configuration sources.

    Populated automatically from configuration sources using the ``auth_client``
    prefix via pico-ioc's ``@configured`` decorator.

    Attributes:
        enabled: Whether authentication middleware is active.
        issuer: Expected JWT issuer (``iss`` claim).
        audience: Expected JWT audience (``aud`` claim).
        jwks_ttl_seconds: How long to cache the JWKS key set (seconds).
        jwks_endpoint: URL to fetch JWKS from. Defaults to ``{issuer}/api/v1/auth/jwks``.
    """

    enabled: bool = True
    issuer: str = ""
    audience: str = ""
    jwks_ttl_seconds: int = 300
    jwks_endpoint: str = ""
    accepted_algorithms: tuple[str, ...] = ("RS256",)
    # ── Revocation denylist (jti) ────────────────────────────────
    # Endpoint the validator polls to refresh its local cache of
    # revoked JWT IDs. Empty disables the check entirely (back to
    # signature-only validation — safe default for setups that
    # haven't wired the issuer's revoke endpoint).
    # ``revocation_ttl_seconds`` is the worst-case window between
    # an operator clicking Revoke and validators actually
    # rejecting the token. Lower = snappier, higher = fewer
    # round-trips. JWKS rotation remains the instant-kill path.
    revocation_endpoint: str = ""
    revocation_ttl_seconds: int = 15
    # SECURITY: on a revocation-fetch error, fail closed by default — treat the
    # token as unable-to-confirm (revoked) rather than serving a stale denylist
    # (fail-open). Set True to restore the old fail-open behaviour if availability
    # matters more than promptly honouring revocations.
    revocation_fail_open: bool = False
    # If the revocation endpoint requires a Bearer token (the
    # default — pico-server-auth gates it behind role=service),
    # the auth-client uses this token for the poll. Empty falls
    # back to anonymous (works for dev / unauthenticated setups).
    revocation_bearer: str = ""

class AuthClientError(Exception):
    """Base error for all pico-client-auth exceptions."""

class AuthConfigurationError(AuthClientError):
    """Authentication configuration is missing or invalid."""

class InsufficientPermissionsError(AuthClientError):
    """The authenticated user lacks the required role(s)."""

class MissingTokenError(AuthClientError):
    """No Bearer token found in the Authorization header."""

class TokenExpiredError(AuthClientError):
    """The JWT token has expired."""

class TokenInvalidError(AuthClientError):
    """The JWT token is malformed or has an invalid signature."""

@component
class JWKSClient:
    """Fetches and caches the JSON Web Key Set from the auth server.

    Supports automatic cache refresh when a key ID (``kid``) is not found
    (handles key rotation) and TTL-based expiration.
    """

    def __init__(self, settings: AuthClientSettings):
        self._settings = settings
        self._keys: dict = {}
        self._fetched_at: float = 0.0
        self._endpoint = settings.jwks_endpoint or f"{settings.issuer.rstrip('/')}/api/v1/auth/jwks"

    def _is_expired(self) -> bool:
        return (time.monotonic() - self._fetched_at) >= self._settings.jwks_ttl_seconds

    async def _fetch_keys(self) -> None:
        logger.debug("Fetching JWKS from %s", self._endpoint)
        # SECURITY: require TLS (https) for the JWKS endpoint and bound the
        # request with a timeout so a hung server can't stall validation.
        _require_https(self._endpoint)
        async with httpx.AsyncClient(timeout=5.0) as client:
            response = await client.get(self._endpoint)
            response.raise_for_status()
            data = response.json()
        self._keys = {k["kid"]: k for k in data.get("keys", [])}
        self._fetched_at = time.monotonic()

    async def get_key(self, kid: str) -> dict:
        """Return the JWK for the given key ID, fetching/refreshing as needed.

        Args:
            kid: The ``kid`` header value from the JWT.

        Returns:
            The JWK dict matching the requested key ID.

        Raises:
            KeyError: If the key ID is not found even after a refresh.
        """
        if self._is_expired() or not self._keys:
            await self._fetch_keys()

        if kid in self._keys:
            return self._keys[kid]

        # Key not found — might be key rotation, force refresh once
        await self._fetch_keys()
        if kid in self._keys:
            return self._keys[kid]

        raise KeyError(f"Key ID '{kid}' not found in JWKS")

@dataclass(frozen=True)
class TokenClaims:
    """Immutable representation of the essential JWT claims.

    Attributes:
        sub: Subject identifier (user ID).
        email: User email address.
        role: Primary role claim from the token.
        org_id: Organisation identifier.
        jti: Unique token identifier (JWT ID).
    """

    sub: str
    email: str
    role: str
    org_id: str
    jti: str
    groups: tuple[str, ...] = ()

@runtime_checkable
class RoleResolver(Protocol):
    """Protocol for resolving user roles from JWT claims.

    Implement this protocol to customise how roles are extracted
    (e.g. from a ``roles`` array claim, from an external service, etc.).
    Register your implementation as a ``@component`` and it will
    automatically replace the default resolver.
    """

    async def resolve(self, claims: TokenClaims, raw_claims: dict) -> list[str]: ...

class SecurityContext:
    """Singleton-style accessor for the current request's authentication state.

    All methods are static; the underlying storage uses ``ContextVar`` so
    each async task / thread has its own isolated copy.
    """

    @staticmethod
    def get() -> Optional[TokenClaims]:
        """Return the current claims, or ``None`` if not authenticated."""
        return _claims_var.get()

    @staticmethod
    def require() -> TokenClaims:
        """Return the current claims, raising if not authenticated.

        Raises:
            MissingTokenError: If no authentication context is set.
        """
        claims = _claims_var.get()
        if claims is None:
            raise MissingTokenError("No authenticated user in SecurityContext")
        return claims

    @staticmethod
    def get_roles() -> list[str]:
        """Return the resolved roles for the current request."""
        return list(_roles_var.get())

    @staticmethod
    def has_role(role: str) -> bool:
        """Check whether the current user has the given role."""
        return role in _roles_var.get()

    @staticmethod
    def require_role(*roles: str) -> None:
        """Assert that the current user has at least one of the given roles.

        Raises:
            InsufficientPermissionsError: If none of the roles match.
        """
        current = set(_roles_var.get())
        if not current.intersection(roles):
            raise InsufficientPermissionsError(f"Required one of {roles}, but user has {sorted(current)}")

    @staticmethod
    def get_groups() -> tuple[str, ...]:
        """Return the group IDs for the current request."""
        return _groups_var.get()

    @staticmethod
    def has_group(group_id: str) -> bool:
        """Check whether the current user belongs to the given group."""
        return group_id in _groups_var.get()

    @staticmethod
    def require_group(*group_ids: str) -> None:
        """Assert that the current user belongs to at least one of the given groups.

        Raises:
            InsufficientPermissionsError: If none of the groups match.
        """
        current = set(_groups_var.get())
        if not current.intersection(group_ids):
            raise InsufficientPermissionsError(f"Required one of groups {group_ids}, but user has {sorted(current)}")

    @staticmethod
    def set(claims: TokenClaims, roles: list[str]) -> None:
        """Populate the security context (called by the middleware)."""
        _claims_var.set(claims)
        _roles_var.set(roles)
        _groups_var.set(claims.groups)

    @staticmethod
    def clear() -> None:
        """Clear the security context (called by the middleware in ``finally``)."""
        _claims_var.set(None)
        _roles_var.set([])
        _groups_var.set(())