"""Integrity protection for the multi-round-trip `requestState` (MCP 2026-07-28).

The spec requires servers to treat the client-echoed `requestState` as

attacker-controlled: `RequestStateBoundary` seals every outgoing value and

verifies every inbound echo, so handlers only ever see plaintext they minted.

"""

from __future__ import annotations

import base64

import hashlib

import hmac

import json

import logging

import math

import os

import time

from collections.abc import Callable, Mapping, Sequence

from dataclasses import replace

from typing import Any, NoReturn, Protocol, cast

from cryptography.exceptions import InvalidTag

from cryptography.hazmat.primitives.ciphers.aead import AESGCM

from cryptography.hazmat.primitives.hashes import SHA256

from cryptography.hazmat.primitives.kdf.hkdf import HKDF

from mcp_types import INTERNAL_ERROR, INVALID_PARAMS

from mcp_types.methods import INPUT_REQUIRED_METHODS, is_input_required

from mcp.server.auth.middleware.auth_context import get_access_token

from mcp.server.auth.provider import principal_components

from mcp.server.context import CallNext, HandlerResult, ServerRequestContext

from mcp.shared.exceptions import MCPError

__all__ = [

"AESGCMRequestStateCodec",

"InvalidRequestState",

"RequestStateBoundary",

"RequestStateCodec",

"RequestStateSecurity",

"authenticated_principal",

]

logger = logging.getLogger(__name__)

class InvalidRequestState(Exception):

"""A sealed `requestState` token failed verification.

The message is a log-only reason code; the boundary never puts it on the wire.

"""

class RequestStateCodec(Protocol):

"""Authenticated crypto over the framework's request-state envelope.

The framework stamps and re-verifies every envelope claim (expiry, request

binding, principal); a codec only provides integrity and, ideally,

confidentiality (a sign-only codec leaves the payload client-readable).

Requirements: `unseal(seal(payload))` round-trips, and `unseal` raises

`InvalidRequestState` for any token it did not mint unmodified; tokens

never name their algorithm (version with a format prefix bound under the

authentication tag, RFC 8725); comparisons are constant-time. Both methods

are synchronous, so cache key material rather than calling a KMS per token.

"""

def seal(self, payload: bytes) -> str:

"""Return an opaque URL-safe token protecting `payload`."""

...

def unseal(self, token: str) -> bytes:

"""Reverse `seal`.

Raises:

InvalidRequestState: Malformed, unauthentic, or unknown-key token.

"""

...

def authenticated_principal(ctx: ServerRequestContext[Any, Any]) -> str | None:

"""Default principal binding: the authenticated (client, issuer, subject) identity.

Uses the same components session ownership uses, so two users of one OAuth

client are distinct principals whenever the token verifier supplies a

subject, and the binding degrades to the client identity when it does not.

Returns `None` (state not principal-bound) on unauthenticated transports.

"""

token = get_access_token()

if token is None:

return None

return compact_json(principal_components(token))

class RequestStateSecurity:

"""Policy for protecting `requestState`: codec, TTL, principal, audience.

Exactly one of `keys` or `codec`:

RequestStateSecurity(keys=[secret]) # built-in AES-256-GCM

RequestStateSecurity(codec=MyKmsCodec()) # bring your own crypto

RequestStateSecurity.ephemeral() # process-local key

`keys` is the rotation ring: `keys[0]` seals, every key unseals.

Zero-downtime rotation, each phase fully rolled out before the next:

`keys=[old, new]`, then `keys=[new, old]`, then `keys=[new]` after one TTL.

The boundary enforces expiry, request binding, audience, and principal for

every codec, fail-closed in both directions. `audience=None` defers to the

boundary's `default_audience` (`MCPServer` passes its server name).

"""

codec: RequestStateCodec

ttl: float

bind_principal: Callable[[ServerRequestContext[Any, Any]], str | None] | None

audience: str | None

def __init__(

self,

*,

keys: Sequence[bytes | bytearray | str] | None = None,

codec: RequestStateCodec | None = None,

ttl: float = 600.0,

bind_principal: Callable[[ServerRequestContext[Any, Any]], str | None] | None = authenticated_principal,

audience: str | None = None,

) -> None:

if (keys is None) == (codec is None):

raise ValueError("RequestStateSecurity takes exactly one of keys= or codec=")

if not (math.isfinite(ttl) and ttl > 0):

raise ValueError(f"request-state ttl must be a positive finite number, got {ttl!r}")

if keys is not None:

self.codec = AESGCMRequestStateCodec(keys)

else:

assert codec is not None

self.codec = codec

self.ttl = ttl

self.bind_principal = bind_principal

self.audience = audience

@classmethod

def ephemeral(cls, *, ttl: float = 600.0, audience: str | None = None) -> RequestStateSecurity:

"""Protection under a key generated now and held only by this process.

This is the policy `MCPServer` installs when `request_state_security=`

is omitted; call it yourself on the lowlevel tier or to set `ttl`/

`audience`. Suits single-process deployments (stdio, one HTTP worker):

state minted before a restart or by another worker is rejected.

Multi-instance deployments must share a key via `keys=[...]`.

"""

return cls(keys=[os.urandom(32)], ttl=ttl, audience=audience)

_KDF_INFO = b"mcp/request-state/v1/aes-256-gcm"

_KID_INFO = b"mcp/request-state/v1/kid:"

_TOKEN_PREFIX = "v1."

_KID_LEN = 4

_NONCE_LEN = 12

def compact_json(value: Any, *, sort_keys: bool = False) -> str:

"""Canonical JSON for everything the state path digests or seals.

ASCII output keeps the encode total: a lone surrogate in client-supplied

text escapes instead of raising. Anything consuming this must parse with

stdlib `json.loads`, which accepts those escapes (pydantic's JSON parser

does not).

"""

return json.dumps(value, sort_keys=sort_keys, separators=(",", ":"))

def _b64u(data: bytes) -> str:

return base64.urlsafe_b64encode(data).decode().rstrip("=")

def _b64u_decode(text: str) -> bytes:

"""Strict inverse of `_b64u`: only the canonical unpadded encoding decodes."""

raw = base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))

if _b64u(raw) != text:

raise ValueError("non-canonical base64url")

return raw

def _derive_key(secret: bytes) -> bytes:

"""Stretch an operator secret (>= 32 bytes, any format) into the AES-256 key."""

return HKDF(algorithm=SHA256(), length=32, salt=None, info=_KDF_INFO).derive(secret)

class AESGCMRequestStateCodec:

"""Built-in codec: AES-256-GCM under key(s) derived with HKDF-SHA256.

Tokens are encrypted, not merely signed, so clients cannot read the state.

`keys[0]` seals; all keys unseal (rotation, see `RequestStateSecurity`).

Each token carries a 4-byte non-secret key fingerprint for an O(1) ring

lookup, and the "v1." prefix and fingerprint are bound into the GCM

associated data, so a token cannot be replayed into another format version

or ring slot. Key bytes are copied at construction.

"""

def __init__(self, keys: Sequence[bytes | bytearray | str]) -> None:

for i, key in enumerate(cast("Sequence[object]", keys)):

if not isinstance(key, bytes | bytearray | str):

# Never coerce: bytes(32) would silently build an all-zero key.

raise TypeError(

f"request-state keys must be bytes, bytearray, or str; keys[{i}] is {type(key).__name__}"

)

material = [k.encode() if isinstance(k, str) else bytes(k) for k in keys]

if not material:

raise ValueError("AESGCMRequestStateCodec requires at least one key")

for i, k in enumerate(material):

if len(k) < 32:

raise ValueError(

f"request-state keys must be at least 32 bytes of secret randomness; "

f"keys[{i}] is {len(k)} bytes. "

'Generate one with: python -c "import secrets; print(secrets.token_hex(32))"'

)

self._ring: dict[bytes, AESGCM] = {}

self._mint_kid = b""

for i, secret in enumerate(material):

key = _derive_key(secret)

kid = hashlib.sha256(_KID_INFO + key).digest()[:_KID_LEN]

if kid in self._ring:

raise ValueError(f"keys[{i}] duplicates an earlier ring key")

self._ring[kid] = AESGCM(key)

if i == 0:

self._mint_kid = kid

def seal(self, payload: bytes) -> str:

kid = self._mint_kid

nonce = os.urandom(_NONCE_LEN)

sealed = self._ring[kid].encrypt(nonce, payload, _TOKEN_PREFIX.encode() + kid)

return _TOKEN_PREFIX + _b64u(kid + nonce + sealed)

def unseal(self, token: str) -> bytes:

if not token.startswith(_TOKEN_PREFIX):

raise InvalidRequestState("malformed")

try:

raw = _b64u_decode(token[len(_TOKEN_PREFIX) :])

except ValueError as exc:

raise InvalidRequestState("malformed") from exc

if len(raw) < _KID_LEN + _NONCE_LEN + 16:

raise InvalidRequestState("malformed")

kid, nonce, sealed = raw[:_KID_LEN], raw[_KID_LEN : _KID_LEN + _NONCE_LEN], raw[_KID_LEN + _NONCE_LEN :]

aead = self._ring.get(kid)

if aead is None:

raise InvalidRequestState("unknown key")

try:

return aead.decrypt(nonce, sealed, _TOKEN_PREFIX.encode() + kid)

except InvalidTag:

raise InvalidRequestState("seal") from None

# The multi-round-trip carriers: the only methods whose results may carry `requestState`.

_MRTR_METHODS = INPUT_REQUIRED_METHODS

_ENVELOPE_VERSION = 1

_FUTURE_SKEW = 60.0

_PRINCIPAL_LABEL = b"mcp/request-state/principal:"

_RoundBinding = tuple[str, str, str | None]

"""The (target, args-digest, principal) one round's envelope binds, computed once per round."""

def _reject(method: str, reason: str) -> NoReturn:

"""Refuse a round: frozen wire error, real reason to the server log only."""

logger.warning("requestState rejected on %s: %s", method, reason)

raise MCPError(

code=INVALID_PARAMS,

message="Invalid or expired requestState",

data={"reason": "invalid_request_state"},

)

def _request_identity(method: str, params: Mapping[str, Any] | None) -> tuple[str, str]:

"""Salient (target, args-digest) for the request a token binds to.

Per-method allowlist, never a denylist: a future wire field cannot silently join the digest.

"""

p: Mapping[str, Any] = params or {}

args: dict[str, Any] = {}

if method == "resources/read":

target = str(p.get("uri", ""))

else:

target, args = str(p.get("name", "")), p.get("arguments") or args

return target, _b64u(hashlib.sha256(compact_json(args, sort_keys=True).encode()).digest()[:16])

def _principal_claim(principal: str) -> str:

salt = os.urandom(8)

tag = hashlib.sha256(_PRINCIPAL_LABEL + salt + _principal_bytes(principal)).digest()[:16]

return _b64u(salt + tag)

def _principal_matches(claim: str, principal: str) -> bool:

try:

raw = _b64u_decode(claim)

except ValueError:

return False

# A wrong-length claim never matches: compare_digest handles mismatched sizes.

expected = hashlib.sha256(_PRINCIPAL_LABEL + raw[:8] + _principal_bytes(principal)).digest()[:16]

return hmac.compare_digest(raw[8:], expected)

def _principal_bytes(principal: str) -> bytes:

# The digest input is one-way and never decoded, so surrogatepass keeps it total.

return principal.encode("utf-8", "surrogatepass")

def _bound_principal(

security: RequestStateSecurity,

ctx: ServerRequestContext[Any, Any],

fail: Callable[[str], NoReturn],

) -> str | None:

"""Run `bind_principal` under the deny-on-error discipline, in one place for both directions.

`fail` converts a failure into the calling direction's wire shape: the

frozen rejection when verifying, the sanitized internal error when sealing.

"""

try:

principal = security.bind_principal(ctx) if security.bind_principal is not None else None

except Exception: # deny-on-error: a raising principal binding must fail closed

logger.exception("bind_principal raised while processing requestState on %s", ctx.method)

fail("principal binding error")

# The declared return type is str | None, but a user callback can ignore it.

if principal is not None and not isinstance(cast("object", principal), str):

fail(f"bind_principal returned {type(principal).__name__}, expected str or None")

return principal

class RequestStateBoundary:

"""Server middleware sealing/unsealing `requestState` at the wire boundary.

Acts only on the multi-round-trip carriers (tools/call, prompts/get,

resources/read); every other method passes through untouched.

Inbound state is verified (codec unseal plus claims check) and replaced

with the plaintext the server minted before any interceptor or handler

runs; failure answers -32602 with the frozen message "Invalid or expired

requestState", the real reason going to the server log only. Outbound, an

`input_required` result carrying `requestState` is sealed in a fresh

claims envelope; handlers and resolvers never call the codec.

`default_audience` seeds the audience claim when the policy sets none, and

must be stated explicitly: it is the service identity that stops state

minted by another service sharing the same keys. `MCPServer` installs this

middleware with its server name by default (under an ephemeral policy

unless `request_state_security=` supplies one); lowlevel `Server` users

append one to `server.middleware`, passing their server's name (or `None`

to deliberately leave tokens audience-free).

"""

def __init__(self, security: RequestStateSecurity, *, default_audience: str | None) -> None:

self._security = security

self._audience = security.audience if security.audience is not None else default_audience

async def __call__(self, ctx: ServerRequestContext[Any, Any], call_next: CallNext) -> HandlerResult:

if ctx.method not in _MRTR_METHODS:

return await call_next(ctx)

binding: _RoundBinding | None = None

if ctx.params is not None and ctx.params.get("requestState") is not None:

# An explicit JSON null counts as absent: stripping the field is already in any client's power.

plaintext, binding = self._unseal(ctx)

ctx = replace(ctx, params={**ctx.params, "requestState": plaintext})

result = await call_next(ctx)

return self._seal_result(ctx, result, binding)

def _unseal(self, ctx: ServerRequestContext[Any, Any]) -> tuple[str, _RoundBinding]:

assert ctx.params is not None

wire = ctx.params["requestState"]

if not isinstance(wire, str):

_reject(ctx.method, "non-string requestState")

security = self._security

try:

payload = security.codec.unseal(wire)

except InvalidRequestState as exc:

_reject(ctx.method, str(exc))

except Exception: # deny-on-error: a buggy custom codec must fail closed

logger.exception("requestState codec raised during unseal on %s", ctx.method)

_reject(ctx.method, "codec error")

try:

claims = json.loads(payload)

version, iat, exp, inner = claims["v"], claims["iat"], claims["exp"], claims["s"]

except (ValueError, KeyError, TypeError):

_reject(ctx.method, "malformed")

if version != _ENVELOPE_VERSION or not isinstance(inner, str):

_reject(ctx.method, "malformed")

now = time.time()

# Accept-conditions are stated positively so a NaN claim fails the comparison and rejects.

if not isinstance(iat, int | float) or not (iat <= now + _FUTURE_SKEW):

_reject(ctx.method, "minted in the future")

if not isinstance(exp, int | float) or not (now < exp):

_reject(ctx.method, "expired")

target, args_digest = _request_identity(ctx.method, ctx.params)

if claims.get("m") != ctx.method or claims.get("t") != target or claims.get("a") != args_digest:

_reject(ctx.method, "request binding")

if claims.get("aud") != self._audience:

_reject(ctx.method, "audience")

def fail_verify(reason: str) -> NoReturn:

_reject(ctx.method, reason)

principal = _bound_principal(security, ctx, fail_verify)

claim = claims.get("p")

if (claim is None) != (principal is None):

_reject(ctx.method, "principal drift")

if claim is not None and principal is not None:

if not isinstance(claim, str) or not _principal_matches(claim, principal):

_reject(ctx.method, "principal")

return inner, (target, args_digest, principal)

def _seal_result(

self, ctx: ServerRequestContext[Any, Any], result: HandlerResult, binding: _RoundBinding | None

) -> HandlerResult:

# Spec-path results arrive as wire mappings; a short-circuiting middleware may return a model.

if not is_input_required(result):

return result

state = result.get("requestState") if isinstance(result, Mapping) else result.request_state

if state is None:

return result

if isinstance(result, Mapping):

if not isinstance(state, str):

# Only a short-circuiting middleware can put a non-string here; nothing to seal.

return result

return {**result, "requestState": self._seal(ctx, state, binding)}

return result.model_copy(update={"request_state": self._seal(ctx, state, binding)})

def _seal(self, ctx: ServerRequestContext[Any, Any], state: str, binding: _RoundBinding | None = None) -> str:

security = self._security

if binding is None:

def fail_seal(reason: str) -> NoReturn:

logger.error("refusing to seal requestState on %s: %s", ctx.method, reason)

raise MCPError(code=INTERNAL_ERROR, message="Internal error")

target, args_digest = _request_identity(ctx.method, ctx.params)

binding = (target, args_digest, _bound_principal(security, ctx, fail_seal))

target, args_digest, principal = binding

now = time.time()

claims: dict[str, Any] = {

"v": _ENVELOPE_VERSION,

"iat": now,

"exp": now + security.ttl,

"m": ctx.method,

"t": target,

"a": args_digest,

"s": state,

}

if self._audience is not None:

claims["aud"] = self._audience

if principal is not None:

claims["p"] = _principal_claim(principal)

payload = compact_json(claims).encode()

try:

return security.codec.seal(payload)

except Exception: # deny-on-error: a raising custom codec must not leak its failure

logger.exception("requestState codec raised during seal on %s", ctx.method)

raise MCPError(code=INTERNAL_ERROR, message="Internal error") from None