inbound

def encode_header_value(value: str) -> str:
    """Wrap `value` in the `=?base64?...?=` sentinel when it would not survive an HTTP field round-trip.

    Plain printable ASCII without leading/trailing whitespace passes verbatim;
    anything else (control chars, non-ASCII, edge whitespace, or a value that
    already looks like the sentinel) is base64-wrapped so the receiver can
    recover the exact bytes.
    """
    if _HEADER_SAFE.fullmatch(value) and value == value.strip() and not _B64_SENTINEL.fullmatch(value):
        return value
    return f"=?base64?{base64.b64encode(value.encode('utf-8')).decode('ascii')}?="

def decode_header_value(value: str | None) -> str | None:
    """Inverse of :func:`encode_header_value`.

    Returns the value verbatim unless it carries the `=?base64?...?=` sentinel,
    in which case the payload is decoded as UTF-8. A malformed sentinel (bad
    base64, non-canonical base64, or bad UTF-8) yields `None` so a corrupt
    header never matches a body value by accident. `None` in → `None` out so
    callers can pass `headers.get(...)` directly.
    """
    if value is None:
        return None
    m = _B64_SENTINEL.fullmatch(value)
    if m is None:
        return value
    payload = m.group("payload")
    try:
        decoded = base64.b64decode(payload, validate=True)
    except binascii.Error:
        return None
    # Reject non-canonical base64 (e.g. non-zero trailing bits), which
    # `validate=True` tolerates; the encoder only ever emits canonical form.
    if base64.b64encode(decoded).decode("ascii") != payload:
        return None
    try:
        return decoded.decode("utf-8")
    except UnicodeDecodeError:
        return None

def find_invalid_x_mcp_header(input_schema: Any) -> str | None:
    """Return a reason string if any `x-mcp-header` annotation in `input_schema` is invalid; else `None`.

    Walks every JSON Schema 2020-12 schema position. An annotation is valid
    only when it sits on a property statically reachable from the root via a
    chain of pure `properties` keys, names a non-empty RFC 9110 token, is on
    an integer/string/boolean property, and is case-insensitively unique
    across the whole schema. A `None` / non-mapping schema has no schema
    positions and returns `None`.
    """
    seen: dict[str, str] = {}
    for path, schema in _walk_schema_positions(input_schema):
        if X_MCP_HEADER_KEY not in schema:
            continue
        if not path:  # None (off the pure-properties chain) or () (the root itself)
            return f"{X_MCP_HEADER_KEY} found at a schema position not reachable via a pure `properties` chain"
        where = ".".join(path)
        header = schema[X_MCP_HEADER_KEY]
        # Wrong type and malformed value are distinct failures with distinct messages: the
        # non-str arm returns before any interpolation, because `repr` of an arbitrary
        # schema value is not total (a large `int` exceeds `sys.get_int_max_str_digits`).
        if not isinstance(header, str):
            return f"property {where!r}: {X_MCP_HEADER_KEY} must be a string, not {type(header).__name__}"
        if not _RFC9110_TOKEN.fullmatch(header):
            return f"property {where!r}: {X_MCP_HEADER_KEY} {header!r} is not an RFC 9110 token"
        prop_type = schema.get("type")
        if not isinstance(prop_type, str):
            return (
                f"property {where!r}: {X_MCP_HEADER_KEY} is only permitted on "
                f"integer/string/boolean properties (the type keyword is {type(prop_type).__name__}, not a string)"
            )
        if prop_type not in _X_MCP_HEADER_PRIMITIVE_TYPES:
            return (
                f"property {where!r}: {X_MCP_HEADER_KEY} is only permitted on "
                f"integer/string/boolean properties (got {prop_type!r})"
            )
        lower = header.lower()
        if lower in seen:
            return f"{X_MCP_HEADER_KEY} {header!r} on property {where!r} duplicates property {seen[lower]!r}"
        seen[lower] = where
    return None

def x_mcp_header_map(input_schema: Any) -> dict[tuple[str, ...], str]:
    """Map each property carrying a valid `x-mcp-header` to its annotation token, keyed by property path.

    The key is the chain of `properties` keys from the schema root to the
    annotated property; a top-level property has a one-element path, a nested
    one a longer path. Call only on a schema that
    :func:`find_invalid_x_mcp_header` accepts; an invalid schema yields an
    undefined subset.
    """
    return {path: token for path, token, _ in _annotated_positions(input_schema)}

def mcp_param_headers(header_map: Mapping[tuple[str, ...], str], arguments: Mapping[str, Any]) -> dict[str, str]:
    """Build the `Mcp-Param-*` headers a `tools/call` mirrors from its arguments.

    For each `(path, token)` in `header_map`, read the value at that property
    path in `arguments` and, when it is present and not `None`, emit
    `Mcp-Param-<token>` carrying it: `bool` as `true`/`false`, other scalars via
    `str`, each passed through :func:`encode_header_value` so a non-token value
    is base64-wrapped. A path that hits a missing key or a non-mapping node is
    skipped, matching the spec's "omit the header when no value is present",
    as is a value with no header rendering.
    """
    headers: dict[str, str] = {}
    for path, token in header_map.items():
        value = _value_at_path(arguments, path)
        if value is None or (rendered := _render_header_scalar(value)) is None:
            continue
        headers[f"{MCP_PARAM_HEADER_PREFIX}{token}"] = encode_header_value(rendered)
    return headers

def has_envelope_intent(params: Mapping[str, Any] | None) -> bool:
    """Whether `params._meta` carries any modern per-request envelope key: presence,
    not validity (a bare `progressToken` is not envelope evidence)."""
    if not params:
        return False
    meta = params.get("_meta")
    return isinstance(meta, Mapping) and not _ENVELOPE_KEYS.isdisjoint(cast("Mapping[str, Any]", meta))

def parse_envelope(params: Mapping[str, Any] | None) -> ModernEnvelope | InboundLadderRejection | None:
    """Parse the per-request modern envelope out of `params._meta`.

    `None` when `_meta` carries no envelope key; an INVALID_PARAMS rejection
    naming the missing key(s) when a claim is present but the required pair
    (protocol version, client capabilities) is incomplete; else the envelope
    (unchecked). Client info is optional (SHOULD-include, spec PR #3002), so
    absent reads as `None`.
    """
    if not has_envelope_intent(params):
        return None
    # `has_envelope_intent` proved `params._meta` is a mapping carrying a claim.
    meta = cast("Mapping[str, Any]", (params or {})["_meta"])
    if missing := [key for key in (PROTOCOL_VERSION_META_KEY, CLIENT_CAPABILITIES_META_KEY) if key not in meta]:
        return InboundLadderRejection(
            code=INVALID_PARAMS,
            message=f"params._meta is missing the required envelope key(s): {', '.join(missing)}",
        )
    return ModernEnvelope(
        protocol_version=meta[PROTOCOL_VERSION_META_KEY],
        client_info=meta.get(CLIENT_INFO_META_KEY),
        client_capabilities=meta[CLIENT_CAPABILITIES_META_KEY],
    )

def check_supported_version(
    envelope: ModernEnvelope, supported_modern_versions: Sequence[str] = MODERN_PROTOCOL_VERSIONS
) -> InboundModernRoute | InboundLadderRejection:
    """Promote a parsed envelope to a route: non-string version → INVALID_PARAMS;
    unserved → UNSUPPORTED_PROTOCOL_VERSION with `{"supported": [...], "requested": ...}`."""
    if not isinstance(envelope.protocol_version, str):
        return InboundLadderRejection(
            code=INVALID_PARAMS, message="the protocol-version envelope value must be a string"
        )
    if envelope.protocol_version not in supported_modern_versions:
        return InboundLadderRejection(
            code=UNSUPPORTED_PROTOCOL_VERSION,
            message="Unsupported protocol version",
            data=UnsupportedProtocolVersionErrorData(
                supported=list(supported_modern_versions), requested=envelope.protocol_version
            ).model_dump(mode="json"),
        )
    return InboundModernRoute(
        protocol_version=envelope.protocol_version,
        client_info=envelope.client_info,
        client_capabilities=envelope.client_capabilities,
    )

def find_duplicated_routing_header(headers: Iterable[tuple[str, str]]) -> str | None:
    """Name of a routing header supplied more than once in raw header lines, or `None`.

    Takes raw `(name, value)` pairs — a folded mapping hides duplicates. A
    duplicate is rejected because first-copy and last-copy readers would
    disagree. `Mcp-Param-*` duplicates are :func:`validate_mcp_param_headers`'s job.
    """
    seen: set[str] = set()
    for name, _ in headers:
        key = name.lower()
        if key in _ROUTING_HEADER_NAMES:
            if key in seen:
                return key
            seen.add(key)
    return None

def classify_inbound_request(
    body: Mapping[str, Any],
    *,
    headers: Mapping[str, str] | None = None,
    supported_modern_versions: Sequence[str] = MODERN_PROTOCOL_VERSIONS,
) -> InboundModernRoute | InboundLadderRejection:
    """Run the modern-protocol validation ladder over a decoded JSON-RPC body.

    Rungs, in order; first failure wins:

    0. `body.method` is not `initialize` → else
       :data:`~mcp_types.jsonrpc.UNSUPPORTED_PROTOCOL_VERSION` naming the served
       versions (:func:`_handshake_refusal`).
    1. :func:`parse_envelope`: `params._meta` is a mapping carrying the required
       envelope pair (protocol version, client capabilities) → else
       :data:`~mcp_types.jsonrpc.INVALID_PARAMS`, naming the missing key(s)
       (basic/index.mdx "Per-request protocol fields"). Client info is
       optional (SHOULD-include, spec PR #3002); absent reads as `None`.
    2. When `headers` is given: `MCP-Protocol-Version`, `Mcp-Method`, and (for
       :data:`NAME_BEARING_METHODS`) `Mcp-Name` agree with the body → else
       :data:`~mcp_types.jsonrpc.HEADER_MISMATCH`. Runs before the
       supported-version rung so a client that disagrees with itself is told
       so, rather than told the body's version is unsupported.
    3. :func:`check_supported_version` → non-string values are
       :data:`~mcp_types.jsonrpc.INVALID_PARAMS` (a shape defect, not a
       negotiation outcome), else
       :data:`~mcp_types.jsonrpc.UNSUPPORTED_PROTOCOL_VERSION` with
       `data = {"supported": [...], "requested": <value>}`.

    Method existence is *not* a rung: kernel dispatch owns that decision so
    custom-registered methods route and the answer lives in one place.

    Args:
        body: The decoded JSON-RPC request mapping.
        headers: Transport headers keyed by lowercase name, or `None` to
            skip the header rung.
        supported_modern_versions: Modern protocol revisions this server
            accepts on the per-request-envelope path.
    """
    raw_params = body.get("params")
    params = cast("Mapping[str, Any]", raw_params) if isinstance(raw_params, Mapping) else None
    if body.get("method") == "initialize":
        return _handshake_refusal(params, supported_modern_versions)
    envelope = parse_envelope(params)
    if envelope is None:
        return InboundLadderRejection(code=INVALID_PARAMS, message=_ENVELOPE_REQUIRED_MESSAGE)
    if isinstance(envelope, InboundLadderRejection):
        return envelope
    if headers is not None:
        version_header = headers.get(MCP_PROTOCOL_VERSION_HEADER)
        # Presence is checked explicitly: a null body version would otherwise
        # slip the equality check (None == None) and mask the absent header.
        if version_header is None or version_header != envelope.protocol_version:
            return InboundLadderRejection(
                code=HEADER_MISMATCH,
                message=f"{MCP_PROTOCOL_VERSION_HEADER} header does not match the request envelope's protocol version",
            )
        method: Any = body.get("method")
        if headers.get(MCP_METHOD_HEADER) != method:
            return InboundLadderRejection(
                code=HEADER_MISMATCH,
                message=f"{MCP_METHOD_HEADER} header does not match the request body's method",
            )
        name_key = NAME_BEARING_METHODS.get(method)
        if name_key is not None and params is not None:
            body_value = params.get(name_key)
            if body_value is not None and decode_header_value(headers.get(MCP_NAME_HEADER)) != body_value:
                return InboundLadderRejection(
                    code=HEADER_MISMATCH,
                    message=f"{MCP_NAME_HEADER} header does not match the request body's {name_key!r} parameter",
                )
    return check_supported_version(envelope, supported_modern_versions)

def validate_mcp_param_headers(
    input_schema: Any,
    arguments: Mapping[str, Any],
    headers: Mapping[str, str],
) -> InboundLadderRejection | None:
    """Compare a `tools/call` request's `Mcp-Param-*` headers against its body arguments.

    Each annotated property's header and argument must agree: present together
    and equal after sentinel decoding, or absent together (`null` counts as
    absent). Returns the first failure as a `HEADER_MISMATCH` rejection, else `None`.

    A header whose argument is absent or unrenderable is deliberately rejected:
    the spec's purpose clause is exactly an intermediary routing on a value the
    body never carried. A duplicated recognized header is rejected — first-copy
    and last-copy readers would disagree. A schema :func:`find_invalid_x_mcp_header`
    rejects validates nothing: conforming clients drop the tool and emit no headers.
    """
    if find_invalid_x_mcp_header(input_schema) is not None:
        return None
    folded: dict[str, str] = {}
    duplicated: set[str] = set()
    for name, value in headers.items():
        key = name.lower()
        if key in folded:
            duplicated.add(key)
        folded[key] = value
    for path, token, schema in _annotated_positions(input_schema):
        header_name = f"{MCP_PARAM_HEADER_PREFIX}{token}"
        key = header_name.lower()
        raw = folded.get(key)
        value = _value_at_path(arguments, path)
        argument = ".".join(path)
        if raw is not None and key in duplicated:
            return InboundLadderRejection(
                code=HEADER_MISMATCH,
                message=f"{header_name} header appears more than once",
            )
        if value is None:
            if raw is not None:
                return InboundLadderRejection(
                    code=HEADER_MISMATCH,
                    message=f"{header_name} header is present but the request body's {argument!r} argument is absent",
                )
            continue
        rendered = _render_header_scalar(value)
        if rendered is None:
            # Unrenderable value: a conforming client omitted the header, so one claiming it can never match.
            if raw is not None:
                return InboundLadderRejection(
                    code=HEADER_MISMATCH,
                    message=f"{header_name} header does not match the request body's {argument!r} argument",
                )
            continue
        if raw is None:
            return InboundLadderRejection(
                code=HEADER_MISMATCH,
                message=f"{header_name} header is missing but the request body's {argument!r} argument is present",
            )
        decoded = decode_header_value(raw)
        if decoded is None:
            return InboundLadderRejection(
                code=HEADER_MISMATCH,
                message=f"{header_name} header carries a malformed base64 sentinel value",
            )
        if not _mcp_param_value_matches(schema.get("type"), value, rendered, decoded):
            return InboundLadderRejection(
                code=HEADER_MISMATCH,
                message=f"{header_name} header does not match the request body's {argument!r} argument",
            )
    return None