diff --git a/README.md b/README.md index c57ce55..1af1034 100644 --- a/README.md +++ b/README.md @@ -226,6 +226,8 @@ The SDK reads standard settings from the environment or `.env` files: | Environment Variable | Description | |---|---| +| `HOST` | Bind address for HTTP/SSE (default: `127.0.0.1`). Set `0.0.0.0` to listen on all interfaces (containers). | +| `TRUSTED_PROXIES` / `MCP_TRUSTED_PROXIES` | Comma-separated IPs, CIDRs, or hostnames allowed to send `X-Forwarded-Host` / `X-Forwarded-Proto`. Unset means those headers are ignored. `X-Forwarded-For` is never used to decide trust. | | `PORT` / `MCP_SERVER_PORT` | The port to bind for HTTP/SSE transport (default: `3000`). Overridden by `nitrostack-py --port`. | | `WIDGETS_DEV_PORT` | Widget Next.js port (default: `3001`). Overridden by `nitrostack-py --widget`. | | `MCP_TRANSPORT_TYPE` | Transport selection: `stdio`, `http`, or `dual` (combining stdio + HTTP/SSE). | @@ -233,7 +235,8 @@ The SDK reads standard settings from the environment or `.env` files: | `MCP_MAX_SESSIONS` | Cap on concurrent Streamable HTTP sessions; new sessions beyond the cap get an HTTP `429`. Unset = unlimited. | | `MCP_SESSION_TIMEOUT_MS` | Idle timeout (ms) for stateful HTTP sessions; sessions with no activity for this long are terminated automatically. Unset = no timeout. | | `MCP_GRACEFUL_SHUTDOWN_TIMEOUT_MS` | How long (ms) the HTTP transport waits for in-flight requests to finish when shutting down (default: `10000`). | -| `MCP_STATELESS` | Set to `true` to run the HTTP transport in stateless mode: every request gets a fresh context with no session id and no `initialize` handshake required. | +| `NITRO_MCP_PROTOCOL_VERSION` | Protocol era (case-insensitive): `auto` / `both` / `dual` / `dual-spec` (default when unset or unknown), `modern` / `latest` / `2026` / `2026-07-28`, or `legacy` / `2025` / `2025-06-18` / `2025-11-25`. `auto` is not the same as `modern`. Wins over `ServerConfig.protocol_era`. | +| `MCP_STATELESS` | Explicit override: `true` forces `modern` (stateless HTTP), `false` forces `legacy` (sessionful). Wins over `NITRO_MCP_PROTOCOL_VERSION` and `ServerConfig.protocol_era`. | | `MCP_ALLOWED_HOSTS` / `MCP_ALLOWED_ORIGINS` | Comma-separated allow-lists for DNS-rebinding protection, used only when CORS is disabled. | | `NITROSTACK_LOG_FILE` | Destination file for logs (default: `nitrostack.log`). | | `NITROSTACK_LOG_LEVEL` | Log level (`DEBUG`, `INFO`, `WARNING`, `ERROR`). | @@ -260,6 +263,8 @@ Example: server = ServerConfig(name="my-server", transport_type="http", max_sessions=100, session_timeout_ms=1_800_000) ``` +`ServerConfig.protocol_era` is used only when `MCP_STATELESS` and `NITRO_MCP_PROTOCOL_VERSION` are both unset. Unknown tokens become `auto`, same as an unknown env value. + --- ## Developing & Testing diff --git a/nitrostack/__init__.py b/nitrostack/__init__.py index 37df45d..9f1c3d5 100644 --- a/nitrostack/__init__.py +++ b/nitrostack/__init__.py @@ -84,6 +84,17 @@ OAuthService, generate_www_authenticate_header, ) +from nitrostack.auth.cimd import ( + CimdFetchError, + CimdValidationError, + is_blocked_ip, + resolve_cimd, + validate_client_identifier_url, +) +from nitrostack.auth.oauth_security import ( + AuthorizationIssuerMismatchError, + validate_authorization_iss, +) from nitrostack.auth.pkce import ( generate_code_challenge, generate_code_verifier, @@ -118,9 +129,13 @@ from nitrostack.testing import ( NitroTestingModule, ) -from nitrostack.testing import ( - NitroTestingModule, -) +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.protocol.errors import JsonRpcErrorCode +from nitrostack.protocol.mrtr import InputRequest, InputRequiredResult, accepted_content, input_required +from nitrostack.tasks import InMemoryTaskStore, TaskAccessContext, TaskStore +from nitrostack.tasks.authorization import check_task_access, extract_task_access_context +from nitrostack.runtime import StatelessInvariants, assert_stateless_headers +from nitrostack.transports import wrap_stateless_transport, StatelessIngressPipeline __all__ = [ @@ -165,6 +180,13 @@ "OAuthModule", "OAuthService", "generate_www_authenticate_header", + "validate_client_identifier_url", + "resolve_cimd", + "is_blocked_ip", + "CimdValidationError", + "CimdFetchError", + "validate_authorization_iss", + "AuthorizationIssuerMismatchError", "generate_code_challenge", "generate_code_verifier", "generate_pkce_params", @@ -207,4 +229,19 @@ "RESOURCE_MIME_TYPE_MCP_APP", "RESOURCE_MIME_TYPE_OPENAI", "NitroTestingModule", + "MODERN_PROTOCOL_VERSION", + "JsonRpcErrorCode", + "InputRequest", + "InputRequiredResult", + "accepted_content", + "input_required", + "TaskAccessContext", + "TaskStore", + "InMemoryTaskStore", + "check_task_access", + "extract_task_access_context", + "StatelessInvariants", + "assert_stateless_headers", + "wrap_stateless_transport", + "StatelessIngressPipeline", ] diff --git a/nitrostack/auth/cimd.py b/nitrostack/auth/cimd.py new file mode 100644 index 0000000..f9aa7a3 --- /dev/null +++ b/nitrostack/auth/cimd.py @@ -0,0 +1,392 @@ +"""Client ID Metadata Document (CIMD) resolution with SSRF defenses.""" + +from __future__ import annotations + +import asyncio +import http.client +import ipaddress +import json +import socket +import ssl +import urllib.error +import urllib.parse +import urllib.request +from typing import Any, Optional + +from nitrostack.protocol.constants import MAX_CIMD_BYTES + +CIMD_FETCH_TIMEOUT_SEC = 5.0 +_LOOPBACK_HOSTS = frozenset({"localhost", "127.0.0.1", "::1"}) + +_BLOCKED_IPV4_NETWORKS = tuple( + ipaddress.ip_network(cidr) + for cidr in ( + "0.0.0.0/8", + "10.0.0.0/8", + "100.64.0.0/10", + "127.0.0.0/8", + "169.254.0.0/16", + "172.16.0.0/12", + "192.0.0.0/24", + "192.168.0.0/16", + "198.18.0.0/15", + "198.51.100.0/24", + "203.0.113.0/24", + "224.0.0.0/4", + "240.0.0.0/4", + ) +) + +_BLOCKED_IPV6_NETWORKS = tuple( + ipaddress.ip_network(cidr) + for cidr in ( + "::1/128", + "fc00::/7", + "fe80::/10", + "ff00::/8", + "100::/64", + "2001:db8::/32", + ) +) + + +class CimdValidationError(ValueError): + """Raised when a client identifier URL or metadata document is invalid.""" + + +class CimdFetchError(ValueError): + """Raised when a CIMD document cannot be fetched safely.""" + + +def validate_client_identifier_url(client_id_url: str, *, allow_loopback: bool = False) -> str: + """ + Validate a CIMD ``client_id`` URL before any network fetch. + + Returns the normalized URL string on success. + """ + if not client_id_url or not isinstance(client_id_url, str): + raise CimdValidationError("client_id must be a non-empty URL string") + + parsed = urllib.parse.urlparse(client_id_url.strip()) + scheme = (parsed.scheme or "").lower() + hostname = (parsed.hostname or "").lower() + + if scheme == "https": + pass + elif scheme == "http" and allow_loopback and hostname in _LOOPBACK_HOSTS: + pass + else: + raise CimdValidationError("client_id must use https:// (or http:// on loopback when allowed)") + + if parsed.username or parsed.password: + raise CimdValidationError("client_id URL must not contain userinfo") + + if parsed.fragment: + raise CimdValidationError("client_id URL must not contain a fragment") + + path = parsed.path or "" + if not path or path == "/": + raise CimdValidationError("client_id URL must include a non-root path") + + segments = [segment for segment in path.split("/") if segment] + if any(segment in {".", ".."} for segment in segments): + raise CimdValidationError("client_id URL path must not contain '.' or '..' segments") + + if not hostname: + raise CimdValidationError("client_id URL must include a hostname") + + return client_id_url.strip() + + +def is_blocked_ip(ip_str: str) -> bool: + """Return True when ``ip_str`` resolves to a RFC 6890 special-use address.""" + try: + ip = ipaddress.ip_address(ip_str) + except ValueError: + return True + + if isinstance(ip, ipaddress.IPv4Address): + return any(ip in network for network in _BLOCKED_IPV4_NETWORKS) + + if ip.ipv4_mapped is not None: + return is_blocked_ip(str(ip.ipv4_mapped)) + + if any(ip in network for network in _BLOCKED_IPV6_NETWORKS): + return True + + return ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_multicast or ip.is_reserved + + +def _pick_pinned_ip(resolved_ips: set[str]) -> str: + ipv4 = sorted(ip for ip in resolved_ips if ":" not in ip) + if ipv4: + return ipv4[0] + return sorted(resolved_ips)[0] + + +def cimd_peer_is_acceptable(peer: str, pinned_ip: str) -> bool: + """Require the direct socket peer to match the pinned IP. + + A configured trusted proxy hop is the only exception; untrusted peers + cannot stand in for the destination. + """ + if peer == pinned_ip: + return not is_blocked_ip(peer) + from nitrostack.transports.proxy import peer_is_trusted + + return peer_is_trusted(peer) + + +def request_host_for_cimd( + headers: dict[str, str], + peer: Optional[str] = None, + trusted: Optional[Any] = None, +) -> Optional[str]: + """Inbound host used for CIMD pins. Forwarded host requires a trusted peer.""" + from nitrostack.transports.proxy import request_host_for_cimd as resolve_host + + return resolve_host(headers, peer, trusted=trusted) + + +def cimd_host_matches_request( + cimd_url: str, + headers: dict[str, str], + peer: Optional[str] = None, + trusted: Optional[Any] = None, +) -> bool: + """True when the CIMD URL host equals the trusted-proxy-aware request host.""" + from nitrostack.transports.proxy import cimd_url_matches_request_host + + return cimd_url_matches_request_host(cimd_url, headers, peer, trusted=trusted) + + +async def assert_safe_fetch_target(url_str: str, *, allow_loopback: bool = False) -> Optional[str]: + """DNS pre-resolution and IP range filtering. Returns the pinned destination IP.""" + validate_client_identifier_url(url_str, allow_loopback=allow_loopback) + parsed = urllib.parse.urlparse(url_str) + hostname = parsed.hostname + if not hostname: + raise CimdFetchError(f"Invalid hostname in URL: {url_str}") + + if allow_loopback and hostname.lower() in _LOOPBACK_HOSTS: + if hostname.lower() == "::1": + return "::1" + return "127.0.0.1" + + loop = asyncio.get_running_loop() + try: + addr_info = await loop.run_in_executor( + None, + lambda: socket.getaddrinfo(hostname, None, type=socket.SOCK_STREAM), + ) + except socket.gaierror as exc: + raise CimdFetchError(f"DNS resolution failed for {hostname}: {exc}") from exc + + resolved_ips = {info[4][0] for info in addr_info if info and info[4]} + if not resolved_ips: + raise CimdFetchError(f"DNS resolution returned no addresses for {hostname}") + + for ip in resolved_ips: + if is_blocked_ip(ip): + raise CimdFetchError(f"Destination {hostname} resolved to blocked IP: {ip}") + + return _pick_pinned_ip(resolved_ips) + + +class _PinnedHTTPConnection(http.client.HTTPConnection): + def __init__(self, hostname: str, pinned_ip: str, port: Optional[int] = None, **kwargs: Any) -> None: + super().__init__(hostname, port=port, **kwargs) + self._pinned_ip = pinned_ip + + def connect(self) -> None: + self.sock = socket.create_connection((self._pinned_ip, self.port), self.timeout) + peer = self.sock.getpeername()[0] + if not cimd_peer_is_acceptable(peer, self._pinned_ip): + self.sock.close() + raise CimdFetchError(f"Peer address {peer} is not the pinned safe IP") + + +class _PinnedHTTPSConnection(http.client.HTTPSConnection): + def __init__(self, hostname: str, pinned_ip: str, port: Optional[int] = None, **kwargs: Any) -> None: + super().__init__(hostname, port=port, **kwargs) + self._pinned_ip = pinned_ip + + def connect(self) -> None: + self.sock = socket.create_connection((self._pinned_ip, self.port), self.timeout) + peer = self.sock.getpeername()[0] + if not cimd_peer_is_acceptable(peer, self._pinned_ip): + self.sock.close() + raise CimdFetchError(f"Peer address {peer} is not the pinned safe IP") + context = self._context if getattr(self, "_context", None) else ssl.create_default_context() + self.sock = context.wrap_socket(self.sock, server_hostname=self.host) + + +class _PinnedHTTPHandler(urllib.request.HTTPHandler): + def __init__(self, hostname: str, pinned_ip: str) -> None: + super().__init__() + self._hostname = hostname + self._pinned_ip = pinned_ip + + def http_open(self, req: urllib.request.Request): + return self.do_open( + lambda host, **kwargs: _PinnedHTTPConnection(self._hostname, self._pinned_ip, **kwargs), + req, + ) + + +class _PinnedHTTPSHandler(urllib.request.HTTPSHandler): + def __init__(self, hostname: str, pinned_ip: str) -> None: + super().__init__() + self._hostname = hostname + self._pinned_ip = pinned_ip + + def https_open(self, req: urllib.request.Request): + return self.do_open( + lambda host, **kwargs: _PinnedHTTPSConnection(self._hostname, self._pinned_ip, **kwargs), + req, + ) + + +class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: ANN001 + raise urllib.error.HTTPError( + req.full_url, + code, + "HTTP redirects are not allowed for CIMD fetch", + headers, + fp, + ) + + +def _fetch_cimd_bytes( + url_str: str, + *, + timeout_sec: float, + pinned_ip: Optional[str] = None, +) -> bytes: + parsed = urllib.parse.urlparse(url_str) + hostname = parsed.hostname + handlers: list[urllib.request.BaseHandler] = [_NoRedirectHandler()] + if pinned_ip and hostname: + if (parsed.scheme or "").lower() == "http": + handlers.append(_PinnedHTTPHandler(hostname, pinned_ip)) + else: + handlers.append(_PinnedHTTPSHandler(hostname, pinned_ip)) + + opener = urllib.request.build_opener(*handlers) + request = urllib.request.Request(url_str, headers={"Accept": "application/json"}, method="GET") + + try: + with opener.open(request, timeout=timeout_sec) as response: + content_length = response.headers.get("Content-Length") + if content_length is not None: + try: + if int(content_length) > MAX_CIMD_BYTES: + raise CimdFetchError(f"CIMD exceeds maximum size of {MAX_CIMD_BYTES} bytes") + except ValueError as exc: + raise CimdFetchError("Invalid Content-Length header") from exc + + chunks: list[bytes] = [] + total = 0 + while True: + chunk = response.read(MAX_CIMD_BYTES - total + 1) + if not chunk: + break + total += len(chunk) + if total > MAX_CIMD_BYTES: + raise CimdFetchError(f"CIMD exceeds maximum size of {MAX_CIMD_BYTES} bytes") + chunks.append(chunk) + return b"".join(chunks) + except CimdFetchError: + raise + except urllib.error.HTTPError as exc: + if 300 <= exc.code < 400: + raise CimdFetchError("HTTP redirects are not allowed for CIMD fetch") from exc + raise CimdFetchError(f"CIMD fetch failed with HTTP {exc.code}") from exc + except urllib.error.URLError as exc: + raise CimdFetchError(f"CIMD fetch failed: {exc.reason}") from exc + except OSError as exc: + raise CimdFetchError(f"CIMD fetch failed: {exc}") from exc + + +def _validate_cimd_document(doc: Any, fetched_url: str) -> dict[str, Any]: + if not isinstance(doc, dict): + raise CimdValidationError("CIMD document must be a JSON object") + if doc.get("client_id") != fetched_url: + raise CimdValidationError( + f"CIMD client_id '{doc.get('client_id')}' does not match URL '{fetched_url}'" + ) + return doc + + +async def resolve_cimd( + client_id_url: str, + *, + allow_loopback: bool = False, + timeout_sec: float = CIMD_FETCH_TIMEOUT_SEC, +) -> dict[str, Any]: + """ + Fetch and validate a Client ID Metadata Document. + + Applies URL validation, DNS/IP filtering, redirect blocking, timeout, + payload size bounding, and anti-impersonation ``client_id`` checks. + """ + normalized = validate_client_identifier_url(client_id_url, allow_loopback=allow_loopback) + pinned_ip = await assert_safe_fetch_target(normalized, allow_loopback=allow_loopback) + + body = await asyncio.to_thread( + _fetch_cimd_bytes, + normalized, + timeout_sec=timeout_sec, + pinned_ip=pinned_ip, + ) + if len(body) > MAX_CIMD_BYTES: + raise CimdFetchError(f"CIMD exceeds maximum size of {MAX_CIMD_BYTES} bytes") + try: + document = json.loads(body.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise CimdValidationError("CIMD document must be valid UTF-8 JSON") from exc + + return _validate_cimd_document(document, normalized) + + +def looks_like_cimd_url(value: Any) -> bool: + """Return True when ``value`` is an http(s) client identifier URL.""" + return isinstance(value, str) and value.startswith(("https://", "http://")) + + +def resolve_cimd_sync( + client_id_url: str, + *, + allow_loopback: bool = False, + timeout_sec: float = CIMD_FETCH_TIMEOUT_SEC, +) -> dict[str, Any]: + """Synchronous wrapper for registration / other non-async callers.""" + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run( + resolve_cimd(client_id_url, allow_loopback=allow_loopback, timeout_sec=timeout_sec) + ) + raise CimdFetchError("CIMD resolution cannot run nested on a running event loop") + + +def looks_like_cimd_url(value: Any) -> bool: + """Return True when ``value`` is an http(s) client identifier URL.""" + return isinstance(value, str) and value.startswith(("https://", "http://")) + + +def resolve_cimd_sync( + client_id_url: str, + *, + allow_loopback: bool = False, + timeout_sec: float = CIMD_FETCH_TIMEOUT_SEC, +) -> dict[str, Any]: + """Synchronous wrapper for registration / other non-async callers.""" + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run( + resolve_cimd(client_id_url, allow_loopback=allow_loopback, timeout_sec=timeout_sec) + ) + raise CimdFetchError("CIMD resolution cannot run nested on a running event loop") diff --git a/nitrostack/auth/oauth.py b/nitrostack/auth/oauth.py index feb6d28..b6b3358 100644 --- a/nitrostack/auth/oauth.py +++ b/nitrostack/auth/oauth.py @@ -10,11 +10,13 @@ from nitrostack.core.module import module from nitrostack.core.di import DIContainer from nitrostack.auth.oauth_module import ( + apply_cimd_to_registration_body, build_authorization_server_metadata, build_protected_resource_metadata, build_registration_response, is_client_registration_enabled, ) +from nitrostack.auth.cimd import CimdFetchError, CimdValidationError from nitrostack.core.errors import AudienceMismatchError, ConfigurationError, TokenInactiveError @@ -174,17 +176,42 @@ def do_OPTIONS(self): self.send_header("Connection", "close") self.end_headers() + def _proxy_headers(self) -> Dict[str, str]: + return {key: value for key, value in self.headers.items()} + + def _proxy_peer(self) -> Optional[str]: + return self.client_address[0] if self.client_address else None + + def _trusted_public_origin(self) -> Optional[str]: + from nitrostack.transports.proxy import public_origin, trusted_forwarded_host + + headers = self._proxy_headers() + peer = self._proxy_peer() + if not trusted_forwarded_host(headers, peer): + return None + return public_origin( + headers, + peer, + fallback_host=f"localhost:{service_instance.discovery_port}", + fallback_proto="http", + ) + def do_GET(self): if self.path == "/.well-known/oauth-protected-resource": _write_json(self, 200, build_protected_resource_metadata(service_instance)) elif self.path == "/.well-known/oauth-authorization-server": + origin = self._trusted_public_origin() registration_endpoint = ( registration_path if is_client_registration_enabled(service_instance) else None ) _write_json( self, 200, - build_authorization_server_metadata(service_instance, registration_endpoint), + build_authorization_server_metadata( + service_instance, + registration_endpoint, + public_origin=origin, + ), ) else: _write_empty(self, 404) @@ -209,6 +236,23 @@ def do_POST(self): except Exception: body = {} + try: + body = apply_cimd_to_registration_body( + body, + headers=self._proxy_headers(), + peer=self._proxy_peer(), + ) + except (CimdValidationError, CimdFetchError) as exc: + _write_json( + self, + 400, + { + "error": "invalid_client_metadata", + "error_description": str(exc), + }, + ) + return + _write_json(self, 200, build_registration_response(service_instance, body)) def run_server(): diff --git a/nitrostack/auth/oauth_module.py b/nitrostack/auth/oauth_module.py index 31f8bab..b2585c4 100644 --- a/nitrostack/auth/oauth_module.py +++ b/nitrostack/auth/oauth_module.py @@ -20,12 +20,20 @@ import time from typing import Any, Dict, Optional, TYPE_CHECKING +from nitrostack.auth.cimd import ( + looks_like_cimd_url, + resolve_cimd_sync, +) + if TYPE_CHECKING: from nitrostack.auth.oauth import OAuthService def build_authorization_server_metadata( - service: "OAuthService", registration_endpoint: Optional[str] = None + service: "OAuthService", + registration_endpoint: Optional[str] = None, + *, + public_origin: Optional[str] = None, ) -> Dict[str, Any]: """ Build an RFC 8414 Authorization Server Metadata document. @@ -34,9 +42,15 @@ def build_authorization_server_metadata( document describes the *external* IdP configured via `authorization_servers`/ `token_introspection_endpoint`/`jwks_uri` — it does not mean nitrostack serves these endpoints itself. + + ``public_origin`` is the trusted-proxy-aware request origin. It is used only + as a last-resort issuer fallback and to make a relative registration path + absolute. """ issuer = service.issuer or ( - service.authorization_servers[0] if service.authorization_servers else "http://localhost" + service.authorization_servers[0] + if service.authorization_servers + else (public_origin or "http://localhost") ) auth_server_base = service.authorization_servers[0] if service.authorization_servers else issuer @@ -52,7 +66,10 @@ def build_authorization_server_metadata( "code_challenge_methods_supported": ["S256"], } if registration_endpoint: - metadata["registration_endpoint"] = registration_endpoint + if public_origin and registration_endpoint.startswith("/"): + metadata["registration_endpoint"] = f"{public_origin.rstrip('/')}{registration_endpoint}" + else: + metadata["registration_endpoint"] = registration_endpoint return metadata @@ -62,6 +79,7 @@ def build_protected_resource_metadata(service: "OAuthService") -> Dict[str, Any] "resource": service.resource_uri, "authorization_servers": service.authorization_servers, "scopes_supported": service.scopes_supported, + "bearer_methods_supported": ["header"], } @@ -72,10 +90,47 @@ def is_client_registration_enabled(service: "OAuthService") -> bool: Requires BOTH an explicit opt-in (`enable_client_registration`, from config or `OAUTH_ENABLE_CLIENT_REGISTRATION=true`) AND a configured client id — never a literal default. Without a configured client id there is nothing to hand back. + + Deprecated on MCP 2026-07-28 in favor of Client ID Metadata Documents (CIMD). """ return bool(service.enable_client_registration and service.static_client_id) +def apply_cimd_to_registration_body( + body: Optional[Dict[str, Any]] = None, + *, + headers: Optional[Dict[str, str]] = None, + peer: Optional[str] = None, +) -> Dict[str, Any]: + """ + When registration includes a CIMD ``client_id`` URL, fetch and validate it. + + Returns the body unchanged when ``client_id`` is not a URL. Raises + ``CimdValidationError`` / ``CimdFetchError`` on a failed CIMD fetch. + + ``headers`` / ``peer`` pin the inbound request host. Forwarded host is + honored only when the peer is trusted, so an untrusted + ``X-Forwarded-Host`` cannot rebind the CIMD host comparison. + """ + from nitrostack.auth.cimd import cimd_host_matches_request, request_host_for_cimd + + payload = dict(body or {}) + client_id = payload.get("client_id") + if not looks_like_cimd_url(client_id): + return payload + document = resolve_cimd_sync(str(client_id)) + payload["client_id"] = document["client_id"] + payload["_cimd"] = document + if headers is not None: + payload["_cimd_request_host"] = request_host_for_cimd(headers, peer) + payload["_cimd_host_matches_request"] = cimd_host_matches_request( + str(document["client_id"]), + headers, + peer, + ) + return payload + + def build_registration_response(service: "OAuthService", body: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """ Build an RFC 7591 client-registration response. diff --git a/nitrostack/auth/oauth_security.py b/nitrostack/auth/oauth_security.py new file mode 100644 index 0000000..4c0bc1d --- /dev/null +++ b/nitrostack/auth/oauth_security.py @@ -0,0 +1,28 @@ +"""OAuth 2.1 security helpers.""" + +from __future__ import annotations + + +class AuthorizationIssuerMismatchError(ValueError): + """Raised when RFC 9207 ``iss`` does not match the expected authorization server.""" + + +def validate_authorization_iss(response_iss: str | None, expected_issuer: str) -> None: + """ + Validate RFC 9207 ``iss`` anti-mixup parameter. + + Authorization responses MUST include ``iss`` and it MUST match the intended + authorization server's issuer identifier. + """ + if not expected_issuer or not expected_issuer.strip(): + raise ValueError("expected_issuer is required") + + if not response_iss or not str(response_iss).strip(): + raise AuthorizationIssuerMismatchError("Authorization response missing required iss parameter") + + normalized_expected = expected_issuer.rstrip("/") + normalized_actual = str(response_iss).strip().rstrip("/") + if normalized_actual != normalized_expected: + raise AuthorizationIssuerMismatchError( + f"Authorization iss mismatch: expected '{normalized_expected}', got '{normalized_actual}'" + ) diff --git a/nitrostack/auth/request.py b/nitrostack/auth/request.py new file mode 100644 index 0000000..b3c22c4 --- /dev/null +++ b/nitrostack/auth/request.py @@ -0,0 +1,157 @@ +"""Resolve handler identity from HTTP headers and the request envelope. + +Precedence (first present token is the only candidate): +1. HTTP ``Authorization`` Bearer +2. Spec envelope ``io.modelcontextprotocol/auth`` token +3. ``_meta.authorization`` Bearer (transports without HTTP headers) + +A present token that fails verification yields empty identity. Unsigned +``userId`` / ``tenantId`` fields are never used. +""" + +from __future__ import annotations + +from typing import Any, Mapping, Optional + +from nitrostack.core.context import AuthContext +from nitrostack.protocol.meta import MCP_META_PREFIX, flatten_request_meta_object + + +def bearer_token_from_header(value: Any) -> Optional[str]: + if not isinstance(value, str): + return None + stripped = value.strip() + if stripped.startswith("Bearer "): + token = stripped[len("Bearer ") :].strip() + return token or None + return None + + +def _header_value(headers: Mapping[str, str], name: str) -> Optional[str]: + target = name.lower() + for key, value in headers.items(): + if key.lower() == target: + return value + return None + + +def envelope_auth_slot(raw_meta: Mapping[str, Any] | None) -> Any: + """Return the spec auth slot. Identity fields inside it are not trusted.""" + if not raw_meta: + return None + prefixed = f"{MCP_META_PREFIX}auth" + if prefixed in raw_meta: + return raw_meta[prefixed] + return raw_meta.get("auth") + + +def bearer_token_from_envelope_auth(auth_slot: Any) -> Optional[str]: + """Extract a Bearer or explicit token from the spec auth slot only.""" + if isinstance(auth_slot, str): + return bearer_token_from_header(auth_slot) + if not isinstance(auth_slot, dict): + return None + header = auth_slot.get("authorization") or auth_slot.get("Authorization") + token = bearer_token_from_header(header) + if token: + return token + for key in ("token", "accessToken"): + value = auth_slot.get(key) + if isinstance(value, str) and value.strip(): + return bearer_token_from_header(value) or value.strip() + return None + + +def authorization_token_from_headers(headers: Mapping[str, str] | None) -> Optional[str]: + if not headers: + return None + return bearer_token_from_header( + _header_value(headers, "authorization") or _header_value(headers, "Authorization") + ) + + +def authorization_token_from_meta(raw_meta: Mapping[str, Any] | None) -> Optional[str]: + if not raw_meta: + return None + token = bearer_token_from_envelope_auth(envelope_auth_slot(raw_meta)) + if token: + return token + token = bearer_token_from_header( + raw_meta.get("authorization") or raw_meta.get("Authorization") + ) + if token: + return token + headers = raw_meta.get("headers") + if isinstance(headers, dict): + return bearer_token_from_header( + headers.get("authorization") or headers.get("Authorization") + ) + return None + + +def authorization_token_from_request(rc: Any) -> Optional[str]: + """Return the chosen Bearer token, or None when no candidate is present.""" + if rc is None: + return None + request = getattr(rc, "request", None) + headers_obj = getattr(request, "headers", None) if request is not None else None + if headers_obj is not None: + try: + header_map = {str(key): str(value) for key, value in headers_obj.items()} + except Exception: + header_map = {} + header_token = authorization_token_from_headers(header_map) + if header_token: + return header_token + + return authorization_token_from_meta(flatten_request_meta_object(getattr(rc, "meta", None))) + + +def verify_bearer_payload(token: str) -> Optional[dict[str, Any]]: + try: + from nitrostack.auth.jwt import JWTService + from nitrostack.core.di import DIContainer + + payload = DIContainer.get_instance().resolve(JWTService).verify_token(token) + except Exception: + return None + return payload if isinstance(payload, dict) else None + + +def tenant_from_claims(claims: Mapping[str, Any]) -> Optional[str]: + for key in ("tenant_id", "tenantId", "org_id", "orgId"): + value = claims.get(key) + if isinstance(value, str) and value.strip(): + return value + return None + + +def auth_context_from_payload(payload: Mapping[str, Any]) -> AuthContext: + raw_aud = payload.get("aud") + aud = raw_aud if isinstance(raw_aud, list) else ([raw_aud] if raw_aud else None) + scopes = payload.get("scopes") or payload.get("scope") or [] + if isinstance(scopes, str): + scopes = [item for item in scopes.split(" ") if item] + subject = payload.get("sub") + return AuthContext( + subject=subject if isinstance(subject, str) and subject.strip() else None, + scopes=list(scopes) if isinstance(scopes, list) else [], + client_id=payload.get("client_id"), + exp=payload.get("exp"), + iat=payload.get("iat"), + iss=payload.get("iss"), + aud=aud, + claims=dict(payload), + token_payload=dict(payload), + ) + + +def auth_context_from_request(rc: Any) -> Optional[AuthContext]: + """Verified JWT identity, or None when missing or invalid.""" + token = authorization_token_from_request(rc) + if not token: + return None + payload = verify_bearer_payload(token) + if payload is None: + return None + return auth_context_from_payload(payload) diff --git a/nitrostack/core/additional_decorators.py b/nitrostack/core/additional_decorators.py index 14df053..aa30a47 100644 --- a/nitrostack/core/additional_decorators.py +++ b/nitrostack/core/additional_decorators.py @@ -16,6 +16,7 @@ def cache(ttl: int = 60): """ def decorator(func: Callable): cache_store: Dict[Tuple[Any, ...], Tuple[Any, float]] = {} + setattr(func, "_mcp_cache_ttl", ttl) @wraps(func) async def wrapper(*args, **kwargs): @@ -47,6 +48,7 @@ async def wrapper(*args, **kwargs): return result copy_mcp_attributes(func, wrapper) + setattr(wrapper, "_mcp_cache_ttl", ttl) return wrapper return decorator diff --git a/nitrostack/core/app.py b/nitrostack/core/app.py index 0b3da71..a0bc83e 100644 --- a/nitrostack/core/app.py +++ b/nitrostack/core/app.py @@ -9,12 +9,21 @@ import logging from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Callable, Dict, List, Literal, Optional, Pattern, Set, Tuple, Type +from typing import Any, Callable, Dict, List, Literal, Mapping, Optional, Pattern, Set, Tuple, Type import mcp.types as types -from mcp.server.lowlevel.server import request_ctx +from mcp import MCPError +from mcp.server.context import ServerRequestContext from mcp.server.lowlevel.helper_types import ReadResourceContents -from mcp.server.stdio import stdio_server +from mcp.types import ( + CallToolRequestParams, + GetPromptRequestParams, + PaginatedRequestParams, + ReadResourceRequestParams, + RequestParams, + SubscribeRequestParams, + UnsubscribeRequestParams, +) from pydantic import BaseModel, create_model from nitrostack.core.context import ExecutionContext, TaskContext @@ -34,6 +43,55 @@ from nitrostack.core.additional_decorators import HealthCheckRegistry from nitrostack.core.task import TaskManager, TaskStatus from nitrostack.events.event_emitter import EventEmitter +from nitrostack.protocol.schema import ( + gate_registered_schema, + normalize_input_schema, + normalize_output_schema, +) +from nitrostack.protocol.resources import extract_template_param_names, uri_template_to_pattern +from nitrostack.protocol.version import ( + MODERN_PROTOCOL_VERSION, + EraSource, + ProtocolEra, + protocol_version_for_era, + resolve_http_engine, + resolve_protocol_era_resolution, + wire_mode_for_era, +) +from nitrostack.protocol.mrtr import InputRequiredResult, split_mrtr_from_arguments +from nitrostack.protocol.cache_hints import ( + build_list_endpoint_cache_hint_meta, + resolve_resource_cache_hint_meta, + resolve_tool_cache_hint_meta, +) +from nitrostack.auth.request import ( + auth_context_from_request, + bearer_token_from_envelope_auth, + envelope_auth_slot, +) +from nitrostack.protocol.meta import ( + bind_request_envelope, + flatten_request_meta_object, + strip_tool_arguments, +) +from nitrostack.protocol.observability import TraceContext, extract_trace_context +from nitrostack.runtime.correlation import InFlightRegistry, new_correlation_id +from nitrostack.transports.headers import ( + extract_mcp_param_headers, + extract_mcp_scope_headers, + merge_mcp_param_headers, +) +from nitrostack.protocol.deprecated import ( + deprecated_method_message, + rejects_deprecated_method, +) +from nitrostack.protocol.tasks import ( + DEFAULT_TASK_TTL_MS, + task_support_forbidden_message, + task_support_required_message, +) +from nitrostack.runtime.request_ctx import ServerResult, request_ctx +from nitrostack.tasks.authorization import extract_task_access_context from nitrostack.widgets.component import Component, find_project_root, load_widget_html, parse_widget_options from nitrostack.widgets.mcp_meta import build_call_tool_result_meta, build_tool_list_meta, resource_read_contents_meta from nitrostack.widgets.route_templates import build_missing_widget @@ -43,24 +101,54 @@ logger = logging.getLogger(__name__) +def _iso_timestamp(value: Any) -> str: + """Official Task timestamps are ISO-8601 strings.""" + if hasattr(value, "isoformat"): + text = value.isoformat() + if text.endswith("+00:00"): + return text[:-6] + "Z" + return text + return str(value) + + +class GetTaskResult(types.GetTaskResult): + """Official GetTaskResult plus NitroStack-owned ``result`` / ``error`` slots.""" + + result: Any = None + error: Any = None + + def resolve_http_port() -> int: """MCP HTTP/dual bind port. Defaults to 3000; 3001 is reserved for widgets.""" return int(os.environ.get("PORT") or os.environ.get("MCP_SERVER_PORT") or DEFAULT_HTTP_PORT) +def resolve_http_host() -> str: + """Bind host. ``HOST`` matches the TypeScript templates; default is ``127.0.0.1``.""" + return (os.environ.get("HOST") or "127.0.0.1").strip() or "127.0.0.1" + + @dataclass class ServerConfig: name: str version: str = "1.0.0" transport_type: Optional[Literal["stdio", "http", "dual"]] = None - # Streamable HTTP options (Phase 3). Each can also be set via env var at - # `start()` time (`MCP_STATELESS`, `MCP_MAX_SESSIONS`, `MCP_SESSION_TIMEOUT_MS`); - # the env var wins if both are set, matching the existing `transport_type`/ - # `MCP_TRANSPORT_TYPE` precedence below. + protocol_version: str = MODERN_PROTOCOL_VERSION + # Era fallback when MCP_STATELESS and NITRO_MCP_PROTOCOL_VERSION are unset. + # Same tokens as the env var (modern/latest/2026/2026-07-28, auto/both/dual/ + # dual-spec, legacy/2025/2025-06-18/2025-11-25). None means default era (`auto`). + protocol_era: Optional[str] = None + # Streamable HTTP options. Era resolution at `start()` / `get_combined_app()`: + # `MCP_STATELESS`, then `NITRO_MCP_PROTOCOL_VERSION`, then `protocol_era`, + # then `auto`. Also: `ENABLE_CORS`, `MCP_MAX_SESSIONS`, + # `MCP_SESSION_TIMEOUT_MS`, `MCP_TRANSPORT_TYPE`. + # `MCP_STATELESS` wins over the protocol-era mapping. Unset era is `auto` + # and does not force this flag; only `modern` sets stateless HTTP. stateless: bool = False max_sessions: Optional[int] = None session_timeout_ms: Optional[int] = None json_response: bool = False + extensions: Optional[Dict[str, str]] = None def mcp_app(module: Type, server: ServerConfig): @@ -159,14 +247,14 @@ def inspector_friendly_schema(node: Any) -> Any: def tool_json_schema(schema_spec: Any) -> Optional[Dict[str, Any]]: - """Convert a Pydantic model class or JSON-schema dict to Inspector-friendly schema.""" + """Convert a Pydantic model class or JSON-schema dict to MCP 2026-07-28 outputSchema.""" if schema_spec is None: return None if isinstance(schema_spec, dict): - return inspector_friendly_schema(schema_spec) + return normalize_output_schema(inspector_friendly_schema(schema_spec)) model = get_pydantic_model(schema_spec) if model is not None: - return inspector_friendly_schema(model.model_json_schema()) + return normalize_output_schema(inspector_friendly_schema(model.model_json_schema())) return None @@ -242,6 +330,132 @@ class _PromptEntry: _AUTH_META_KEYS = ("authorization", "x-api-key", "token", "_oauth", "headers") +def _request_meta_from_ctx(rc: Any) -> Dict[str, Any]: + """Flatten MCP request ``_meta`` from the low-level request context.""" + if rc is None: + return {} + + raw_meta = getattr(rc, "meta", None) + if raw_meta is None: + return {} + data = flatten_request_meta_object(raw_meta) + if data: + return data + for key in _AUTH_META_KEYS: + value = getattr(raw_meta, key, None) + if value is not None: + data[key] = value + return data + + +def _trace_context_from_request_ctx(rc: Any) -> TraceContext | None: + return extract_trace_context(_request_meta_from_ctx(rc)) + + +def _read_contents_to_result(uri: str, items: List[ReadResourceContents]) -> types.ReadResourceResult: + import base64 + + contents: List[Any] = [] + for item in items: + mime = item.mime_type or "text/plain" + extra = {"_meta": item.meta} if getattr(item, "meta", None) else {} + if isinstance(item.content, bytes): + contents.append( + types.BlobResourceContents( + uri=uri, + mimeType=mime, + blob=base64.b64encode(item.content).decode("ascii"), + **extra, + ) + ) + else: + contents.append( + types.TextResourceContents( + uri=uri, mimeType=mime, text=str(item.content), **extra + ) + ) + return types.ReadResourceResult(contents=contents) + + +def _http_headers_from_request_ctx(rc: Any) -> Dict[str, str]: + if rc is None: + return {} + request = getattr(rc, "request", None) + raw = getattr(request, "headers", None) if request is not None else None + if raw is None: + session = getattr(rc, "session", None) + transport = getattr(session, "transport", None) if session is not None else None + raw = getattr(transport, "headers", None) if transport is not None else None + if raw is None: + raw = getattr(rc, "headers", None) + if raw is None: + return {} + try: + return {str(key): str(value) for key, value in raw.items()} + except Exception: + return {} + + +def _apply_request_envelope(ctx: ExecutionContext, rc: Any) -> None: + http_headers = _http_headers_from_request_ctx(rc) + envelope = bind_request_envelope( + raw_meta=_request_meta_from_ctx(rc), + mcp_headers=extract_mcp_scope_headers(http_headers), + ) + ctx.rpc_meta = envelope.meta + ctx.mcp_headers = dict(envelope.mcp_headers) + ctx.mcp_param_headers = extract_mcp_param_headers(http_headers) + ctx.protocol_version = envelope.protocol_version + ctx.auth = auth_context_from_request(rc) + if ctx.trace is None: + ctx.trace = extract_trace_context(envelope.meta.raw) + if ctx.jsonrpc_id is None and rc is not None: + ctx.jsonrpc_id = getattr(rc, "request_id", None) + if rc is not None and getattr(rc, "correlation_id", None) is None: + try: + rc.correlation_id = ctx.correlation_id + except Exception: + pass + + +def _bind_correlation(rc: Any) -> tuple[str, Any]: + """Allocate a correlation id. Never reuse the client JSON-RPC ``id`` as the key.""" + jsonrpc_id = getattr(rc, "request_id", None) if rc is not None else None + existing = getattr(rc, "correlation_id", None) if rc is not None else None + if existing: + return str(existing), jsonrpc_id + correlation_id = new_correlation_id() + if rc is not None: + try: + rc.correlation_id = correlation_id + except Exception: + pass + return correlation_id, jsonrpc_id + + +def _tool_arguments_with_mcp_params( + arguments: Dict[str, Any], + param_headers: Mapping[str, str], + input_model: Type[BaseModel], +) -> Dict[str, Any]: + """Fill missing tool fields from ``Mcp-Param-*``. Does not rewrite ``name``.""" + payload = dict(arguments or {}) + inner = payload.get("input") + looks_wrapped = ( + isinstance(inner, dict) + and set(payload.keys()) <= {"input"} + and "input" not in input_model.model_fields + ) + allowed = set(input_model.model_fields) + if looks_wrapped: + return { + "input": merge_mcp_param_headers( + inner, param_headers, allowed_fields=allowed + ) + } + return merge_mcp_param_headers(payload, param_headers, allowed_fields=allowed) + + def _auth_metadata_from_request_ctx(rc: Any) -> Dict[str, Any]: """Copy host-sent auth slots from MCP request ``_meta`` into ExecutionContext. @@ -255,26 +469,7 @@ def _auth_metadata_from_request_ctx(rc: Any) -> Dict[str, Any]: if rc is None: return extra - raw_meta = getattr(rc, "meta", None) - data: Dict[str, Any] = {} - if raw_meta is not None: - extra_fields = getattr(raw_meta, "model_extra", None) or getattr(raw_meta, "__pydantic_extra__", None) - if isinstance(extra_fields, dict): - data.update(extra_fields) - if hasattr(raw_meta, "model_dump"): - try: - dumped = raw_meta.model_dump(exclude_none=True) - if isinstance(dumped, dict): - data.update(dumped) - except Exception: - pass - elif isinstance(raw_meta, dict): - data.update(raw_meta) - else: - for key in _AUTH_META_KEYS: - value = getattr(raw_meta, key, None) - if value is not None: - data[key] = value + data = _request_meta_from_ctx(rc) auth = data.get("authorization") or data.get("Authorization") if isinstance(auth, str) and auth.strip(): @@ -300,6 +495,10 @@ def _auth_metadata_from_request_ctx(rc: Any) -> Dict[str, Any]: if isinstance(header_key, str) and header_key.strip(): extra["x-api-key"] = header_key + envelope_token = bearer_token_from_envelope_auth(envelope_auth_slot(data)) + if envelope_token: + extra["authorization"] = f"Bearer {envelope_token}" + request = getattr(rc, "request", None) headers_obj = getattr(request, "headers", None) if request is not None else None if headers_obj is not None: @@ -327,10 +526,14 @@ def __init__(self, app_class: Type): raise ValueError("Invalid application class. Must be decorated with @mcp_app or @module.") self.mcp_server: Optional[NitroStackMcpServer] = None + era_resolution = resolve_protocol_era_resolution( + config_value=self.server_config.protocol_era + ) + self.protocol_era: ProtocolEra = era_resolution.era + self.protocol_era_source: EraSource = era_resolution.source - # nitrostack owns these registries directly (no FastMCP-managed tool/resource - # manager in between) so that any number of low-level `Server` instances can be - # wired against the same registered tools/resources/prompts (see + # nitrostack owns these registries so any number of low-level `Server` + # instances can be wired against the same tools/resources/prompts (see # `create_configured_mcp_server`). self._tools: Dict[str, _ToolEntry] = {} self._resources: Dict[str, _ResourceEntry] = {} @@ -338,11 +541,12 @@ def __init__(self, app_class: Type): self._prompts: Dict[str, _PromptEntry] = {} self._initial_tools: List[Tuple[Any, Callable, ToolConfig]] = [] self.task_manager = TaskManager() + self._in_flight = InFlightRegistry() self._bootstrap() def _bootstrap(self) -> None: - # 1. Construct the low-level server directly (no FastMCP) + # 1. Construct the low-level server directly. self.mcp_server = NitroStackMcpServer( name=self.server_config.name, version=self.server_config.version, @@ -355,19 +559,26 @@ def _bootstrap(self) -> None: container = DIContainer.get_instance() self._assert_declared_dependencies(resolved_modules, container) - # Instantiate all providers and controllers to populate container + # Instantiate all providers and controllers to populate container. + # Scan both: providers with @tool/@resource/@prompt must stay discoverable. + module_instances: List[Any] = [] + seen_instance_ids: Set[int] = set() for mod in resolved_modules: mod_config = getattr(mod, "_mcp_module_config", None) if mod_config: - # Register & Resolve all providers for provider in mod_config.providers: - container.resolve(provider) - # Register & Resolve all controllers + instance = container.resolve(provider) + if id(instance) not in seen_instance_ids: + seen_instance_ids.add(id(instance)) + module_instances.append(instance) for controller in mod_config.controllers: - container.resolve(controller) + instance = container.resolve(controller) + if id(instance) not in seen_instance_ids: + seen_instance_ids.add(id(instance)) + module_instances.append(instance) - # 3. Discover decorated methods on all instances in the container - for token, instance in list(container._instances.items()): + # 3. Discover decorated methods on resolved module providers and controllers + for instance in module_instances: # Scan members of this instance for name, member in inspect.getmembers(instance): # Discover Tools @@ -450,6 +661,12 @@ def _resolve_modules(self, module_class: Type, resolved_modules: Set[Type]) -> N # ------------------------------------------------------------------ def _register_tool(self, instance: Any, method: Callable, tool_config: ToolConfig) -> None: + gate_registered_schema( + tool_config.input_schema, name=f"tool {tool_config.name!r} input" + ) + gate_registered_schema( + tool_config.output_schema, name=f"tool {tool_config.name!r} output" + ) input_model = get_pydantic_model(tool_config.input_schema) entry = _ToolEntry(config=tool_config, input_model=input_model, instance=instance, method=method) @@ -509,21 +726,26 @@ async def widget_resource_handler(context: ExecutionContext) -> str: ) def _register_resource(self, instance: Any, method: Callable, resource_config: ResourceConfig) -> None: - param_names = re.findall(r"\{([^}]+)\}", resource_config.uri) + gate_registered_schema( + getattr(resource_config, "schema", None) + or (resource_config.metadata or {}).get("schema"), + name=f"resource {resource_config.uri!r}", + ) + param_names = extract_template_param_names(resource_config.uri) entry = _ResourceEntry(config=resource_config, instance=instance, method=method, param_names=param_names) if param_names: - # Build a matching regex from the URI template, e.g. "a://b/{id}" -> - # "^a://b/(?P[^/]+)$", preserving the existing template-matching semantics. - regex_str = re.escape(resource_config.uri) - for pname in param_names: - regex_str = regex_str.replace(re.escape("{" + pname + "}"), f"(?P<{pname}>[^/]+)") - entry.pattern = re.compile(f"^{regex_str}$") + entry.pattern = uri_template_to_pattern(resource_config.uri) self._resource_templates.append(entry) else: self._resources[resource_config.uri] = entry def _register_prompt(self, instance: Any, method: Callable, prompt_config: PromptConfig) -> None: + for argument in prompt_config.arguments or []: + gate_registered_schema( + getattr(argument, "schema", None), + name=f"prompt {prompt_config.name!r} argument {argument.name!r}", + ) self._prompts[prompt_config.name] = _PromptEntry(config=prompt_config, instance=instance, method=method) def _register_health_resource(self) -> None: @@ -543,46 +765,202 @@ async def health_status_resource(context: ExecutionContext) -> str: # Protocol handler wiring (owned low-level `mcp.server.lowlevel.Server`) # ------------------------------------------------------------------ - def _setup_handlers(self, server: NitroStackMcpServer) -> None: - @server.list_tools() - async def _list_tools() -> List[types.Tool]: - return [self._build_tool_definition(entry) for entry in self._tools.values()] + def _advertise_tasks_extension(self) -> bool: + return any( + entry.config.task_support in ("optional", "required") + for entry in self._tools.values() + ) - @server.call_tool(validate_input=False) - async def _call_tool(name: str, arguments: Optional[Dict[str, Any]]): - return await self._call_tool(name, arguments or {}) + def _custom_extensions(self) -> Optional[Dict[str, str]]: + extensions = getattr(self.server_config, "extensions", None) + if not extensions: + return None + return dict(extensions) - @server.list_resources() - async def _list_resources() -> List[types.Resource]: - return [self._build_resource_definition(entry) for entry in self._resources.values()] + def handle_server_discover(self, protocol_version: Optional[str] = None) -> Dict[str, Any]: + """Single ``server/discover`` result for the mounted HTTP engine.""" + from nitrostack.protocol.discovery import build_discover_result - @server.list_resource_templates() - async def _list_resource_templates() -> List[types.ResourceTemplate]: - return [self._build_resource_template_definition(entry) for entry in self._resource_templates] + has_widgets = any( + getattr(entry, "component", None) is not None + for entry in getattr(self, "_tools", {}).values() + ) + version = protocol_version or protocol_version_for_era( + getattr(self, "protocol_era", None), + self.server_config.protocol_version, + ) + return build_discover_result( + server_name=self.server_config.name, + server_version=self.server_config.version, + protocol_version=version, + advertise_tasks=self._advertise_tasks_extension(), + advertise_app=has_widgets, + custom_extensions=self._custom_extensions(), + ) - @server.read_resource() - async def _read_resource(uri: Any): - return await self._read_resource(str(uri)) + def handle_sessionless_initialize(self, requested_version: Optional[str] = None) -> Dict[str, Any]: + """Answer 2025 ``initialize`` on the sessionless ``auto`` path. No session.""" + from nitrostack.protocol.discovery import build_sessionless_initialize_result + + has_widgets = any( + getattr(entry, "component", None) is not None + for entry in getattr(self, "_tools", {}).values() + ) + version = protocol_version_for_era( + getattr(self, "protocol_era", None), + self.server_config.protocol_version, + ) + return build_sessionless_initialize_result( + server_name=self.server_config.name, + server_version=self.server_config.version, + requested_version=requested_version, + protocol_version=version, + advertise_tasks=self._advertise_tasks_extension(), + advertise_app=has_widgets, + custom_extensions=self._custom_extensions(), + ) - @server.subscribe_resource() - async def _subscribe_resource(uri: Any) -> None: - if str(uri) not in self._resources: - raise ResourceNotFoundError(str(uri)) + def _list_endpoint_cache_meta(self) -> Dict[str, Any]: + return build_list_endpoint_cache_hint_meta() - @server.unsubscribe_resource() - async def _unsubscribe_resource(uri: Any) -> None: - return None + def _setup_handlers(self, server: NitroStackMcpServer) -> None: + async def _list_tools(_ctx: ServerRequestContext, _params: Optional[PaginatedRequestParams]): + return types.ListToolsResult( + tools=[self._build_tool_definition(entry) for entry in self._tools.values()], + _meta=self._list_endpoint_cache_meta(), + ) + + async def _call_tool(ctx: ServerRequestContext, params: CallToolRequestParams): + token = request_ctx.set(ctx) + try: + result = await self._call_tool( + params.name, params.arguments or {}, task=params.task + ) + if isinstance(result, types.CreateTaskResult): + task = result.task + return types.CallToolResult( + content=[types.TextContent(type="text", text=task.task_id)], + structuredContent={ + "resultType": "task", + "task": task.model_dump(by_alias=True), + }, + isError=False, + ) + return result + finally: + request_ctx.reset(token) + + async def _list_resources(_ctx: ServerRequestContext, _params: Optional[PaginatedRequestParams]): + return types.ListResourcesResult( + resources=[self._build_resource_definition(entry) for entry in self._resources.values()], + _meta=self._list_endpoint_cache_meta(), + ) + + async def _list_resource_templates(_ctx: ServerRequestContext, _params: Optional[PaginatedRequestParams]): + return types.ListResourceTemplatesResult( + resource_templates=[ + self._build_resource_template_definition(entry) for entry in self._resource_templates + ], + ) + + async def _read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams): + token = request_ctx.set(ctx) + try: + uri = str(params.uri) + return _read_contents_to_result(uri, await self._read_resource(uri)) + finally: + request_ctx.reset(token) + + async def _subscribe_resource(_ctx: ServerRequestContext, params: SubscribeRequestParams): + if rejects_deprecated_method("resources/subscribe", self.protocol_era): + message = deprecated_method_message("resources/subscribe") + raise MCPError(types.METHOD_NOT_FOUND, message or "Not supported") + if str(params.uri) not in self._resources: + raise ResourceNotFoundError(str(params.uri)) + return types.EmptyResult() + + async def _unsubscribe_resource(_ctx: ServerRequestContext, _params: UnsubscribeRequestParams): + return types.EmptyResult() + + async def _list_prompts(_ctx: ServerRequestContext, _params: Optional[PaginatedRequestParams]): + return types.ListPromptsResult( + prompts=[self._build_prompt_definition(entry) for entry in self._prompts.values()], + _meta=self._list_endpoint_cache_meta(), + ) - @server.list_prompts() - async def _list_prompts() -> List[types.Prompt]: - return [self._build_prompt_definition(entry) for entry in self._prompts.values()] + async def _get_prompt(ctx: ServerRequestContext, params: GetPromptRequestParams): + token = request_ctx.set(ctx) + try: + arguments = params.arguments or {} + return await self._get_prompt( + params.name, + {str(key): str(value) for key, value in arguments.items()}, + ) + finally: + request_ctx.reset(token) + + async def _discover(ctx: ServerRequestContext, _params: Optional[RequestParams]): + return self.handle_server_discover(getattr(ctx, "protocol_version", None)) + + server.add_request_handler("server/discover", RequestParams, _discover) + server.add_request_handler("tools/list", PaginatedRequestParams, _list_tools) + server.add_request_handler("tools/call", CallToolRequestParams, _call_tool) + server.add_request_handler("resources/list", PaginatedRequestParams, _list_resources) + server.add_request_handler( + "resources/templates/list", PaginatedRequestParams, _list_resource_templates + ) + server.add_request_handler("resources/read", ReadResourceRequestParams, _read_resource) + server.add_request_handler("resources/subscribe", SubscribeRequestParams, _subscribe_resource) + server.add_request_handler("resources/unsubscribe", UnsubscribeRequestParams, _unsubscribe_resource) + server.add_request_handler("prompts/list", PaginatedRequestParams, _list_prompts) + server.add_request_handler("prompts/get", GetPromptRequestParams, _get_prompt) + + async def handle_list_tools(req=None): + return ServerResult(await _list_tools(None, None)) + + async def handle_call_tool(req=None): + params = getattr(req, "params", None) if req is not None else None + name = getattr(params, "name", "") if params is not None else "" + arguments = (getattr(params, "arguments", None) or {}) if params is not None else {} + task = getattr(params, "task", None) if params is not None else None + return ServerResult(await self._call_tool(name, arguments, task=task)) + + async def handle_list_resources(req=None): + return ServerResult(await _list_resources(None, None)) + + async def handle_list_resource_templates(req=None): + return ServerResult(await _list_resource_templates(None, None)) + + async def handle_read_resource(req): + params = getattr(req, "params", req) + uri = str(params.uri) + return ServerResult(_read_contents_to_result(uri, await self._read_resource(uri))) + + async def handle_list_prompts(req=None): + return ServerResult(await _list_prompts(None, None)) + + async def handle_get_prompt(req): + params = getattr(req, "params", req) + arguments = getattr(params, "arguments", None) or {} + return ServerResult( + await self._get_prompt( + params.name, + {str(key): str(value) for key, value in arguments.items()}, + ) + ) - @server.get_prompt() - async def _get_prompt(name: str, arguments: Optional[Dict[str, str]]) -> types.GetPromptResult: - return await self._get_prompt(name, arguments or {}) + server.request_handlers[types.ListToolsRequest] = handle_list_tools + server.request_handlers[types.CallToolRequest] = handle_call_tool + server.request_handlers[types.ListResourcesRequest] = handle_list_resources + server.request_handlers[types.ListResourceTemplatesRequest] = handle_list_resource_templates + server.request_handlers[types.ReadResourceRequest] = handle_read_resource + server.request_handlers[types.ListPromptsRequest] = handle_list_prompts + server.request_handlers[types.GetPromptRequest] = handle_get_prompt self._register_task_handlers(server) self._register_initialized_handler(server) + server.discover_handler = self.handle_server_discover + server.initialize_handler = self.handle_sessionless_initialize def create_configured_mcp_server(self) -> NitroStackMcpServer: """ @@ -612,7 +990,7 @@ def _tool_input_schema(self, input_model: Type[BaseModel]) -> Dict[str, Any]: schema = {} schema.setdefault("type", "object") schema.setdefault("properties", {}) - return schema + return normalize_input_schema(schema) def _build_tool_definition(self, entry: _ToolEntry) -> types.Tool: cfg = entry.config @@ -643,6 +1021,10 @@ def _build_tool_definition(self, entry: _ToolEntry) -> types.Tool: "description": cfg.examples.description, } + cache_meta = resolve_tool_cache_hint_meta(cfg, entry.method) + if cache_meta: + meta.update(cache_meta) + if is_openai_mode(): meta["openai/type"] = "function" meta["openai/function"] = { @@ -682,14 +1064,21 @@ def _build_tool_definition(self, entry: _ToolEntry) -> types.Tool: def _build_resource_definition(self, entry: _ResourceEntry) -> types.Resource: cfg = entry.config - return types.Resource( - uri=cfg.uri, - name=cfg.name, - title=cfg.title, - description=cfg.description, - mimeType=cfg.mime_type, - size=cfg.size, - ) + meta = dict(cfg.metadata or {}) + cache_meta = resolve_resource_cache_hint_meta(cfg) + if cache_meta: + meta.update(cache_meta) + resource_kwargs: Dict[str, Any] = { + "uri": cfg.uri, + "name": cfg.name, + "title": cfg.title, + "description": cfg.description, + "mimeType": cfg.mime_type, + "size": cfg.size, + } + if meta: + resource_kwargs["_meta"] = meta + return types.Resource(**resource_kwargs) def _build_resource_template_definition(self, entry: _ResourceEntry) -> types.ResourceTemplate: cfg = entry.config @@ -728,6 +1117,25 @@ def _to_call_tool_result( component: Optional[Component] = None, context: Optional[ExecutionContext] = None, ) -> types.CallToolResult: + if isinstance(result, InputRequiredResult): + wire = result.to_wire_dict() + return types.CallToolResult( + content=[types.TextContent(type="text", text=wire.get("message", "Input required"))], + structuredContent=wire, + isError=False, + ) + mrtr_result = None + if isinstance(result, dict) and result.get("resultType") == "input_required": + from nitrostack.protocol.mrtr import coerce_input_required_result + + mrtr_result = coerce_input_required_result(result) + if mrtr_result is not None: + wire = mrtr_result.to_wire_dict() + return types.CallToolResult( + content=[types.TextContent(type="text", text=wire.get("message", "Input required"))], + structuredContent=wire, + isError=False, + ) if isinstance(result, types.CallToolResult): if component is not None: result = result.model_copy( @@ -817,7 +1225,7 @@ def _widget_result_content(self, structured: Dict[str, Any], component: Optional ) return content - async def _call_tool(self, name: str, arguments: Dict[str, Any]): + async def _call_tool(self, name: str, arguments: Dict[str, Any], task: Any = None): entry = self._tools.get(name) if entry is None: return types.CallToolResult( @@ -826,48 +1234,93 @@ async def _call_tool(self, name: str, arguments: Dict[str, Any]): ) cfg = entry.config + tool_arguments, input_responses, request_state = split_mrtr_from_arguments( + strip_tool_arguments(arguments) + ) + rc = request_ctx.get(None) + tool_arguments = _tool_arguments_with_mcp_params( + tool_arguments, + extract_mcp_param_headers(_http_headers_from_request_ctx(rc)), + entry.input_model, + ) # Pydantic validates after accepting either Inspector top-level fields # or the older `{input: {...}}` wrap. Low-level jsonschema is off # (`validate_input=False`) so the wrap is not rejected against the # published top-level inputSchema. - input_instance = parse_tool_input(entry.input_model, arguments) + input_instance = parse_tool_input(entry.input_model, tool_arguments) guards, middleware, interceptors, pipes, filters = self._pipeline_stages(entry.method) - # Detect task-augmented invocation via the request context's public - # `experimental.task_metadata` field (populated by the low-level server - # from `req.params.task`) rather than reaching into private state. - task_metadata = None + # Task metadata: explicit ``params.task``, then the 1.x experimental + # slot, then the official request context ``params`` mapping. + task_metadata = task session = None progress_token = None - rc = request_ctx.get(None) if rc is not None: - if getattr(rc, "experimental", None) is not None: + if task_metadata is None and getattr(rc, "experimental", None) is not None: task_metadata = rc.experimental.task_metadata - if getattr(rc, "meta", None) is not None: - progress_token = rc.meta.progressToken + params = getattr(rc, "params", None) + if task_metadata is None: + task_metadata = getattr(params, "task", None) + if task_metadata is None and isinstance(params, Mapping): + task_metadata = params.get("task") + meta = getattr(rc, "meta", None) + if isinstance(meta, Mapping): + progress_token = meta.get("progress_token") or meta.get("progressToken") + elif meta is not None: + progress_token = getattr(meta, "progress_token", None) or getattr( + meta, "progressToken", None + ) session = getattr(rc, "session", None) auth_meta = _auth_metadata_from_request_ctx(rc) + trace = _trace_context_from_request_ctx(rc) - is_task = (task_metadata is not None) or (cfg.task_support == "required") - if cfg.task_support == "forbidden": - is_task = False + is_task = task_metadata is not None + if cfg.task_support == "forbidden" and task_metadata is not None: + raise MCPError( + types.METHOD_NOT_FOUND, + task_support_forbidden_message(cfg.name), + ) + if cfg.task_support == "required" and task_metadata is None: + raise MCPError( + types.INVALID_REQUEST, + task_support_required_message(cfg.name), + ) if is_task: - ttl = task_metadata.ttl if task_metadata and task_metadata.ttl is not None else 300 - task = self.task_manager.create_task(ttl_seconds=ttl) + ttl_ms = ( + task_metadata.ttl + if task_metadata and task_metadata.ttl is not None + else DEFAULT_TASK_TTL_MS + ) + task_access = extract_task_access_context(rc) + task = await self.task_manager.create_task( + ttl_ms=ttl_ms, + tool_name=cfg.name, + owner_id=task_access.user_id if task_access else None, + tenant_id=task_access.tenant_id if task_access else None, + session_id=task_access.session_id if task_access else None, + ) task_id = task.id + correlation_id, jsonrpc_id = _bind_correlation(rc) async def background_execution(): task_ctx = ExecutionContext( - request_id=str(uuid.uuid4()), + request_id=correlation_id, + correlation_id=correlation_id, + jsonrpc_id=jsonrpc_id, tool_name=cfg.name, metadata={"input": input_instance, **auth_meta}, + input_responses=input_responses, + request_state=request_state, + trace=trace, ) + _apply_request_envelope(task_ctx, rc) task_ctx.task = TaskContext( task_id, self.task_manager, session=session, progress_token=progress_token, + correlation_id=correlation_id, ) try: result = await run_pipeline( @@ -884,12 +1337,19 @@ async def background_execution(): param_name="input", param_type=entry.input_model, ) - self.task_manager.complete_task( - task_id, self._to_call_tool_result(result, entry.component, task_ctx) - ) + if isinstance(result, InputRequiredResult): + await self.task_manager.require_input( + task_id, + result.to_wire_dict(), + progress=result.message or "Additional input required", + ) + else: + await self.task_manager.complete_task( + task_id, self._to_call_tool_result(result, entry.component, task_ctx) + ) except Exception as e: try: - self.task_manager.fail_task(task_id, e) + await self.task_manager.fail_task(task_id, e) except (TaskAlreadyTerminalError, TaskExpiredError): # Cancelled/expired while running — leave terminal state as-is. pass @@ -897,7 +1357,19 @@ async def background_execution(): asyncio.create_task(background_execution()) return types.CreateTaskResult(task=self._task_data_to_mcp_task(task)) - ctx = ExecutionContext(request_id=str(uuid.uuid4()), tool_name=cfg.name, metadata={"input": input_instance, **auth_meta}) + correlation_id, jsonrpc_id = _bind_correlation(rc) + ctx = ExecutionContext( + request_id=correlation_id, + correlation_id=correlation_id, + jsonrpc_id=jsonrpc_id, + tool_name=cfg.name, + metadata={"input": input_instance, **auth_meta}, + input_responses=input_responses, + request_state=request_state, + trace=trace, + ) + _apply_request_envelope(ctx, rc) + ticket = self._in_flight.register(correlation_id, jsonrpc_id=jsonrpc_id) try: result = await run_pipeline( handler=entry.method, @@ -919,6 +1391,13 @@ async def background_execution(): content=[types.TextContent(type="text", text=str(exc))], isError=True, ) + finally: + self._in_flight.discard(correlation_id) + if ticket.cancel_requested.is_set(): + return types.CallToolResult( + content=[types.TextContent(type="text", text="Request was cancelled.")], + isError=True, + ) return self._to_call_tool_result(result, entry.component, ctx) async def _read_resource(self, uri: str) -> List[ReadResourceContents]: @@ -947,6 +1426,7 @@ async def _read_resource(self, uri: str) -> List[ReadResourceContents]: cfg = entry.config ctx = ExecutionContext(request_id=str(uuid.uuid4()), metadata=dict(path_kwargs)) + _apply_request_envelope(ctx, request_ctx.get(None)) guards, middleware, interceptors, pipes, filters = self._pipeline_stages(entry.method) result = await run_pipeline( @@ -985,7 +1465,15 @@ async def _get_prompt(self, name: str, arguments: Dict[str, str]) -> types.GetPr cfg = entry.config args_dict = dict(arguments or {}) + rc = request_ctx.get(None) + allowed = {arg.name for arg in cfg.arguments} if cfg.arguments else None + args_dict = merge_mcp_param_headers( + args_dict, + extract_mcp_param_headers(_http_headers_from_request_ctx(rc)), + allowed_fields=allowed, + ) ctx = ExecutionContext(request_id=str(uuid.uuid4()), metadata=args_dict) + _apply_request_envelope(ctx, rc) guards, middleware, interceptors, pipes, filters = self._pipeline_stages(entry.method) raw_messages = await run_pipeline( @@ -1013,93 +1501,114 @@ async def _get_prompt(self, name: str, arguments: Dict[str, str]) -> types.GetPr return types.GetPromptResult(description=cfg.description, messages=messages) # ------------------------------------------------------------------ - # Task subsystem — registered directly on the low-level server's public - # `request_handlers`/`notification_handlers` dicts (no FastMCP reach-through). + # Task subsystem — registered on the low-level server's public + # `request_handlers`/`notification_handlers` dicts. # ------------------------------------------------------------------ def _task_data_to_mcp_task(self, task) -> types.Task: """Map TaskData to MCP Task. EXPIRED is not an MCP wire status — surface as error.""" if task.status == TaskStatus.EXPIRED: - raise types.McpError( - types.ErrorData( - code=types.INVALID_PARAMS, - message=f"Task {task.id} has expired", - ) - ) + raise MCPError(types.INVALID_PARAMS, f"Task {task.id} has expired") return types.Task( taskId=task.id, status=task.status.value, statusMessage=task.progress or "", - createdAt=task.created_at, - lastUpdatedAt=task.last_updated_at or task.created_at, - ttl=task.ttl_seconds if task.ttl_seconds is not None else 0, + createdAt=_iso_timestamp(task.created_at), + lastUpdatedAt=_iso_timestamp(task.last_updated_at or task.created_at), + ttl=task.ttl if task.ttl is not None else 0, pollInterval=task.poll_interval, ) + def _serialize_task_result_payload(self, result: Any) -> Any: + if isinstance(result, types.CallToolResult): + return result.model_dump(by_alias=True, exclude_none=True) + if isinstance(result, BaseModel): + return result.model_dump() + return result + + def _build_get_task_result(self, task) -> types.GetTaskResult: + mcp_task = self._task_data_to_mcp_task(task) + payload: Dict[str, Any] = { + "taskId": mcp_task.task_id, + "status": mcp_task.status, + "statusMessage": mcp_task.status_message, + "createdAt": mcp_task.created_at, + "lastUpdatedAt": mcp_task.last_updated_at, + "ttl": mcp_task.ttl, + "pollInterval": mcp_task.poll_interval, + } + if task.status == TaskStatus.COMPLETED and task.result is not None: + payload["result"] = self._serialize_task_result_payload(task.result) + elif task.status == TaskStatus.FAILED and task.error is not None: + payload["error"] = {"message": str(task.error)} + elif task.status == TaskStatus.INPUT_REQUIRED and task.result is not None: + payload["result"] = task.result + return GetTaskResult(**payload) + def _register_task_handlers(self, server: NitroStackMcpServer) -> None: async def handle_list_tasks(req): + if rejects_deprecated_method("tasks/list", self.protocol_era): + message = deprecated_method_message("tasks/list") + raise MCPError(types.METHOD_NOT_FOUND, message or "Not supported") tasks_list = [] - for t in self.task_manager.list_tasks(): + access = extract_task_access_context(request_ctx.get(None)) + params = getattr(req, "params", req) + cursor = getattr(params, "cursor", None) if params is not None else None + tasks_page, next_cursor = await self.task_manager.list_tasks_page( + access_context=access, + cursor=cursor, + ) + for t in tasks_page: if t.status == TaskStatus.EXPIRED: continue tasks_list.append(self._task_data_to_mcp_task(t)) - return types.ListTasksResult(tasks=tasks_list, nextCursor=None) + return types.ListTasksResult(tasks=tasks_list, nextCursor=next_cursor) async def handle_get_task(req): - task_id = req.params.taskId + params = getattr(req, "params", req) + task_id = getattr(params, "task_id", None) or getattr(params, "taskId", None) + access = extract_task_access_context(request_ctx.get(None)) try: - t = self.task_manager.get_task(task_id) + t = await self.task_manager.get_task(task_id, access_context=access) except TaskNotFoundError: - raise types.McpError( - types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") - ) - mcp_task = self._task_data_to_mcp_task(t) - return types.GetTaskResult( - taskId=mcp_task.taskId, - status=mcp_task.status, - statusMessage=mcp_task.statusMessage, - createdAt=mcp_task.createdAt, - lastUpdatedAt=mcp_task.lastUpdatedAt, - ttl=mcp_task.ttl, - pollInterval=mcp_task.pollInterval, - ) + raise MCPError(types.INVALID_PARAMS, f"Task {task_id} not found") + return self._build_get_task_result(t) async def handle_cancel_task(req): - task_id = req.params.taskId + params = getattr(req, "params", req) + task_id = getattr(params, "task_id", None) or getattr(params, "taskId", None) + access = extract_task_access_context(request_ctx.get(None)) try: - self.task_manager.cancel_task(task_id) - t = self.task_manager.get_task(task_id) + await self.task_manager.cancel_task(task_id, access_context=access) + t = await self.task_manager.get_task(task_id, access_context=access) except TaskNotFoundError: - raise types.McpError( - types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") - ) + raise MCPError(types.INVALID_PARAMS, f"Task {task_id} not found") except TaskExpiredError: - raise types.McpError( - types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} has expired") - ) + raise MCPError(types.INVALID_PARAMS, f"Task {task_id} has expired") except TaskAlreadyTerminalError as e: - raise types.McpError( - types.ErrorData(code=types.INVALID_PARAMS, message=str(e)) - ) + raise MCPError(types.INVALID_PARAMS, str(e)) mcp_task = self._task_data_to_mcp_task(t) return types.CancelTaskResult( - taskId=mcp_task.taskId, + taskId=mcp_task.task_id, status=mcp_task.status, - statusMessage=mcp_task.statusMessage, - createdAt=mcp_task.createdAt, - lastUpdatedAt=mcp_task.lastUpdatedAt, + statusMessage=mcp_task.status_message, + createdAt=mcp_task.created_at, + lastUpdatedAt=mcp_task.last_updated_at, ttl=mcp_task.ttl, - pollInterval=mcp_task.pollInterval, + pollInterval=mcp_task.poll_interval, ) async def handle_get_task_payload(req): - task_id = req.params.taskId + if rejects_deprecated_method("tasks/result", self.protocol_era): + message = deprecated_method_message("tasks/result") + raise MCPError(types.METHOD_NOT_FOUND, message or "Not supported") + params = getattr(req, "params", req) + task_id = getattr(params, "task_id", None) or getattr(params, "taskId", None) + access = extract_task_access_context(request_ctx.get(None)) try: - t = await self.task_manager.wait_until_done(task_id) + t = await self.task_manager.wait_until_done(task_id, access_context=access) except TaskNotFoundError: - raise types.McpError( - types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") - ) + raise MCPError(types.INVALID_PARAMS, f"Task {task_id} not found") if t.status == TaskStatus.COMPLETED: return t.result if t.status == TaskStatus.CANCELLED: @@ -1117,6 +1626,38 @@ async def handle_get_task_payload(req): isError=True, ) + async def on_list_tasks(ctx: ServerRequestContext, params: Optional[PaginatedRequestParams]): + token = request_ctx.set(ctx) + try: + return await handle_list_tasks(types.ListTasksRequest(params=params)) + finally: + request_ctx.reset(token) + + async def on_get_task(ctx: ServerRequestContext, params: types.GetTaskRequestParams): + token = request_ctx.set(ctx) + try: + return await handle_get_task(types.GetTaskRequest(params=params)) + finally: + request_ctx.reset(token) + + async def on_cancel_task(ctx: ServerRequestContext, params: types.CancelTaskRequestParams): + token = request_ctx.set(ctx) + try: + return await handle_cancel_task(types.CancelTaskRequest(params=params)) + finally: + request_ctx.reset(token) + + async def on_get_task_payload(ctx: ServerRequestContext, params: types.GetTaskPayloadRequestParams): + token = request_ctx.set(ctx) + try: + return await handle_get_task_payload(types.GetTaskPayloadRequest(params=params)) + finally: + request_ctx.reset(token) + + server.add_request_handler("tasks/list", PaginatedRequestParams, on_list_tasks) + server.add_request_handler("tasks/get", types.GetTaskRequestParams, on_get_task) + server.add_request_handler("tasks/cancel", types.CancelTaskRequestParams, on_cancel_task) + server.add_request_handler("tasks/result", types.GetTaskPayloadRequestParams, on_get_task_payload) server.request_handlers[types.ListTasksRequest] = handle_list_tasks server.request_handlers[types.GetTaskRequest] = handle_get_task server.request_handlers[types.CancelTaskRequest] = handle_cancel_task @@ -1138,6 +1679,7 @@ async def handle_initialized(notification: types.InitializedNotification): tool_name=config.name, metadata={}, ) + _apply_request_envelope(ctx, request_ctx.get(None)) guards, middleware, interceptors, pipes, filters = self._pipeline_stages(method) await run_pipeline( @@ -1158,18 +1700,38 @@ async def handle_initialized(notification: types.InitializedNotification): sys.stderr.write(f"Error executing initial tool '{config.name}': {e}\n") sys.stderr.flush() + async def on_initialized(ctx: ServerRequestContext, _params: Optional[types.NotificationParams]): + token = request_ctx.set(ctx) + try: + await handle_initialized(types.InitializedNotification()) + finally: + request_ctx.reset(token) + + server.add_notification_handler( + "notifications/initialized", types.NotificationParams, on_initialized + ) server.notification_handlers[types.InitializedNotification] = handle_initialized # ------------------------------------------------------------------ # Transports # ------------------------------------------------------------------ + def _apply_protocol_era(self) -> ProtocolEra: + """Re-read env/config, store the era, and log how it was chosen.""" + resolution = resolve_protocol_era_resolution( + config_value=self.server_config.protocol_era + ) + self.protocol_era = resolution.era + self.protocol_era_source = resolution.source + logger.info(resolution.log_line()) + return resolution.era + def get_combined_app( self, *, max_sessions: Optional[int] = None, session_idle_timeout: Optional[float] = None, - enable_cors: bool = True, + enable_cors: Optional[bool] = None, stateless: Optional[bool] = None, json_response: Optional[bool] = None, ) -> Any: @@ -1179,22 +1741,71 @@ def get_combined_app( (`/mcp/health`) endpoints. See `nitrostack.transports.http.build_http_app` for the full behavior (session cap, CORS, DNS-rebinding protection). - Any argument left as `None` falls back to this app's `ServerConfig`. + Any argument left as ``None`` falls back to env + (``MCP_STATELESS``, ``NITRO_MCP_PROTOCOL_VERSION``, ``ENABLE_CORS``) + then this app's ``ServerConfig.protocol_era``. Unset protocol era is ``auto``. """ from nitrostack.transports.http import build_http_app - return build_http_app( + era = self._apply_protocol_era() + wire_mode = wire_mode_for_era(era) + # Sessionful 1.x only when era is legacy. auto/modern stay sessionless. + http_engine = resolve_http_engine(era, stateless=stateless) + effective_stateless = http_engine == "sessionless" + + if enable_cors is None: + env_cors = self._env_bool("ENABLE_CORS") + enable_cors = True if env_cors is None else env_cors + + http_app = build_http_app( self, max_sessions=max_sessions if max_sessions is not None else self.server_config.max_sessions, session_idle_timeout=session_idle_timeout, enable_cors=enable_cors, - stateless=self.server_config.stateless if stateless is None else stateless, + stateless=effective_stateless, json_response=self.server_config.json_response if json_response is None else json_response, + protocol_era=era, + wire_mode=wire_mode, + http_engine=http_engine, ) + if http_engine == "sessionless": + from nitrostack.transports.middleware import wrap_stateless_transport + + def _discover_handler(_request): + return self.handle_server_discover() + + def _initialize_handler(request): + params = getattr(request, "params", None) or {} + requested = params.get("protocolVersion") if isinstance(params, dict) else None + return self.handle_sessionless_initialize( + requested if isinstance(requested, str) else None + ) + + http_app = wrap_stateless_transport( + http_app, + server_name=self.server_config.name, + server_version=self.server_config.version, + protocol_version=protocol_version_for_era(era, self.server_config.protocol_version), + advertise_tasks=self._advertise_tasks_extension(), + advertise_app=any( + getattr(entry, "component", None) is not None + for entry in getattr(self, "_tools", {}).values() + ), + custom_extensions=self._custom_extensions(), + wire_mode=wire_mode, + protocol_era=era, + enable_cors=enable_cors, + discover_handler=_discover_handler, + initialize_handler=_initialize_handler, + ) + + return http_app + async def _run_stdio(self) -> None: - async with stdio_server() as (read_stream, write_stream): - await self.mcp_server.run(read_stream, write_stream, self.mcp_server.create_initialization_options()) + from nitrostack.transports.stdio import run_stdio + + await run_stdio(self.mcp_server, self.protocol_era) @staticmethod def _env_int(name: str) -> Optional[int]: @@ -1228,10 +1839,8 @@ async def start(self) -> None: transport = os.environ.get("MCP_TRANSPORT_TYPE") or self.server_config.transport_type node_env = os.environ.get("NODE_ENV", "development") port = resolve_http_port() + host = resolve_http_host() - stateless = self._env_bool("MCP_STATELESS") - if stateless is None: - stateless = self.server_config.stateless json_response = self._env_bool("MCP_JSON_RESPONSE") if json_response is None: json_response = self.server_config.json_response @@ -1245,12 +1854,11 @@ async def start(self) -> None: app = self.get_combined_app( max_sessions=max_sessions, session_idle_timeout=session_idle_timeout, - stateless=stateless, json_response=json_response, ) config = uvicorn.Config( app, - host="0.0.0.0", + host=host, port=port, log_level="info", timeout_graceful_shutdown=graceful_timeout_ms / 1000, @@ -1263,18 +1871,18 @@ async def start(self) -> None: app = self.get_combined_app( max_sessions=max_sessions, session_idle_timeout=session_idle_timeout, - stateless=stateless, json_response=json_response, ) await run_dual( self, app, - host="0.0.0.0", + host=host, port=port, graceful_timeout=graceful_timeout_ms / 1000, ) else: # Default Stdio + self._apply_protocol_era() from nitrostack.transports.stdio import safe_stdio_transport with safe_stdio_transport(): await self._run_stdio() diff --git a/nitrostack/core/context.py b/nitrostack/core/context.py index 817b431..42cf129 100644 --- a/nitrostack/core/context.py +++ b/nitrostack/core/context.py @@ -2,10 +2,14 @@ import os import sys from dataclasses import dataclass, field -from typing import Any, Protocol, List, Dict, Optional +from typing import TYPE_CHECKING, Any, Protocol, List, Dict, Optional from nitrostack.core.errors import TaskCancelledError +if TYPE_CHECKING: + from nitrostack.protocol.meta import RequestMeta + from nitrostack.protocol.observability import TraceContext + # Logger protocol used by ExecutionContext (Section 13) class Logger(Protocol): def debug(self, message: str, meta: dict | None = None) -> None: ... @@ -22,7 +26,11 @@ def __init__(self, log_file: Optional[str] = None, name: str = "nitrostack"): self.logger = logging.getLogger(name) # Read log level from environment - level_str = os.environ.get("NITROSTACK_LOG_LEVEL", "DEBUG").upper() + level_str = ( + os.environ.get("NITROSTACK_LOG_LEVEL") + or os.environ.get("NITRO_LOG_LEVEL") + or "DEBUG" + ).upper() level = getattr(logging, level_str, logging.DEBUG) self.logger.setLevel(level) @@ -110,6 +118,7 @@ def __init__( *, session: Any = None, progress_token: Any = None, + correlation_id: Any = None, ): self.task_id = task_id self.progress_message: str = "" @@ -117,17 +126,17 @@ def __init__( self._task_manager = task_manager self._session = session self._progress_token = progress_token + self._correlation_id = correlation_id self._progress_count = 0 def update_progress(self, message: str) -> None: self.progress_message = message manager = self._task_manager if manager is not None: - try: - manager.update_progress(self.task_id, message) - except Exception: - # Task may already be terminal/expired — ignore for handler ergonomics. - pass + updater = getattr(manager, "update_progress_sync", None) + if updater is None: + raise RuntimeError("Task manager does not support synchronous progress updates") + updater(self.task_id, message) self._push_progress_notification(message) def _push_progress_notification(self, message: str) -> None: @@ -141,6 +150,7 @@ def _push_progress_notification(self, message: str) -> None: progress_token=self._progress_token, progress=self._progress_count, message=message, + related_request_id=self._correlation_id, ) ) except Exception: @@ -154,16 +164,16 @@ def cancel(self) -> None: manager = self._task_manager if manager is None: return - try: - manager.cancel_task(self.task_id) - except Exception: - pass + canceller = getattr(manager, "cancel_task_sync", None) + if canceller is None: + raise RuntimeError("Task manager does not support synchronous cancel") + canceller(self.task_id) def throw_if_cancelled(self) -> None: manager = self._task_manager if manager is not None: try: - if manager.is_task_cancelled(self.task_id): + if manager.is_task_cancelled_sync(self.task_id): self.is_cancelled = True except Exception: pass @@ -173,8 +183,28 @@ def throw_if_cancelled(self) -> None: @dataclass class ExecutionContext: request_id: str + correlation_id: str | None = None + jsonrpc_id: Any = None tool_name: str | None = None logger: Logger = field(default_factory=lambda: FileLogger()) metadata: dict = field(default_factory=dict) auth: AuthContext | None = None task: TaskContext | None = None + input_responses: Dict[str, Any] = field(default_factory=dict) + request_state: Optional[Dict[str, Any]] = None + trace: "TraceContext | None" = None + protocol_version: Optional[str] = None + rpc_meta: Optional["RequestMeta"] = None + mcp_headers: Dict[str, str] = field(default_factory=dict) + mcp_param_headers: Dict[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.correlation_id: + self.correlation_id = self.request_id + + @property + def user(self) -> Optional[str]: + """Verified identity only. Unsigned ``_meta.userId`` is never used.""" + if self.auth is None: + return None + return self.auth.subject diff --git a/nitrostack/core/mcp_server.py b/nitrostack/core/mcp_server.py index baec5d7..b383ff2 100644 --- a/nitrostack/core/mcp_server.py +++ b/nitrostack/core/mcp_server.py @@ -1,40 +1,70 @@ """ Low-level MCP server ownership for nitrostack. -Builds directly on the Python `mcp` SDK's low-level `Server` class: nitrostack -owns registration/dispatch and only declares the protocol-level capabilities -it actually implements. +Builds on official ``mcp`` 2.x ``Server``: nitrostack owns registration and +declares only the capabilities it implements. ``/mcp`` is the v2 +``streamable_http_app()`` for every era. """ -from typing import Any, Dict, Optional +from typing import Any, Callable, Dict, Optional import mcp.types as types from mcp.server.lowlevel import Server as LowLevelServer from mcp.server.lowlevel.server import NotificationOptions +from mcp.server.subscriptions import InMemorySubscriptionBus, ListenHandler class NitroStackMcpServer(LowLevelServer): """ - Thin subclass of the official low-level MCP `Server`. + Thin subclass of the official low-level MCP ``Server``. - The base `Server.get_capabilities()` only advertises a capability when a - handler for the corresponding request type has been registered, and it - always reports `resources.subscribe=False`. This subclass declares: - `listChanged=True` for tools/resources/prompts, `resources.subscribe=True` - (nitrostack always registers subscribe/unsubscribe handlers), and the - `tasks` capability whenever nitrostack's task subsystem handlers are - registered on this server. + One instance owns tool, resource, and prompt registration. The HTTP factory + mounts that same server on ``/mcp`` via ``streamable_http_app()``. + Sessionful 1.x is not used; v2 serves both protocol eras. """ def __init__(self, name: str, version: Optional[str] = None): - super().__init__(name=name, version=version) + self.subscription_bus = InMemorySubscriptionBus() + self.listen_handler = ListenHandler(self.subscription_bus) + super().__init__( + name=name, + version=version or "", + on_subscriptions_listen=self.listen_handler, + ) self.has_task_support: bool = False + self.http_engine: Optional[str] = None + self.sessionful: bool = False + self.discover_handler: Optional[Callable[[], dict[str, Any]]] = None + self.initialize_handler: Optional[Callable[[Optional[str]], dict[str, Any]]] = None + # In-process tests still look up handlers by request type. + self.request_handlers: Dict[Any, Callable[..., Any]] = {} + self.notification_handlers: Dict[Any, Callable[..., Any]] = {} + + def handle_server_discover(self) -> dict[str, Any]: + """Answer ``server/discover`` from this server instance.""" + if self.discover_handler is None: + raise RuntimeError("server/discover is not configured on this server") + return self.discover_handler() + + def handle_sessionless_initialize(self, requested_version: Optional[str] = None) -> dict[str, Any]: + """Answer sessionless ``initialize`` from this server instance.""" + if self.initialize_handler is None: + raise RuntimeError("sessionless initialize is not configured on this server") + return self.initialize_handler(requested_version) def get_capabilities( self, - notification_options: NotificationOptions, - experimental_capabilities: Dict[str, Dict[str, Any]], + notification_options: Optional[NotificationOptions] = None, + experimental_capabilities: Optional[Dict[str, Dict[str, Any]]] = None, + extensions: Optional[Dict[str, Dict[str, Any]]] = None, + *, + protocol_version: Optional[str] = None, ) -> types.ServerCapabilities: - caps = super().get_capabilities(notification_options, experimental_capabilities) + caps = super().get_capabilities( + notification_options, + experimental_capabilities, + extensions, + protocol_version=protocol_version, + ) if caps.resources is not None: caps.resources.subscribe = True @@ -51,8 +81,6 @@ def get_capabilities( return caps def create_initialization_options(self, **kwargs: Any): - # Match TS's static `listChanged: true` for tools/resources/prompts by - # default, unless the caller explicitly supplies notification_options. kwargs.setdefault( "notification_options", NotificationOptions( diff --git a/nitrostack/core/task.py b/nitrostack/core/task.py index 57c292a..d1ec1b2 100644 --- a/nitrostack/core/task.py +++ b/nitrostack/core/task.py @@ -1,26 +1,19 @@ """ -MCP Task state machine (Phase 1). +MCP task state machine. -Validated in-memory task store used by McpApplication for MCP Tasks. +``TaskManager`` owns lifecycle logic; persistence is delegated to a pluggable +``TaskStore`` (default: ``InMemoryTaskStore``). -State machine (ASCII):: - - [*] --> WORKING - WORKING --> COMPLETED - WORKING --> FAILED - WORKING --> CANCELLED - WORKING --> EXPIRED (lazy TTL check-on-read) - COMPLETED / FAILED / CANCELLED / EXPIRED are terminal — no further transitions - -TTL: optional ``ttl_seconds`` sets ``expires_at``. Expiration is evaluated lazily -on read/mutate (no background thread). An expired non-terminal task is transitioned -to ``EXPIRED``. +TTL eviction invariants: +- ``working`` / ``input_required`` tasks are never evicted. +- TTL countdown starts only after a terminal transition. +- ``cleanup_expired(now_ms)`` removes terminal tasks where + ``(now_ms - lastUpdatedAt) > ttl_ms``. """ from __future__ import annotations import asyncio -import datetime import uuid from dataclasses import dataclass, field, replace from enum import Enum @@ -32,18 +25,27 @@ TaskExpiredError, TaskNotFoundError, ) +from nitrostack.protocol.tasks import DEFAULT_POLL_INTERVAL_MS, ttl_seconds_to_ms +from nitrostack.tasks.memory import InMemoryTaskStore +from nitrostack.tasks.store import TaskStore +from nitrostack.tasks.types import TaskEntry, TaskWireData, datetime_to_ms, utc_now +from nitrostack.tasks.types import TaskAccessContext +from nitrostack.tasks.authorization import check_task_access, list_task_wire_data_for_context class TaskStatus(Enum): - """Task lifecycle statuses required by Phase 1.""" + """Task lifecycle statuses required by MCP 2026-07-28.""" WORKING = "working" + INPUT_REQUIRED = "input_required" COMPLETED = "completed" FAILED = "failed" CANCELLED = "cancelled" EXPIRED = "expired" +ACTIVE_STATUSES = frozenset({TaskStatus.WORKING, TaskStatus.INPUT_REQUIRED}) + TERMINAL_STATUSES = frozenset( { TaskStatus.COMPLETED, @@ -59,17 +61,21 @@ def is_terminal_status(status: TaskStatus) -> bool: return status in TERMINAL_STATUSES +def _status_from_wire(value: str) -> TaskStatus: + return TaskStatus(value) + + +def _status_to_wire(status: TaskStatus) -> str: + return status.value + + @dataclass class TaskData: """ Snapshot of a task's protocol-visible and result state. ``progress`` holds the latest progress/status message. ``result`` / ``error`` - are populated on successful completion or failure respectively. Tasks created - without a TTL have ``expires_at is None`` and never expire. - - Compatibility aliases (``task_id``, ``status_message``, ``ttl``) mirror the - previous task-entry attribute names used by callers. + are populated on successful completion or failure respectively. """ id: str @@ -77,13 +83,12 @@ class TaskData: progress: Optional[str] = None result: Any = None error: Any = None - created_at: datetime.datetime = field( - default_factory=lambda: datetime.datetime.now(datetime.timezone.utc) - ) - expires_at: Optional[datetime.datetime] = None - last_updated_at: Optional[datetime.datetime] = None + created_at: Any = field(default_factory=utc_now) + expires_at: Optional[Any] = None + last_updated_at: Optional[Any] = None ttl_seconds: Optional[int] = None - poll_interval: int = 5 + ttl_ms: Optional[int] = None + poll_interval: int = DEFAULT_POLL_INTERVAL_MS def __post_init__(self) -> None: if self.last_updated_at is None: @@ -99,211 +104,370 @@ def status_message(self) -> Optional[str]: @property def ttl(self) -> Optional[int]: - return self.ttl_seconds + return self.ttl_ms if self.ttl_ms is not None else ttl_seconds_to_ms(self.ttl_seconds) @dataclass -class _TaskEntry: - """Internal store entry (not part of the public API).""" +class _RuntimeTaskHandle: + """Local-only execution primitives (not persisted to distributed stores).""" - data: TaskData done_event: asyncio.Event = field(default_factory=asyncio.Event) + cancelled: bool = False class TaskManager: """ - In-memory task store with validated transitions, result storage, and lazy TTL. + Task lifecycle manager backed by a pluggable ``TaskStore``. - Typical flow:: - - manager = TaskManager() - task = manager.create_task(ttl_seconds=60) - manager.update_progress(task.id, "halfway") - manager.complete_task(task.id, {"ok": True}) - assert manager.get_task(task.id).result == {"ok": True} + Runtime wait/cancel handles remain process-local even when using Redis or + PostgreSQL persistence. """ - def __init__(self) -> None: - self._tasks: Dict[str, _TaskEntry] = {} + def __init__(self, store: Optional[TaskStore] = None) -> None: + self._store = store or InMemoryTaskStore() + self._runtime: Dict[str, _RuntimeTaskHandle] = {} - def create_task( + async def create_task( self, ttl_seconds: Optional[int] = None, *, + ttl_ms: Optional[int] = None, task_id: Optional[str] = None, + poll_interval_ms: int = DEFAULT_POLL_INTERVAL_MS, + tool_name: Optional[str] = None, + owner_id: Optional[str] = None, + tenant_id: Optional[str] = None, + session_id: Optional[str] = None, ) -> TaskData: - """ - Create a new task in ``WORKING`` status. - - ``ttl_seconds=None`` means the task never expires. ``task_id`` is optional - and intended for compatibility callers that supply their own ID. - """ - now = datetime.datetime.now(datetime.timezone.utc) + """Create a new task in ``WORKING`` status.""" + now = utc_now() resolved_id = task_id or f"task_{uuid.uuid4().hex[:12]}" - if resolved_id in self._tasks: + if await self._store.has(resolved_id): raise ValueError(f"Task {resolved_id} already exists") - expires_at = None - if ttl_seconds is not None: - expires_at = now + datetime.timedelta(seconds=ttl_seconds) - - data = TaskData( - id=resolved_id, - status=TaskStatus.WORKING, - progress="Task started", + + resolved_ttl_seconds = ttl_seconds + resolved_ttl_ms = ttl_ms + if resolved_ttl_ms is not None and resolved_ttl_seconds is None: + resolved_ttl_seconds = max(1, int(resolved_ttl_ms / 1000)) + elif resolved_ttl_seconds is not None and resolved_ttl_ms is None: + resolved_ttl_ms = ttl_seconds_to_ms(resolved_ttl_seconds) + + wire = TaskWireData( + task_id=resolved_id, + status="working", + status_message="Task created", created_at=now, last_updated_at=now, - expires_at=expires_at, - ttl_seconds=ttl_seconds, - poll_interval=5, + ttl_ms=resolved_ttl_ms, + poll_interval_ms=poll_interval_ms, + owner_id=owner_id, + tenant_id=tenant_id, + session_id=session_id, ) - self._tasks[resolved_id] = _TaskEntry(data=data) - return self._snapshot(data) - - def get_task(self, task_id: str) -> TaskData: - """ - Return a snapshot of the task. - - Missing IDs raise ``TaskNotFoundError``. Expired non-terminal tasks are - lazily transitioned to ``EXPIRED`` before the snapshot is returned. - """ - entry = self._get_entry(task_id) - self._maybe_expire(entry) - return self._snapshot(entry.data) - - def update_progress(self, task_id: str, progress: Any) -> None: - """ - Update progress for a ``WORKING`` task. - - Raises ``TaskAlreadyTerminalError`` if the task is already terminal - (including after lazy expiration). - """ - entry = self._get_entry(task_id) - self._maybe_expire(entry) - if is_terminal_status(entry.data.status): - if entry.data.status == TaskStatus.EXPIRED: + entry = TaskEntry( + task_id=resolved_id, + data=wire, + status="working", + tool_name=tool_name, + owner_id=owner_id, + tenant_id=tenant_id, + session_id=session_id, + ) + await self._store.set(resolved_id, entry) + self._runtime[resolved_id] = _RuntimeTaskHandle() + return self._snapshot_from_entry(entry) + + async def get_task( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> TaskData: + """Return a snapshot of the task or raise ``TaskNotFoundError``.""" + entry = await self._require_entry(task_id, access_context=access_context) + return self._snapshot_from_entry(entry) + + def update_progress_sync(self, task_id: str, progress: Any) -> None: + """Synchronous progress update for ``TaskContext`` (same-loop safe).""" + entry = self._require_entry_sync(task_id) + status = _status_from_wire(entry.status) + if is_terminal_status(status): + if status == TaskStatus.EXPIRED: raise TaskExpiredError(task_id) - raise TaskAlreadyTerminalError(task_id, entry.data.status) - entry.data.progress = progress - entry.data.last_updated_at = datetime.datetime.now(datetime.timezone.utc) - - def complete_task(self, task_id: str, result: Any) -> None: - """Transition ``WORKING`` → ``COMPLETED`` and store ``result``.""" - entry = self._get_entry(task_id) - self._maybe_expire(entry) - self._require_working_for_transition(entry, TaskStatus.COMPLETED) - entry.data.result = result - entry.data.error = None - entry.data.progress = "Task completed successfully" - self._set_status(entry, TaskStatus.COMPLETED) - - def fail_task(self, task_id: str, error: Any) -> None: - """Transition ``WORKING`` → ``FAILED`` and store ``error``.""" - entry = self._get_entry(task_id) - self._maybe_expire(entry) - self._require_working_for_transition(entry, TaskStatus.FAILED) - entry.data.error = error - entry.data.progress = f"Task failed: {error}" - self._set_status(entry, TaskStatus.FAILED) - - def cancel_task(self, task_id: str) -> None: - """Transition ``WORKING`` → ``CANCELLED``.""" - entry = self._get_entry(task_id) - self._maybe_expire(entry) - if is_terminal_status(entry.data.status): - if entry.data.status == TaskStatus.EXPIRED: + raise TaskAlreadyTerminalError(task_id, status) + now = utc_now() + entry.data.status_message = str(progress) + entry.data.last_updated_at = now + entry.status = entry.data.status + self._store_set_sync(task_id, entry) + + async def update_progress(self, task_id: str, progress: Any) -> None: + """Update progress for an active task.""" + self.update_progress_sync(task_id, progress) + + async def require_input( + self, + task_id: str, + pause_payload: Any, + *, + progress: str = "Additional input required", + ) -> None: + """Transition an active task to ``input_required``.""" + entry = await self._require_entry(task_id) + status = _status_from_wire(entry.status) + if is_terminal_status(status): + if status == TaskStatus.EXPIRED: raise TaskExpiredError(task_id) - raise TaskAlreadyTerminalError(task_id, entry.data.status) - entry.data.progress = "Task cancelled by client" - self._set_status(entry, TaskStatus.CANCELLED) - - def list_tasks(self) -> List[TaskData]: - """Return snapshots for all known tasks (applies lazy expiration).""" - snapshots: List[TaskData] = [] - for entry in list(self._tasks.values()): - self._maybe_expire(entry) - snapshots.append(self._snapshot(entry.data)) - return snapshots + raise TaskAlreadyTerminalError(task_id, status) + if status not in ACTIVE_STATUSES: + raise InvalidTaskTransitionError(status, TaskStatus.INPUT_REQUIRED) + now = utc_now() + entry.result = pause_payload + entry.status = "input_required" + entry.data.status = "input_required" + entry.data.status_message = progress + entry.data.last_updated_at = now + await self._store.set(task_id, entry) + + async def resume_task(self, task_id: str, *, progress: str = "Resuming task") -> None: + """Transition ``input_required`` back to ``working``.""" + entry = await self._require_entry(task_id) + status = _status_from_wire(entry.status) + if status != TaskStatus.INPUT_REQUIRED: + raise InvalidTaskTransitionError(status, TaskStatus.WORKING) + now = utc_now() + entry.status = "working" + entry.data.status = "working" + entry.data.status_message = progress + entry.data.last_updated_at = now + await self._store.set(task_id, entry) + + async def complete_task(self, task_id: str, result: Any) -> None: + """Transition an active task to ``completed``.""" + entry = await self._require_entry(task_id) + self._require_active_for_transition(entry, TaskStatus.COMPLETED) + entry.result = result + entry.error = None + entry.status = "completed" + entry.data.status = "completed" + entry.data.status_message = "Task completed successfully" + entry.data.last_updated_at = utc_now() + await self._store.set(task_id, entry) + self._signal_done(task_id) + + async def fail_task(self, task_id: str, error: Any) -> None: + """Transition an active task to ``failed``.""" + entry = await self._require_entry(task_id) + self._require_active_for_transition(entry, TaskStatus.FAILED) + entry.error = {"message": str(error)} + entry.status = "failed" + entry.data.status = "failed" + entry.data.status_message = f"Task failed: {error}" + entry.data.last_updated_at = utc_now() + await self._store.set(task_id, entry) + self._signal_done(task_id) + + def cancel_task_sync( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> None: + """Synchronous cancel for ``TaskContext`` (same-loop safe).""" + entry = self._require_entry_sync(task_id, access_context=access_context) + status = _status_from_wire(entry.status) + if is_terminal_status(status): + if status == TaskStatus.EXPIRED: + raise TaskExpiredError(task_id) + raise TaskAlreadyTerminalError(task_id, status) + entry.status = "cancelled" + entry.data.status = "cancelled" + entry.data.status_message = "Task cancelled by client" + entry.data.last_updated_at = utc_now() + self._store_set_sync(task_id, entry) + handle = self._runtime.setdefault(task_id, _RuntimeTaskHandle()) + handle.cancelled = True + self._signal_done(task_id) + + async def cancel_task( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> None: + """Transition an active task to ``cancelled``.""" + self.cancel_task_sync(task_id, access_context=access_context) - def has_task(self, task_id: str) -> bool: - """Return True if a task with ``task_id`` exists in the store.""" - return task_id in self._tasks + async def list_tasks( + self, + *, + access_context: Optional[TaskAccessContext] = None, + cursor: Optional[str] = None, + limit: int = 50, + ) -> List[TaskData]: + """Return task snapshots filtered by caller access context.""" + snapshots, _ = await self.list_tasks_page( + access_context=access_context, + cursor=cursor, + limit=limit, + ) + return snapshots - def is_task_cancelled(self, task_id: str) -> bool: - """Return True if the task exists and is in ``CANCELLED`` status.""" - if task_id not in self._tasks: + async def list_tasks_page( + self, + *, + access_context: Optional[TaskAccessContext] = None, + cursor: Optional[str] = None, + limit: int = 50, + ) -> tuple[List[TaskData], Optional[str]]: + """Return a filtered, paginated task page and optional next cursor.""" + entries = await self._store.list() + page, next_cursor = list_task_wire_data_for_context( + entries, + access_context, + cursor=cursor, + limit=limit, + ) + by_id = {entry.task_id: entry for entry in entries} + snapshots = [ + self._snapshot_from_entry(by_id[wire.task_id]) + for wire in page + if wire.task_id in by_id + ] + return snapshots, next_cursor + + async def has_task(self, task_id: str) -> bool: + return await self._store.has(task_id) + + async def is_task_cancelled(self, task_id: str) -> bool: + handle = self._runtime.get(task_id) + if handle is not None and handle.cancelled: + return True + if not await self._store.has(task_id): return False - entry = self._tasks[task_id] - self._maybe_expire(entry) - return entry.data.status == TaskStatus.CANCELLED - - async def wait_until_done(self, task_id: str) -> TaskData: - """ - Block until the task reaches a terminal state, then return a snapshot. - - Applies lazy expiration before waiting when the task is still working. - """ - entry = self._get_entry(task_id) - self._maybe_expire(entry) - if not is_terminal_status(entry.data.status): - await entry.done_event.wait() - # Re-fetch: status may have changed while waiting. - entry = self._get_entry(task_id) - return self._snapshot(entry.data) - - def get_result(self, task_id: str) -> Any: - """ - Return the stored result for a completed task. - - Raises ``TaskNotFoundError``, ``TaskExpiredError``, or - ``InvalidTaskTransitionError`` if the task is not completed. - """ - data = self.get_task(task_id) + entry = await self._store.get(task_id) + return entry is not None and entry.status == "cancelled" + + def is_task_cancelled_sync(self, task_id: str) -> bool: + """Best-effort synchronous cancel probe for ``TaskContext.throw_if_cancelled``.""" + handle = self._runtime.get(task_id) + if handle is not None and handle.cancelled: + return True + return False + + async def wait_until_done( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> TaskData: + """Block until the task reaches a terminal state.""" + entry = await self._require_entry(task_id, access_context=access_context) + status = _status_from_wire(entry.status) + if not is_terminal_status(status): + handle = self._runtime.setdefault(task_id, _RuntimeTaskHandle()) + await handle.done_event.wait() + entry = await self._require_entry(task_id, access_context=access_context) + return self._snapshot_from_entry(entry) + + async def get_result( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> Any: + """Return the stored result for a completed task.""" + data = await self.get_task(task_id, access_context=access_context) if data.status == TaskStatus.EXPIRED: raise TaskExpiredError(task_id) if data.status != TaskStatus.COMPLETED: raise InvalidTaskTransitionError(data.status, TaskStatus.COMPLETED) return data.result + async def cleanup_expired(self, now_ms: Optional[int] = None) -> int: + """Evict terminal tasks whose post-completion TTL has elapsed.""" + resolved_now = now_ms if now_ms is not None else datetime_to_ms(utc_now()) + evicted = await self._store.cleanup_expired(resolved_now) + for task_id in list(self._runtime.keys()): + if not await self._store.has(task_id): + self._runtime.pop(task_id, None) + return evicted + + async def destroy(self) -> None: + await self._store.destroy() + self._runtime.clear() + # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ - def _get_entry(self, task_id: str) -> _TaskEntry: - entry = self._tasks.get(task_id) + def _store_get_sync(self, task_id: str) -> Optional[TaskEntry]: + getter = getattr(self._store, "get_sync", None) + if getter is None: + raise RuntimeError("Task store does not support synchronous reads") + return getter(task_id) + + def _store_set_sync(self, task_id: str, entry: TaskEntry) -> None: + setter = getattr(self._store, "set_sync", None) + if setter is None: + raise RuntimeError("Task store does not support synchronous writes") + setter(task_id, entry) + + def _require_entry_sync( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> TaskEntry: + entry = self._store_get_sync(task_id) if entry is None: raise TaskNotFoundError(task_id) + check_task_access(entry, access_context) return entry - def _maybe_expire(self, entry: _TaskEntry) -> None: - if entry.data.expires_at is None: - return - if is_terminal_status(entry.data.status): - return - now = datetime.datetime.now(datetime.timezone.utc) - if now >= entry.data.expires_at: - entry.data.status = TaskStatus.EXPIRED - entry.data.progress = "Task expired" - entry.data.last_updated_at = now - entry.done_event.set() - - def _require_working_for_transition( - self, entry: _TaskEntry, to_status: TaskStatus - ) -> None: - current = entry.data.status + async def _require_entry( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> TaskEntry: + entry = await self._store.get(task_id) + if entry is None: + raise TaskNotFoundError(task_id) + check_task_access(entry, access_context) + return entry + + def _require_active_for_transition(self, entry: TaskEntry, to_status: TaskStatus) -> None: + current = _status_from_wire(entry.status) if current == TaskStatus.EXPIRED: - raise TaskExpiredError(entry.data.id) + raise TaskExpiredError(entry.task_id) if is_terminal_status(current): - raise TaskAlreadyTerminalError(entry.data.id, current) - if current != TaskStatus.WORKING: + raise TaskAlreadyTerminalError(entry.task_id, current) + if current not in ACTIVE_STATUSES: raise InvalidTaskTransitionError(current, to_status) - def _set_status(self, entry: _TaskEntry, status: TaskStatus) -> None: - entry.data.status = status - entry.data.last_updated_at = datetime.datetime.now(datetime.timezone.utc) - if is_terminal_status(status): - entry.done_event.set() + def _signal_done(self, task_id: str) -> None: + handle = self._runtime.setdefault(task_id, _RuntimeTaskHandle()) + handle.done_event.set() + + @staticmethod + def _snapshot_from_entry(entry: TaskEntry) -> TaskData: + ttl_ms = entry.data.ttl_ms + ttl_seconds = max(1, int(ttl_ms / 1000)) if ttl_ms is not None else None + return TaskData( + id=entry.task_id, + status=_status_from_wire(entry.status), + progress=entry.data.status_message, + result=entry.result, + error=entry.error, + created_at=entry.data.created_at, + last_updated_at=entry.data.last_updated_at, + expires_at=None, + ttl_seconds=ttl_seconds, + ttl_ms=ttl_ms, + poll_interval=entry.data.poll_interval_ms, + ) @staticmethod def _snapshot(data: TaskData) -> TaskData: - """Return a shallow copy so callers cannot mutate internal state.""" return replace(data) diff --git a/nitrostack/protocol/__init__.py b/nitrostack/protocol/__init__.py new file mode 100644 index 0000000..86dad4b --- /dev/null +++ b/nitrostack/protocol/__init__.py @@ -0,0 +1,208 @@ +"""MCP 2.0 protocol layer — version constants, extensions, and wire invariants.""" + +from nitrostack.protocol.constants import ( + LEGACY_SESSION_HEADER, + MAX_CIMD_BYTES, + MAX_SCHEMA_DEPTH, +) +from nitrostack.protocol.contracts import ( + build_cache_hint_meta, + build_prompt_get_result, + build_resource_blob_content, + build_resource_text_content, +) +from nitrostack.protocol.deprecated import deprecated_method_message, rejects_deprecated_method +from nitrostack.protocol.method_contract import ( + DEPRECATED_MODERN_METHODS, + MODERN_METHOD_CONTRACTS, + MethodContract, + contract_for, + mcp_method_is_required, + mcp_name_field, + mcp_name_is_required, +) +from nitrostack.protocol.discovery import ( + DISCOVER_RESULT_TYPE, + INITIALIZE_METHOD, + SERVER_DISCOVER_METHOD, + build_discover_result, + build_sessionless_initialize_result, +) +from nitrostack.protocol.errors import ERROR_CODE_MESSAGES, JsonRpcErrorCode +from nitrostack.protocol.extensions import MCPExtensionId +from nitrostack.protocol.jsonrpc import ( + JsonRpcWireError, + build_ping_response, + build_tool_error_result, + jsonrpc_error, + jsonrpc_success, + map_exception_to_jsonrpc, + parse_jsonrpc_request, +) +from nitrostack.protocol.layers import RuntimeLayer +from nitrostack.protocol.cache_hints import ( + build_list_endpoint_cache_hint_meta, + resolve_resource_cache_hint_meta, + resolve_tool_cache_hint_meta, +) +from nitrostack.protocol.observability import TraceContext, extract_trace_context, trace_context_from_request_meta +from nitrostack.protocol.meta import ( + RequestEnvelope, + RequestMeta, + bind_request_envelope, + envelope_identity_is_ignored, + envelope_protocol_version, + extract_request_meta, + flatten_request_meta_object, + split_params_and_meta, + strip_tool_arguments, +) +from nitrostack.protocol.mrtr import ( + InputRequest, + InputRequiredResult, + accepted_content, + build_input_required_jsonrpc_result, + input_required, + split_mrtr_tool_params, +) +from nitrostack.protocol.resources import resolve_resource_uri, uri_template_to_pattern +from nitrostack.protocol.schema import ( + JSON_SCHEMA_2020_12_URI, + UnsupportedJsonSchemaError, + assert_json_schema_2020_12, + bound_schema_depth, + gate_registered_schema, + normalize_input_schema, + normalize_output_schema, +) +from nitrostack.protocol.tasks import ( + DEFAULT_POLL_INTERVAL_MS, + DEFAULT_TASK_TTL_MS, + RESULT_TYPE_TASK, + build_task_create_jsonrpc_result, + task_support_forbidden_message, + task_support_required_message, + ttl_ms_to_seconds, + ttl_seconds_to_ms, +) +from nitrostack.protocol.version import ( + LEGACY_PROTOCOL_VERSION, + MODERN_PROTOCOL_VERSION, + PROTOCOL_ERA_ENV, + STATELESS_OVERRIDE_ENV, + SUPPORTED_PROTOCOL_VERSIONS, + ProtocolEra, + WireMode, + HttpEngine, + http_engine_for_era, + resolve_http_engine, + accepts_sessionless_initialize, + rejects_legacy_initialize, + needs_modern_engine, + needs_sessionful_engine, + protocol_era_for_wire_mode, + protocol_version_for_era, + EraSource, + ProtocolEraResolution, + resolve_protocol_era, + resolve_protocol_era_resolution, + stateless_for_era, + supported_protocol_versions_for_era, + wire_mode_for_era, +) + +__all__ = [ + "LEGACY_PROTOCOL_VERSION", + "MODERN_PROTOCOL_VERSION", + "PROTOCOL_ERA_ENV", + "STATELESS_OVERRIDE_ENV", + "SUPPORTED_PROTOCOL_VERSIONS", + "ProtocolEra", + "WireMode", + "HttpEngine", + "http_engine_for_era", + "resolve_http_engine", + "accepts_sessionless_initialize", + "rejects_legacy_initialize", + "needs_modern_engine", + "needs_sessionful_engine", + "protocol_era_for_wire_mode", + "protocol_version_for_era", + "EraSource", + "ProtocolEraResolution", + "resolve_protocol_era", + "resolve_protocol_era_resolution", + "stateless_for_era", + "supported_protocol_versions_for_era", + "wire_mode_for_era", + "MCPExtensionId", + "RuntimeLayer", + "LEGACY_SESSION_HEADER", + "MAX_CIMD_BYTES", + "MAX_SCHEMA_DEPTH", + "JsonRpcErrorCode", + "ERROR_CODE_MESSAGES", + "RequestMeta", + "RequestEnvelope", + "bind_request_envelope", + "envelope_identity_is_ignored", + "envelope_protocol_version", + "extract_request_meta", + "flatten_request_meta_object", + "split_params_and_meta", + "strip_tool_arguments", + "deprecated_method_message", + "rejects_deprecated_method", + "DEPRECATED_MODERN_METHODS", + "MODERN_METHOD_CONTRACTS", + "MethodContract", + "contract_for", + "mcp_method_is_required", + "mcp_name_field", + "mcp_name_is_required", + "DISCOVER_RESULT_TYPE", + "INITIALIZE_METHOD", + "SERVER_DISCOVER_METHOD", + "build_sessionless_initialize_result", + "build_discover_result", + "parse_jsonrpc_request", + "build_ping_response", + "jsonrpc_success", + "jsonrpc_error", + "build_tool_error_result", + "map_exception_to_jsonrpc", + "JsonRpcWireError", + "JSON_SCHEMA_2020_12_URI", + "UnsupportedJsonSchemaError", + "assert_json_schema_2020_12", + "bound_schema_depth", + "gate_registered_schema", + "normalize_input_schema", + "normalize_output_schema", + "resolve_resource_uri", + "uri_template_to_pattern", + "build_cache_hint_meta", + "build_resource_text_content", + "build_resource_blob_content", + "build_prompt_get_result", + "InputRequest", + "InputRequiredResult", + "accepted_content", + "input_required", + "split_mrtr_tool_params", + "build_input_required_jsonrpc_result", + "RESULT_TYPE_TASK", + "DEFAULT_TASK_TTL_MS", + "DEFAULT_POLL_INTERVAL_MS", + "build_task_create_jsonrpc_result", + "ttl_ms_to_seconds", + "ttl_seconds_to_ms", + "task_support_forbidden_message", + "task_support_required_message", + "TraceContext", + "extract_trace_context", + "trace_context_from_request_meta", + "build_list_endpoint_cache_hint_meta", + "resolve_tool_cache_hint_meta", + "resolve_resource_cache_hint_meta", +] diff --git a/nitrostack/protocol/cache_hints.py b/nitrostack/protocol/cache_hints.py new file mode 100644 index 0000000..7e4b151 --- /dev/null +++ b/nitrostack/protocol/cache_hints.py @@ -0,0 +1,78 @@ +"""Cache hint resolution for tools, resources, and list endpoints.""" + +from __future__ import annotations + +from typing import Any, Callable, Literal, Optional + +from nitrostack.protocol.contracts import MCP_CACHE_HINT_KEY, build_cache_hint_meta + +DEFAULT_LIST_CACHE_TTL_MS = 60_000 + + +def resolve_tool_cache_hint_meta( + tool_config: Any, + method: Optional[Callable] = None, +) -> Optional[dict[str, Any]]: + """ + Resolve tool ``_meta`` cache hints. + + Priority: explicit ``metadata['cacheHint']`` / wire key, then ``@cache(ttl=...)``. + """ + metadata = getattr(tool_config, "metadata", None) or {} + explicit = metadata.get("cacheHint") or metadata.get(MCP_CACHE_HINT_KEY) + if isinstance(explicit, dict): + ttl_ms = explicit.get("ttlMs") + scope = explicit.get("cacheScope", "private") + if isinstance(ttl_ms, int) and ttl_ms >= 0: + return build_cache_hint_meta(ttl_ms, cache_scope=scope if scope in ("public", "private") else "private") + + cache_hint = getattr(tool_config, "cache_hint", None) + if isinstance(cache_hint, dict): + ttl_ms = cache_hint.get("ttlMs") + scope = cache_hint.get("cacheScope", "private") + if isinstance(ttl_ms, int) and ttl_ms >= 0: + return build_cache_hint_meta(ttl_ms, cache_scope=scope if scope in ("public", "private") else "private") + + if method is not None: + ttl_seconds = getattr(method, "_mcp_cache_ttl", None) + if isinstance(ttl_seconds, (int, float)) and ttl_seconds >= 0: + scope = getattr(method, "_mcp_cache_scope", "private") + return build_cache_hint_meta( + int(ttl_seconds * 1000), + cache_scope=scope if scope in ("public", "private") else "private", + ) + + return None + + +def resolve_resource_cache_hint_meta(resource_config: Any) -> Optional[dict[str, Any]]: + """Resolve resource cache hints from ``cacheHint`` or ``cacheMaxAge`` metadata.""" + metadata = getattr(resource_config, "metadata", None) or {} + explicit = metadata.get("cacheHint") or metadata.get(MCP_CACHE_HINT_KEY) + if isinstance(explicit, dict): + ttl_ms = explicit.get("ttlMs") + scope = explicit.get("cacheScope", "private") + if isinstance(ttl_ms, int) and ttl_ms >= 0: + return build_cache_hint_meta(ttl_ms, cache_scope=scope if scope in ("public", "private") else "private") + + cache_max_age = metadata.get("cacheMaxAge") + if isinstance(cache_max_age, (int, float)) and cache_max_age >= 0: + return build_cache_hint_meta(int(cache_max_age * 1000), cache_scope="private") + + cache_hint = getattr(resource_config, "cache_hint", None) + if isinstance(cache_hint, dict): + ttl_ms = cache_hint.get("ttlMs") + scope = cache_hint.get("cacheScope", "private") + if isinstance(ttl_ms, int) and ttl_ms >= 0: + return build_cache_hint_meta(ttl_ms, cache_scope=scope if scope in ("public", "private") else "private") + + return None + + +def build_list_endpoint_cache_hint_meta( + *, + ttl_ms: int = DEFAULT_LIST_CACHE_TTL_MS, + cache_scope: Literal["public", "private"] = "private", +) -> dict[str, Any]: + """Cache hint for ``tools/list``, ``resources/list``, and ``prompts/list`` responses.""" + return build_cache_hint_meta(ttl_ms, cache_scope=cache_scope) diff --git a/nitrostack/protocol/constants.py b/nitrostack/protocol/constants.py new file mode 100644 index 0000000..3441a54 --- /dev/null +++ b/nitrostack/protocol/constants.py @@ -0,0 +1,10 @@ +"""Cross-cutting security and validation limits.""" + +# RFC 6890 CIMD resolver payload cap +MAX_CIMD_BYTES = 5120 + +# JSON Schema depth bounding for DoS protection (SEP-2106) +MAX_SCHEMA_DEPTH = 64 + +# Legacy session header that MUST NOT be emitted in stateless mode +LEGACY_SESSION_HEADER = "Mcp-Session-Id" diff --git a/nitrostack/protocol/contracts.py b/nitrostack/protocol/contracts.py new file mode 100644 index 0000000..a6cfdea --- /dev/null +++ b/nitrostack/protocol/contracts.py @@ -0,0 +1,80 @@ +"""Wire contract builders for tools, resources, and prompts.""" + +from __future__ import annotations + +from typing import Any, Literal, Optional + +MCP_CACHE_HINT_KEY = "io.modelcontextprotocol/cacheHint" + + +def build_cache_hint_meta( + ttl_ms: int, + *, + cache_scope: Literal["public", "private"] = "private", +) -> dict[str, Any]: + """Build tool/resource ``_meta`` cache hint (SEP-2549).""" + return { + MCP_CACHE_HINT_KEY: { + "ttlMs": ttl_ms, + "cacheScope": cache_scope, + } + } + + +def merge_contract_meta(*parts: Optional[dict[str, Any]]) -> dict[str, Any]: + """Merge multiple ``_meta`` dicts for wire definitions.""" + merged: dict[str, Any] = {} + for part in parts: + if part: + merged.update(part) + return merged + + +def build_resource_text_content( + uri: str, + text: str, + *, + mime_type: str = "application/json", +) -> dict[str, Any]: + """Build a text ``resources/read`` content item.""" + return { + "uri": uri, + "mimeType": mime_type, + "text": text, + } + + +def build_resource_blob_content( + uri: str, + blob: str, + *, + mime_type: str = "application/octet-stream", +) -> dict[str, Any]: + """Build a base64 blob ``resources/read`` content item.""" + return { + "uri": uri, + "mimeType": mime_type, + "blob": blob, + } + + +def build_prompt_text_message(role: str, text: str) -> dict[str, Any]: + """Build a single prompt message entry.""" + return { + "role": role, + "content": { + "type": "text", + "text": text, + }, + } + + +def build_prompt_get_result( + description: str, + messages: list[dict[str, Any]], +) -> dict[str, Any]: + """Build the ``prompts/get`` result payload.""" + return { + "description": description, + "messages": messages, + } diff --git a/nitrostack/protocol/deprecated.py b/nitrostack/protocol/deprecated.py new file mode 100644 index 0000000..264ac71 --- /dev/null +++ b/nitrostack/protocol/deprecated.py @@ -0,0 +1,13 @@ +"""Deprecated MCP methods rejected on the 2026-07-28 wire.""" + +from nitrostack.protocol.method_contract import ( + DEPRECATED_MODERN_METHODS, + deprecated_method_message, + rejects_deprecated_method, +) + +__all__ = [ + "DEPRECATED_MODERN_METHODS", + "deprecated_method_message", + "rejects_deprecated_method", +] diff --git a/nitrostack/protocol/discovery.py b/nitrostack/protocol/discovery.py new file mode 100644 index 0000000..d15fc0a --- /dev/null +++ b/nitrostack/protocol/discovery.py @@ -0,0 +1,104 @@ +"""server/discover capability negotiation. + +Single payload builder for the HTTP engine. The ingress pipeline must not +construct a second discover result. +""" + +from __future__ import annotations + +from typing import Any, Literal, Optional + +from nitrostack.protocol.cache_hints import DEFAULT_LIST_CACHE_TTL_MS, build_list_endpoint_cache_hint_meta +from nitrostack.protocol.extensions import MCPExtensionId +from nitrostack.protocol.version import ( + LEGACY_PROTOCOL_VERSION, + MODERN_PROTOCOL_VERSION, + SUPPORTED_PROTOCOL_VERSIONS, +) + +SERVER_DISCOVER_METHOD = "server/discover" +DISCOVER_RESULT_TYPE = "complete" +INITIALIZE_METHOD = "initialize" +INITIALIZED_NOTIFICATION = "notifications/initialized" +_SESSIONLESS_INITIALIZE_VERSIONS = frozenset( + {LEGACY_PROTOCOL_VERSION, MODERN_PROTOCOL_VERSION, "2025-11-25"} +) + + +def build_discover_result( + *, + server_name: str, + server_version: str, + protocol_version: str = MODERN_PROTOCOL_VERSION, + supported_versions: Optional[tuple[str, ...]] = None, + advertise_tasks: bool = True, + advertise_app: bool = False, + custom_extensions: Optional[dict[str, str]] = None, + tools_list_changed: bool = True, + resources_list_changed: bool = True, + prompts_list_changed: bool = True, + ttl_ms: int = DEFAULT_LIST_CACHE_TTL_MS, + cache_scope: Literal["public", "private"] = "private", +) -> dict[str, Any]: + """Build the ``server/discover`` result for the mounted HTTP engine.""" + versions = list(supported_versions or SUPPORTED_PROTOCOL_VERSIONS) + extensions: dict[str, dict[str, str]] = {} + if advertise_app: + extensions[MCPExtensionId.APP.value] = {"version": protocol_version} + if advertise_tasks: + extensions[MCPExtensionId.TASKS.value] = {"version": protocol_version} + for extension_id, version in (custom_extensions or {}).items(): + if extension_id and version: + extensions[str(extension_id)] = {"version": str(version)} + + capabilities: dict[str, Any] = { + "tools": {"listChanged": tools_list_changed}, + "resources": {"subscribe": False, "listChanged": resources_list_changed}, + "prompts": {"listChanged": prompts_list_changed}, + } + if extensions: + capabilities["extensions"] = extensions + + return { + "resultType": DISCOVER_RESULT_TYPE, + "protocolVersion": protocol_version, + "supportedVersions": versions, + "serverInfo": {"name": server_name, "version": server_version}, + "capabilities": capabilities, + "ttlMs": ttl_ms, + "cacheScope": cache_scope, + "_meta": build_list_endpoint_cache_hint_meta(ttl_ms=ttl_ms, cache_scope=cache_scope), + } + + +def build_sessionless_initialize_result( + *, + server_name: str, + server_version: str, + requested_version: Optional[str] = None, + protocol_version: str = MODERN_PROTOCOL_VERSION, + advertise_tasks: bool = True, + advertise_app: bool = False, + custom_extensions: Optional[dict[str, str]] = None, +) -> dict[str, Any]: + """ + Initialize-shaped result for era ``auto`` with no session. + + Official v2 ``legacy: 'stateless'`` is not mounted yet. This adapter + answers 2025 ``initialize`` on the same engine as ``server/discover``. + """ + requested = (requested_version or "").strip() + negotiated = requested if requested in _SESSIONLESS_INITIALIZE_VERSIONS else protocol_version + discovered = build_discover_result( + server_name=server_name, + server_version=server_version, + protocol_version=protocol_version, + advertise_tasks=advertise_tasks, + advertise_app=advertise_app, + custom_extensions=custom_extensions, + ) + return { + "protocolVersion": negotiated, + "capabilities": discovered["capabilities"], + "serverInfo": discovered["serverInfo"], + } diff --git a/nitrostack/protocol/errors.py b/nitrostack/protocol/errors.py new file mode 100644 index 0000000..09d4bb0 --- /dev/null +++ b/nitrostack/protocol/errors.py @@ -0,0 +1,26 @@ +"""Standard JSON-RPC error codes for MCP 2026-07-28.""" + +from __future__ import annotations + +from enum import IntEnum + + +class JsonRpcErrorCode(IntEnum): + PARSE_ERROR = -32700 + INVALID_REQUEST = -32600 + METHOD_NOT_FOUND = -32601 + INVALID_PARAMS = -32602 + INTERNAL_ERROR = -32603 + HEADER_BODY_MISMATCH = -32020 + UNSUPPORTED_PROTOCOL_VERSION = -32022 + + +ERROR_CODE_MESSAGES: dict[JsonRpcErrorCode, str] = { + JsonRpcErrorCode.PARSE_ERROR: "Parse error", + JsonRpcErrorCode.INVALID_REQUEST: "Invalid Request", + JsonRpcErrorCode.METHOD_NOT_FOUND: "Method not found", + JsonRpcErrorCode.INVALID_PARAMS: "Invalid params", + JsonRpcErrorCode.INTERNAL_ERROR: "Internal error", + JsonRpcErrorCode.HEADER_BODY_MISMATCH: "Header/body mismatch", + JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION: "Unsupported protocol version", +} diff --git a/nitrostack/protocol/extensions.py b/nitrostack/protocol/extensions.py new file mode 100644 index 0000000..43d0125 --- /dev/null +++ b/nitrostack/protocol/extensions.py @@ -0,0 +1,10 @@ +"""Named extension identifiers (SEP-2133) advertised via server/discover.""" + +from enum import Enum + + +class MCPExtensionId(str, Enum): + """Canonical MCP 2.0 extension keys from the capabilities.extensions map.""" + + APP = "io.modelcontextprotocol/app" + TASKS = "io.modelcontextprotocol/tasks" diff --git a/nitrostack/protocol/jsonrpc.py b/nitrostack/protocol/jsonrpc.py new file mode 100644 index 0000000..24239e2 --- /dev/null +++ b/nitrostack/protocol/jsonrpc.py @@ -0,0 +1,282 @@ +"""JSON-RPC 2.0 wire protocol.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from collections.abc import Collection +from typing import Any, Optional + +from nitrostack.protocol.errors import JsonRpcErrorCode, ERROR_CODE_MESSAGES +from nitrostack.protocol.meta import RequestMeta, split_params_and_meta + +JSONRPC_VERSION = "2.0" + +# Backward-compatible aliases +PARSE_ERROR = int(JsonRpcErrorCode.PARSE_ERROR) +HEADER_BODY_MISMATCH = int(JsonRpcErrorCode.HEADER_BODY_MISMATCH) +UNSUPPORTED_PROTOCOL_VERSION = int(JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION) + + +@dataclass(frozen=True) +class JsonRpcRequest: + """Parsed JSON-RPC 2.0 request object.""" + + id: Any + method: str + params: dict[str, Any] + meta: RequestMeta + + +class JsonRpcWireError(Exception): + """Base class for JSON-RPC wire failures.""" + + code: JsonRpcErrorCode + message: str + data: Any = None + + def __init__( + self, + code: JsonRpcErrorCode, + message: Optional[str] = None, + data: Any = None, + ) -> None: + self.code = code + self.message = message or ERROR_CODE_MESSAGES.get(code, "Error") + self.data = data + super().__init__(self.message) + + def to_response(self, request_id: Any) -> dict[str, Any]: + return jsonrpc_error(request_id, int(self.code), self.message, self.data) + + +class JsonRpcParseError(JsonRpcWireError): + def __init__(self, message: str = "Parse error") -> None: + super().__init__(JsonRpcErrorCode.PARSE_ERROR, message) + + +class InvalidRequestError(JsonRpcWireError): + def __init__(self, message: str = "Invalid Request", data: Any = None) -> None: + super().__init__(JsonRpcErrorCode.INVALID_REQUEST, message, data) + + +class MethodNotFoundError(JsonRpcWireError): + def __init__(self, method: str) -> None: + super().__init__( + JsonRpcErrorCode.METHOD_NOT_FOUND, + f"Method not found: {method}", + ) + + +class InvalidParamsError(JsonRpcWireError): + def __init__(self, message: str = "Invalid params", data: Any = None) -> None: + super().__init__(JsonRpcErrorCode.INVALID_PARAMS, message, data) + + +class InternalError(JsonRpcWireError): + def __init__(self, message: str = "Internal error", data: Any = None) -> None: + super().__init__(JsonRpcErrorCode.INTERNAL_ERROR, message, data) + + +class HeaderBodyMismatchError(JsonRpcWireError): + def __init__(self, message: str = "Header/body mismatch") -> None: + super().__init__(JsonRpcErrorCode.HEADER_BODY_MISMATCH, message) + + +class UnsupportedProtocolVersionError(JsonRpcWireError): + def __init__(self, message: Optional[str] = None) -> None: + super().__init__( + JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION, + message or ERROR_CODE_MESSAGES[JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION], + ) + + +def parse_jsonrpc_request(raw_body: bytes) -> JsonRpcRequest: + """Parse and validate a JSON-RPC 2.0 request from HTTP POST body.""" + try: + payload = json.loads(raw_body.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise JsonRpcParseError(str(exc)) from exc + + if not isinstance(payload, dict): + raise InvalidRequestError("Request must be a JSON object") + if payload.get("jsonrpc") != JSONRPC_VERSION: + raise InvalidRequestError("jsonrpc must be '2.0'") + if "method" not in payload or not isinstance(payload["method"], str): + raise InvalidRequestError("method is required and must be a string") + + params = payload.get("params") or {} + if not isinstance(params, dict): + raise InvalidRequestError("params must be an object when present") + + business_params, meta = split_params_and_meta(params) + + return JsonRpcRequest( + id=payload.get("id"), + method=payload["method"], + params=business_params, + meta=meta, + ) + + +def jsonrpc_method_from_body(raw_body: Optional[bytes]) -> Optional[str]: + """JSON-RPC method from a POST body, or None when the body is not a request.""" + if not raw_body: + return None + try: + return parse_jsonrpc_request(raw_body).method + except (JsonRpcParseError, JsonRpcWireError): + return None + + +def validate_header_body_method(header_method: Optional[str], body_method: str) -> None: + """SEP-2243: reject when Mcp-Method header mirrors a different JSON-RPC method.""" + if header_method is None: + return + if header_method != body_method: + raise HeaderBodyMismatchError( + f"Mcp-Method header '{header_method}' does not match body method '{body_method}'" + ) + + +def validate_header_body_name( + header_name: Optional[str], + body_name: Optional[str], +) -> None: + """SEP-2243: reject when Mcp-Name header mirrors a different resource/tool name.""" + if header_name is None or body_name is None: + return + if header_name != body_name: + raise HeaderBodyMismatchError( + f"Mcp-Name header '{header_name}' does not match body name '{body_name}'" + ) + + +REQUIRED_MCP_NAME_MESSAGE = "Mcp-Name header is required" +REQUIRED_MCP_METHOD_MESSAGE = "Mcp-Method header is required" + + +def validate_required_header_body( + header_value: Optional[str], + body_value: Optional[str], + *, + header_name: str, + body_label: str, + required_message: str, +) -> None: + """Require a header and match it to the JSON-RPC body field exactly.""" + header = header_value.strip() if isinstance(header_value, str) else "" + body = body_value.strip() if isinstance(body_value, str) else "" + if not header or not body: + raise HeaderBodyMismatchError(required_message) + if header != body: + raise HeaderBodyMismatchError( + f"{header_name} header '{header_value}' does not match body {body_label} '{body_value}'" + ) + + +def validate_required_mcp_name( + header_name: Optional[str], + body_name: Optional[str], + *, + body_label: str = "name", +) -> None: + """Require ``Mcp-Name`` and match the body name or URI exactly.""" + validate_required_header_body( + header_name, + body_name, + header_name="Mcp-Name", + body_label=body_label, + required_message=REQUIRED_MCP_NAME_MESSAGE, + ) + + +def validate_required_mcp_method( + header_method: Optional[str], + body_method: Optional[str], +) -> None: + """Require ``Mcp-Method`` and match JSON-RPC ``method`` exactly.""" + validate_required_header_body( + header_method, + body_method, + header_name="Mcp-Method", + body_label="method", + required_message=REQUIRED_MCP_METHOD_MESSAGE, + ) + + +def validate_protocol_version_header_meta( + header_version: Optional[str], + meta_version: Optional[str], +) -> None: + """Reject when header and envelope protocol versions are both set and differ.""" + header = header_version.strip() if isinstance(header_version, str) else "" + meta = meta_version.strip() if isinstance(meta_version, str) else "" + if header and meta and header != meta: + raise HeaderBodyMismatchError( + f"MCP-Protocol-Version header '{header_version}' does not match " + f"envelope protocol version '{meta_version}'" + ) + + +def validate_supported_protocol_version( + version: Optional[str], + supported: Collection[str], +) -> None: + """Reject a present protocol version that is not in the era's supported set.""" + if isinstance(version, str) and version.strip() and version.strip() not in supported: + raise UnsupportedProtocolVersionError() + + +def jsonrpc_success(request_id: Any, result: Any) -> dict[str, Any]: + return {"jsonrpc": JSONRPC_VERSION, "id": request_id, "result": result} + + +def jsonrpc_error(request_id: Any, code: int, message: str, data: Any = None) -> dict[str, Any]: + error: dict[str, Any] = {"code": code, "message": message} + if data is not None: + error["data"] = data + return {"jsonrpc": JSONRPC_VERSION, "id": request_id, "error": error} + + +def build_ping_response(request_id: Any) -> dict[str, Any]: + """Ping fast path.""" + return jsonrpc_success(request_id, {}) + + +def build_tool_error_result(message: str, *, text_type: str = "text") -> dict[str, Any]: + """ + Tool business failure — JSON-RPC success with isError: true. + NOT a top-level JSON-RPC error response. + """ + return { + "content": [{"type": text_type, "text": message}], + "isError": True, + } + + +def wrap_tool_success_result(content: list[dict[str, Any]]) -> dict[str, Any]: + """Standard successful tool result payload.""" + return {"content": content, "isError": False} + + +def map_exception_to_jsonrpc(exc: Exception, request_id: Any) -> dict[str, Any]: + """Map SDK exceptions to JSON-RPC wire responses.""" + if isinstance(exc, JsonRpcWireError): + return exc.to_response(request_id) + + from nitrostack.core.errors import ( + PromptNotFoundError, + ResourceNotFoundError, + ToolExecutionError, + ValidationError, + ) + + if isinstance(exc, (ResourceNotFoundError, PromptNotFoundError, ValidationError)): + # SEP-2164: missing resources return -32602 + return InvalidParamsError(str(exc)).to_response(request_id) + + if isinstance(exc, ToolExecutionError): + return jsonrpc_success(request_id, build_tool_error_result(str(exc))) + + return InternalError(str(exc)).to_response(request_id) diff --git a/nitrostack/protocol/layers.py b/nitrostack/protocol/layers.py new file mode 100644 index 0000000..7f30a65 --- /dev/null +++ b/nitrostack/protocol/layers.py @@ -0,0 +1,12 @@ +"""Runtime layer identifiers for the MCP 2026-07-28 architecture.""" + +from enum import Enum + + +class RuntimeLayer(str, Enum): + """Decoupled layers of the NitroStack stateless MCP runtime.""" + + TRANSPORT_SECURITY = "transport_security" + PROTOCOL_DISPATCHER = "protocol_dispatcher" + REGISTRIES = "registries" + TASK_MANAGEMENT = "task_management" diff --git a/nitrostack/protocol/meta.py b/nitrostack/protocol/meta.py new file mode 100644 index 0000000..a068e30 --- /dev/null +++ b/nitrostack/protocol/meta.py @@ -0,0 +1,178 @@ +"""Request _meta envelope parsing.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping, Optional + +MCP_META_PREFIX = "io.modelcontextprotocol/" +# Unsigned envelope keys that must not become handler identity. +_IDENTITY_META_KEYS = frozenset( + {"userid", "user_id", "user", "tenantid", "tenant_id", "tenant"} +) +_PROTOCOL_VERSION_HEADER = "mcp-protocol-version" + + +@dataclass(frozen=True) +class RequestMeta: + """Parsed _meta envelope lifted into execution context.""" + + protocol_version: Optional[str] = None + client_info: Optional[dict[str, Any]] = None + client_capabilities: Optional[dict[str, Any]] = None + traceparent: Optional[str] = None + tracestate: Optional[str] = None + baggage: Optional[str] = None + trace: Optional[dict[str, Any]] = None + auth: Optional[Any] = None + raw: dict[str, Any] = field(default_factory=dict) + + +def _meta_get(meta: dict[str, Any], bare: str, prefixed: str) -> Optional[Any]: + if prefixed in meta: + return meta[prefixed] + return meta.get(bare) + + +def extract_request_meta(params: dict[str, Any]) -> RequestMeta: + """Parse _meta from JSON-RPC params (reverse-DNS and bare keys).""" + raw = params.get("_meta") + if not isinstance(raw, dict): + return RequestMeta() + + protocol_version = _meta_get(raw, "protocolVersion", f"{MCP_META_PREFIX}protocolVersion") + if not isinstance(protocol_version, str): + mcp = raw.get("mcp") + if isinstance(mcp, dict): + nested = mcp.get("protocolVersion") + protocol_version = nested if isinstance(nested, str) else None + client_info = _meta_get(raw, "clientInfo", f"{MCP_META_PREFIX}clientInfo") + client_capabilities = _meta_get( + raw, "clientCapabilities", f"{MCP_META_PREFIX}clientCapabilities" + ) + traceparent = _meta_get(raw, "traceparent", f"{MCP_META_PREFIX}traceparent") + tracestate = _meta_get(raw, "tracestate", f"{MCP_META_PREFIX}tracestate") + baggage = _meta_get(raw, "baggage", f"{MCP_META_PREFIX}baggage") + trace = _meta_get(raw, "trace", f"{MCP_META_PREFIX}trace") + auth = _meta_get(raw, "auth", f"{MCP_META_PREFIX}auth") + + return RequestMeta( + protocol_version=protocol_version if isinstance(protocol_version, str) else None, + client_info=client_info if isinstance(client_info, dict) else None, + client_capabilities=client_capabilities if isinstance(client_capabilities, dict) else None, + traceparent=traceparent if isinstance(traceparent, str) else None, + tracestate=tracestate if isinstance(tracestate, str) else None, + baggage=baggage if isinstance(baggage, str) else None, + trace=trace if isinstance(trace, dict) else None, + auth=auth, + raw=dict(raw), + ) + + +def split_params_and_meta(params: dict[str, Any]) -> tuple[dict[str, Any], RequestMeta]: + """Separate business params from the _meta envelope.""" + meta = extract_request_meta(params) + business = {key: value for key, value in params.items() if key != "_meta"} + return business, meta + + +def _is_envelope_argument_key(key: Any) -> bool: + name = str(key) + return name == "_meta" or name.startswith(MCP_META_PREFIX) + + +def strip_tool_arguments(arguments: Optional[Mapping[str, Any]]) -> dict[str, Any]: + """Copy ``tools/call`` arguments without protocol envelope keys. + + Drops ``_meta`` and ``io.modelcontextprotocol/*`` so guards, pipes, and + user handlers never receive envelope slots as input. Nested user values + are left unchanged. The request envelope stays on ``ExecutionContext``. + """ + if not arguments: + return {} + cleaned = { + key: value + for key, value in arguments.items() + if not _is_envelope_argument_key(key) + } + inner = cleaned.get("input") + if isinstance(inner, Mapping): + cleaned["input"] = { + key: value + for key, value in inner.items() + if not _is_envelope_argument_key(key) + } + return cleaned + + +@dataclass(frozen=True) +class RequestEnvelope: + """Allowed envelope mapping for handler context. Identity is not included.""" + + meta: RequestMeta + mcp_headers: dict[str, str] + protocol_version: Optional[str] + + +def bind_request_envelope( + raw_meta: Optional[dict[str, Any]] = None, + mcp_headers: Optional[Mapping[str, str]] = None, +) -> RequestEnvelope: + """ + Map JSON-RPC ``_meta`` and MCP headers onto handler context fields. + + Unsigned ``userId`` / ``tenantId`` remain on ``meta.raw`` only and are not + treated as identity. + """ + meta = extract_request_meta({"_meta": raw_meta or {}}) + headers = {key: value for key, value in (mcp_headers or {}).items()} + header_version = None + for key, value in headers.items(): + if key.lower() == _PROTOCOL_VERSION_HEADER: + header_version = value + break + return RequestEnvelope( + meta=meta, + mcp_headers=headers, + protocol_version=header_version or meta.protocol_version, + ) + + +def envelope_identity_is_ignored(raw_meta: dict[str, Any]) -> bool: + """True when the envelope contains unsigned identity keys.""" + return any(str(key).lower() in _IDENTITY_META_KEYS for key in raw_meta) + + +def envelope_protocol_version(meta: RequestMeta) -> Optional[str]: + """Protocol version from ``_meta.mcp.protocolVersion``, then other envelope keys.""" + mcp = meta.raw.get("mcp") if meta.raw else None + if isinstance(mcp, dict): + nested = mcp.get("protocolVersion") + if isinstance(nested, str) and nested.strip(): + return nested.strip() + if isinstance(meta.protocol_version, str) and meta.protocol_version.strip(): + return meta.protocol_version.strip() + return None + + +def flatten_request_meta_object(raw_meta: Any) -> dict[str, Any]: + """Flatten an MCP ``_meta`` object, including Pydantic extras.""" + if raw_meta is None: + return {} + data: dict[str, Any] = {} + extra = getattr(raw_meta, "model_extra", None) or getattr(raw_meta, "__pydantic_extra__", None) + if isinstance(extra, dict): + data.update(extra) + if hasattr(raw_meta, "model_dump"): + try: + dumped = raw_meta.model_dump(exclude_none=True) + if isinstance(dumped, dict): + nested_extra = dumped.pop("__pydantic_extra__", None) + if isinstance(nested_extra, dict): + data.update(nested_extra) + data.update(dumped) + except Exception: + pass + elif isinstance(raw_meta, dict): + data.update(raw_meta) + return data diff --git a/nitrostack/protocol/method_contract.py b/nitrostack/protocol/method_contract.py new file mode 100644 index 0000000..b9142ab --- /dev/null +++ b/nitrostack/protocol/method_contract.py @@ -0,0 +1,135 @@ +"""SEP-2243 method contracts for the 2026 method surface. + +Official v2 owns this table once mounted. Until then the sidecar validator +reads one row per method: required ``Mcp-Method``, optional ``Mcp-Name`` field, +and whether ``auto`` requires the method header. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +from nitrostack.protocol.version import ProtocolEra, WireMode + +NAME_FIELD_NAME = "name" +NAME_FIELD_URI = "uri" + + +@dataclass(frozen=True) +class MethodContract: + """Header contract for one JSON-RPC method.""" + + method: str + name_field: Optional[str] = None + # ``auto`` keeps handshake / ping / discover optional (2025 clients). + # ``modern`` requires ``Mcp-Method`` on every JSON-RPC POST. + requires_method_header: bool = True + + +def _contract( + method: str, + *, + name_field: Optional[str] = None, + requires_method_header: bool = True, +) -> MethodContract: + return MethodContract( + method=method, + name_field=name_field, + requires_method_header=requires_method_header, + ) + + +MODERN_METHOD_CONTRACTS: tuple[MethodContract, ...] = ( + _contract("ping", requires_method_header=False), + _contract("initialize", requires_method_header=False), + _contract("notifications/initialized", requires_method_header=False), + _contract("notifications/cancelled"), + _contract("notifications/progress"), + _contract("server/discover", requires_method_header=False), + _contract("subscriptions/listen"), + _contract("notifications/subscriptions/acknowledged"), + _contract("tools/list", requires_method_header=False), + _contract("tools/call", name_field=NAME_FIELD_NAME, requires_method_header=False), + _contract("notifications/tools/list_changed"), + _contract("resources/list", requires_method_header=False), + _contract("resources/templates/list", requires_method_header=False), + _contract("resources/read", name_field=NAME_FIELD_URI), + _contract("resources/subscribe", name_field=NAME_FIELD_URI), + _contract("resources/unsubscribe", name_field=NAME_FIELD_URI), + _contract("notifications/resources/list_changed"), + _contract("notifications/resources/updated"), + _contract("prompts/list", requires_method_header=False), + _contract("prompts/get", name_field=NAME_FIELD_NAME), + _contract("notifications/prompts/list_changed"), + _contract("completion/complete"), + _contract("tasks/get"), + _contract("tasks/cancel"), + _contract("tasks/result"), + _contract("tasks/list"), + _contract("logging/setLevel"), +) + +_CONTRACTS_BY_METHOD: dict[str, MethodContract] = { + row.method: row for row in MODERN_METHOD_CONTRACTS +} + +NAME_SCOPED_METHODS: frozenset[str] = frozenset( + row.method for row in MODERN_METHOD_CONTRACTS if row.name_field +) + +# 2025 methods removed from the 2026-07-28 wire. Handshake methods +# (``initialize``, ``notifications/initialized``) are not in this table: +# ``modern`` rejects them as method-not-found; ``auto`` still answers them. +DEPRECATED_MODERN_METHODS: dict[str, str] = { + "tasks/result": "Method 'tasks/result' is not supported in MCP 2026-07-28; use 'tasks/get'.", + "tasks/list": "Method 'tasks/list' is not supported in modern stateless MCP 2026-07-28.", + "resources/subscribe": ( + "Method 'resources/subscribe' is not supported in stateless MCP 2026-07-28; " + "use SSE subscriptions/listen." + ), + "logging/setLevel": ( + "Method 'logging/setLevel' is not supported in stateless MCP 2026-07-28; " + "configure logging at the host level." + ), +} + + +def deprecated_method_message(method: str) -> Optional[str]: + """Return the modern-wire rejection text when ``method`` is retired.""" + return DEPRECATED_MODERN_METHODS.get(method) + + +def rejects_deprecated_method(method: str, era: ProtocolEra) -> bool: + """True when era ``modern`` must answer ``method`` as not found.""" + return era == "modern" and method in DEPRECATED_MODERN_METHODS + + +def contract_for(method: str) -> Optional[MethodContract]: + """Return the table row for ``method``, or ``None`` when unlisted.""" + return _CONTRACTS_BY_METHOD.get(method) + + +def mcp_name_field(method: str) -> Optional[str]: + """Body field mirrored by ``Mcp-Name``, if this method is name-scoped.""" + row = contract_for(method) + return row.name_field if row is not None else None + + +def mcp_name_is_required(method: str) -> bool: + """``Mcp-Name`` is required on ``tools/call``, ``resources/read``, ``prompts/get``.""" + return mcp_name_field(method) is not None and method in { + "tools/call", + "resources/read", + "prompts/get", + } + + +def mcp_method_is_required(method: str, wire_mode: WireMode) -> bool: + """``modern`` requires ``Mcp-Method`` on every POST. ``auto`` follows the table.""" + if wire_mode == "reject": + return True + row = contract_for(method) + if row is None: + return False + return row.requires_method_header diff --git a/nitrostack/protocol/mrtr.py b/nitrostack/protocol/mrtr.py new file mode 100644 index 0000000..49f51e4 --- /dev/null +++ b/nitrostack/protocol/mrtr.py @@ -0,0 +1,153 @@ +"""Multi Round-Trip Request (MRTR) helpers (SEP-2322).""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal, Optional + +InputRequestKind = Literal["form", "url"] +RESULT_TYPE_INPUT_REQUIRED = "input_required" + +DEFAULT_INPUT_REQUIRED_MESSAGE = "Additional input needed to complete this operation" + + +@dataclass +class InputRequest: + """One elicitation prompt in an MRTR exchange.""" + + id: str + message: Optional[str] = None + schema: Optional[dict[str, Any]] = None + kind: InputRequestKind = "form" + url: Optional[str] = None + + def to_wire_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = {"id": self.id, "kind": self.kind} + if self.message is not None: + payload["message"] = self.message + if self.schema is not None: + payload["schema"] = self.schema + if self.url is not None: + payload["url"] = self.url + return payload + + +@dataclass +class InputRequiredResult: + """Wire result pausing tool execution until the client supplies input.""" + + input_requests: list[InputRequest] + request_state: dict[str, Any] + message: Optional[str] = None + result_type: str = RESULT_TYPE_INPUT_REQUIRED + + def to_wire_dict(self) -> dict[str, Any]: + return { + "resultType": self.result_type, + "message": self.message or DEFAULT_INPUT_REQUIRED_MESSAGE, + "inputRequests": [req.to_wire_dict() for req in self.input_requests], + "requestState": self.request_state, + } + + +def accepted_content(input_responses: Optional[dict[str, Any]], request_id: str) -> Any: + """ + Return the client's answer for ``request_id``, or ``None`` if not yet supplied. + + Handler primitive for MRTR elicitation. + """ + if not input_responses or request_id not in input_responses: + return None + return input_responses[request_id] + + +def input_required( + requests: list[InputRequest], + request_state: dict[str, Any], + message: Optional[str] = None, +) -> InputRequiredResult: + """Halt tool execution and request additional client input.""" + if not requests: + raise ValueError("input_required requires at least one InputRequest") + return InputRequiredResult( + input_requests=list(requests), + request_state=dict(request_state), + message=message, + ) + + +def split_mrtr_tool_params(params: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any], Optional[dict[str, Any]]]: + """ + Extract MRTR resume fields from a ``tools/call`` params object. + + Returns ``(arguments, input_responses, request_state)``. + """ + raw = dict(params or {}) + input_responses = raw.pop("inputResponses", None) or {} + request_state = raw.pop("requestState", None) + arguments = raw.pop("arguments", None) or {} + + if not isinstance(arguments, dict): + arguments = {} + if not isinstance(input_responses, dict): + input_responses = {} + + # Clients may echo MRTR fields inside ``arguments`` on resume. + if "inputResponses" in arguments: + nested = arguments.pop("inputResponses") or {} + if isinstance(nested, dict): + input_responses = {**input_responses, **nested} + if "requestState" in arguments: + nested_state = arguments.pop("requestState") + if isinstance(nested_state, dict): + request_state = nested_state + + return arguments, input_responses, request_state + + +def split_mrtr_from_arguments(arguments: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any], Optional[dict[str, Any]]]: + """Extract MRTR resume fields when only the arguments dict is available.""" + return split_mrtr_tool_params({"arguments": arguments}) + + +def is_input_required_result(value: Any) -> bool: + """True when ``value`` is an InputRequiredResult or matching wire dict.""" + if isinstance(value, InputRequiredResult): + return True + return isinstance(value, dict) and value.get("resultType") == RESULT_TYPE_INPUT_REQUIRED + + +def coerce_input_required_result(value: Any) -> Optional[InputRequiredResult]: + """Convert a handler return value or wire dict into InputRequiredResult.""" + if isinstance(value, InputRequiredResult): + return value + if not isinstance(value, dict) or value.get("resultType") != RESULT_TYPE_INPUT_REQUIRED: + return None + + requests = [] + for item in value.get("inputRequests") or []: + if not isinstance(item, dict) or "id" not in item: + continue + requests.append( + InputRequest( + id=item["id"], + message=item.get("message"), + schema=item.get("schema"), + kind=item.get("kind", "form"), + url=item.get("url"), + ) + ) + return InputRequiredResult( + input_requests=requests, + request_state=dict(value.get("requestState") or {}), + message=value.get("message"), + ) + + +def build_input_required_jsonrpc_result(request_id: Any, result: InputRequiredResult) -> dict[str, Any]: + """Build a JSON-RPC success envelope for an MRTR pause response.""" + return { + "jsonrpc": "2.0", + "id": request_id, + "result": result.to_wire_dict(), + } diff --git a/nitrostack/protocol/observability.py b/nitrostack/protocol/observability.py new file mode 100644 index 0000000..29f1050 --- /dev/null +++ b/nitrostack/protocol/observability.py @@ -0,0 +1,50 @@ +"""W3C Trace Context extraction for MCP _meta envelopes.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional + +from nitrostack.protocol.meta import MCP_META_PREFIX, RequestMeta + + +@dataclass(frozen=True) +class TraceContext: + """Distributed tracing metadata propagated through MCP _meta.""" + + traceparent: str | None = None + tracestate: str | None = None + baggage: str | None = None + + +def _meta_get(meta: dict[str, Any], bare: str, prefixed: str) -> Optional[str]: + value = meta.get(prefixed) + if value is None: + value = meta.get(bare) + return value if isinstance(value, str) and value.strip() else None + + +def extract_trace_context(meta: dict[str, Any] | None) -> TraceContext | None: + """Extract W3C trace fields from a ``_meta`` dict.""" + if not meta: + return None + + traceparent = _meta_get(meta, "traceparent", f"{MCP_META_PREFIX}traceparent") + tracestate = _meta_get(meta, "tracestate", f"{MCP_META_PREFIX}tracestate") + baggage = _meta_get(meta, "baggage", f"{MCP_META_PREFIX}baggage") + + if not traceparent and not tracestate and not baggage: + return None + + return TraceContext(traceparent=traceparent, tracestate=tracestate, baggage=baggage) + + +def trace_context_from_request_meta(meta: RequestMeta) -> TraceContext | None: + """Build ``TraceContext`` from parsed ``RequestMeta``.""" + if not meta.traceparent and not meta.tracestate and not meta.baggage: + return None + return TraceContext( + traceparent=meta.traceparent, + tracestate=meta.tracestate, + baggage=meta.baggage, + ) diff --git a/nitrostack/protocol/resources.py b/nitrostack/protocol/resources.py new file mode 100644 index 0000000..1eb904c --- /dev/null +++ b/nitrostack/protocol/resources.py @@ -0,0 +1,63 @@ +"""Resource URI resolution for static resources and templates.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any, Mapping, Optional, Pattern, Sequence + +_TEMPLATE_PARAM_RE = re.compile(r"\{([^}]+)\}") + + +def extract_template_param_names(uri_template: str) -> list[str]: + """Extract `{param}` names from a URI template.""" + return _TEMPLATE_PARAM_RE.findall(uri_template) + + +def uri_template_to_pattern(uri_template: str) -> Pattern[str]: + """ + Compile a URI template to a regex with named groups. + + Example: ``mcp://customers/{customerId}/orders`` matches concrete URIs and + captures ``customerId``. + """ + param_names = extract_template_param_names(uri_template) + regex_str = re.escape(uri_template) + for pname in param_names: + regex_str = regex_str.replace(re.escape("{" + pname + "}"), f"(?P<{pname}>[^/]+)") + return re.compile(f"^{regex_str}$") + + +@dataclass(frozen=True) +class ResourceMatch: + """Result of resolving a resource URI against the registry.""" + + entry: Any + path_params: dict[str, str] + matched_via_template: bool = False + + +def resolve_resource_uri( + uri: str, + static_resources: Mapping[str, Any], + templates: Sequence[tuple[Pattern[str], Any]], +) -> Optional[ResourceMatch]: + """ + Resolve a resource URI: + 1. Exact static map lookup + 2. URI template pattern match + """ + static_entry = static_resources.get(uri) + if static_entry is not None: + return ResourceMatch(entry=static_entry, path_params={}) + + for pattern, entry in templates: + match = pattern.match(uri) + if match: + return ResourceMatch( + entry=entry, + path_params=match.groupdict(), + matched_via_template=True, + ) + + return None diff --git a/nitrostack/protocol/schema.py b/nitrostack/protocol/schema.py new file mode 100644 index 0000000..65b71cd --- /dev/null +++ b/nitrostack/protocol/schema.py @@ -0,0 +1,230 @@ +"""JSON Schema 2020-12 normalization, dialect gate, and depth bounding.""" + +from __future__ import annotations + +from typing import Any, Optional + +from nitrostack.protocol.constants import MAX_SCHEMA_DEPTH + +JSON_SCHEMA_2020_12_URI = "https://json-schema.org/draft/2020-12/schema" +_ACCEPTED_DIALECT_URIS = frozenset( + { + JSON_SCHEMA_2020_12_URI, + "http://json-schema.org/draft/2020-12/schema", + f"{JSON_SCHEMA_2020_12_URI}#", + "http://json-schema.org/draft/2020-12/schema#", + } +) + +_COMPOSITION_KEYS = frozenset({"allOf", "anyOf", "oneOf"}) +_MAP_SCHEMA_KEYS = frozenset({"properties", "patternProperties"}) +_SINGLE_SCHEMA_KEYS = frozenset( + { + "additionalProperties", + "contains", + "propertyNames", + "if", + "then", + "else", + "not", + } +) + + +class UnsupportedJsonSchemaError(ValueError): + """Raised when a registered schema is not JSON Schema 2020-12.""" + + +def bound_schema_depth( + node: Any, + *, + max_depth: int = MAX_SCHEMA_DEPTH, + depth: int = 0, +) -> Any: + """Collapse schema nodes beyond ``max_depth`` to permissive ``{}`` (DoS defense).""" + if depth >= max_depth: + return {} + + if isinstance(node, list): + return [bound_schema_depth(item, max_depth=max_depth, depth=depth + 1) for item in node] + + if not isinstance(node, dict): + return node + + result: dict[str, Any] = {} + for key, value in node.items(): + if key in _COMPOSITION_KEYS and isinstance(value, list): + result[key] = [ + bound_schema_depth(item, max_depth=max_depth, depth=depth + 1) for item in value + ] + elif key == "$defs" and isinstance(value, dict): + result[key] = { + def_key: bound_schema_depth(def_val, max_depth=max_depth, depth=depth + 1) + for def_key, def_val in value.items() + } + elif key in _MAP_SCHEMA_KEYS and isinstance(value, dict): + result[key] = { + nested_key: bound_schema_depth(nested_val, max_depth=max_depth, depth=depth + 1) + for nested_key, nested_val in value.items() + } + elif key == "items": + if isinstance(value, list): + raise UnsupportedJsonSchemaError( + "JSON Schema 2020-12 'items' must be a single schema; " + "use 'prefixItems' for tuple validation" + ) + result[key] = bound_schema_depth(value, max_depth=max_depth, depth=depth + 1) + elif key == "prefixItems": + if isinstance(value, list): + result[key] = [ + bound_schema_depth(item, max_depth=max_depth, depth=depth + 1) for item in value + ] + elif isinstance(value, dict): + result[key] = bound_schema_depth(value, max_depth=max_depth, depth=depth + 1) + else: + result[key] = value + elif key in _SINGLE_SCHEMA_KEYS: + if isinstance(value, dict): + result[key] = bound_schema_depth(value, max_depth=max_depth, depth=depth + 1) + else: + result[key] = bound_schema_depth(value, max_depth=max_depth, depth=depth + 1) + elif key == "$ref": + result[key] = value + elif isinstance(value, (dict, list)): + result[key] = bound_schema_depth(value, max_depth=max_depth, depth=depth + 1) + else: + result[key] = value + return result + + +def _dialect_uri(value: Any) -> str: + return str(value).strip() + + +def is_json_schema_2020_12_dialect(schema_uri: Optional[str]) -> bool: + """True when ``$schema`` is absent (treated as 2020-12) or a 2020-12 URI.""" + if schema_uri is None: + return True + uri = _dialect_uri(schema_uri) + if not uri: + return True + return uri in _ACCEPTED_DIALECT_URIS + + +def _assert_items_not_tuple(node: Any, *, name: str, path: str = "$") -> None: + if isinstance(node, list): + for index, item in enumerate(node): + _assert_items_not_tuple(item, name=name, path=f"{path}[{index}]") + return + if not isinstance(node, dict): + return + items = node.get("items") + if isinstance(items, list): + raise UnsupportedJsonSchemaError( + f"{name} at {path}.items is a tuple list; JSON Schema 2020-12 " + "requires a single schema for items (use prefixItems for tuples)" + ) + for key, value in node.items(): + if key in _COMPOSITION_KEYS and isinstance(value, list): + for index, item in enumerate(value): + _assert_items_not_tuple(item, name=name, path=f"{path}.{key}[{index}]") + elif key in {"$defs", "definitions"} and isinstance(value, dict): + for def_key, def_val in value.items(): + _assert_items_not_tuple(def_val, name=name, path=f"{path}.{key}.{def_key}") + elif key in _MAP_SCHEMA_KEYS and isinstance(value, dict): + for nested_key, nested_val in value.items(): + _assert_items_not_tuple( + nested_val, name=name, path=f"{path}.{key}.{nested_key}" + ) + elif key == "prefixItems" and isinstance(value, list): + for index, item in enumerate(value): + _assert_items_not_tuple(item, name=name, path=f"{path}.prefixItems[{index}]") + elif key in _SINGLE_SCHEMA_KEYS or key == "items": + _assert_items_not_tuple(value, name=name, path=f"{path}.{key}") + elif key != "enum" and isinstance(value, (dict, list)): + _assert_items_not_tuple(value, name=name, path=f"{path}.{key}") + + +def _assert_nested_dialects(node: Any, *, name: str, path: str = "$") -> None: + if isinstance(node, list): + for index, item in enumerate(node): + _assert_nested_dialects(item, name=name, path=f"{path}[{index}]") + return + if not isinstance(node, dict): + return + dialect = node.get("$schema") + if dialect is not None and not is_json_schema_2020_12_dialect(_dialect_uri(dialect)): + raise UnsupportedJsonSchemaError( + f"{name} at {path} uses unsupported JSON Schema dialect {dialect!r}; " + f"required dialect is {JSON_SCHEMA_2020_12_URI}" + ) + for key, value in node.items(): + if key == "enum": + continue + if isinstance(value, (dict, list)): + _assert_nested_dialects(value, name=name, path=f"{path}.{key}") + + +def assert_json_schema_2020_12(schema: Any, *, name: str = "schema") -> None: + """Reject unsupported drafts and tuple-style ``items`` lists. + + A missing ``$schema`` is treated as JSON Schema 2020-12. + """ + if schema is None: + return + if not isinstance(schema, dict): + return + _assert_nested_dialects(schema, name=name) + _assert_items_not_tuple(schema, name=name) + + +def gate_registered_schema(schema: Any, *, name: str) -> None: + """Registration-time dialect gate for a dict or Pydantic model.""" + if schema is None: + return + if isinstance(schema, dict): + assert_json_schema_2020_12(schema, name=name) + return + model_schema = getattr(schema, "model_json_schema", None) + if callable(model_schema): + assert_json_schema_2020_12(model_schema(), name=name) + + +def normalize_input_schema(schema: dict[str, Any]) -> dict[str, Any]: + """ + Normalize tool inputSchema for MCP 2026-07-28. + + Root MUST be ``type: object``; ``$schema`` is set to JSON Schema 2020-12. + Unsupported drafts and tuple-style ``items`` fail instead of being rewritten. + """ + assert_json_schema_2020_12(schema, name="inputSchema") + bounded = bound_schema_depth(schema) + if not isinstance(bounded, dict): + bounded = {} + bounded = dict(bounded) + bounded["type"] = "object" + bounded.setdefault("properties", {}) + bounded["$schema"] = JSON_SCHEMA_2020_12_URI + return bounded + + +def normalize_output_schema(schema: dict[str, Any]) -> dict[str, Any]: + """Normalize tool outputSchema — unrestricted root type, bounded depth.""" + assert_json_schema_2020_12(schema, name="outputSchema") + bounded = bound_schema_depth(schema) + if not isinstance(bounded, dict): + bounded = {"type": "object"} + bounded = dict(bounded) + bounded["$schema"] = JSON_SCHEMA_2020_12_URI + return bounded + + +def normalize_json_schema( + schema: dict[str, Any], + *, + require_object_root: bool = False, +) -> dict[str, Any]: + """Generic schema normalizer used by tool input/output builders.""" + if require_object_root: + return normalize_input_schema(schema) + return normalize_output_schema(schema) diff --git a/nitrostack/protocol/tasks.py b/nitrostack/protocol/tasks.py new file mode 100644 index 0000000..e1d654f --- /dev/null +++ b/nitrostack/protocol/tasks.py @@ -0,0 +1,43 @@ +"""MCP Tasks protocol helpers for 2026-07-28.""" + +from __future__ import annotations + +from typing import Any, Optional + +RESULT_TYPE_TASK = "task" +DEFAULT_TASK_TTL_MS = 300_000 +DEFAULT_POLL_INTERVAL_MS = 2_000 + + +def ttl_ms_to_seconds(ttl_ms: Optional[int]) -> Optional[int]: + """Convert wire TTL (milliseconds) to internal expiry seconds.""" + if ttl_ms is None: + return None + return max(1, int(ttl_ms / 1000)) + + +def ttl_seconds_to_ms(ttl_seconds: Optional[int]) -> Optional[int]: + """Convert internal TTL seconds to wire milliseconds.""" + if ttl_seconds is None: + return None + return int(ttl_seconds * 1000) + + +def build_task_create_jsonrpc_result(request_id: Any, task: dict[str, Any]) -> dict[str, Any]: + """Build task-augmented ``tools/call`` success envelope.""" + return { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "resultType": RESULT_TYPE_TASK, + "task": task, + }, + } + + +def task_support_forbidden_message(tool_name: str) -> str: + return f"Tool '{tool_name}' does not support task augmentation" + + +def task_support_required_message(tool_name: str) -> str: + return f"Task augmentation required for tool '{tool_name}'" diff --git a/nitrostack/protocol/version.py b/nitrostack/protocol/version.py new file mode 100644 index 0000000..24cf3f0 --- /dev/null +++ b/nitrostack/protocol/version.py @@ -0,0 +1,239 @@ +"""MCP protocol version identifiers and protocol-era selection.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Literal, Optional + +MODERN_PROTOCOL_VERSION = "2026-07-28" +LEGACY_PROTOCOL_VERSION = "2025-06-18" + +SUPPORTED_PROTOCOL_VERSIONS: tuple[str, ...] = (MODERN_PROTOCOL_VERSION,) + +PROTOCOL_ERA_ENV = "NITRO_MCP_PROTOCOL_VERSION" +STATELESS_OVERRIDE_ENV = "MCP_STATELESS" + +ProtocolEra = Literal["legacy", "modern", "auto"] +EraSource = Literal["mcp_stateless", "env", "config", "default"] +# How the HTTP factory should treat 2025-shaped traffic for this era. +# ``stateless`` here is the dual-spec fallback (sessionless initialize), not +# the 1.x ``StreamableHTTPSessionManager(stateless=True)`` flag. +WireMode = Literal["sessionful", "stateless", "reject"] +# Which /mcp HTTP engine the era factory mounts. Official mcp 2.x replaces +# the sessionless 1.x manager when that dependency is installed. +HttpEngine = Literal["sessionless", "sessionful"] + +_AUTO_ALIASES = frozenset({"auto", "both", "dual", "dual-spec"}) +_MODERN_ALIASES = frozenset({"modern", "latest", "2026", MODERN_PROTOCOL_VERSION}) +_LEGACY_ALIASES = frozenset({"legacy", "2025", "2025-11-25", LEGACY_PROTOCOL_VERSION}) +_TRUE_TOKENS = frozenset({"1", "true", "yes", "on"}) +_FALSE_TOKENS = frozenset({"0", "false", "no", "off"}) + + +def _parse_bool_token(raw: Optional[str]) -> Optional[bool]: + if raw is None: + return None + value = raw.strip().lower() + if not value: + return None + if value in _TRUE_TOKENS: + return True + if value in _FALSE_TOKENS: + return False + return None + + +def _era_token(raw: Optional[str]) -> str: + return (raw or "").strip().lower() + + +def _era_from_token(value: str) -> ProtocolEra: + if not value or value in _AUTO_ALIASES: + return "auto" + if value in _MODERN_ALIASES: + return "modern" + if value in _LEGACY_ALIASES: + return "legacy" + return "auto" + + +@dataclass(frozen=True) +class ProtocolEraResolution: + """Resolved era plus which input selected it.""" + + era: ProtocolEra + source: EraSource + + def log_line(self) -> str: + return f"protocol era={self.era} (source={self.source})" + + +def resolve_protocol_era_resolution( + raw: Optional[str] = None, + *, + stateless_override: Optional[str] = None, + config_value: Optional[str] = None, +) -> ProtocolEraResolution: + """ + Resolve the active protocol era and record how it was chosen. + + Precedence: + 1. ``MCP_STATELESS`` (explicit boolean) → source ``mcp_stateless`` + 2. ``NITRO_MCP_PROTOCOL_VERSION`` (or the ``raw`` argument) → ``env`` + 3. ``ServerConfig.protocol_era`` (``config_value``) → ``config`` + 4. ``auto`` → ``default`` + + Tokens (case-insensitive, trimmed): ``modern`` / ``latest`` / ``2026`` / + ``2026-07-28``; ``auto`` / ``both`` / ``dual`` / ``dual-spec``; ``legacy`` / + ``2025`` / ``2025-06-18`` / ``2025-11-25``. Unknown tokens resolve to + ``auto``. + + ``auto`` is not ``modern``. ``modern`` is stateless-only; ``auto`` is the + dual-spec era and does not force the 1.x ``stateless=True`` transport flag. + """ + override = ( + stateless_override + if stateless_override is not None + else os.environ.get(STATELESS_OVERRIDE_ENV) + ) + flag = _parse_bool_token(override) + if flag is True: + return ProtocolEraResolution("modern", "mcp_stateless") + if flag is False: + return ProtocolEraResolution("legacy", "mcp_stateless") + + if raw is not None: + value = _era_token(raw) + source: EraSource = "env" if value else "default" + return ProtocolEraResolution(_era_from_token(value), source) + + value = _era_token(os.environ.get(PROTOCOL_ERA_ENV)) + if value: + return ProtocolEraResolution(_era_from_token(value), "env") + value = _era_token(config_value) + if value: + return ProtocolEraResolution(_era_from_token(value), "config") + return ProtocolEraResolution("auto", "default") + + +def resolve_protocol_era( + raw: Optional[str] = None, + *, + stateless_override: Optional[str] = None, + config_value: Optional[str] = None, +) -> ProtocolEra: + """Resolve the active protocol era. See ``resolve_protocol_era_resolution``.""" + return resolve_protocol_era_resolution( + raw, + stateless_override=stateless_override, + config_value=config_value, + ).era + + +def supported_protocol_versions_for_era(era: ProtocolEra) -> frozenset[str]: + """Dated protocol versions this era accepts on the wire header or envelope.""" + legacy_versions = frozenset({LEGACY_PROTOCOL_VERSION, "2025-11-25"}) + if era == "modern": + return frozenset({MODERN_PROTOCOL_VERSION}) + if era == "legacy": + return legacy_versions + return frozenset({MODERN_PROTOCOL_VERSION, *legacy_versions}) + + +def protocol_era_for_wire_mode(wire_mode: WireMode) -> ProtocolEra: + """Map dual-spec wire mode onto the era that owns its supported versions.""" + if wire_mode == "reject": + return "modern" + if wire_mode == "sessionful": + return "legacy" + return "auto" + + +def protocol_version_for_era(era: Optional[ProtocolEra], fallback: str = MODERN_PROTOCOL_VERSION) -> str: + if era == "legacy": + return LEGACY_PROTOCOL_VERSION + if era in ("modern", "auto"): + return MODERN_PROTOCOL_VERSION + return fallback + + +def stateless_for_era(era: Optional[ProtocolEra]) -> Optional[bool]: + """ + Map era onto the 1.x Streamable HTTP ``stateless`` flag. + + ``modern`` → True, ``legacy`` → False, ``auto`` → None so the HTTP factory + does not treat dual-spec as modern-only. + """ + if era == "modern": + return True + if era == "legacy": + return False + return None + + +def wire_mode_for_era(era: ProtocolEra) -> WireMode: + """ + Dual-spec policy for an era. + + * ``legacy`` — sessionful 2025 wire only + * ``auto`` — accept 2025 ``initialize`` without a session (official v2 fallback) + * ``modern`` — reject 2025 sessionful wire + """ + if era == "modern": + return "reject" + if era == "auto": + return "stateless" + return "sessionful" + + +def accepts_sessionless_initialize(era: ProtocolEra) -> bool: + """True when era ``auto`` answers 2025 ``initialize`` without a session.""" + return era == "auto" + + +def rejects_legacy_initialize(era: ProtocolEra) -> bool: + """True when era ``modern`` rejects 2025 ``initialize`` / ``initialized``.""" + return era == "modern" + + +def needs_modern_engine(era: ProtocolEra) -> bool: + """True when ``/mcp`` should be the official 2026 engine (``modern`` or ``auto``).""" + return era in ("modern", "auto") + + +def needs_sessionful_engine(era: ProtocolEra) -> bool: + """True only for ``legacy``. ``auto`` does not mount a second session manager.""" + return era == "legacy" + + +def http_engine_for_era(era: ProtocolEra) -> HttpEngine: + """ + Select the /mcp HTTP engine for an era. + + ``legacy`` uses the sessionful 1.x manager. ``modern`` and ``auto`` use the + sessionless /mcp path (one engine; official mcp 2.x when mounted). + """ + if needs_sessionful_engine(era): + return "sessionful" + return "sessionless" + + +def resolve_http_engine( + era: ProtocolEra, + *, + http_engine: Optional[HttpEngine] = None, + stateless: Optional[bool] = None, +) -> HttpEngine: + """ + Sessionful 1.x is mounted only when era is ``legacy``. + + ``auto`` and ``modern`` stay sessionless even if a caller passes + ``stateless=False`` or ``http_engine='sessionful'``. ``legacy`` is + sessionful unless ``stateless`` is True. + """ + if not needs_sessionful_engine(era): + return "sessionless" + if stateless is True or http_engine == "sessionless": + return "sessionless" + return "sessionful" diff --git a/nitrostack/runtime/__init__.py b/nitrostack/runtime/__init__.py new file mode 100644 index 0000000..6dabc41 --- /dev/null +++ b/nitrostack/runtime/__init__.py @@ -0,0 +1,23 @@ +"""Stateless runtime policies and invariants for MCP 2026-07-28.""" + +from nitrostack.runtime.correlation import InFlightRegistry, InFlightTicket, new_correlation_id +from nitrostack.runtime.stateless import ( + StatelessInvariants, + assert_stateless_headers, + has_incoming_session_id, + is_unsupported_protocol_version, + request_protocol_version, + sessionless_rejects_incoming_session_id, +) + +__all__ = [ + "InFlightRegistry", + "InFlightTicket", + "new_correlation_id", + "StatelessInvariants", + "assert_stateless_headers", + "has_incoming_session_id", + "is_unsupported_protocol_version", + "request_protocol_version", + "sessionless_rejects_incoming_session_id", +] diff --git a/nitrostack/runtime/acceptance.py b/nitrostack/runtime/acceptance.py new file mode 100644 index 0000000..17992ed --- /dev/null +++ b/nitrostack/runtime/acceptance.py @@ -0,0 +1,204 @@ +"""MCP 2026-07-28 feature-area and acceptance-criteria registry.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum +from typing import Iterable + +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION + + +class ProtocolArea(IntEnum): + """Major MCP 2026-07-28 implementation areas.""" + + CORE_JSONRPC = 1 + STATELESS_HTTP = 2 + REGISTRIES_SCHEMA = 3 + ASYNC_TASKS = 4 + MULTI_TENANT = 5 + OAUTH_CIMD = 6 + MRTR_OBSERVABILITY = 7 + + +@dataclass(frozen=True) +class ProtocolDeliverable: + area: ProtocolArea + item: str + test_module: str + + +@dataclass(frozen=True) +class AcceptanceCriterion: + key: str + description: str + test_module: str + + +PROTOCOL_DELIVERABLES: tuple[ProtocolDeliverable, ...] = ( + ProtocolDeliverable( + ProtocolArea.CORE_JSONRPC, + "JSON-RPC error code hierarchy (-32700..-32603, -32020)", + "tests/test_mcp20_jsonrpc_wire.py", + ), + ProtocolDeliverable( + ProtocolArea.CORE_JSONRPC, + "SEP-2164 missing resource returns -32602", + "tests/test_mcp20_contracts.py", + ), + ProtocolDeliverable( + ProtocolArea.CORE_JSONRPC, + "Reject tasks/result and tasks/list on modern wire", + "tests/test_mcp20_tasks.py", + ), + ProtocolDeliverable( + ProtocolArea.CORE_JSONRPC, + "Deprecated methods share one policy on every modern route", + "tests/test_mcp20_deprecated.py", + ), + ProtocolDeliverable( + ProtocolArea.CORE_JSONRPC, + "Strip envelope keys from tool arguments before handlers", + "tests/test_mcp20_tool_args.py", + ), + ProtocolDeliverable( + ProtocolArea.STATELESS_HTTP, + "POST /mcp stateless ingress without Mcp-Session-Id", + "tests/test_mcp20_stateless_http.py", + ), + ProtocolDeliverable( + ProtocolArea.STATELESS_HTTP, + "server/discover and ping fast path", + "tests/test_mcp20_stateless_http.py", + ), + ProtocolDeliverable( + ProtocolArea.STATELESS_HTTP, + "CORS headers for MCP preflight", + "tests/test_mcp20_stateless_http.py", + ), + ProtocolDeliverable( + ProtocolArea.STATELESS_HTTP, + "GET /subscriptions/listen SSE attach on modern and auto", + "tests/test_mcp20_subscriptions.py", + ), + ProtocolDeliverable( + ProtocolArea.STATELESS_HTTP, + "Concurrent identical JSON-RPC ids stay isolated", + "tests/test_mcp20_concurrency.py", + ), + ProtocolDeliverable( + ProtocolArea.STATELESS_HTTP, + "Official mcp 2.x stdio follows the active era", + "tests/test_mcp20_stdio.py", + ), + ProtocolDeliverable( + ProtocolArea.REGISTRIES_SCHEMA, + "JSON Schema 2020-12 with depth bounding (max 64)", + "tests/test_mcp20_contracts.py", + ), + ProtocolDeliverable( + ProtocolArea.REGISTRIES_SCHEMA, + "Resource URI templates and resolution", + "tests/test_mcp20_contracts.py", + ), + ProtocolDeliverable( + ProtocolArea.ASYNC_TASKS, + "Task state machine and tasks/get embedded result", + "tests/test_mcp20_tasks.py", + ), + ProtocolDeliverable( + ProtocolArea.ASYNC_TASKS, + "Pluggable TaskStore and terminal-only TTL eviction", + "tests/test_mcp20_task_store.py", + ), + ProtocolDeliverable( + ProtocolArea.MULTI_TENANT, + "TaskAccessContext isolation and anti-enumeration", + "tests/test_mcp20_task_authorization.py", + ), + ProtocolDeliverable( + ProtocolArea.OAUTH_CIMD, + "CIMD resolver with SSRF defenses and RFC 9207 iss", + "tests/test_mcp20_oauth_cimd.py", + ), + ProtocolDeliverable( + ProtocolArea.MRTR_OBSERVABILITY, + "MRTR input_required elicitation helpers", + "tests/test_mcp20_mrtr.py", + ), + ProtocolDeliverable( + ProtocolArea.MRTR_OBSERVABILITY, + "Extensions map, trace context, and cache hints", + "tests/test_mcp20_extensions_cache_observability.py", + ), +) + +ACCEPTANCE_CRITERIA: tuple[AcceptanceCriterion, ...] = ( + AcceptanceCriterion( + "stateless", + "No Mcp-Session-Id emitted; stateless POST /mcp works", + "tests/test_mcp20_blueprint.py", + ), + AcceptanceCriterion( + "task_conformance", + "Task-augmented tools/call returns immediate TaskData", + "tests/test_mcp20_tasks.py", + ), + AcceptanceCriterion( + "task_cancellation", + "tasks/cancel marks cancelled; terminal cancel returns -32602", + "tests/test_mcp20_tasks.py", + ), + AcceptanceCriterion( + "ttl_safety", + "Active tasks never evicted; terminal TTL from lastUpdatedAt", + "tests/test_mcp20_task_store.py", + ), + AcceptanceCriterion( + "multi_tenant", + "Cross-tenant access returns TaskNotFoundError / -32602", + "tests/test_mcp20_task_authorization.py", + ), + AcceptanceCriterion( + "ssrf_security", + "CIMD blocks special-use IPs and oversized payloads", + "tests/test_mcp20_oauth_cimd.py", + ), + AcceptanceCriterion( + "error_parity", + "JSON-RPC codes match specification including SEP-2164", + "tests/test_mcp20_jsonrpc_wire.py", + ), + AcceptanceCriterion( + "automated_tests", + "pytest suite covers all seven protocol areas", + "tests/test_mcp20_acceptance.py", + ), +) + +MCP20_TEST_MODULES: tuple[str, ...] = tuple( + sorted({d.test_module for d in PROTOCOL_DELIVERABLES} | {c.test_module for c in ACCEPTANCE_CRITERIA}) +) + + +def deliverables_for_area(area: ProtocolArea) -> tuple[ProtocolDeliverable, ...]: + return tuple(item for item in PROTOCOL_DELIVERABLES if item.area == area) + + +def protocol_coverage_complete() -> bool: + """True when every protocol area has at least one mapped deliverable.""" + return all(deliverables_for_area(area) for area in ProtocolArea) + + +def acceptance_criteria_registered() -> bool: + return len(ACCEPTANCE_CRITERIA) >= 8 + + +def modern_protocol_target() -> str: + return MODERN_PROTOCOL_VERSION + + +def iter_area_summary() -> Iterable[str]: + for area in ProtocolArea: + items = deliverables_for_area(area) + yield f"Area {area.value}: {len(items)} deliverable(s)" diff --git a/nitrostack/runtime/conformance.py b/nitrostack/runtime/conformance.py new file mode 100644 index 0000000..dc352b0 --- /dev/null +++ b/nitrostack/runtime/conformance.py @@ -0,0 +1,98 @@ +"""MCP 2026-07-28 implementation blueprint conformance registry.""" + +from __future__ import annotations + +from enum import Enum +from importlib import import_module +from pathlib import Path +from typing import Iterable + +from nitrostack.protocol.errors import JsonRpcErrorCode +from nitrostack.protocol.method_contract import DEPRECATED_MODERN_METHODS as _DEPRECATED_MESSAGES +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION + +_PACKAGE_ROOT = Path(__file__).resolve().parent.parent + + +class ConformanceArea(str, Enum): + """Five verification areas for the MCP 2026-07-28 blueprint.""" + + STATELESS_HTTP = "stateless_http" + JSONRPC = "jsonrpc" + TASKS = "tasks" + MULTI_TENANT = "multi_tenant" + CIMD_SSRF = "cimd_ssrf" + + +BLUEPRINT_CONFORMANCE_AREAS: dict[ConformanceArea, str] = { + ConformanceArea.STATELESS_HTTP: ( + "Stateless POST /mcp, server/discover on 2026-07-28, no Mcp-Session-Id, CORS" + ), + ConformanceArea.JSONRPC: ( + "Standard JSON-RPC error codes; deprecated methods rejected on modern wire" + ), + ConformanceArea.TASKS: ( + "Task-augmented tools/call, tasks/get lifecycle, cancel, terminal-only TTL eviction" + ), + ConformanceArea.MULTI_TENANT: ( + "Cross-tenant task access raises TaskNotFoundError without enumeration" + ), + ConformanceArea.CIMD_SSRF: ( + "CIMD fetch blocks special-use IPs, redirects, and payloads over 5 KiB" + ), +} + +# Canonical module paths implementing the recommended package layout. +REQUIRED_BLUEPRINT_MODULES: tuple[str, ...] = ( + "nitrostack.core.app", + "nitrostack.core.context", + "nitrostack.core.decorators", + "nitrostack.core.errors", + "nitrostack.core.task", + "nitrostack.protocol.version", + "nitrostack.protocol.schema", + "nitrostack.protocol.mrtr", + "nitrostack.protocol.cache_hints", + "nitrostack.protocol.observability", + "nitrostack.tasks.store", + "nitrostack.tasks.memory", + "nitrostack.auth.cimd", + "nitrostack.auth.oauth", + "nitrostack.transports.http", + "nitrostack.transports.sse", + "nitrostack.transports.stdio", +) + +MODERN_JSONRPC_ERROR_CODES: frozenset[int] = frozenset( + { + JsonRpcErrorCode.PARSE_ERROR, + JsonRpcErrorCode.INVALID_REQUEST, + JsonRpcErrorCode.METHOD_NOT_FOUND, + JsonRpcErrorCode.INVALID_PARAMS, + JsonRpcErrorCode.INTERNAL_ERROR, + } +) + +DEPRECATED_MODERN_METHODS: frozenset[str] = frozenset(_DEPRECATED_MESSAGES) + + +def verify_package_layout(modules: Iterable[str] | None = None) -> list[str]: + """Return import errors for any required blueprint module that cannot load.""" + missing: list[str] = [] + for module_path in modules or REQUIRED_BLUEPRINT_MODULES: + try: + import_module(module_path) + except Exception as exc: # pragma: no cover - surfaced by tests + missing.append(f"{module_path}: {exc}") + return missing + + +def assert_blueprint_layout() -> None: + """Raise ``ImportError`` when a required blueprint module is unavailable.""" + errors = verify_package_layout() + if errors: + raise ImportError("Blueprint package layout incomplete:\n" + "\n".join(errors)) + + +def protocol_version_matches_blueprint() -> bool: + return MODERN_PROTOCOL_VERSION == "2026-07-28" diff --git a/nitrostack/runtime/correlation.py b/nitrostack/runtime/correlation.py new file mode 100644 index 0000000..05ed212 --- /dev/null +++ b/nitrostack/runtime/correlation.py @@ -0,0 +1,61 @@ +"""Per-request correlation for overlapping JSON-RPC calls. + +Client JSON-RPC ``id`` values are echoed on the wire only. In-flight work, +progress, and cancel target an internal correlation id so two concurrent +POSTs that reuse ``id: 1`` stay isolated. +""" + +from __future__ import annotations + +import asyncio +import uuid +from dataclasses import dataclass, field +from typing import Any, Iterator, Optional + + +def new_correlation_id() -> str: + return str(uuid.uuid4()) + + +@dataclass +class InFlightTicket: + """One in-flight handler, keyed only by ``correlation_id``.""" + + correlation_id: str + jsonrpc_id: Any = None + cancel_requested: asyncio.Event = field(default_factory=asyncio.Event) + + def request_cancel(self) -> None: + self.cancel_requested.set() + + +class InFlightRegistry: + """In-flight map that never indexes by the client JSON-RPC ``id``.""" + + def __init__(self) -> None: + self._tickets: dict[str, InFlightTicket] = {} + + def register(self, correlation_id: str, *, jsonrpc_id: Any = None) -> InFlightTicket: + ticket = InFlightTicket(correlation_id=correlation_id, jsonrpc_id=jsonrpc_id) + self._tickets[correlation_id] = ticket + return ticket + + def get(self, correlation_id: str) -> Optional[InFlightTicket]: + return self._tickets.get(correlation_id) + + def discard(self, correlation_id: str) -> None: + self._tickets.pop(correlation_id, None) + + def cancel(self, correlation_id: str) -> bool: + """Cancel one ticket. Same client ``id`` on another ticket is untouched.""" + ticket = self._tickets.get(correlation_id) + if ticket is None: + return False + ticket.request_cancel() + return True + + def __len__(self) -> int: + return len(self._tickets) + + def __iter__(self) -> Iterator[InFlightTicket]: + return iter(list(self._tickets.values())) diff --git a/nitrostack/runtime/request_ctx.py b/nitrostack/runtime/request_ctx.py new file mode 100644 index 0000000..485e849 --- /dev/null +++ b/nitrostack/runtime/request_ctx.py @@ -0,0 +1,65 @@ +"""Per-request context for official mcp 2.x (replaces 1.x ``request_ctx``).""" + +from __future__ import annotations + +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Any, Mapping, Optional + +request_ctx: ContextVar[Any] = ContextVar("nitrostack_request_ctx", default=None) + + +@dataclass +class Experimental: + """Task metadata previously hung off the 1.x experimental request slot.""" + + task_metadata: Any = None + + +@dataclass +class RequestContext: + """Test and in-process stand-in for the 1.x low-level request context.""" + + request_id: Any = None + correlation_id: Any = None + meta: Any = None + session: Any = None + lifespan_context: Any = None + experimental: Optional[Experimental] = None + request: Any = None + + +class RequestParamsMeta: + """Stand-in for 1.x ``types.RequestParams.Meta`` extra fields.""" + + def __init__(self, __pydantic_extra__: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> None: + extra = dict(__pydantic_extra__ or {}) + extra.update(kwargs) + self.__pydantic_extra__ = extra + for key, value in extra.items(): + setattr(self, key, value) + + @classmethod + def model_validate(cls, data: Any) -> "RequestParamsMeta": + if isinstance(data, cls): + return data + if data is None: + return cls() + if isinstance(data, Mapping): + return cls(**dict(data)) + extra = getattr(data, "__pydantic_extra__", None) + if isinstance(extra, Mapping): + return cls(__pydantic_extra__=extra) + return cls() + + +class ServerResult: + """1.x ``ServerResult`` wrapper: in-process callers still read ``.root``.""" + + def __init__(self, root: Any) -> None: + self.root = root + + +def bind_request_ctx(ctx: Any): + """Push ``ctx`` onto the request ContextVar and return the reset token.""" + return request_ctx.set(ctx) diff --git a/nitrostack/runtime/stateless.py b/nitrostack/runtime/stateless.py new file mode 100644 index 0000000..cb849db --- /dev/null +++ b/nitrostack/runtime/stateless.py @@ -0,0 +1,74 @@ +"""Stateless HTTP invariants for MCP 2026-07-28.""" + +from dataclasses import dataclass +from typing import Mapping, Optional + +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER +from nitrostack.protocol.version import ProtocolEra, supported_protocol_versions_for_era + + +@dataclass(frozen=True) +class StatelessInvariants: + """Architectural invariants for MCP 2026-07-28 stateless transport.""" + + emit_session_headers: bool = False + require_initialize_handshake: bool = False + allow_session_stickiness: bool = False + + +DEFAULT_STATELESS_INVARIANTS = StatelessInvariants() + + +def has_incoming_session_id(request_headers: Mapping[str, str]) -> bool: + """True when the client sent a non-empty ``Mcp-Session-Id``.""" + target = LEGACY_SESSION_HEADER.lower() + for key, value in request_headers.items(): + if key.lower() == target and str(value).strip(): + return True + return False + + +def sessionless_strips_incoming_session_id(wire_mode: str) -> bool: + """``modern`` and ``auto`` ignore client ``Mcp-Session-Id``; ``legacy`` keeps it.""" + return wire_mode != "sessionful" + + +def sessionless_rejects_incoming_session_id(wire_mode: str) -> bool: + """Deprecated alias for ``sessionless_strips_incoming_session_id``.""" + return sessionless_strips_incoming_session_id(wire_mode) + + +def request_protocol_version( + header_version: Optional[str], + envelope_version: Optional[str] = None, +) -> Optional[str]: + """ + Resolve the request protocol version. + + The header wins when present. Otherwise ``_meta.mcp.protocolVersion`` is + used. When neither is present the request proceeds (legacy default). The + header is not required on ``modern`` in this sidecar. + """ + for raw in (header_version, envelope_version): + if isinstance(raw, str) and raw.strip(): + return raw.strip() + return None + + +def is_unsupported_protocol_version( + version: Optional[str], + era: ProtocolEra, +) -> bool: + """True when a version is present and not in the era's supported set.""" + if version is None or not str(version).strip(): + return False + return version.strip() not in supported_protocol_versions_for_era(era) + + +def assert_stateless_headers(response_headers: dict[str, str]) -> None: + """Raise if a response violates stateless wire rules (no Mcp-Session-Id).""" + for key in response_headers: + if key.lower() == LEGACY_SESSION_HEADER.lower(): + raise ValueError( + f"Stateless MCP MUST NOT emit '{LEGACY_SESSION_HEADER}' on responses" + ) diff --git a/nitrostack/tasks/__init__.py b/nitrostack/tasks/__init__.py new file mode 100644 index 0000000..d208aa9 --- /dev/null +++ b/nitrostack/tasks/__init__.py @@ -0,0 +1,18 @@ +"""Async MCP task subsystem with pluggable persistence.""" + +from nitrostack.tasks.memory import InMemoryTaskStore +from nitrostack.tasks.authorization import check_task_access, extract_task_access_context, list_task_wire_data_for_context +from nitrostack.tasks.store import TaskStore +from nitrostack.tasks.types import TaskAccessContext, TaskEntry, TaskStatus, TaskWireData + +__all__ = [ + "TaskStore", + "InMemoryTaskStore", + "TaskAccessContext", + "TaskEntry", + "TaskStatus", + "TaskWireData", + "check_task_access", + "extract_task_access_context", + "list_task_wire_data_for_context", +] diff --git a/nitrostack/tasks/authorization.py b/nitrostack/tasks/authorization.py new file mode 100644 index 0000000..e16cd29 --- /dev/null +++ b/nitrostack/tasks/authorization.py @@ -0,0 +1,136 @@ +"""Task authorization and multi-tenant isolation.""" + +from __future__ import annotations + +from typing import List, Optional, Tuple + +from nitrostack.auth.request import ( + authorization_token_from_request, + tenant_from_claims, + verify_bearer_payload, +) +from nitrostack.core.errors import TaskNotFoundError +from nitrostack.tasks.types import TaskAccessContext, TaskEntry, TaskWireData + + +def check_task_access(entry: TaskEntry, context: TaskAccessContext | None) -> None: + """ + Enforce tenant/user/session isolation for task access. + + ``context is None`` is the internal, non-HTTP path (TaskManager progress + / complete / fail). HTTP handlers must pass a ``TaskAccessContext`` — + possibly empty — so missing identity cannot read a scoped task. + + When the task was created with owner/tenant/session, every set dimension + must match. A missing dimension on the caller is a mismatch. + + Raises ``TaskNotFoundError`` (not Forbidden) on mismatch — anti-enumeration. + """ + if context is None: + return + + task_id = entry.data.task_id + + if entry.tenant_id and entry.tenant_id != context.tenant_id: + raise TaskNotFoundError(task_id) + + if entry.owner_id and entry.owner_id != context.user_id: + raise TaskNotFoundError(task_id) + + if entry.session_id and entry.session_id != context.session_id: + raise TaskNotFoundError(task_id) + + +def entry_matches_access_context(entry: TaskEntry, context: TaskAccessContext | None) -> bool: + """Return True when ``entry`` is visible to ``context`` (list filtering).""" + if context is None: + return True + try: + check_task_access(entry, context) + return True + except TaskNotFoundError: + return False + + +def list_task_wire_data_for_context( + entries: List[TaskEntry], + context: TaskAccessContext | None, + *, + cursor: str | None = None, + limit: int = 50, +) -> Tuple[List[TaskWireData], Optional[str]]: + """ + Filter, sort, and paginate task entries for a caller context. + """ + filtered = [entry for entry in entries if entry_matches_access_context(entry, context)] + filtered.sort(key=lambda entry: entry.data.created_at, reverse=True) + + start_index = 0 + if cursor: + cursor_index = next( + (index for index, entry in enumerate(filtered) if entry.task_id == cursor), + None, + ) + if cursor_index is None: + raise TaskNotFoundError(cursor) + start_index = cursor_index + 1 + + page_entries = filtered[start_index : start_index + limit] + page = [entry.data for entry in page_entries] + next_cursor = page[-1].task_id if (start_index + limit) < len(filtered) else None + return page, next_cursor + + +def _tenant_from_claims(claims: dict[str, Any]) -> Optional[str]: + return tenant_from_claims(claims) + + +def _authorization_from_request_context(rc: Any) -> Optional[str]: + """HTTP Authorization, then envelope auth, then ``_meta`` Bearer.""" + return authorization_token_from_request(rc) + + +def _session_id_from_request_context(rc: Any) -> Optional[str]: + session = getattr(rc, "session", None) + if session is None: + return None + session_id = getattr(session, "id", None) or getattr(session, "session_id", None) + if session_id is None: + return None + return str(session_id) + + +def extract_task_access_context(rc: Any) -> Optional[TaskAccessContext]: + """ + Build ``TaskAccessContext`` from an MCP request context. + + Identity comes from a verified JWT: HTTP ``Authorization``, then the + spec envelope auth slot, then a Bearer token in ``_meta``. Unsigned + ``userId`` / ``tenantId`` fields are never used. If a Bearer token is + present and verification fails, returns an empty context so scoped + tasks are denied. + + ``None`` is reserved for callers with no request context (internal path). + """ + if rc is None: + return None + + user_id: Optional[str] = None + tenant_id: Optional[str] = None + session_id = _session_id_from_request_context(rc) + token = _authorization_from_request_context(rc) + + if token: + payload = verify_bearer_payload(token) + if payload is None: + return TaskAccessContext() + subject = payload.get("sub") + if isinstance(subject, str) and subject.strip(): + user_id = subject + tenant_id = _tenant_from_claims(payload) + + return TaskAccessContext( + user_id=str(user_id) if user_id else None, + tenant_id=str(tenant_id) if tenant_id else None, + session_id=str(session_id) if session_id else None, + ) diff --git a/nitrostack/tasks/eviction.py b/nitrostack/tasks/eviction.py new file mode 100644 index 0000000..5b488d1 --- /dev/null +++ b/nitrostack/tasks/eviction.py @@ -0,0 +1,25 @@ +"""Terminal-only TTL eviction rules.""" + +from __future__ import annotations + +from nitrostack.tasks.types import TERMINAL_TASK_STATUSES, TaskEntry + + +def is_terminal_eviction_eligible(status: str) -> bool: + """Only completed, failed, and cancelled tasks may be evicted.""" + return status in TERMINAL_TASK_STATUSES + + +def should_evict_terminal_task(entry: TaskEntry, now_ms: int) -> bool: + """ + Return True when a terminal task exceeded its post-completion TTL. + + Active tasks (``working``, ``input_required``) are never evicted. + """ + if not is_terminal_eviction_eligible(entry.status): + return False + ttl_ms = entry.data.ttl_ms + if ttl_ms is None: + return False + last_updated_ms = entry.data.last_updated_at_ms + return (now_ms - last_updated_ms) > ttl_ms diff --git a/nitrostack/tasks/memory.py b/nitrostack/tasks/memory.py new file mode 100644 index 0000000..3353df1 --- /dev/null +++ b/nitrostack/tasks/memory.py @@ -0,0 +1,49 @@ +"""In-memory TaskStore for development and single-replica deployments.""" + +from __future__ import annotations + +from typing import Dict, List, Optional + +from nitrostack.tasks.eviction import should_evict_terminal_task +from nitrostack.tasks.store import TaskStore +from nitrostack.tasks.types import TaskEntry + + +class InMemoryTaskStore(TaskStore): + """Process-local task persistence with terminal-only TTL eviction.""" + + def __init__(self) -> None: + self._entries: Dict[str, TaskEntry] = {} + + def get_sync(self, task_id: str) -> Optional[TaskEntry]: + return self._entries.get(task_id) + + def set_sync(self, task_id: str, entry: TaskEntry) -> None: + self._entries[task_id] = entry + + async def get(self, task_id: str) -> Optional[TaskEntry]: + return self.get_sync(task_id) + + async def set(self, task_id: str, entry: TaskEntry) -> None: + self.set_sync(task_id, entry) + + async def delete(self, task_id: str) -> bool: + return self._entries.pop(task_id, None) is not None + + async def has(self, task_id: str) -> bool: + return task_id in self._entries + + async def list(self) -> List[TaskEntry]: + return list(self._entries.values()) + + async def cleanup_expired(self, now_ms: int) -> int: + evicted = 0 + for task_id in list(self._entries.keys()): + entry = self._entries[task_id] + if should_evict_terminal_task(entry, now_ms): + del self._entries[task_id] + evicted += 1 + return evicted + + async def destroy(self) -> None: + self._entries.clear() diff --git a/nitrostack/tasks/notify.py b/nitrostack/tasks/notify.py new file mode 100644 index 0000000..fc81677 --- /dev/null +++ b/nitrostack/tasks/notify.py @@ -0,0 +1,194 @@ +"""Fan-out for ``notifications/tasks/status`` across live transports.""" + +from __future__ import annotations + +import logging +from typing import Any, Callable, Optional, Protocol + +from nitrostack.tasks.authorization import entry_matches_access_context +from nitrostack.tasks.types import TaskAccessContext, TaskEntry + +logger = logging.getLogger(__name__) + +TASK_STATUS_METHOD = "notifications/tasks/status" + + +def build_task_status_notification( + task_id: str, + status: str, + *, + status_message: Optional[str] = None, +) -> dict[str, Any]: + """Wire notification for a task status or progress change.""" + params: dict[str, Any] = {"taskId": task_id, "status": status} + if status_message is not None: + params["statusMessage"] = status_message + return {"method": TASK_STATUS_METHOD, "params": params} + + +class TaskStatusSink(Protocol): + """One connected transport that can receive task status.""" + + access: Optional[TaskAccessContext] + session_id: Optional[str] + task_id: Optional[str] + + async def send(self, notification: dict[str, Any]) -> None: + """Deliver one status notification. Must not raise to the router.""" + ... + + +class CallbackTaskSink: + """Test and in-process sink that records or forwards notifications.""" + + def __init__( + self, + callback: Callable[[dict[str, Any]], Any], + *, + access: Optional[TaskAccessContext] = None, + session_id: Optional[str] = None, + task_id: Optional[str] = None, + ) -> None: + self._callback = callback + self.access = access + self.session_id = session_id + self.task_id = task_id + + async def send(self, notification: dict[str, Any]) -> None: + result = self._callback(notification) + if hasattr(result, "__await__"): + await result + + +class QueueTaskSink: + """Push notifications onto an asyncio queue (listen / SSE attach).""" + + def __init__( + self, + put: Callable[[dict[str, Any]], None], + *, + access: Optional[TaskAccessContext] = None, + session_id: Optional[str] = None, + task_id: Optional[str] = None, + ) -> None: + self._put = put + self.access = access + self.session_id = session_id + self.task_id = task_id + + async def send(self, notification: dict[str, Any]) -> None: + self._put(notification) + + +class StreamTaskSink: + """Write a JSON-RPC notification onto an official stream write side.""" + + def __init__( + self, + write_stream: Any, + *, + access: Optional[TaskAccessContext] = None, + session_id: Optional[str] = None, + task_id: Optional[str] = None, + ) -> None: + self._write_stream = write_stream + self.access = access + self.session_id = session_id + self.task_id = task_id + + async def send(self, notification: dict[str, Any]) -> None: + from mcp.shared.message import SessionMessage + from mcp_types import jsonrpc_message_adapter + + message = jsonrpc_message_adapter.validate_python( + { + "jsonrpc": "2.0", + "method": notification["method"], + "params": notification.get("params") or {}, + } + ) + await self._write_stream.send(SessionMessage(message)) + + +class SessionTaskSink: + """Notify the originating MCP session (stdio or legacy SSE).""" + + def __init__( + self, + session: Any, + *, + access: Optional[TaskAccessContext] = None, + session_id: Optional[str] = None, + task_id: Optional[str] = None, + ) -> None: + self._session = session + self.access = access + self.session_id = session_id + self.task_id = task_id + + async def send(self, notification: dict[str, Any]) -> None: + session = self._session + if session is None: + return + method = notification["method"] + params = notification.get("params") or {} + sender = getattr(session, "send_notification", None) + if sender is not None: + await sender(method, params) + return + outbound = getattr(session, "outbound", None) + notify = getattr(outbound, "notify", None) if outbound is not None else None + if notify is not None: + await notify(method, params) + + +def _sink_accepts(sink: TaskStatusSink, entry: TaskEntry) -> bool: + task_id = getattr(sink, "task_id", None) + if task_id is not None and task_id != entry.task_id: + return False + session_id = getattr(sink, "session_id", None) + if session_id is not None and entry.session_id and session_id != entry.session_id: + return False + access = getattr(sink, "access", None) + if access is not None and not entry_matches_access_context(entry, access): + return False + return True + + +class TaskStatusRouter: + """Deliver task status to every live channel. A failed send is ignored.""" + + def __init__(self) -> None: + self._sinks: dict[object, TaskStatusSink] = {} + + def register(self, sink: TaskStatusSink) -> Callable[[], None]: + token = object() + self._sinks[token] = sink + + def unsubscribe() -> None: + self._sinks.pop(token, None) + + return unsubscribe + + async def notify_task_status( + self, + entry: TaskEntry, + *, + status: Optional[str] = None, + status_message: Optional[str] = None, + ) -> None: + """Fan out one status change. Never raises to the task lifecycle.""" + notification = build_task_status_notification( + entry.task_id, + status if status is not None else entry.status, + status_message=status_message + if status_message is not None + else entry.data.status_message, + ) + for sink in list(self._sinks.values()): + if not _sink_accepts(sink, entry): + continue + try: + await sink.send(notification) + except Exception: + logger.exception("task status notify failed; continuing") diff --git a/nitrostack/tasks/store.py b/nitrostack/tasks/store.py new file mode 100644 index 0000000..6b7a463 --- /dev/null +++ b/nitrostack/tasks/store.py @@ -0,0 +1,37 @@ +"""Pluggable task persistence interface.""" + +from abc import ABC, abstractmethod +from typing import List, Optional + +from nitrostack.tasks.types import TaskEntry + + +class TaskStore(ABC): + """Abstract storage adapter decoupling TaskManager from persistence backend.""" + + @abstractmethod + async def get(self, task_id: str) -> Optional[TaskEntry]: + """Fetch task entry by ID.""" + + @abstractmethod + async def set(self, task_id: str, entry: TaskEntry) -> None: + """Save or update task entry atomically.""" + + @abstractmethod + async def delete(self, task_id: str) -> bool: + """Remove task entry. Returns True if it existed.""" + + @abstractmethod + async def has(self, task_id: str) -> bool: + """Check existence without deserializing the full payload.""" + + @abstractmethod + async def list(self) -> List[TaskEntry]: + """Return all stored entries (for pagination and TTL sweep).""" + + @abstractmethod + async def cleanup_expired(self, now_ms: int) -> int: + """Evict expired terminal tasks. Returns count evicted.""" + + async def destroy(self) -> None: + """Release connections, timers, and background resources.""" diff --git a/nitrostack/tasks/types.py b/nitrostack/tasks/types.py new file mode 100644 index 0000000..f923d55 --- /dev/null +++ b/nitrostack/tasks/types.py @@ -0,0 +1,79 @@ +"""Task subsystem shared types.""" + +from __future__ import annotations + +import datetime +from dataclasses import dataclass, field +from typing import Any, Literal, Optional + +from pydantic import BaseModel, ConfigDict, Field + +TaskStatus = Literal["working", "input_required", "completed", "failed", "cancelled"] + +TERMINAL_TASK_STATUSES: frozenset[str] = frozenset({"completed", "failed", "cancelled"}) + + +def utc_now() -> datetime.datetime: + return datetime.datetime.now(datetime.timezone.utc) + + +def datetime_to_ms(value: datetime.datetime) -> int: + if value.tzinfo is None: + value = value.replace(tzinfo=datetime.timezone.utc) + return int(value.timestamp() * 1000) + + +class TaskAccessContext(BaseModel): + """Identity metadata for multi-tenant task authorization.""" + + model_config = ConfigDict(populate_by_name=True) + + user_id: Optional[str] = Field(default=None, alias="userId") + tenant_id: Optional[str] = Field(default=None, alias="tenantId") + session_id: Optional[str] = Field(default=None, alias="sessionId") + + +@dataclass +class TaskWireData: + """Protocol-visible task metadata.""" + + task_id: str + status: TaskStatus = "working" + status_message: Optional[str] = None + created_at: datetime.datetime = field(default_factory=utc_now) + last_updated_at: datetime.datetime = field(default_factory=utc_now) + ttl_ms: Optional[int] = None + poll_interval_ms: int = 2000 + owner_id: Optional[str] = None + tenant_id: Optional[str] = None + session_id: Optional[str] = None + + def __post_init__(self) -> None: + if self.last_updated_at is None: + self.last_updated_at = self.created_at + + @property + def last_updated_at_ms(self) -> int: + return datetime_to_ms(self.last_updated_at) + + +@dataclass +class TaskEntry: + """Internal server task state persisted by ``TaskStore``.""" + + task_id: str + data: TaskWireData + status: TaskStatus = "working" + result: Any = None + error: Optional[dict[str, Any]] = None + tool_name: Optional[str] = None + owner_id: Optional[str] = None + tenant_id: Optional[str] = None + session_id: Optional[str] = None + metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.data.task_id != self.task_id: + self.data.task_id = self.task_id + if self.status != self.data.status: + self.data.status = self.status diff --git a/nitrostack/testing/__init__.py b/nitrostack/testing/__init__.py index 7899d4b..0db12b6 100644 --- a/nitrostack/testing/__init__.py +++ b/nitrostack/testing/__init__.py @@ -11,9 +11,8 @@ class NitroTestingModule: In-process test harness for testing NitroStack applications without spinning up real transports or subprocesses (Section 15). - Dispatches directly through the owned low-level `mcp.server.lowlevel.Server`'s - registered `request_handlers` (the same dict the real stdio/HTTP transports use), - rather than any FastMCP-specific convenience method. + Dispatches through the owned low-level `mcp.server.lowlevel.Server` + `request_handlers` (the same dict the real stdio/HTTP transports use). """ @classmethod async def create(cls, app_module: Type) -> "NitroTestingModule": diff --git a/nitrostack/transports/__init__.py b/nitrostack/transports/__init__.py new file mode 100644 index 0000000..6cf62b1 --- /dev/null +++ b/nitrostack/transports/__init__.py @@ -0,0 +1,30 @@ +"""MCP transport adapters (stdio, stateless HTTP).""" + +from nitrostack.transports.dispatch import ( + DispatchStage, + IngressContext, + StatelessIngressPipeline, + is_task_wire_interception, +) +from nitrostack.transports.middleware import StatelessTransportMiddleware, wrap_stateless_transport +from nitrostack.transports.sse import format_sse_message, sse_connect_headers, sse_notification +from nitrostack.transports.subscriptions import ( + http_listen_requires_auth, + listen_auth_error, + subscriptions_listen_endpoint, +) + +__all__ = [ + "DispatchStage", + "IngressContext", + "StatelessIngressPipeline", + "StatelessTransportMiddleware", + "wrap_stateless_transport", + "is_task_wire_interception", + "format_sse_message", + "sse_connect_headers", + "sse_notification", + "http_listen_requires_auth", + "listen_auth_error", + "subscriptions_listen_endpoint", +] diff --git a/nitrostack/transports/cors.py b/nitrostack/transports/cors.py new file mode 100644 index 0000000..39059da --- /dev/null +++ b/nitrostack/transports/cors.py @@ -0,0 +1,92 @@ +"""CORS configuration for stateless MCP HTTP.""" + +from __future__ import annotations + +import os +from typing import Mapping, Optional, Sequence + +from nitrostack.transports.headers import ( + CORS_ALLOW_HEADERS, + CORS_ALLOW_METHODS, + CORS_EXPOSE_HEADERS, + HEADER_ACCESS_CONTROL_ALLOW_HEADERS, + HEADER_ACCESS_CONTROL_ALLOW_METHODS, + HEADER_ACCESS_CONTROL_ALLOW_ORIGIN, + HEADER_ACCESS_CONTROL_EXPOSE_HEADERS, + HEADER_MCP_PARAM_PREFIX, + get_header, +) + + +def configured_cors_origins() -> tuple[str, ...]: + """Comma-separated allowlist from ``MCP_CORS_ALLOWED_ORIGINS``.""" + raw = os.environ.get("MCP_CORS_ALLOWED_ORIGINS", "") + return tuple(item.strip() for item in raw.split(",") if item.strip()) + + +def resolve_allowed_origin( + origin: Optional[str] = None, + *, + allow_origin: str = "*", + allowed_origins: Optional[Sequence[str]] = None, +) -> str: + """ + Choose ``Access-Control-Allow-Origin`` without reflecting arbitrary Origins. + + An explicit allowlist (argument or ``MCP_CORS_ALLOWED_ORIGINS``) is required + before a request Origin is echoed. Otherwise the configured ``allow_origin`` + default (``*``) is used. + """ + allowlist = ( + tuple(allowed_origins) if allowed_origins is not None else configured_cors_origins() + ) + if allowlist: + if origin and origin in allowlist: + return origin + if allow_origin != "*" and allow_origin in allowlist: + return allow_origin + return allowlist[0] + return allow_origin + + +def build_cors_headers( + origin: Optional[str] = None, + *, + allow_origin: str = "*", + allowed_origins: Optional[Sequence[str]] = None, +) -> dict[str, str]: + """Build CORS headers for MCP browser clients (SEP-2243 & SEP-2575).""" + resolved_origin = resolve_allowed_origin( + origin, + allow_origin=allow_origin, + allowed_origins=allowed_origins, + ) + return { + HEADER_ACCESS_CONTROL_ALLOW_ORIGIN: resolved_origin, + HEADER_ACCESS_CONTROL_ALLOW_METHODS: CORS_ALLOW_METHODS, + HEADER_ACCESS_CONTROL_ALLOW_HEADERS: CORS_ALLOW_HEADERS, + HEADER_ACCESS_CONTROL_EXPOSE_HEADERS: CORS_EXPOSE_HEADERS, + } + + +def requested_mcp_param_headers(request_headers: Mapping[str, str]) -> tuple[str, ...]: + """Exact ``Mcp-Param-*`` names from ``Access-Control-Request-Headers``.""" + raw = get_header(request_headers, "Access-Control-Request-Headers") or "" + prefix = HEADER_MCP_PARAM_PREFIX.lower() + echoed: list[str] = [] + for item in raw.split(","): + name = item.strip() + if name.lower().startswith(prefix) and len(name) > len(HEADER_MCP_PARAM_PREFIX): + echoed.append(name) + return tuple(echoed) + + +def cors_preflight_response_headers(request_headers: Mapping[str, str]) -> dict[str, str]: + """Headers for OPTIONS preflight — HTTP 204 No Content.""" + origin = get_header(request_headers, "Origin") + headers = build_cors_headers(origin=origin) + extra = requested_mcp_param_headers(request_headers) + if extra: + current = headers[HEADER_ACCESS_CONTROL_ALLOW_HEADERS] + headers[HEADER_ACCESS_CONTROL_ALLOW_HEADERS] = current + ", " + ", ".join(extra) + return headers diff --git a/nitrostack/transports/dispatch.py b/nitrostack/transports/dispatch.py new file mode 100644 index 0000000..f4005b9 --- /dev/null +++ b/nitrostack/transports/dispatch.py @@ -0,0 +1,558 @@ +"""Stateless HTTP ingress pipeline.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from enum import Enum +from collections.abc import Awaitable +from typing import Any, Callable, Optional, Union + +from nitrostack.protocol.deprecated import ( + deprecated_method_message, + rejects_deprecated_method, +) +from nitrostack.protocol.discovery import ( + INITIALIZE_METHOD, + INITIALIZED_NOTIFICATION, + SERVER_DISCOVER_METHOD, + build_sessionless_initialize_result, +) +from nitrostack.protocol.errors import ERROR_CODE_MESSAGES, JsonRpcErrorCode +from nitrostack.protocol.jsonrpc import ( + HeaderBodyMismatchError, + InvalidRequestError, + JsonRpcParseError, + JsonRpcRequest, + JsonRpcWireError, + MethodNotFoundError, + build_ping_response, + jsonrpc_error, + jsonrpc_success, + parse_jsonrpc_request, + validate_header_body_method, + validate_header_body_name, + UnsupportedProtocolVersionError, + validate_protocol_version_header_meta, + validate_required_mcp_method, + validate_required_mcp_name, + validate_supported_protocol_version, +) +from nitrostack.protocol.method_contract import ( + mcp_method_is_required, + mcp_name_field, + mcp_name_is_required, +) +from nitrostack.protocol.meta import envelope_protocol_version +from nitrostack.protocol.version import ( + LEGACY_PROTOCOL_VERSION, + ProtocolEra, + WireMode, + accepts_sessionless_initialize, + protocol_era_for_wire_mode, + rejects_legacy_initialize, + supported_protocol_versions_for_era, +) +from nitrostack.runtime.stateless import ( + is_unsupported_protocol_version, + request_protocol_version, + sessionless_strips_incoming_session_id, +) +from nitrostack.transports.headers import ( + HEADER_MCP_METHOD, + HEADER_MCP_NAME, + HEADER_MCP_PROTOCOL_VERSION, + first_oversized_mcp_param, + get_header, + strip_legacy_session_headers, +) + +TaskDispatchHandler = Callable[[JsonRpcRequest], Awaitable[Optional[dict[str, Any]]]] +RegistryDispatchHandler = Callable[[JsonRpcRequest], Awaitable[Optional[dict[str, Any]]]] +DiscoverHandler = Callable[[JsonRpcRequest], Union[Awaitable[dict[str, Any]], dict[str, Any]]] +InitializeHandler = Callable[[JsonRpcRequest], Union[Awaitable[dict[str, Any]], dict[str, Any]]] + + +class DispatchStage(str, Enum): + """Six-step stateless HTTP ingress lifecycle.""" + + CORS_SECURITY = "cors_security" + BODY_PARSING = "body_parsing" + PING_FAST_PATH = "ping_fast_path" + TASK_INTERCEPTION = "task_interception" + REGISTRY_DISPATCH = "registry_dispatch" + RESPONSE_SERIALIZATION = "response_serialization" + + +TASK_METHOD_PREFIX = "tasks/" +TOOLS_CALL_METHOD = "tools/call" +PING_METHOD = "ping" +LEGACY_HANDSHAKE_METHODS = frozenset({"initialize", "notifications/initialized"}) + + +def is_header_only_ping(raw_body: bytes, request_headers: dict[str, str]) -> bool: + """True when ``Mcp-Method: ping`` and the body is empty or not JSON-RPC. + + A parsed JSON-RPC body is never header-only: header/body match still applies. + Other ``Mcp-Method`` values are not answered from the header alone. + """ + header_method = get_header(request_headers, HEADER_MCP_METHOD) + if header_method is None or header_method.strip() != PING_METHOD: + return False + if not raw_body or not raw_body.strip(): + return True + try: + parse_jsonrpc_request(raw_body) + except (JsonRpcParseError, JsonRpcWireError): + return True + return False + + +@dataclass +class IngressContext: + server_name: str + server_version: str + protocol_version: str + advertise_tasks: bool = True + advertise_app: bool = False + custom_extensions: Optional[dict[str, str]] = None + wire_mode: WireMode = "stateless" + protocol_era: Optional[ProtocolEra] = None + + def resolved_era(self) -> ProtocolEra: + if self.protocol_era is not None: + return self.protocol_era + return protocol_era_for_wire_mode(self.wire_mode) + + def accepts_sessionless_initialize(self) -> bool: + """True when this engine answers 2025 ``initialize`` without a session.""" + return accepts_sessionless_initialize(self.resolved_era()) + + def rejects_legacy_initialize(self) -> bool: + """True when this engine rejects 2025 ``initialize`` / ``initialized``.""" + return rejects_legacy_initialize(self.resolved_era()) + + +def prepare_sessionless_request_headers( + request_headers: dict[str, str], + wire_mode: WireMode, +) -> dict[str, str]: + """ + Drop client ``Mcp-Session-Id`` on sessionless engines (``modern`` / ``auto``). + + Obsolete session headers must not become a dependency; they are ignored rather + than rejected so stateless clients stay unambiguous. + """ + if not sessionless_strips_incoming_session_id(wire_mode): + return request_headers + return strip_legacy_session_headers(request_headers) + + +def reject_legacy_handshake( + request: JsonRpcRequest, + era: ProtocolEra, +) -> Optional[tuple[int, dict[str, Any]]]: + """Era ``modern`` answers ``initialize`` / ``initialized`` as method-not-found.""" + return reject_legacy_handshake_method(request.method, request.id, era) + + +def reject_legacy_handshake_method( + method: str, + request_id: Any, + era: ProtocolEra, +) -> Optional[tuple[int, dict[str, Any]]]: + """Era ``modern`` answers handshake methods as method-not-found.""" + if not rejects_legacy_initialize(era): + return None + if method not in LEGACY_HANDSHAKE_METHODS: + return None + return 200, MethodNotFoundError(method).to_response(request_id) + + +def reject_deprecated_method( + request: JsonRpcRequest, + era: ProtocolEra, +) -> Optional[tuple[int, dict[str, Any]]]: + """Era ``modern`` answers retired 2025 methods as method-not-found.""" + return reject_deprecated_method_name(request.method, request.id, era) + + +def reject_deprecated_method_name( + method: str, + request_id: Any, + era: ProtocolEra, +) -> Optional[tuple[int, dict[str, Any]]]: + """Same retired-method error used on POST, GET, and replay.""" + if not rejects_deprecated_method(method, era): + return None + message = deprecated_method_message(method) + return 200, jsonrpc_error( + request_id, + int(JsonRpcErrorCode.METHOD_NOT_FOUND), + message or f"Method not found: {method}", + ) + + +def reject_modern_method_policy( + method: str, + request_id: Any, + era: ProtocolEra, +) -> Optional[tuple[int, dict[str, Any]]]: + """Handshake then retired-method policy. ``auto`` does not error here.""" + handshake = reject_legacy_handshake_method(method, request_id, era) + if handshake is not None: + return handshake + return reject_deprecated_method_name(method, request_id, era) + + +def reject_legacy_wire( + request: JsonRpcRequest, + request_headers: dict[str, str], + era: ProtocolEra, +) -> Optional[tuple[int, dict[str, Any]]]: + """ + Era ``modern`` fails closed on 2025 handshake and 2025 protocol versions. + + Official v2 ``legacy: 'reject'`` is not mounted yet; this is the sidecar + stand-in. Handshake methods are method-not-found before header contracts. + Incoming session ids are rejected earlier. + """ + handshake = reject_legacy_handshake(request, era) + if handshake is not None: + return handshake + + if not rejects_legacy_initialize(era): + return None + + header_version = get_header(request_headers, HEADER_MCP_PROTOCOL_VERSION) + body_version = request.params.get("protocolVersion") + if header_version == LEGACY_PROTOCOL_VERSION or body_version == LEGACY_PROTOCOL_VERSION: + return 400, jsonrpc_error( + request.id, + int(JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION), + ERROR_CODE_MESSAGES[JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION], + ) + + return None + + +def is_task_wire_interception(method: str, params: dict[str, Any]) -> bool: + """Ingress step 4 — route to the task subsystem when matched.""" + if method.startswith(TASK_METHOD_PREFIX): + return True + return method == TOOLS_CALL_METHOD and bool(params.get("task")) + + +def reject_required_mcp_name( + request: JsonRpcRequest, + request_headers: dict[str, str], + wire_mode: WireMode, +) -> Optional[tuple[int, dict[str, Any]]]: + """Require or cross-check ``Mcp-Name`` on name-scoped methods.""" + if not mcp_name_is_required(request.method): + return None + field = mcp_name_field(request.method) or "name" + header_name = get_header(request_headers, HEADER_MCP_NAME) + body_name = request.params.get(field) + body_value = body_name if isinstance(body_name, str) else None + try: + if wire_mode == "stateless": + if header_name is not None: + validate_header_body_name(header_name, body_value) + elif not (body_value and body_value.strip()): + raise InvalidRequestError( + f"{field!r} is required in params for {request.method!r}" + ) + else: + validate_required_mcp_name( + header_name, + body_value, + body_label=field, + ) + except HeaderBodyMismatchError as exc: + return 400, exc.to_response(request.id) + except InvalidRequestError as exc: + return 400, exc.to_response(request.id) + return None + + +def reject_required_mcp_method( + request: JsonRpcRequest, + request_headers: dict[str, str], + wire_mode: WireMode, +) -> Optional[tuple[int, dict[str, Any]]]: + """Require or optionally cross-check ``Mcp-Method`` against the JSON-RPC method.""" + header_method = get_header(request_headers, HEADER_MCP_METHOD) + try: + if mcp_method_is_required(request.method, wire_mode): + validate_required_mcp_method(header_method, request.method) + else: + validate_header_body_method(header_method, request.method) + except HeaderBodyMismatchError as exc: + return 400, exc.to_response(request.id) + return None + + +def reject_protocol_version_mismatch( + request: JsonRpcRequest, + request_headers: dict[str, str], +) -> Optional[tuple[int, dict[str, Any]]]: + """Reject when header and ``_meta.mcp.protocolVersion`` both exist and differ.""" + header_version = get_header(request_headers, HEADER_MCP_PROTOCOL_VERSION) + meta_version = envelope_protocol_version(request.meta) + try: + validate_protocol_version_header_meta(header_version, meta_version) + except HeaderBodyMismatchError as exc: + return 400, exc.to_response(request.id) + return None + + +def reject_unsupported_protocol_version( + request: JsonRpcRequest, + request_headers: dict[str, str], + era: ProtocolEra, +) -> Optional[tuple[int, dict[str, Any]]]: + """ + Reject a present protocol version that the era does not support. + + Absent header and envelope versions are allowed (legacy default). The + header is not required on ``modern`` in this sidecar. + """ + header_version = get_header(request_headers, HEADER_MCP_PROTOCOL_VERSION) + meta_version = envelope_protocol_version(request.meta) + version = request_protocol_version(header_version, meta_version) + if not is_unsupported_protocol_version(version, era): + return None + try: + validate_supported_protocol_version(version, supported_protocol_versions_for_era(era)) + except UnsupportedProtocolVersionError as exc: + return 400, exc.to_response(request.id) + return None + + +class StatelessIngressPipeline: + """ + Deterministic JSON-RPC pre-dispatch for stateless POST /mcp. + + Handles ping and deprecated-method rejection inline. + ``server/discover`` is forwarded to the HTTP engine handler when provided. + Task and tool methods always return None so ``TaskManager`` plus the + low-level MCP server remain the only production task path. + """ + + def __init__( + self, + context: IngressContext, + *, + task_handler: Optional[TaskDispatchHandler] = None, + registry_handler: Optional[RegistryDispatchHandler] = None, + discover_handler: Optional[DiscoverHandler] = None, + initialize_handler: Optional[InitializeHandler] = None, + ) -> None: + self._context = context + self._task_handler = task_handler + self._registry_handler = registry_handler + self._discover_handler = discover_handler + self._initialize_handler = initialize_handler + + def reject_tools_call_mcp_name( + self, + raw_body: bytes, + request_headers: dict[str, str], + ) -> Optional[tuple[int, dict[str, Any]]]: + """Replay-path ``Mcp-Name`` check for name-scoped methods.""" + try: + request = parse_jsonrpc_request(raw_body) + except (JsonRpcParseError, JsonRpcWireError): + return None + return reject_required_mcp_name( + request, request_headers, self._context.wire_mode + ) + + def reject_jsonrpc_mcp_method( + self, + raw_body: bytes, + request_headers: dict[str, str], + ) -> Optional[tuple[int, dict[str, Any]]]: + """Replay-path ``Mcp-Method`` check for JSON-RPC POST.""" + try: + request = parse_jsonrpc_request(raw_body) + except (JsonRpcParseError, JsonRpcWireError): + return None + return reject_required_mcp_method(request, request_headers, self._context.wire_mode) + + def reject_protocol_version_cross_check( + self, + raw_body: bytes, + request_headers: dict[str, str], + ) -> Optional[tuple[int, dict[str, Any]]]: + """Replay-path header vs envelope protocol version check.""" + try: + request = parse_jsonrpc_request(raw_body) + except (JsonRpcParseError, JsonRpcWireError): + return None + return reject_protocol_version_mismatch(request, request_headers) + + def reject_unsupported_protocol_version_header( + self, + raw_body: bytes, + request_headers: dict[str, str], + ) -> Optional[tuple[int, dict[str, Any]]]: + """Replay-path unsupported protocol version check.""" + try: + request = parse_jsonrpc_request(raw_body) + except (JsonRpcParseError, JsonRpcWireError): + return None + return reject_unsupported_protocol_version( + request, request_headers, self._context.resolved_era() + ) + + def reject_method_policy( + self, + raw_body: bytes, + request_headers: dict[str, str], + ) -> Optional[tuple[int, dict[str, Any]]]: + """Handshake and retired-method policy for POST, GET, and replay.""" + era = self._context.resolved_era() + try: + request = parse_jsonrpc_request(raw_body) + except (JsonRpcParseError, JsonRpcWireError): + header_method = get_header(request_headers, HEADER_MCP_METHOD) + if header_method is None: + return None + return reject_modern_method_policy(header_method.strip(), None, era) + return reject_modern_method_policy(request.method, request.id, era) + + def response_protocol_version(self) -> str: + """Advertised version used when the request does not name a supported one.""" + return self._context.protocol_version + + def response_supported_versions(self) -> frozenset[str]: + return supported_protocol_versions_for_era(self._context.resolved_era()) + + async def handle_post( + self, + raw_body: bytes, + request_headers: dict[str, str], + ) -> Optional[tuple[int, dict[str, Any]]]: + """ + Run ingress steps 2–5. Returns None to delegate to the underlying MCP app. + Step 1 (CORS) is handled by transport middleware. + """ + request_headers = prepare_sessionless_request_headers( + request_headers, self._context.wire_mode + ) + + if is_header_only_ping(raw_body, request_headers): + return 200, build_ping_response(None) + + try: + request = parse_jsonrpc_request(raw_body) + except JsonRpcParseError as exc: + return 400, exc.to_response(None) + except JsonRpcWireError as exc: + return 400, exc.to_response(None) + + rejected_handshake = reject_legacy_handshake( + request, self._context.resolved_era() + ) + if rejected_handshake is not None: + return rejected_handshake + + required_method = reject_required_mcp_method( + request, request_headers, self._context.wire_mode + ) + if required_method is not None: + return required_method + + required_name = reject_required_mcp_name( + request, request_headers, self._context.wire_mode + ) + if required_name is not None: + return required_name + + oversized = first_oversized_mcp_param(request_headers) + if oversized is not None: + return 400, InvalidRequestError( + "Mcp-Param header exceeds the maximum size" + ).to_response(request.id) + + header_name = get_header(request_headers, HEADER_MCP_NAME) + name_field = mcp_name_field(request.method) + body_name = request.params.get(name_field) if name_field else ( + request.params.get("name") or request.params.get("uri") + ) + if not mcp_name_is_required(request.method) and isinstance(body_name, str): + try: + validate_header_body_name(header_name, body_name) + except HeaderBodyMismatchError as exc: + return 400, exc.to_response(request.id) + + version_mismatch = reject_protocol_version_mismatch(request, request_headers) + if version_mismatch is not None: + return version_mismatch + + unsupported = reject_unsupported_protocol_version( + request, request_headers, self._context.resolved_era() + ) + if unsupported is not None: + return unsupported + + rejected = reject_legacy_wire(request, request_headers, self._context.resolved_era()) + if rejected is not None: + return rejected + + deprecated = reject_deprecated_method(request, self._context.resolved_era()) + if deprecated is not None: + return deprecated + + if request.method == PING_METHOD: + return 200, build_ping_response(request.id) + + if self._context.accepts_sessionless_initialize(): + if request.method == INITIALIZE_METHOD: + if self._initialize_handler is not None: + result = self._initialize_handler(request) + if isinstance(result, Awaitable): + result = await result + else: + requested = request.params.get("protocolVersion") + result = build_sessionless_initialize_result( + server_name=self._context.server_name, + server_version=self._context.server_version, + requested_version=requested if isinstance(requested, str) else None, + protocol_version=self._context.protocol_version, + advertise_tasks=self._context.advertise_tasks, + advertise_app=self._context.advertise_app, + custom_extensions=self._context.custom_extensions, + ) + return 200, jsonrpc_success(request.id, result) + if request.method == INITIALIZED_NOTIFICATION: + return 202, {} + + if request.method == SERVER_DISCOVER_METHOD: + if self._discover_handler is None: + return None + result = self._discover_handler(request) + if isinstance(result, Awaitable): + result = await result + return 200, jsonrpc_success(request.id, result) + + if is_task_wire_interception(request.method, request.params): + if self._task_handler is not None: + response = await self._task_handler(request) + if response is not None: + return 200, response + return None + + if self._registry_handler is not None: + response = await self._registry_handler(request) + if response is not None: + return 200, response + + return None + + @staticmethod + def serialize_response(jsonrpc_response: dict[str, Any]) -> bytes: + """Step 6 — JSON-RPC response serialization.""" + return json.dumps(jsonrpc_response).encode("utf-8") diff --git a/nitrostack/transports/dual.py b/nitrostack/transports/dual.py index 6e52f66..794c0c7 100644 --- a/nitrostack/transports/dual.py +++ b/nitrostack/transports/dual.py @@ -37,7 +37,7 @@ async def run_dual( mcp_app: "McpApplication", http_app: Any, *, - host: str = "0.0.0.0", + host: str = "127.0.0.1", port: int = 3000, graceful_timeout: float = 10.0, ) -> None: diff --git a/nitrostack/transports/headers.py b/nitrostack/transports/headers.py new file mode 100644 index 0000000..4f32766 --- /dev/null +++ b/nitrostack/transports/headers.py @@ -0,0 +1,261 @@ +"""MCP 2.0 HTTP header names and builders.""" + +from __future__ import annotations + +from collections.abc import Collection +from typing import Any, Mapping, Optional + +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION + +# Request headers +HEADER_CONTENT_TYPE = "Content-Type" +HEADER_MCP_PROTOCOL_VERSION = "MCP-Protocol-Version" +HEADER_MCP_METHOD = "Mcp-Method" +HEADER_MCP_NAME = "Mcp-Name" +HEADER_MCP_PARAM_PREFIX = "Mcp-Param-" +MAX_MCP_PARAM_VALUE_BYTES = 4096 +_MCP_PARAM_RESERVED = frozenset({"name", "method", "uri"}) +HEADER_AUTHORIZATION = "Authorization" +HEADER_LAST_EVENT_ID = "Last-Event-ID" + +# Response headers +HEADER_VARY = "Vary" + +# CORS headers +HEADER_ACCESS_CONTROL_ALLOW_ORIGIN = "Access-Control-Allow-Origin" +HEADER_ACCESS_CONTROL_ALLOW_METHODS = "Access-Control-Allow-Methods" +HEADER_ACCESS_CONTROL_ALLOW_HEADERS = "Access-Control-Allow-Headers" +HEADER_ACCESS_CONTROL_EXPOSE_HEADERS = "Access-Control-Expose-Headers" + +# SSE anti-buffering +HEADER_X_ACCEL_BUFFERING = "X-Accel-Buffering" +HEADER_CACHE_CONTROL = "Cache-Control" + +MCP_JSON_CONTENT_TYPE = "application/json" +MCP_SSE_CONTENT_TYPE = "text/event-stream" + +CORS_ALLOW_METHODS = "GET, POST, DELETE, OPTIONS" +CORS_ALLOW_HEADER_NAMES: tuple[str, ...] = ( + "Content-Type", + "Accept", + "Authorization", + "MCP-Protocol-Version", + "Mcp-Method", + "Mcp-Name", + "Last-Event-ID", + "Mcp-Session-Id", +) +CORS_ALLOW_HEADERS = ", ".join(CORS_ALLOW_HEADER_NAMES) +CORS_EXPOSE_HEADER_NAMES: tuple[str, ...] = ( + "MCP-Protocol-Version", + "Mcp-Method", + "Mcp-Name", +) +CORS_EXPOSE_HEADERS = ", ".join(CORS_EXPOSE_HEADER_NAMES) + +MCP_HTTP_PATH = "/mcp" +SSE_SUBSCRIPTIONS_PATH = "/subscriptions/listen" + + +def normalize_headers(headers: Mapping[str, str]) -> dict[str, str]: + """Case-insensitive header map keyed by original casing where possible.""" + return {k: v for k, v in headers.items()} + + +def get_header(headers: Mapping[str, str], name: str) -> Optional[str]: + """Fetch a header value case-insensitively.""" + target = name.lower() + for key, value in headers.items(): + if key.lower() == target: + return value + return None + + +def strip_legacy_session_headers(headers: dict[str, str]) -> dict[str, str]: + """Remove legacy Mcp-Session-Id from incoming or outgoing headers.""" + return { + key: value + for key, value in headers.items() + if key.lower() != LEGACY_SESSION_HEADER.lower() + } + + +def strip_legacy_session_headers_asgi( + headers: list[tuple[bytes, bytes]], +) -> list[tuple[bytes, bytes]]: + """Remove ``Mcp-Session-Id`` from an ASGI header list.""" + target = LEGACY_SESSION_HEADER.lower().encode("latin-1") + return [(key, value) for key, value in headers if key.lower() != target] + + +def scope_without_session_headers(scope: dict) -> dict: + """Copy an ASGI scope with incoming session headers removed.""" + copied = dict(scope) + copied["headers"] = strip_legacy_session_headers_asgi(list(scope.get("headers") or [])) + return copied + + +def decode_asgi_headers(headers: list[tuple[bytes, bytes]]) -> dict[str, str]: + """Decode an ASGI header list to a text map.""" + return {key.decode("latin-1"): value.decode("latin-1") for key, value in headers} + + +def snapshot_validated_asgi_headers( + headers: list[tuple[bytes, bytes]], +) -> tuple[tuple[bytes, bytes], ...]: + """Immutable header snapshot after ingress validators. Session ids are dropped.""" + return tuple(strip_legacy_session_headers_asgi(list(headers))) + + +def scope_with_header_snapshot( + scope: dict, + snapshot: tuple[tuple[bytes, bytes], ...] | list[tuple[bytes, bytes]], +) -> dict: + """Copy an ASGI scope pinned to a validated header snapshot. + + Replay and the inner engine read this snapshot, not a later forged header set. + """ + copied = dict(scope) + copied["headers"] = list(snapshot) + return copied + + +def handled_protocol_version( + request_headers: Mapping[str, str], + *, + fallback: str, + supported: Optional[Collection[str]] = None, +) -> str: + """Protocol version that handled the request. Unsupported client values are ignored.""" + header = get_header(request_headers, HEADER_MCP_PROTOCOL_VERSION) + if isinstance(header, str) and header.strip(): + value = header.strip() + if supported is None or value in supported: + return value + return fallback + + +def build_mcp_response_headers( + *, + content_type: str = MCP_JSON_CONTENT_TYPE, + protocol_version: str = MODERN_PROTOCOL_VERSION, + method: Optional[str] = None, + extra: Optional[Mapping[str, str]] = None, +) -> dict[str, str]: + """Standard MCP response headers. Echo values win over inner-app extras.""" + headers = { + HEADER_CONTENT_TYPE: content_type, + HEADER_VARY: "Origin", + } + if extra: + headers.update(extra) + headers[HEADER_MCP_PROTOCOL_VERSION] = protocol_version + if method: + headers[HEADER_MCP_METHOD] = method + return strip_legacy_session_headers(headers) + + +def build_mcp_echo_headers( + request_headers: Mapping[str, str], + *, + protocol_version: str, + method: Optional[str] = None, + supported_versions: Optional[Collection[str]] = None, + content_type: str = MCP_JSON_CONTENT_TYPE, + extra: Optional[Mapping[str, str]] = None, +) -> dict[str, str]: + """Echo the handled protocol version and method on every ``/mcp`` response.""" + echo_method = method or get_header(request_headers, HEADER_MCP_METHOD) + return build_mcp_response_headers( + content_type=content_type, + protocol_version=handled_protocol_version( + request_headers, fallback=protocol_version, supported=supported_versions + ), + method=echo_method, + extra=extra, + ) + + +def build_sse_stream_headers( + protocol_version: str = MODERN_PROTOCOL_VERSION, +) -> dict[str, str]: + """SSE stream headers including proxy buffering guards.""" + return build_mcp_response_headers( + content_type=MCP_SSE_CONTENT_TYPE, + protocol_version=protocol_version, + extra={ + HEADER_X_ACCEL_BUFFERING: "no", + HEADER_CACHE_CONTROL: "no-transform", + }, + ) + + +def extract_mcp_param_headers(headers: Mapping[str, str]) -> dict[str, str]: + """Extract ``Mcp-Param-*`` values keyed by the header suffix.""" + params: dict[str, str] = {} + prefix = HEADER_MCP_PARAM_PREFIX.lower() + for key, value in headers.items(): + lower = key.lower() + if not lower.startswith(prefix): + continue + param_name = key[len(prefix) :] + if param_name: + params[param_name] = value + return params + + +def first_oversized_mcp_param(headers: Mapping[str, str]) -> Optional[str]: + """Return the first ``Mcp-Param-*`` suffix whose value exceeds the size cap.""" + for name, value in extract_mcp_param_headers(headers).items(): + if len(str(value).encode("utf-8")) > MAX_MCP_PARAM_VALUE_BYTES: + return name + return None + + +def merge_mcp_param_headers( + arguments: Mapping[str, Any], + param_headers: Mapping[str, str], + *, + allowed_fields: Optional[set[str]] = None, + reserved: frozenset[str] = _MCP_PARAM_RESERVED, +) -> dict[str, Any]: + """ + Copy ``arguments`` and fill missing keys from ``Mcp-Param-*``. + + Existing JSON-RPC values win. ``name`` / ``method`` / ``uri`` are never + taken from headers. When ``allowed_fields`` is set, unknown suffixes are + ignored (schema-declared params only). + """ + merged = dict(arguments) + reserved_lower = {item.lower() for item in reserved} + field_map = ( + {field.lower(): field for field in allowed_fields} + if allowed_fields is not None + else None + ) + for raw_name, value in param_headers.items(): + if not raw_name or raw_name.lower() in reserved_lower: + continue + dest = field_map.get(raw_name.lower()) if field_map is not None else raw_name + if dest is None: + continue + current = merged.get(dest) + if dest in merged and current is not None and current != "": + continue + merged[dest] = value + return merged + + +def extract_mcp_scope_headers(headers: Mapping[str, str]) -> dict[str, str]: + """Copy MCP scoped request headers. Authorization is not included.""" + scoped: dict[str, str] = {} + for name in (HEADER_MCP_PROTOCOL_VERSION, HEADER_MCP_METHOD, HEADER_MCP_NAME): + value = get_header(headers, name) + if value: + scoped[name] = value + prefix = HEADER_MCP_PARAM_PREFIX.lower() + for key, value in headers.items(): + if key.lower().startswith(prefix) and key.lower() != prefix: + scoped[key] = value + return scoped diff --git a/nitrostack/transports/http.py b/nitrostack/transports/http.py index ad73564..a343747 100644 --- a/nitrostack/transports/http.py +++ b/nitrostack/transports/http.py @@ -46,7 +46,23 @@ from mcp.server.sse import SseServerTransport from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from mcp.server.transport_security import TransportSecuritySettings -from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS + +from nitrostack.protocol.version import ( + HttpEngine, + ProtocolEra, + WireMode, + protocol_version_for_era, + resolve_http_engine, +) +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER +from nitrostack.transports.headers import ( + CORS_ALLOW_HEADER_NAMES, + CORS_EXPOSE_HEADER_NAMES, + SSE_SUBSCRIPTIONS_PATH, + strip_legacy_session_headers_asgi, +) +from nitrostack.transports.subscriptions import subscriptions_listen_endpoint +from nitrostack.transports.proxy import public_url_for_request if TYPE_CHECKING: from nitrostack.core.app import McpApplication @@ -55,16 +71,8 @@ DEFAULT_ENDPOINT = "/mcp" -# CORS headers for browser-based MCP clients (Inspector). -CORS_ALLOW_HEADERS = [ - "Content-Type", - "Accept", - "Authorization", - "Mcp-Session-Id", - "MCP-Protocol-Version", - "Last-Event-ID", -] -CORS_EXPOSE_HEADERS = ["Mcp-Session-Id"] +# Shared 2026 allow-headers. Session id is exposed only on the sessionful engine. +CORS_ALLOW_HEADERS = list(CORS_ALLOW_HEADER_NAMES) # The Accept value `StreamableHTTPServerTransport` requires: it needs # `application/json` on POST and `text/event-stream` on both POST and GET. @@ -81,11 +89,16 @@ def _server_meta(mcp_app: "McpApplication") -> Dict[str, str]: } -def _landing_html(name: str, version: str, endpoint: str) -> str: +def _landing_html( + name: str, + version: str, + endpoint: str, + public_mcp_url: Optional[str] = None, +) -> str: safe_name = html.escape(name) safe_version = html.escape(version) - mcp_path = html.escape(endpoint.rstrip("/") or "/mcp") - health_path = f"{mcp_path}/health" + mcp_path = html.escape(public_mcp_url or (endpoint.rstrip("/") or "/mcp")) + health_path = html.escape(f"{(endpoint.rstrip('/') or '/mcp')}/health") return f""" @@ -289,25 +302,21 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: class HeaderCompatMiddleware: """ - Normalize `Accept` and `MCP-Protocol-Version` on the Streamable HTTP mount - so tolerable client quirks don't turn into a failed connection. - - `StreamableHTTPServerTransport` matches Accept media types with - `str.startswith`, so it does not honour wildcards: a client sending - `Accept: */*` (or no Accept at all, which RFC 9110 also defines as - accepting anything) is rejected with `406` even though it accepts - everything the transport can send. It also rejects a request with `400` when - `MCP-Protocol-Version` names a version it doesn't know, which breaks a - client that advertises a spec release newer than the installed `mcp` SDK - even though the session itself negotiated a version both sides support. - - Both rejections happen before the JSON-RPC layer, so the client sees a - stream that opens and closes with no response on it and no explanation. - Requests that already satisfy the transport pass through untouched. + Normalize ``Accept`` on the Streamable HTTP mount so wildcard or absent + Accept does not 406. Session ids are stripped only when + ``drop_session_headers`` is set. + + ``MCP-Protocol-Version`` is never deleted. Duplicate casings are collapsed + to one ``mcp-protocol-version`` entry and the client value is kept so later + version checks see what the client sent. + + Stack order on the HTTP app: CORS → this middleware (preserve) → handler. + Sidecar version checks run on the combined app outside this mount. """ - def __init__(self, app: ASGIApp) -> None: + def __init__(self, app: ASGIApp, *, drop_session_headers: bool = False) -> None: self.app = app + self.drop_session_headers = drop_session_headers @staticmethod def _normalize_accept(value: Optional[str]) -> Optional[str]: @@ -325,44 +334,40 @@ def _normalize_accept(value: Optional[str]) -> Optional[str]: return MCP_ACCEPT return None + @staticmethod + def _canonicalize_protocol_version(headers: List[Any]) -> List[Any]: + """Keep the protocol version value; emit one lowercase header name.""" + version_values: List[bytes] = [] + kept: List[Any] = [] + for key, value in headers: + if key.lower() == b"mcp-protocol-version": + version_values.append(value) + else: + kept.append((key, value)) + if not version_values: + return headers + kept.append((b"mcp-protocol-version", version_values[0])) + return kept + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": await self.app(scope, receive, send) return headers: List[Any] = list(scope.get("headers") or []) + if self.drop_session_headers: + headers = strip_legacy_session_headers_asgi(headers) + headers = self._canonicalize_protocol_version(headers) + raw_accept = next((value for key, value in headers if key.lower() == b"accept"), None) accept = self._normalize_accept(raw_accept.decode("latin-1") if raw_accept is not None else None) - - raw_version = next( - (value for key, value in headers if key.lower() == b"mcp-protocol-version"), - None, - ) - drop_version = raw_version is not None and raw_version.decode("latin-1") not in SUPPORTED_PROTOCOL_VERSIONS - - if accept is None and not drop_version: - await self.app(scope, receive, send) - return - - rewritten = [ - (key, value) - for key, value in headers - if not (key.lower() == b"accept" and accept is not None) - and not (key.lower() == b"mcp-protocol-version" and drop_version) - ] if accept is not None: - rewritten.append((b"accept", accept.encode("latin-1"))) + headers = [(key, value) for key, value in headers if key.lower() != b"accept"] + headers.append((b"accept", accept.encode("latin-1"))) logger.debug("Rewrote Accept %r -> %r for %s", raw_accept, accept, scope.get("path")) - if drop_version: - logger.debug( - "Dropped unsupported MCP-Protocol-Version %r for %s (supported: %s)", - raw_version, - scope.get("path"), - ", ".join(SUPPORTED_PROTOCOL_VERSIONS), - ) scope = dict(scope) - scope["headers"] = rewritten + scope["headers"] = headers await self.app(scope, receive, send) @@ -469,6 +474,9 @@ def build_http_app( enable_cors: bool = True, stateless: bool = False, json_response: bool = False, + protocol_era: ProtocolEra = "auto", + wire_mode: WireMode = "stateless", + http_engine: Optional[HttpEngine] = None, ) -> Starlette: """ Build the Starlette app exposing NitroStack's owned low-level server over @@ -497,7 +505,24 @@ def build_http_app( that doesn't implement SSE parsing for POST responses sees the SSE form as a stream that ended without a result. Server-initiated streaming (progress, notifications) is unavailable in this mode. + protocol_era: ``legacy`` / ``modern`` / ``auto``. Distinct from ``stateless``. + wire_mode: Dual-spec policy for 2025 traffic (``sessionful``, ``stateless`` + fallback, or ``reject``). ``auto`` uses ``stateless``; ``modern`` uses + ``reject``. + http_engine: Era factory result. ``sessionless`` for modern/auto, + ``sessionful`` only for legacy. ``auto`` / ``modern`` never start + a sessionful manager, even if ``stateless=False``. """ + http_engine = resolve_http_engine( + protocol_era, http_engine=http_engine, stateless=stateless + ) + stateless = http_engine == "sessionless" + + server = mcp_app.mcp_server + if server is not None: + server.http_engine = http_engine + server.sessionful = http_engine == "sessionful" + security_settings = None if not enable_cors: allowed_hosts = _env_list("MCP_ALLOWED_HOSTS") or ["localhost:*", "127.0.0.1:*"] @@ -507,7 +532,11 @@ def build_http_app( allowed_hosts=allowed_hosts, allowed_origins=allowed_origins, ) + else: + security_settings = TransportSecuritySettings(enable_dns_rebinding_protection=False) + # Official mcp 2.x owns /mcp (streamable HTTP, both protocol eras). + # One manager only; auto/modern are sessionless and legacy is sessionful. session_manager = StreamableHTTPSessionManager( app=mcp_app.mcp_server, stateless=stateless, @@ -528,7 +557,9 @@ async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) -> # Wraps only the Streamable HTTP mount, so `/mcp/health` and the legacy SSE # routes keep their own (correct) content negotiation. - mcp_asgi_app: ASGIApp = HeaderCompatMiddleware(handle_streamable_http) + mcp_asgi_app: ASGIApp = HeaderCompatMiddleware( + handle_streamable_http, drop_session_headers=stateless + ) session_cap: Optional[SessionCapMiddleware] = None if max_sessions and not stateless: session_cap = SessionCapMiddleware( @@ -541,22 +572,40 @@ async def handle_sse(request): await mcp_app.mcp_server.run(streams[0], streams[1], mcp_app.mcp_server.create_initialization_options()) return Response() + def _request_public_mcp_url(request) -> str: + port = os.environ.get("PORT") or os.environ.get("MCP_SERVER_PORT") or "3000" + return public_url_for_request( + request, + path=endpoint.rstrip("/") or "/mcp", + fallback_host=f"localhost:{port}", + ) + async def health_check(request): return JSONResponse( { "status": "ok", "transport": "streamable-http", - "protocolVersion": "2025-06-18", + "protocolVersion": protocol_version_for_era(protocol_era), + "protocolEra": protocol_era, + "statelessCapable": http_engine == "sessionless", "stateless": stateless, "jsonResponse": json_response, "sessions": session_cap.active_session_count if session_cap else None, "uptimeSeconds": round(time.monotonic() - _PROCESS_START, 2), + "publicUrl": _request_public_mcp_url(request), } ) async def root_page(request): meta = _server_meta(mcp_app) - return HTMLResponse(_landing_html(meta["name"], meta["version"], endpoint)) + return HTMLResponse( + _landing_html( + meta["name"], + meta["version"], + endpoint, + public_mcp_url=_request_public_mcp_url(request), + ) + ) async def oauth_not_supported(request): """Inspector DCR posts `/register` when Authentication is on. @@ -564,13 +613,14 @@ async def oauth_not_supported(request): Return JSON (not the HTML 404 page) so the client shows a clear OAuth-off message instead of `Unexpected token '<'`. """ + connect_url = _request_public_mcp_url(request) return JSONResponse( { "error": "invalid_request", "error_description": ( "This MCP server does not use OAuth. In MCP Inspector turn " "Authentication off, then connect with Streamable HTTP to " - f"http://localhost:{os.environ.get('PORT') or os.environ.get('MCP_SERVER_PORT') or '3000'}{endpoint}." + f"{connect_url}." ), }, status_code=404, @@ -610,7 +660,9 @@ async def widgets_preview_call(request): except Exception as exc: logger.exception("Widget preview tool call failed for %s", name) return JSONResponse({"error": str(exc)}, status_code=400) - structured = getattr(result, "structuredContent", None) + structured = getattr(result, "structured_content", None) + if structured is None: + structured = getattr(result, "structuredContent", None) try: html = entry.component.html_with_data(structured) except Exception as exc: @@ -630,7 +682,7 @@ async def widgets_preview_call(request): "structuredContent": structured, "html": html, "resourceUri": entry.component.resource_uri, - "isError": bool(getattr(result, "isError", False)), + "isError": bool(getattr(result, "is_error", None) if getattr(result, "is_error", None) is not None else getattr(result, "isError", False)), } ) @@ -645,6 +697,7 @@ async def json_version(request): "User-Agent": f"NitroStack/{meta['version']}", "webSocketDebuggerUrl": "", "transport": "mcp", + "publicUrl": _request_public_mcp_url(request), "endpoints": { "mcp": endpoint, "sse": "/sse", @@ -704,6 +757,14 @@ async def lifespan(app): Mount(endpoint, app=mcp_asgi_app), Route("/sse", endpoint=handle_sse, methods=["GET"]), ] + if protocol_era != "legacy": + routes.append( + Route( + SSE_SUBSCRIPTIONS_PATH, + endpoint=subscriptions_listen_endpoint(mcp_app), + methods=["GET", "POST"], + ) + ) # Rewrite `/mcp` → `/mcp/` *before* routing so Inspector never sees a 307. # CORS stays outermost so preflight still works on the original path. @@ -711,6 +772,9 @@ async def lifespan(app): Middleware(ExactEndpointSlashMiddleware, endpoint=endpoint), ] if enable_cors: + expose_headers = list(CORS_EXPOSE_HEADER_NAMES) + if http_engine == "sessionful": + expose_headers.append(LEGACY_SESSION_HEADER) middleware.insert( 0, Middleware( @@ -718,7 +782,7 @@ async def lifespan(app): allow_origins=["*"], allow_methods=["GET", "POST", "DELETE", "OPTIONS"], allow_headers=CORS_ALLOW_HEADERS, - expose_headers=CORS_EXPOSE_HEADERS, + expose_headers=expose_headers, ), ) @@ -726,4 +790,14 @@ async def lifespan(app): _enable_trace_logging() middleware.insert(0, Middleware(RequestTraceMiddleware)) - return Starlette(routes=routes, middleware=middleware, lifespan=lifespan) + app = Starlette(routes=routes, middleware=middleware, lifespan=lifespan) + app.state.protocol_era = protocol_era + app.state.wire_mode = wire_mode + app.state.stateless = stateless + app.state.http_engine = http_engine + app.state.sessionful = http_engine == "sessionful" + app.state.session_manager = session_manager + app.state.streamable_http_manager_count = 1 + app.state.subscription_bus = getattr(mcp_app.mcp_server, "subscription_bus", None) + app.state.protocol_version = protocol_version_for_era(protocol_era) + return app diff --git a/nitrostack/transports/middleware.py b/nitrostack/transports/middleware.py new file mode 100644 index 0000000..314f097 --- /dev/null +++ b/nitrostack/transports/middleware.py @@ -0,0 +1,524 @@ +"""Stateless HTTP ASGI middleware.""" + +from __future__ import annotations + +from typing import Any, Callable, Optional + +from nitrostack.protocol.errors import JsonRpcErrorCode +from nitrostack.protocol.jsonrpc import jsonrpc_error, jsonrpc_method_from_body +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION, ProtocolEra, WireMode +from nitrostack.runtime.stateless import assert_stateless_headers +from nitrostack.transports.cors import build_cors_headers, cors_preflight_response_headers +from nitrostack.transports.dispatch import ( + DiscoverHandler, + IngressContext, + InitializeHandler, + StatelessIngressPipeline, +) +from nitrostack.transports.headers import ( + MCP_HTTP_PATH, + build_mcp_echo_headers, + decode_asgi_headers, + get_header, + scope_with_header_snapshot, + scope_without_session_headers, + snapshot_validated_asgi_headers, + strip_legacy_session_headers, +) + +ASGIApp = Callable[..., Any] + +MCP_POST_PATHS = (MCP_HTTP_PATH, f"{MCP_HTTP_PATH}/") + + +class StatelessTransportMiddleware: + """ + ASGI wrapper implementing stateless HTTP transport invariants: + - OPTIONS 204 CORS preflight on MCP paths only + - Legacy session header stripping + - MCP response headers on all responses + - Pre-dispatch for POST /mcp (ping, server/discover) + """ + + def __init__( + self, + app: ASGIApp, + *, + pipeline: Optional[StatelessIngressPipeline] = None, + mcp_paths: tuple[str, ...] = MCP_POST_PATHS, + enable_cors: bool = True, + ) -> None: + self.app = app + self.pipeline = pipeline + self.mcp_paths = mcp_paths + self.enable_cors = enable_cors + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + if scope.get("type") != "http": + await self.app(scope, receive, send) + return + + method = scope.get("method", "GET").upper() + path = scope.get("path", "") + + if method == "OPTIONS" and path in self.mcp_paths and self.enable_cors: + await self._send_options(scope, receive, send) + return + + buffered_body: Optional[bytes] = None + header_snapshot: Optional[tuple[tuple[bytes, bytes], ...]] = None + if method == "POST" and path in self.mcp_paths and self.pipeline is not None: + scope = self._strip_sessionless_scope(scope) + buffered_body = await self._read_body(receive) + handled = await self._try_pre_dispatch(scope, buffered_body, send) + if handled: + return + live_headers = decode_asgi_headers(list(scope.get("headers") or [])) + rejected = self._reject_replay_headers(buffered_body, live_headers) + if rejected is not None: + await self._send_pipeline_response( + scope, send, live_headers, rejected, body=buffered_body + ) + return + header_snapshot = snapshot_validated_asgi_headers(list(scope.get("headers") or [])) + snapshot_headers = decode_asgi_headers(list(header_snapshot)) + rejected = self._reject_replay_headers(buffered_body, snapshot_headers) + if rejected is not None: + await self._send_pipeline_response( + scope, send, snapshot_headers, rejected, body=buffered_body + ) + return + receive = self._replay_receive(buffered_body, receive) + scope = scope_with_header_snapshot(scope, header_snapshot) + + if path in self.mcp_paths and self.pipeline is not None: + scope = self._strip_sessionless_scope(scope) + raw_headers = decode_asgi_headers(list(scope.get("headers") or [])) + if method == "GET": + rejected = self.pipeline.reject_method_policy(b"", raw_headers) + if rejected is not None: + await self._send_pipeline_response( + scope, send, raw_headers, rejected, body=b"" + ) + return + + await self._forward_with_stateless_headers( + scope, receive, send, body=buffered_body + ) + + def _reject_replay_headers( + self, + raw_body: bytes, + request_headers: dict[str, str], + ) -> Optional[tuple[int, dict[str, Any]]]: + """Re-apply method policy, ``-32020``, and ``-32022`` on replay.""" + assert self.pipeline is not None + policy_rejected = self.pipeline.reject_method_policy(raw_body, request_headers) + if policy_rejected is not None: + return policy_rejected + method_rejected = self.pipeline.reject_jsonrpc_mcp_method(raw_body, request_headers) + if method_rejected is not None: + return method_rejected + name_rejected = self.pipeline.reject_tools_call_mcp_name(raw_body, request_headers) + if name_rejected is not None: + return name_rejected + version_rejected = self.pipeline.reject_protocol_version_cross_check( + raw_body, request_headers + ) + if version_rejected is not None: + return version_rejected + return self.pipeline.reject_unsupported_protocol_version_header( + raw_body, request_headers + ) + + async def _send_options(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + headers_list = scope.get("headers") or [] + req_headers = { + k.decode("latin-1"): v.decode("latin-1") for k, v in headers_list + } + cors = cors_preflight_response_headers(req_headers) + response_headers = self._echo_headers(req_headers, extra=cors) + assert_stateless_headers(response_headers) + + await send( + { + "type": "http.response.start", + "status": 204, + "headers": self._encode_headers(response_headers), + } + ) + await send({"type": "http.response.body", "body": b""}) + + async def _try_pre_dispatch( + self, + scope: dict[str, Any], + body: bytes, + send: Any, + ) -> bool: + headers_list = scope.get("headers") or [] + raw_headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in headers_list} + + assert self.pipeline is not None + result = await self.pipeline.handle_post(body, raw_headers) + if result is None: + return False + + await self._send_pipeline_response(scope, send, raw_headers, result, body=body) + return True + + async def _send_pipeline_response( + self, + scope: dict[str, Any], + send: Any, + raw_headers: dict[str, str], + result: tuple[int, dict[str, Any]], + *, + body: Optional[bytes] = None, + ) -> None: + status, jsonrpc_response = result + origin = get_header(raw_headers, "Origin") + cors = build_cors_headers(origin=origin) + response_headers = self._echo_headers(raw_headers, extra=cors, body=body) + assert_stateless_headers(response_headers) + assert self.pipeline is not None + if status == 202 and jsonrpc_response == {}: + payload = b"" + else: + payload = self.pipeline.serialize_response(jsonrpc_response) + await send( + { + "type": "http.response.start", + "status": status, + "headers": self._encode_headers(response_headers), + } + ) + await send({"type": "http.response.body", "body": payload}) + + async def _send_session_id_rejected( + self, + scope: dict[str, Any], + send: Any, + raw_headers: dict[str, str], + ) -> None: + origin = get_header(raw_headers, "Origin") + cors = build_cors_headers(origin=origin) + response_headers = self._echo_headers(raw_headers, extra=cors) + assert_stateless_headers(response_headers) + payload = StatelessIngressPipeline.serialize_response( + jsonrpc_error( + None, + int(JsonRpcErrorCode.INVALID_REQUEST), + "Invalid Request: Mcp-Session-Id is not supported", + ) + ) + await send( + { + "type": "http.response.start", + "status": 400, + "headers": self._encode_headers(response_headers), + } + ) + await send({"type": "http.response.body", "body": payload}) + + async def _forward_with_stateless_headers( + self, + scope: dict[str, Any], + receive: Any, + send: Any, + *, + body: Optional[bytes] = None, + ) -> None: + async def send_wrapper(message: dict[str, Any]) -> None: + if message["type"] == "http.response.start": + raw_headers = { + k.decode("latin-1"): v.decode("latin-1") + for k, v in message.get("headers", []) + } + req_headers = { + k.decode("latin-1"): v.decode("latin-1") + for k, v in (scope.get("headers") or []) + } + inner_headers = { + k: v + for k, v in strip_legacy_session_headers(raw_headers).items() + if k.lower() != "content-type" + } + merged = self._echo_headers( + req_headers, + content_type=raw_headers.get("content-type", "application/json"), + extra={ + **inner_headers, + **build_cors_headers(origin=get_header(req_headers, "Origin")), + }, + body=body, + ) + assert_stateless_headers(merged) + message = { + **message, + "headers": self._encode_headers(merged), + } + await send(message) + + await self.app( + scope_with_header_snapshot( + scope, + snapshot_validated_asgi_headers(list(scope.get("headers") or [])), + ), + receive, + send_wrapper, + ) + + def _strip_sessionless_scope(self, scope: dict[str, Any]) -> dict[str, Any]: + if self.pipeline is None: + return scope + from nitrostack.runtime.stateless import sessionless_strips_incoming_session_id + + if sessionless_strips_incoming_session_id(self.pipeline._context.wire_mode): + return scope_without_session_headers(scope) + return scope + + def _echo_headers( + self, + request_headers: dict[str, str], + *, + extra: Optional[dict[str, str]] = None, + content_type: str = "application/json", + body: Optional[bytes] = None, + ) -> dict[str, str]: + fallback = MODERN_PROTOCOL_VERSION + supported = None + if self.pipeline is not None: + fallback = self.pipeline.response_protocol_version() + supported = self.pipeline.response_supported_versions() + return build_mcp_echo_headers( + request_headers, + protocol_version=fallback, + method=jsonrpc_method_from_body(body), + supported_versions=supported, + content_type=content_type, + extra=extra, + ) + + @staticmethod + async def _read_body(receive: Any) -> bytes: + body = b"" + while True: + message = await receive() + if message["type"] == "http.request": + body += message.get("body", b"") + if not message.get("more_body", False): + break + return body + + @staticmethod + def _replay_receive(body: bytes, original_receive: Any) -> Any: + sent = False + + async def replay() -> dict[str, Any]: + nonlocal sent + if not sent: + sent = True + return {"type": "http.request", "body": body, "more_body": False} + # Body was already buffered. Wait for a real client disconnect + # instead of synthesizing one — Streamable HTTP treats disconnect + # as an abort of the in-flight request. + while True: + message = await original_receive() + if message.get("type") == "http.disconnect": + return message + + return replay + + @staticmethod + def _encode_headers(headers: dict[str, str]) -> list[tuple[bytes, bytes]]: + return [(k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers.items()] + + +class SessionlessHttpGuard: + """ + Transport invariants official mcp 2.x does not own on sessionless ``/mcp``. + + Does not parse or answer ``tools/call``. Forwards those to the v2 app. + """ + + def __init__( + self, + app: ASGIApp, + *, + protocol_era: ProtocolEra = "auto", + enable_cors: bool = True, + discover_handler: Optional[DiscoverHandler] = None, + ) -> None: + self.app = app + self.protocol_era = protocol_era + self.enable_cors = enable_cors + self.discover_handler = discover_handler + wire_mode: WireMode = "reject" if protocol_era == "modern" else "stateless" + self._sender = StatelessTransportMiddleware( + app, + pipeline=StatelessIngressPipeline( + IngressContext( + server_name="", + server_version="", + protocol_version=MODERN_PROTOCOL_VERSION, + wire_mode=wire_mode, + protocol_era=protocol_era, + ), + discover_handler=discover_handler, + ), + enable_cors=enable_cors, + ) + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + if scope.get("type") != "http": + await self.app(scope, receive, send) + return + method = scope.get("method", "").upper() + path = scope.get("path", "") + if method == "OPTIONS" and path in MCP_POST_PATHS and self.enable_cors: + await self._sender._send_options(scope, receive, send) + return + scope = scope_without_session_headers(scope) + raw_headers = decode_asgi_headers(list(scope.get("headers") or [])) + from nitrostack.transports.dispatch import ( + is_header_only_ping, + reject_modern_method_policy, + ) + from nitrostack.transports.headers import HEADER_MCP_METHOD, get_header + if path in MCP_POST_PATHS and method == "GET": + header_method = get_header(raw_headers, HEADER_MCP_METHOD) + if header_method is not None: + rejected = reject_modern_method_policy( + header_method.strip(), None, self.protocol_era + ) + if rejected is not None: + await self._sender._send_pipeline_response( + scope, send, raw_headers, rejected, body=b"" + ) + return + if method != "POST" or path not in MCP_POST_PATHS: + await self.app(scope, receive, send) + return + + body = await StatelessTransportMiddleware._read_body(receive) + if is_header_only_ping(body, raw_headers): + from nitrostack.protocol.jsonrpc import build_ping_response + + await self._sender._send_pipeline_response( + scope, send, raw_headers, (200, build_ping_response(None)), body=body + ) + return + + from nitrostack.protocol.jsonrpc import ( + JsonRpcParseError, + JsonRpcWireError, + parse_jsonrpc_request, + validate_header_body_method, + ) + try: + request = parse_jsonrpc_request(body) + except (JsonRpcParseError, JsonRpcWireError): + await self.app( + scope, + StatelessTransportMiddleware._replay_receive(body, receive), + send, + ) + return + + rejected = reject_modern_method_policy( + request.method, request.id, self.protocol_era + ) + if rejected is not None: + await self._sender._send_pipeline_response( + scope, send, raw_headers, rejected, body=body + ) + return + + header_method = get_header(raw_headers, HEADER_MCP_METHOD) + if header_method is not None: + try: + validate_header_body_method(header_method, request.method) + except Exception as exc: + from nitrostack.protocol.jsonrpc import map_exception_to_jsonrpc + + await self._sender._send_pipeline_response( + scope, + send, + raw_headers, + (400, map_exception_to_jsonrpc(exc, request.id)), + body=body, + ) + return + + if request.method == "server/discover" and self.discover_handler is not None: + from nitrostack.protocol.jsonrpc import jsonrpc_success + + await self._sender._send_pipeline_response( + scope, + send, + raw_headers, + (200, jsonrpc_success(request.id, self.discover_handler())), + body=body, + ) + return + + await self._sender._forward_with_stateless_headers( + scope, + StatelessTransportMiddleware._replay_receive(body, receive), + send, + body=body, + ) + + +def wrap_sessionless_http( + app: ASGIApp, + *, + protocol_era: ProtocolEra = "auto", + enable_cors: bool = True, + discover_handler: Optional[DiscoverHandler] = None, +) -> ASGIApp: + """Sessionless transport guard; official mcp 2.x still owns ``tools/call``.""" + return SessionlessHttpGuard( + app, + protocol_era=protocol_era, + enable_cors=enable_cors, + discover_handler=discover_handler, + ) + + +def wrap_modern_handshake_reject(app: ASGIApp, *, enable_cors: bool = True) -> ASGIApp: + """Reject 2025 ``initialize`` on era ``modern``; leave the v2 app otherwise.""" + return wrap_sessionless_http(app, protocol_era="modern", enable_cors=enable_cors) + + +def wrap_stateless_transport( + app: ASGIApp, + *, + server_name: str, + server_version: str, + protocol_version: str, + advertise_tasks: bool = True, + advertise_app: bool = False, + custom_extensions: Optional[dict[str, str]] = None, + wire_mode: WireMode = "stateless", + protocol_era: Optional[ProtocolEra] = None, + enable_cors: bool = True, + discover_handler: Optional[DiscoverHandler] = None, + initialize_handler: Optional[InitializeHandler] = None, +) -> ASGIApp: + """Wrap an ASGI app with stateless HTTP middleware.""" + pipeline = StatelessIngressPipeline( + IngressContext( + server_name=server_name, + server_version=server_version, + protocol_version=protocol_version, + advertise_tasks=advertise_tasks, + advertise_app=advertise_app, + custom_extensions=custom_extensions, + wire_mode=wire_mode, + protocol_era=protocol_era, + ), + discover_handler=discover_handler, + initialize_handler=initialize_handler, + ) + return StatelessTransportMiddleware(app, pipeline=pipeline, enable_cors=enable_cors) diff --git a/nitrostack/transports/proxy.py b/nitrostack/transports/proxy.py new file mode 100644 index 0000000..e356bb9 --- /dev/null +++ b/nitrostack/transports/proxy.py @@ -0,0 +1,222 @@ +"""Trusted reverse-proxy handling for forwarded host and proto. + +``X-Forwarded-*`` is ignored unless the direct socket peer is on an explicit +allow-list (``TRUSTED_PROXIES`` / ``MCP_TRUSTED_PROXIES``). ``X-Forwarded-For`` +is never used to decide trust. +""" + +from __future__ import annotations + +import ipaddress +import os +from collections.abc import Mapping, Sequence +from typing import Any, Optional +from urllib.parse import urlparse + +from nitrostack.transports.headers import get_header + +TRUSTED_PROXIES_ENV = "TRUSTED_PROXIES" +MCP_TRUSTED_PROXIES_ENV = "MCP_TRUSTED_PROXIES" + + +def configured_trusted_proxies() -> tuple[str, ...]: + """Comma-separated IPs, CIDRs, or hostnames. Empty means trust nobody.""" + raw = os.environ.get(TRUSTED_PROXIES_ENV) or os.environ.get(MCP_TRUSTED_PROXIES_ENV) or "" + return tuple(item.strip() for item in raw.split(",") if item.strip()) + + +def _parse_peer_ip(peer: str) -> Optional[ipaddress.IPv4Address | ipaddress.IPv6Address]: + host = peer.strip() + if host.startswith("["): + end = host.find("]") + host = host[1:end] if end != -1 else host.lstrip("[") + elif host.count(":") == 1: + candidate, _, tail = host.rpartition(":") + if tail.isdigit(): + host = candidate + try: + return ipaddress.ip_address(host) + except ValueError: + return None + + +def peer_is_trusted( + peer: Optional[str], + trusted: Optional[Sequence[str]] = None, +) -> bool: + """True when the direct socket peer is on the configured allow-list.""" + allow = tuple(trusted) if trusted is not None else configured_trusted_proxies() + if not allow or not peer: + return False + ip = _parse_peer_ip(peer) + peer_key = peer.strip().lower() + for item in allow: + token = item.strip() + if not token: + continue + if ip is not None: + try: + if ip in ipaddress.ip_network(token, strict=False): + return True + continue + except ValueError: + pass + if token.lower() == peer_key or token.lower() == (ip and str(ip)): + return True + return False + + +def first_forwarded_value(headers: Mapping[str, str], name: str) -> Optional[str]: + """Left-most value of a comma-separated forwarded header.""" + raw = get_header(headers, name) + if not raw: + return None + value = raw.split(",")[0].strip() + return value or None + + +def trusted_forwarded_host( + headers: Mapping[str, str], + peer: Optional[str], + trusted: Optional[Sequence[str]] = None, +) -> Optional[str]: + """``X-Forwarded-Host`` when the peer is trusted; otherwise ignored.""" + if not peer_is_trusted(peer, trusted): + return None + return first_forwarded_value(headers, "X-Forwarded-Host") + + +def trusted_forwarded_proto( + headers: Mapping[str, str], + peer: Optional[str], + trusted: Optional[Sequence[str]] = None, +) -> Optional[str]: + """``X-Forwarded-Proto`` when the peer is trusted; otherwise ignored.""" + if not peer_is_trusted(peer, trusted): + return None + proto = first_forwarded_value(headers, "X-Forwarded-Proto") + if not proto: + return None + proto = proto.split(";")[0].strip().lower() + if proto in {"http", "https"}: + return proto + return None + + +def hostname_from_host_header(value: str) -> str: + """Hostname from a ``Host`` / ``X-Forwarded-Host`` value (port stripped).""" + host = value.strip() + if not host: + return "" + if host.startswith("["): + end = host.find("]") + return host[1:end].lower() if end != -1 else host.lower() + if host.count(":") == 1: + name, _, tail = host.rpartition(":") + if tail.isdigit(): + return name.lower() + return host.lower() + + +def request_host_for_cimd( + headers: Mapping[str, str], + peer: Optional[str], + trusted: Optional[Sequence[str]] = None, +) -> Optional[str]: + """Inbound host used for CIMD / OAuth host pins. + + Forwarded host is used only when the direct peer is trusted. Untrusted + ``X-Forwarded-Host`` cannot rebind the comparison host. + """ + forwarded = trusted_forwarded_host(headers, peer, trusted) + if forwarded: + return hostname_from_host_header(forwarded) or None + host = get_header(headers, "Host") + if host: + return hostname_from_host_header(host) or None + return None + + +def cimd_url_matches_request_host( + cimd_url: str, + headers: Mapping[str, str], + peer: Optional[str], + trusted: Optional[Sequence[str]] = None, +) -> bool: + """True when the CIMD URL host equals the trusted-proxy-aware request host.""" + expected = (urlparse(cimd_url).hostname or "").lower() + actual = request_host_for_cimd(headers, peer, trusted) + return bool(expected and actual and expected == actual) + + +def public_origin( + headers: Mapping[str, str], + peer: Optional[str], + *, + fallback_host: Optional[str] = None, + fallback_proto: str = "http", + trusted: Optional[Sequence[str]] = None, +) -> str: + """Absolute origin for health, docs, and OAuth URLs.""" + host = trusted_forwarded_host(headers, peer, trusted) or get_header(headers, "Host") or fallback_host + proto = trusted_forwarded_proto(headers, peer, trusted) or fallback_proto or "http" + if not host: + host = "127.0.0.1" + return f"{proto}://{host}" + + +def public_url( + headers: Mapping[str, str], + peer: Optional[str], + path: str = "/mcp", + **kwargs: Any, +) -> str: + """``public_origin`` plus a path (OAuth connect URL, health ``publicUrl``).""" + origin = public_origin(headers, peer, **kwargs).rstrip("/") + if not path.startswith("/"): + path = f"/{path}" + return f"{origin}{path}" + + +def request_peer(request: Any) -> Optional[str]: + """Direct socket peer from a Starlette / ASGI request.""" + client = getattr(request, "client", None) + if client is None: + return None + host = getattr(client, "host", None) + if host: + return str(host) + if isinstance(client, (tuple, list)) and client: + return str(client[0]) + return None + + +def public_origin_for_request( + request: Any, + *, + fallback_host: Optional[str] = None, + fallback_proto: Optional[str] = None, + trusted: Optional[Sequence[str]] = None, +) -> str: + headers = getattr(request, "headers", None) or {} + url = getattr(request, "url", None) + proto = fallback_proto or getattr(url, "scheme", None) or "http" + host = fallback_host or getattr(url, "netloc", None) + return public_origin( + headers, + request_peer(request), + fallback_host=host, + fallback_proto=proto, + trusted=trusted, + ) + + +def public_url_for_request( + request: Any, + path: str = "/mcp", + **kwargs: Any, +) -> str: + origin = public_origin_for_request(request, **kwargs).rstrip("/") + if not path.startswith("/"): + path = f"/{path}" + return f"{origin}{path}" diff --git a/nitrostack/transports/sse.py b/nitrostack/transports/sse.py new file mode 100644 index 0000000..d046435 --- /dev/null +++ b/nitrostack/transports/sse.py @@ -0,0 +1,27 @@ +"""SSE notification bus helpers.""" + +from __future__ import annotations + +import json +from typing import Any + +from nitrostack.transports.headers import build_sse_stream_headers + + +def format_sse_message(payload: dict[str, Any]) -> bytes: + """Format a JSON notification as an SSE event: message.""" + data = json.dumps(payload, separators=(",", ":")) + return f"event: message\ndata: {data}\n\n".encode("utf-8") + + +def sse_notification(method: str, params: dict[str, Any] | None = None) -> bytes: + """Build an SSE frame for notifications/tools/list_changed or tasks/status.""" + body: dict[str, Any] = {"method": method} + if params is not None: + body["params"] = params + return format_sse_message(body) + + +def sse_connect_headers() -> dict[str, str]: + """Headers for GET /mcp Accept: text/event-stream connections.""" + return build_sse_stream_headers() diff --git a/nitrostack/transports/stdio.py b/nitrostack/transports/stdio.py index 5596fa1..d76e580 100644 --- a/nitrostack/transports/stdio.py +++ b/nitrostack/transports/stdio.py @@ -1,6 +1,17 @@ import sys import contextlib -from typing import Generator +from typing import TYPE_CHECKING, Any, Generator + +from mcp.server.runner import serve_loop +from mcp.server.stdio import stdio_server + +from nitrostack.protocol.version import ProtocolEra + +if TYPE_CHECKING: + from mcp.server.lowlevel import Server + from mcp.shared._stream_protocols import ReadStream, WriteStream + from mcp.shared.message import SessionMessage + class SafeStdoutWrapper: """ @@ -32,3 +43,49 @@ def safe_stdio_transport() -> Generator[None, None, None]: yield finally: sys.stdout = original_stdout + + +async def serve_stdio_streams( + server: "Server[Any]", + read_stream: "ReadStream[SessionMessage | Exception]", + write_stream: "WriteStream[SessionMessage]", + era: ProtocolEra, +) -> None: + """Drive official mcp 2.x over an already-open stdio stream pair. + + ``auto`` uses the official dual-era loop (2026 envelope or 2025 handshake). + ``modern`` rejects ``initialize``. ``legacy`` keeps the handshake loop. + """ + init_options = server.create_initialization_options() + if era == "auto": + await server.run(read_stream, write_stream, init_options) + return + + async with server.lifespan(server) as lifespan_state: + try: + if era == "modern": + from mcp.server.runner import _serve_modern_stream + + await _serve_modern_stream( + server, + read_stream, + write_stream, + lifespan_state=lifespan_state, + raise_exceptions=False, + ) + else: + await serve_loop( + server, + read_stream, + write_stream, + lifespan_state=lifespan_state, + init_options=init_options, + ) + finally: + await write_stream.aclose() + + +async def run_stdio(server: "Server[Any]", era: ProtocolEra) -> None: + """Serve official mcp 2.x on process stdin/stdout for the active era.""" + async with stdio_server() as (read_stream, write_stream): + await serve_stdio_streams(server, read_stream, write_stream, era) diff --git a/nitrostack/transports/subscriptions.py b/nitrostack/transports/subscriptions.py new file mode 100644 index 0000000..5f04aca --- /dev/null +++ b/nitrostack/transports/subscriptions.py @@ -0,0 +1,195 @@ +"""HTTP attach point for the official MCP 2026 subscription bus.""" + +from __future__ import annotations + +import asyncio +import uuid +from typing import Any, Mapping, Optional + +from starlette.requests import Request +from starlette.responses import JSONResponse, Response, StreamingResponse + +from nitrostack.auth.jwt import JWTService +from nitrostack.auth.request import authorization_token_from_headers, verify_bearer_payload +from nitrostack.core.di import DIContainer +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.transports.headers import ( + MCP_SSE_CONTENT_TYPE, + build_sse_stream_headers, + get_header, +) +from nitrostack.transports.sse import format_sse_message + +try: + from mcp.shared.subscriptions import ( + SUBSCRIPTION_ID_META_KEY, + event_matches, + event_to_notification, + ) + from mcp_types import SubscriptionFilter +except ImportError: # pragma: no cover + SUBSCRIPTION_ID_META_KEY = "io.modelcontextprotocol/subscriptionId" + event_matches = None # type: ignore[assignment] + event_to_notification = None # type: ignore[assignment] + SubscriptionFilter = None # type: ignore[assignment] + + +def http_listen_requires_auth() -> bool: + """True when ``/mcp`` is credential-gated (JWT registered or OAuth required).""" + from nitrostack.auth.oauth import OAuthService, is_oauth_required + + container = DIContainer.get_instance() + if container.has_value(JWTService): + return True + return container.has_value(OAuthService) and is_oauth_required() + + +def unauthorized_listen_response() -> JSONResponse: + return JSONResponse({"error": "unauthorized"}, status_code=401) + + +def listen_auth_error(headers: Mapping[str, str]) -> Optional[JSONResponse]: + """Deny listen with the same credential rules as a gated ``/mcp`` call.""" + if not http_listen_requires_auth(): + return None + token = authorization_token_from_headers(headers) + if not token: + return unauthorized_listen_response() + if verify_bearer_payload(token) is None: + return unauthorized_listen_response() + return None + + +def _accepts_sse(headers: Mapping[str, str]) -> bool: + accept = (get_header(headers, "Accept") or "*/*").lower() + return "text/event-stream" in accept or "*/*" in accept + + +def _default_filter() -> Any: + if SubscriptionFilter is None: + return None + return SubscriptionFilter( + tools_list_changed=True, + prompts_list_changed=True, + resources_list_changed=True, + ) + + +def _filter_from_body(body: Any) -> Any: + honored = _default_filter() + if SubscriptionFilter is None or not isinstance(body, dict): + return honored + params = body.get("params") if isinstance(body.get("params"), dict) else body + raw = params.get("notifications") if isinstance(params, dict) else None + if not isinstance(raw, dict): + return honored + return SubscriptionFilter( + tools_list_changed=True if raw.get("toolsListChanged") or raw.get("tools_list_changed") else None, + prompts_list_changed=True if raw.get("promptsListChanged") or raw.get("prompts_list_changed") else None, + resources_list_changed=True + if raw.get("resourcesListChanged") or raw.get("resources_list_changed") + else None, + resource_subscriptions=list( + raw.get("resourceSubscriptions") or raw.get("resource_subscriptions") or [] + ) + or None, + ) + + +def _notification_payload(event: Any, meta: dict[str, Any]) -> dict[str, Any]: + if event_to_notification is None: + return {"method": "notifications/tools/list_changed", "params": {"_meta": meta}} + notification = event_to_notification(event, meta) + return notification.model_dump(by_alias=True, exclude_none=True) + + +def _ack_payload(subscription_id: str, honored: Any) -> dict[str, Any]: + notifications: dict[str, Any] = {} + if honored is not None: + if getattr(honored, "tools_list_changed", None): + notifications["toolsListChanged"] = True + if getattr(honored, "prompts_list_changed", None): + notifications["promptsListChanged"] = True + if getattr(honored, "resources_list_changed", None): + notifications["resourcesListChanged"] = True + uris = getattr(honored, "resource_subscriptions", None) + if uris: + notifications["resourceSubscriptions"] = list(uris) + else: + notifications = { + "toolsListChanged": True, + "promptsListChanged": True, + "resourcesListChanged": True, + } + return { + "method": "notifications/subscriptions/acknowledged", + "params": { + "_meta": {SUBSCRIPTION_ID_META_KEY: subscription_id}, + "notifications": notifications, + }, + } + + +def subscriptions_listen_endpoint(mcp_app: Any): + """GET/POST ``/subscriptions/listen``: SSE attach on the official v2 bus.""" + + async def handle(request: Request) -> Response: + denied = listen_auth_error(request.headers) + if denied is not None: + return denied + if not _accepts_sse(request.headers): + return Response(status_code=406) + + body: Any = {} + if request.method == "POST": + try: + body = await request.json() + except Exception: + body = {} + + bus = getattr(getattr(mcp_app, "mcp_server", None), "subscription_bus", None) + if bus is None: + return JSONResponse({"error": "subscriptions are not available"}, status_code=503) + + honored = _filter_from_body(body) + honored_uris = frozenset(getattr(honored, "resource_subscriptions", None) or ()) + subscription_id = str(uuid.uuid4()) + queue: asyncio.Queue[Any] = asyncio.Queue() + + def deliver(event: Any) -> None: + if event_matches is not None and honored is not None: + if not event_matches(honored, honored_uris, event): + return + try: + queue.put_nowait(event) + except Exception: + pass + + unsubscribe = bus.subscribe(deliver) + + async def frames(): + try: + yield format_sse_message(_ack_payload(subscription_id, honored)) + while True: + event = await queue.get() + yield format_sse_message( + _notification_payload( + event, {SUBSCRIPTION_ID_META_KEY: subscription_id} + ) + ) + finally: + unsubscribe() + + return StreamingResponse( + frames(), + media_type=MCP_SSE_CONTENT_TYPE, + headers=build_sse_stream_headers( + protocol_version=getattr( + request.app.state, "protocol_version", MODERN_PROTOCOL_VERSION + ) + if hasattr(request.app, "state") + else MODERN_PROTOCOL_VERSION + ), + ) + + return handle diff --git a/pyproject.toml b/pyproject.toml index 6df79b3..107a3f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ description = "NitroStack Python SDK — A Python-idiomatic port of the NitroSta readme = "README.md" requires-python = ">=3.10" dependencies = [ - "mcp>=1.10.0,<2.0.0", + "mcp>=2,<3", "pydantic>=2.0.0", "starlette>=0.30.0", "uvicorn>=0.20.0", diff --git a/requirements.txt b/requirements.txt index d0f3fc2..ebe86ba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -mcp>=1.10.0,<2.0.0 +mcp>=2,<3 pydantic>=2.0.0 starlette>=0.30.0 uvicorn>=0.20.0 diff --git a/tests/test_flight_booking.py b/tests/test_flight_booking.py index de9b008..d01cd62 100644 --- a/tests/test_flight_booking.py +++ b/tests/test_flight_booking.py @@ -127,8 +127,8 @@ async def run(): ) ) payload = raw.root - assert payload.isError is not True - assert payload.structuredContent["offers"] + assert payload.is_error is not True + assert payload.structured_content["offers"] html = next(b.resource.text for b in payload.content if getattr(b, "type", None) == "resource") assert "JFK" in html assert "LAX" in html @@ -153,7 +153,7 @@ async def run(): ), ) ) - assert resp.root.isError is True + assert resp.root.is_error is True text = resp.root.content[0].text assert "OAuth" in text or "Access denied" in text except PermissionError as exc: diff --git a/tests/test_initial_tool.py b/tests/test_initial_tool.py index e9d494d..bc84708 100644 --- a/tests/test_initial_tool.py +++ b/tests/test_initial_tool.py @@ -46,7 +46,7 @@ async def _test_initial_tool_hook(): # Check that called is False initially assert ExecutionState.called is False - # Retrieve the owned low-level Server directly (no FastMCP wrapper in between) + # Retrieve the owned low-level Server directly. server = harness.app.mcp_server # Get the InitializedNotification handler diff --git a/tests/test_lifecycle_http.py b/tests/test_lifecycle_http.py index 86441b0..c4bb13b 100644 --- a/tests/test_lifecycle_http.py +++ b/tests/test_lifecycle_http.py @@ -122,7 +122,12 @@ class App: ) assert created.status_code == 200, created.text created_body = created.json()["result"] - task_id = created_body.get("task", {}).get("taskId") or created_body.get("taskId") + structured = created_body.get("structuredContent") or {} + task_id = ( + created_body.get("task", {}).get("taskId") + or created_body.get("taskId") + or (structured.get("task") or {}).get("taskId") + ) assert task_id, created_body status = None @@ -143,26 +148,17 @@ class App: break asyncio.run(asyncio.sleep(0.05)) - result = client.post( - "/mcp", - headers=JSON_HEADERS, - json={ - "jsonrpc": "2.0", - "id": 5, - "method": "tasks/result", - "params": {"taskId": task_id}, - }, - ) - assert result.status_code == 200, result.text - payload = result.json()["result"] - assert payload.get("isError") is False - hello = (payload.get("structuredContent") or {}).get("hello") - if hello is None: - hello = payload["content"][0]["text"] - assert "ADA" in hello - else: - assert hello == "ADA" - assert status == "completed" + assert status == "completed" + payload = polled.json()["result"] + assert payload.get("result") is not None + result_payload = payload["result"] + assert result_payload.get("isError") is False + hello = (result_payload.get("structuredContent") or {}).get("hello") + if hello is None: + hello = result_payload["content"][0]["text"] + assert "ADA" in hello + else: + assert hello == "ADA" diff --git a/tests/test_mcp20_acceptance.py b/tests/test_mcp20_acceptance.py new file mode 100644 index 0000000..ffc2cce --- /dev/null +++ b/tests/test_mcp20_acceptance.py @@ -0,0 +1,226 @@ +"""Acceptance-criteria verification for MCP 2026-07-28.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import os +import sys +import time +from pathlib import Path + +import mcp.types as types +import pytest +from nitrostack.runtime.request_ctx import Experimental, RequestContext, RequestParamsMeta, request_ctx +from pydantic import BaseModel, Field + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.core.task import TaskManager, TaskStatus +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.runtime.acceptance import ( + ACCEPTANCE_CRITERIA, + MCP20_TEST_MODULES, + PROTOCOL_DELIVERABLES, + ProtocolArea, + acceptance_criteria_registered, + deliverables_for_area, + iter_area_summary, + modern_protocol_target, + protocol_coverage_complete, +) +from nitrostack.runtime.conformance import assert_blueprint_layout + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class TestProtocolRegistry: + def test_all_areas_have_deliverables(self): + assert protocol_coverage_complete() + for line in iter_area_summary(): + assert "0 deliverable" not in line + + def test_deliverable_count(self): + assert len(PROTOCOL_DELIVERABLES) >= 14 + + def test_acceptance_criteria_registered(self): + assert acceptance_criteria_registered() + keys = {item.key for item in ACCEPTANCE_CRITERIA} + assert keys == { + "stateless", + "task_conformance", + "task_cancellation", + "ttl_safety", + "multi_tenant", + "ssrf_security", + "error_parity", + "automated_tests", + } + + def test_modern_protocol_target(self): + assert modern_protocol_target() == "2026-07-28" + + @pytest.mark.parametrize("area", list(ProtocolArea)) + def test_each_area_maps_to_tests(self, area: ProtocolArea): + items = deliverables_for_area(area) + assert items, f"Area {area} has no deliverables" + for item in items: + assert item.test_module.startswith("tests/test_mcp20_") + + +class TestLayoutAndTestTraceability: + def test_all_mcp20_test_modules_exist(self): + for module_path in MCP20_TEST_MODULES: + path = REPO_ROOT / module_path + assert path.is_file(), module_path + + def test_blueprint_layout_importable(self): + assert_blueprint_layout() + + +class TestAcceptanceTaskConformance: + def test_task_augmented_call_returns_immediately(self): + @injectable() + class AsyncController: + @tool(name="fast_task", description="fast", input_schema=EchoInput, task_support="optional") + async def fast_task(self, input: EchoInput, context: ExecutionContext) -> str: + await asyncio.sleep(0.05) + return input.value + + @module(name="AcceptanceTasks", controllers=[AsyncController]) + class TasksModule: + pass + + @mcp_app(module=TasksModule, server=ServerConfig(name="acceptance-tasks")) + class TasksApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(TasksApp) + token = request_ctx.set( + RequestContext( + request_id="1", + meta=None, + session=None, + lifespan_context=None, + experimental=Experimental(task_metadata=types.TaskMetadata(ttl=60_000)), + ) + ) + try: + started = time.perf_counter() + result = await app._call_tool("fast_task", {"value": "x"}) + elapsed_ms = (time.perf_counter() - started) * 1000 + assert isinstance(result, types.CreateTaskResult) + assert result.task.status == "working" + assert elapsed_ms < 500, f"task handle took {elapsed_ms:.1f}ms" + finally: + request_ctx.reset(token) + + asyncio.run(_run()) + + def test_cancel_on_terminal_task_returns_invalid_params(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task(tool_name="job") + await manager.complete_task(task.id, {"ok": True}) + + @injectable() + class DummyController: + @tool(name="noop_accept", description="noop", input_schema=EchoInput) + async def noop_accept(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="AcceptanceCancel", controllers=[DummyController]) + class CancelModule: + pass + + @mcp_app(module=CancelModule, server=ServerConfig(name="acceptance-cancel")) + class CancelApp: + pass + + from mcp import MCPError as McpError + + app = await McpApplicationFactory.create(CancelApp) + app.task_manager = manager + handler = app.mcp_server.request_handlers[types.CancelTaskRequest] + with pytest.raises(McpError) as exc: + await handler( + types.CancelTaskRequest( + method="tasks/cancel", + params=types.CancelTaskRequestParams(taskId=task.id), + ) + ) + assert exc.value.error.code == types.INVALID_PARAMS + + asyncio.run(_run()) + + def test_active_task_not_evicted_before_terminal(self): + from nitrostack.tasks.eviction import should_evict_terminal_task + from nitrostack.tasks.types import TaskEntry, TaskWireData, datetime_to_ms, utc_now + + now = utc_now() + entry = TaskEntry( + task_id="active", + data=TaskWireData( + task_id="active", + status="working", + created_at=now, + last_updated_at=now, + ttl_ms=1, + ), + status="working", + ) + assert should_evict_terminal_task(entry, datetime_to_ms(now) + 10_000) is False + + def test_completed_task_snapshot_status(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task(tool_name="job") + await manager.cancel_task(task.id) + snapshot = await manager.get_task(task.id) + assert snapshot.status == TaskStatus.CANCELLED + + asyncio.run(_run()) + + +class TestAcceptanceStatelessInvariant: + def test_legacy_session_header_constant(self): + assert LEGACY_SESSION_HEADER == "Mcp-Session-Id" + + def test_server_config_defaults_to_modern_protocol(self): + cfg = ServerConfig(name="x") + assert cfg.protocol_version == MODERN_PROTOCOL_VERSION + + +class TestAutomatedTestCoverageMap: + def test_mcp20_test_modules_importable(self): + for module_path in MCP20_TEST_MODULES: + file_path = REPO_ROOT / module_path + spec = importlib.util.spec_from_file_location( + module_path.replace("/", "_").replace(".", "_"), + file_path, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + def test_full_mcp20_suite_module_count(self): + suite = sorted(REPO_ROOT.glob("tests/test_mcp20_*.py")) + assert len(suite) == 17 diff --git a/tests/test_mcp20_blueprint.py b/tests/test_mcp20_blueprint.py new file mode 100644 index 0000000..ea3b26f --- /dev/null +++ b/tests/test_mcp20_blueprint.py @@ -0,0 +1,413 @@ +"""Blueprint conformance verification.""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from unittest.mock import patch + +import mcp.types as types +import pytest +from nitrostack.runtime.request_ctx import Experimental, RequestContext, RequestParamsMeta, request_ctx +from mcp import MCPError as McpError +from pydantic import BaseModel, Field +from starlette.testclient import TestClient + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.auth.cimd import assert_safe_fetch_target, is_blocked_ip, resolve_cimd +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.core.errors import ResourceNotFoundError, TaskNotFoundError +from nitrostack.core.task import TaskManager +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER, MAX_CIMD_BYTES +from nitrostack.protocol.deprecated import deprecated_method_message +from nitrostack.protocol.discovery import build_discover_result +from nitrostack.protocol.errors import JsonRpcErrorCode +from nitrostack.protocol.extensions import MCPExtensionId +from nitrostack.protocol.jsonrpc import map_exception_to_jsonrpc +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.runtime.conformance import ( + BLUEPRINT_CONFORMANCE_AREAS, + ConformanceArea, + DEPRECATED_MODERN_METHODS, + MODERN_JSONRPC_ERROR_CODES, + REQUIRED_BLUEPRINT_MODULES, + assert_blueprint_layout, + protocol_version_matches_blueprint, + verify_package_layout, +) +from nitrostack.runtime.stateless import assert_stateless_headers +from nitrostack.tasks.eviction import should_evict_terminal_task +from nitrostack.tasks.types import TaskAccessContext, TaskEntry, TaskWireData, datetime_to_ms, utc_now +from nitrostack.tasks.authorization import check_task_access +from nitrostack.transports.cors import cors_preflight_response_headers +from nitrostack.transports.dispatch import IngressContext, StatelessIngressPipeline + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +JSON_HEADERS = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"} + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class TestBlueprintRegistry: + def test_conformance_areas_defined(self): + assert set(BLUEPRINT_CONFORMANCE_AREAS) == set(ConformanceArea) + assert len(BLUEPRINT_CONFORMANCE_AREAS) == 5 + + def test_modern_protocol_version(self): + assert protocol_version_matches_blueprint() + assert MODERN_PROTOCOL_VERSION == "2026-07-28" + + def test_required_modules_import(self): + assert verify_package_layout() == [] + assert_blueprint_layout() + + def test_blueprint_module_count(self): + assert len(REQUIRED_BLUEPRINT_MODULES) >= 17 + + +class TestStatelessHttpConformance: + def test_discover_advertises_modern_protocol_and_extensions(self): + result = build_discover_result( + server_name="blueprint-server", + server_version="1.0.0", + advertise_tasks=True, + ) + assert result["protocolVersion"] == MODERN_PROTOCOL_VERSION + extensions = result["capabilities"]["extensions"] + assert MCPExtensionId.TASKS.value in extensions + assert result["resultType"] == "complete" + assert isinstance(result["ttlMs"], int) and result["ttlMs"] >= 0 + assert result["cacheScope"] in ("public", "private") + + def test_stateless_pipeline_handles_discover_without_session(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION), + discover_handler=lambda _req: build_discover_result( + server_name="srv", + server_version="1.0.0", + protocol_version=MODERN_PROTOCOL_VERSION, + ), + ) + body = json.dumps( + {"jsonrpc": "2.0", "id": "d1", "method": "server/discover", "params": {}} + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + assert resp["result"]["protocolVersion"] == MODERN_PROTOCOL_VERSION + assert_stateless_headers({"Content-Type": "application/json"}) + + asyncio.run(_run()) + + def test_post_mcp_stateless_no_session_header(self): + @injectable() + class PingController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="BlueprintHttp", controllers=[PingController]) + class HttpModule: + pass + + @mcp_app(module=HttpModule, server=ServerConfig(name="blueprint-http", stateless=True)) + class HttpApp: + pass + + app = asyncio.run(McpApplicationFactory.create(HttpApp)) + http_app = app.get_combined_app(stateless=True, json_response=True) + with TestClient(http_app) as client: + resp = client.post( + "/mcp", + headers={**JSON_HEADERS, "Mcp-Method": "tools/call", "Mcp-Name": "echo"}, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + assert resp.status_code == 200 + assert LEGACY_SESSION_HEADER not in {k.title() for k in resp.headers.keys()} + assert resp.json()["result"]["content"][0]["text"] == "ok" + + def test_cors_preflight_includes_mcp_headers(self): + headers = cors_preflight_response_headers({"Origin": "https://app.example.com"}) + assert "Access-Control-Allow-Origin" in headers + allow = headers.get("Access-Control-Allow-Headers", "") + assert "Mcp-Method" in allow + assert "Mcp-Name" in allow + assert "MCP-Protocol-Version" in allow + + +class TestJsonRpcConformance: + @pytest.mark.parametrize( + "code", + sorted(MODERN_JSONRPC_ERROR_CODES), + ) + def test_standard_error_codes_registered(self, code: JsonRpcErrorCode): + assert int(code) in MODERN_JSONRPC_ERROR_CODES + + def test_missing_resource_maps_to_invalid_params(self): + resp = map_exception_to_jsonrpc(ResourceNotFoundError("mcp://missing"), "99") + assert resp["error"]["code"] == JsonRpcErrorCode.INVALID_PARAMS + + @pytest.mark.parametrize("method", sorted(DEPRECATED_MODERN_METHODS)) + def test_deprecated_methods_have_modern_rejection_message(self, method: str): + message = deprecated_method_message(method) + assert "2026-07-28" in message or "tasks/get" in message + + def test_tasks_result_rejected_on_modern_wire(self): + @injectable() + class DummyController: + @tool(name="noop", description="noop", input_schema=EchoInput) + async def noop(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="BlueprintJsonRpc", controllers=[DummyController]) + class JsonRpcModule: + pass + + @mcp_app( + module=JsonRpcModule, + server=ServerConfig(name="jsonrpc", protocol_era="modern"), + ) + class JsonRpcApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(JsonRpcApp) + handler = app.mcp_server.request_handlers[types.GetTaskPayloadRequest] + with pytest.raises(McpError) as exc: + await handler( + types.GetTaskPayloadRequest( + method="tasks/result", + params=types.GetTaskPayloadRequestParams(taskId="missing"), + ) + ) + assert exc.value.error.code == types.METHOD_NOT_FOUND + + asyncio.run(_run()) + + +class TestTaskSubsystemConformance: + def test_task_augmented_call_returns_working_task(self): + @injectable() + class AsyncController: + @tool(name="slow", description="slow", input_schema=EchoInput, task_support="optional") + async def slow(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="BlueprintTasks", controllers=[AsyncController]) + class TasksModule: + pass + + @mcp_app(module=TasksModule, server=ServerConfig(name="tasks")) + class TasksApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(TasksApp) + token = request_ctx.set( + RequestContext( + request_id="1", + meta=None, + session=None, + lifespan_context=None, + experimental=Experimental(task_metadata=types.TaskMetadata(ttl=60_000)), + ) + ) + try: + created = await app._call_tool("slow", {"value": "async"}) + assert isinstance(created, types.CreateTaskResult) + assert created.task.status == "working" + + get_handler = app.mcp_server.request_handlers[types.GetTaskRequest] + snapshot = await get_handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId=created.task.task_id), + ) + ) + assert snapshot.status in ("working", "completed") + finally: + request_ctx.reset(token) + + asyncio.run(_run()) + + def test_tasks_get_embeds_result_after_completion(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task(tool_name="job") + await manager.complete_task(task.id, {"done": True}) + + @injectable() + class DummyController: + @tool(name="noop_bp", description="noop", input_schema=EchoInput) + async def noop_bp(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="BlueprintGet", controllers=[DummyController]) + class GetModule: + pass + + @mcp_app(module=GetModule, server=ServerConfig(name="get")) + class GetApp: + pass + + app = await McpApplicationFactory.create(GetApp) + app.task_manager = manager + handler = app.mcp_server.request_handlers[types.GetTaskRequest] + response = await handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId=task.id), + ) + ) + payload = response.model_dump(by_alias=True) + assert payload["status"] == "completed" + assert payload["result"] is not None + + asyncio.run(_run()) + + def test_cancel_marks_task_cancelled(self): + async def _run(): + from nitrostack.core.task import TaskStatus + + manager = TaskManager() + task = await manager.create_task(tool_name="job") + await manager.cancel_task(task.id) + snapshot = await manager.get_task(task.id) + assert snapshot.status == TaskStatus.CANCELLED + + asyncio.run(_run()) + + def test_terminal_ttl_eviction_uses_last_updated_at(self): + now = utc_now() + wire = TaskWireData( + task_id="evict-me", + status="completed", + created_at=now, + last_updated_at=now, + ttl_ms=500, + ) + entry = TaskEntry(task_id="evict-me", data=wire, status="completed") + before_ttl = datetime_to_ms(now) + 100 + after_ttl = datetime_to_ms(now) + 2_000 + assert should_evict_terminal_task(entry, before_ttl) is False + assert should_evict_terminal_task(entry, after_ttl) is True + + +class TestMultiTenantIsolationConformance: + def test_cross_tenant_access_raises_not_found(self): + entry = TaskEntry( + task_id="iso-1", + data=TaskWireData(task_id="iso-1"), + owner_id="alice", + tenant_id="tenant-a", + ) + foreign = TaskAccessContext(user_id="alice", tenant_id="tenant-b") + with pytest.raises(TaskNotFoundError): + check_task_access(entry, foreign) + + def test_tasks_get_cross_tenant_returns_invalid_params(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task( + tool_name="secret", + owner_id="alice", + tenant_id="tenant-a", + ) + + @injectable() + class DummyController: + @tool(name="noop_iso", description="noop", input_schema=EchoInput) + async def noop_iso(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="BlueprintIso", controllers=[DummyController]) + class IsoModule: + pass + + @mcp_app(module=IsoModule, server=ServerConfig(name="iso")) + class IsoApp: + pass + + app = await McpApplicationFactory.create(IsoApp) + app.task_manager = manager + handler = app.mcp_server.request_handlers[types.GetTaskRequest] + token = request_ctx.set( + RequestContext( + request_id="iso", + meta={"tenantId": "tenant-b", "userId": "alice"}, + session=None, + lifespan_context={}, + ) + ) + try: + with pytest.raises(McpError) as exc: + await handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId=task.id), + ) + ) + assert exc.value.error.code == types.INVALID_PARAMS + finally: + request_ctx.reset(token) + + asyncio.run(_run()) + + +class TestCimdSsrfConformance: + @pytest.mark.parametrize( + "ip", + ["127.0.0.1", "169.254.169.254", "10.0.0.1", "::1", "fc00::1"], + ) + def test_blocks_special_use_ips(self, ip: str): + assert is_blocked_ip(ip) is True + + def test_blocks_private_dns_target(self): + async def _run(): + import socket + + url = "https://metadata.example.com/oauth/client.json" + with patch( + "socket.getaddrinfo", + return_value=[(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.1", 0))], + ): + from nitrostack.auth.cimd import CimdFetchError + + with pytest.raises(CimdFetchError, match="blocked IP"): + await assert_safe_fetch_target(url) + + asyncio.run(_run()) + + def test_rejects_oversized_payload(self): + async def _run(): + from nitrostack.auth.cimd import CimdFetchError + + url = "https://app.example.com/oauth/client-metadata.json" + oversized = b"x" * (MAX_CIMD_BYTES + 1) + + with patch("nitrostack.auth.cimd.assert_safe_fetch_target", return_value=None): + with patch("nitrostack.auth.cimd._fetch_cimd_bytes", return_value=oversized): + with pytest.raises(CimdFetchError, match="maximum size"): + await resolve_cimd(url) + + asyncio.run(_run()) diff --git a/tests/test_mcp20_concurrency.py b/tests/test_mcp20_concurrency.py new file mode 100644 index 0000000..a3a69d5 --- /dev/null +++ b/tests/test_mcp20_concurrency.py @@ -0,0 +1,263 @@ +"""Concurrent identical JSON-RPC ids stay isolated on the sessionless path.""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys + +from pydantic import BaseModel, Field + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.runtime.correlation import InFlightRegistry, new_correlation_id + +CALL_HEADERS = { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-method": "tools/call", + "mcp-protocol-version": "2026-07-28", +} + +CALL_META = { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {}, +} + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def test_registry_cancel_targets_correlation_id_only(): + registry = InFlightRegistry() + first = new_correlation_id() + second = new_correlation_id() + registry.register(first, jsonrpc_id=1) + registry.register(second, jsonrpc_id=1) + assert len(registry) == 2 + assert {ticket.jsonrpc_id for ticket in registry} == {1} + assert registry.cancel(first) is True + assert registry.get(first).cancel_requested.is_set() + assert not registry.get(second).cancel_requested.is_set() + registry.discard(first) + assert len(registry) == 1 + assert registry.get(second) is not None + + +def _header_list(headers: dict[str, str]) -> list[tuple[bytes, bytes]]: + return [(key.lower().encode("latin-1"), value.encode("latin-1")) for key, value in headers.items()] + + +async def _with_lifespan(app, http_call): + started = asyncio.Event() + messages: asyncio.Queue = asyncio.Queue() + await messages.put({"type": "lifespan.startup"}) + + async def receive(): + return await messages.get() + + async def send(message): + if message["type"] == "lifespan.startup.complete": + started.set() + + task = asyncio.create_task(app({"type": "lifespan"}, receive, send)) + await asyncio.wait_for(started.wait(), timeout=2) + try: + return await http_call() + finally: + await messages.put({"type": "lifespan.shutdown"}) + await asyncio.wait_for(task, timeout=2) + + +async def _asgi_json_post(app, payload: dict, headers: dict[str, str]) -> tuple[int, dict]: + status = 0 + chunks: list[bytes] = [] + sent_request = False + finished = asyncio.Event() + body = json.dumps(payload).encode("utf-8") + + async def receive(): + nonlocal sent_request + if not sent_request: + sent_request = True + return {"type": "http.request", "body": body, "more_body": False} + await finished.wait() + return {"type": "http.disconnect"} + + async def send(message): + nonlocal status + if message["type"] == "http.response.start": + status = message["status"] + elif message["type"] == "http.response.body": + chunks.append(message.get("body") or b"") + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/mcp", + "raw_path": b"/mcp", + "query_string": b"", + "headers": _header_list(headers), + "client": ("testclient", 50000), + "server": ("test", 80), + } + try: + await asyncio.wait_for(app(scope, receive, send), timeout=5) + finally: + finished.set() + raw = b"".join(chunks) + parsed = json.loads(raw.decode("utf-8")) if raw else {} + return status, parsed + + +def _call_body(name: str, value: str) -> dict: + return { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": name, + "arguments": {"value": value}, + "_meta": CALL_META, + }, + } + + +def _make_overlap_app(): + started_alpha = asyncio.Event() + started_beta = asyncio.Event() + release = asyncio.Event() + seen: list[str] = [] + + @injectable() + class PairController: + @tool(name="alpha", description="alpha", input_schema=EchoInput) + async def alpha(self, input: EchoInput, context: ExecutionContext) -> str: + seen.append(f"{context.correlation_id}:{context.jsonrpc_id}:{input.value}") + started_alpha.set() + await release.wait() + return f"alpha:{input.value}" + + @tool(name="beta", description="beta", input_schema=EchoInput) + async def beta(self, input: EchoInput, context: ExecutionContext) -> str: + seen.append(f"{context.correlation_id}:{context.jsonrpc_id}:{input.value}") + started_beta.set() + await release.wait() + return f"beta:{input.value}" + + @module(name="overlap-rpc", controllers=[PairController]) + class PairModule: + pass + + @mcp_app(module=PairModule, server=ServerConfig(name="overlap-rpc")) + class App: + pass + + app = asyncio.run(McpApplicationFactory.create(App)) + return app, started_alpha, started_beta, release, seen + + +def _tool_text(payload: dict) -> str: + result = payload.get("result") or {} + content = result.get("content") or [] + if content: + return str(content[0].get("text") or "") + structured = result.get("structuredContent") + if structured is not None: + return str(structured) + return json.dumps(payload) + + +def test_two_overlapping_id_1_calls_return_distinct_results(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + app, started_alpha, started_beta, release, seen = _make_overlap_app() + http_app = app.get_combined_app(json_response=True) + + async def run(): + async def both(): + first = asyncio.create_task( + _asgi_json_post( + http_app, + _call_body("alpha", "one"), + {**CALL_HEADERS, "mcp-name": "alpha"}, + ) + ) + second = asyncio.create_task( + _asgi_json_post( + http_app, + _call_body("beta", "two"), + {**CALL_HEADERS, "mcp-name": "beta"}, + ) + ) + await asyncio.wait_for(started_alpha.wait(), timeout=2) + await asyncio.wait_for(started_beta.wait(), timeout=2) + assert len(app._in_flight) == 2 + assert {ticket.jsonrpc_id for ticket in app._in_flight} == {1} + correlations = [ticket.correlation_id for ticket in app._in_flight] + assert correlations[0] != correlations[1] + release.set() + return await asyncio.gather(first, second) + + return await _with_lifespan(http_app, both) + + (status_a, body_a), (status_b, body_b) = asyncio.run(run()) + assert status_a == 200 and status_b == 200 + assert body_a.get("id") == 1 + assert body_b.get("id") == 1 + texts = {_tool_text(body_a), _tool_text(body_b)} + assert "alpha:one" in texts + assert "beta:two" in texts + assert len({row.split(":")[0] for row in seen}) == 2 + + +def test_cancel_of_one_id_1_call_does_not_cancel_the_other(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + app, started_alpha, started_beta, release, seen = _make_overlap_app() + http_app = app.get_combined_app(json_response=True) + + async def run_cancel(): + async def both(): + first = asyncio.create_task( + _asgi_json_post( + http_app, + _call_body("alpha", "one"), + {**CALL_HEADERS, "mcp-name": "alpha"}, + ) + ) + second = asyncio.create_task( + _asgi_json_post( + http_app, + _call_body("beta", "two"), + {**CALL_HEADERS, "mcp-name": "beta"}, + ) + ) + await asyncio.wait_for(started_alpha.wait(), timeout=2) + await asyncio.wait_for(started_beta.wait(), timeout=2) + alpha_corr = next(row.split(":")[0] for row in seen if row.endswith(":one")) + assert app._in_flight.cancel(alpha_corr) is True + first.cancel() + release.set() + return await second + + return await _with_lifespan(http_app, both) + + status, body = asyncio.run(run_cancel()) + assert status == 200 + assert body.get("id") == 1 + assert "beta:two" in _tool_text(body) diff --git a/tests/test_mcp20_contracts.py b/tests/test_mcp20_contracts.py new file mode 100644 index 0000000..32c24f4 --- /dev/null +++ b/tests/test_mcp20_contracts.py @@ -0,0 +1,285 @@ +"""Tests for MCP 2.0 tool/resource/prompt contracts.""" + +import pytest +from pydantic import BaseModel, Field + +from nitrostack.core.errors import ResourceNotFoundError +from nitrostack.protocol.constants import MAX_SCHEMA_DEPTH +from nitrostack.protocol.contracts import ( + build_cache_hint_meta, + build_prompt_get_result, + build_prompt_text_message, + build_resource_blob_content, + build_resource_text_content, +) +from nitrostack.protocol.jsonrpc import map_exception_to_jsonrpc +from nitrostack.protocol.resources import ( + resolve_resource_uri, + uri_template_to_pattern, +) +from nitrostack.protocol.schema import ( + JSON_SCHEMA_2020_12_URI, + UnsupportedJsonSchemaError, + assert_json_schema_2020_12, + bound_schema_depth, + normalize_input_schema, + normalize_output_schema, +) + + +class _DeepSchema: + @staticmethod + def nested(depth: int) -> dict: + node: dict = {"type": "object", "properties": {"value": {"type": "string"}}} + for _ in range(depth): + node = {"type": "object", "properties": {"child": node}} + return node + + +class TestJsonSchema2020_12: + def test_input_schema_requires_object_root(self): + schema = normalize_input_schema({"type": "string"}) + assert schema["type"] == "object" + assert schema["$schema"] == JSON_SCHEMA_2020_12_URI + assert "properties" in schema + + def test_output_schema_allows_primitive_root(self): + schema = normalize_output_schema({"type": "integer"}) + assert schema["type"] == "integer" + assert schema["$schema"] == JSON_SCHEMA_2020_12_URI + + def test_depth_bounding_collapses_deep_nodes(self): + deep = _DeepSchema.nested(MAX_SCHEMA_DEPTH + 5) + bounded = bound_schema_depth(deep) + # Walk down until we hit the collapsed node + node = bounded + for _ in range(MAX_SCHEMA_DEPTH - 1): + node = node["properties"]["child"] + assert node["properties"]["child"] == {} + + def test_items_object_is_a_single_schema(self): + bounded = bound_schema_depth( + {"type": "array", "items": {"type": "string", "minLength": 1}}, + max_depth=8, + ) + assert bounded["items"] == {"type": "string", "minLength": 1} + + def test_missing_schema_is_treated_as_2020_12(self): + assert_json_schema_2020_12({"type": "object", "properties": {}}) + schema = normalize_input_schema({"type": "object", "properties": {"n": {"type": "integer"}}}) + assert schema["$schema"] == JSON_SCHEMA_2020_12_URI + + def test_draft_04_schema_is_rejected(self): + draft = { + "$schema": "http://json-schema.org/draft-04/schema#", + "type": "object", + "properties": {"n": {"type": "integer"}}, + } + with pytest.raises(UnsupportedJsonSchemaError, match="draft-04"): + assert_json_schema_2020_12(draft, name="tool 'legacy' input") + with pytest.raises(UnsupportedJsonSchemaError, match="draft-04"): + normalize_input_schema(draft) + + def test_items_tuple_list_is_rejected(self): + schema = { + "type": "object", + "properties": { + "pair": { + "type": "array", + "items": [{"type": "string"}, {"type": "integer"}], + } + }, + } + with pytest.raises(UnsupportedJsonSchemaError, match="tuple"): + assert_json_schema_2020_12(schema, name="tool 'pair' input") + with pytest.raises(UnsupportedJsonSchemaError, match="items"): + bound_schema_depth(schema) + + def test_prefix_items_tuple_is_allowed(self): + schema = { + "$schema": JSON_SCHEMA_2020_12_URI, + "type": "object", + "properties": { + "pair": { + "type": "array", + "prefixItems": [{"type": "string"}, {"type": "integer"}], + "items": False, + } + }, + } + assert_json_schema_2020_12(schema) + bounded = bound_schema_depth(schema, max_depth=8) + assert bounded["properties"]["pair"]["prefixItems"][0]["type"] == "string" + + +class TestResourceUriResolution: + def test_static_resource_exact_match(self): + static = {"mcp://telemetry/system_metrics": "entry-a"} + result = resolve_resource_uri("mcp://telemetry/system_metrics", static, []) + assert result is not None + assert result.entry == "entry-a" + assert result.path_params == {} + + def test_template_match_extracts_params(self): + pattern = uri_template_to_pattern("mcp://customers/{customerId}/orders") + templates = [(pattern, "entry-b")] + result = resolve_resource_uri("mcp://customers/acme/orders", {}, templates) + assert result is not None + assert result.matched_via_template is True + assert result.path_params == {"customerId": "acme"} + + def test_missing_resource_maps_to_invalid_params(self): + resp = map_exception_to_jsonrpc(ResourceNotFoundError("mcp://missing"), "1") + assert resp["error"]["code"] == -32602 + + +class TestWireContracts: + def test_cache_hint_meta(self): + meta = build_cache_hint_meta(60000, cache_scope="private") + hint = meta["io.modelcontextprotocol/cacheHint"] + assert hint["ttlMs"] == 60000 + assert hint["cacheScope"] == "private" + + def test_resource_text_content(self): + content = build_resource_text_content( + "mcp://telemetry/system_metrics", + '{"cpuUsage": 14.2}', + ) + assert content["mimeType"] == "application/json" + assert "cpuUsage" in content["text"] + + def test_resource_blob_content(self): + content = build_resource_blob_content( + "mcp://assets/diagram.png", + "iVBORw0KGgo=", + mime_type="image/png", + ) + assert content["blob"] == "iVBORw0KGgo=" + assert content["mimeType"] == "image/png" + + def test_prompt_get_result(self): + messages = [ + build_prompt_text_message("user", "Review this Python code:\ndef foo(): pass") + ] + result = build_prompt_get_result("Senior Code Review Prompt", messages) + assert result["description"] == "Senior Code Review Prompt" + assert result["messages"][0]["content"]["type"] == "text" + + +class TestAppIntegrationSchema: + def test_pydantic_model_produces_modern_input_schema(self): + from nitrostack.core.app import inspector_friendly_schema + + class CalcInput(BaseModel): + a: float = Field(description="First number") + b: float = Field(description="Second number") + + schema = normalize_input_schema(inspector_friendly_schema(CalcInput.model_json_schema())) + assert schema["$schema"] == JSON_SCHEMA_2020_12_URI + assert schema["type"] == "object" + assert "a" in schema["properties"] + + def test_draft_04_tool_fails_registration(self): + import asyncio + + from nitrostack import injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.context import ExecutionContext + from nitrostack.core.di import DIContainer + + draft_schema = { + "$schema": "http://json-schema.org/draft-04/schema#", + "type": "object", + "properties": {"value": {"type": "string"}}, + } + + @injectable() + class DraftController: + @tool(name="legacy_echo", description="echo", input_schema=draft_schema) + async def legacy_echo(self, input, context: ExecutionContext) -> str: + return "ok" + + @module(name="DraftSchema", controllers=[DraftController]) + class DraftModule: + pass + + @mcp_app(module=DraftModule, server=ServerConfig(name="draft-schema")) + class DraftApp: + pass + + DIContainer.reset() + try: + with pytest.raises(UnsupportedJsonSchemaError, match="draft-04"): + asyncio.run(McpApplicationFactory.create(DraftApp)) + finally: + DIContainer.reset() + + def test_valid_2020_12_tool_registers(self): + import asyncio + + from nitrostack import injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.context import ExecutionContext + from nitrostack.core.di import DIContainer + + class EchoIn(BaseModel): + value: str = "" + + @injectable() + class ModernController: + @tool(name="modern_echo", description="echo", input_schema=EchoIn) + async def modern_echo(self, input: EchoIn, context: ExecutionContext) -> str: + return input.value + + @module(name="ModernSchema", controllers=[ModernController]) + class ModernModule: + pass + + @mcp_app(module=ModernModule, server=ServerConfig(name="modern-schema")) + class ModernApp: + pass + + DIContainer.reset() + try: + app = asyncio.run(McpApplicationFactory.create(ModernApp)) + assert "modern_echo" in app._tools + listed = app._tool_input_schema(app._tools["modern_echo"].input_model) + assert listed["$schema"] == JSON_SCHEMA_2020_12_URI + finally: + DIContainer.reset() + + def test_tuple_items_tool_fails_registration(self): + import asyncio + + from nitrostack import injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.context import ExecutionContext + from nitrostack.core.di import DIContainer + + tuple_schema = { + "type": "object", + "properties": { + "pair": {"type": "array", "items": [{"type": "string"}, {"type": "number"}]} + }, + } + + @injectable() + class TupleController: + @tool(name="tuple_echo", description="echo", input_schema=tuple_schema) + async def tuple_echo(self, input, context: ExecutionContext) -> str: + return "ok" + + @module(name="TupleSchema", controllers=[TupleController]) + class TupleModule: + pass + + @mcp_app(module=TupleModule, server=ServerConfig(name="tuple-schema")) + class TupleApp: + pass + + DIContainer.reset() + try: + with pytest.raises(UnsupportedJsonSchemaError, match="items"): + asyncio.run(McpApplicationFactory.create(TupleApp)) + finally: + DIContainer.reset() diff --git a/tests/test_mcp20_deprecated.py b/tests/test_mcp20_deprecated.py new file mode 100644 index 0000000..e06e5e2 --- /dev/null +++ b/tests/test_mcp20_deprecated.py @@ -0,0 +1,306 @@ +"""Deprecated-method policy is the same on every modern entry point.""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys + +import mcp.types as types +import pytest +from mcp import MCPError as McpError +from pydantic import BaseModel, Field +from starlette.testclient import TestClient + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.protocol.deprecated import ( + DEPRECATED_MODERN_METHODS, + deprecated_method_message, + rejects_deprecated_method, +) +from nitrostack.protocol.errors import JsonRpcErrorCode +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.transports.dispatch import IngressContext, StatelessIngressPipeline +from nitrostack.transports.middleware import StatelessTransportMiddleware + +INITIALIZE_BODY = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, +} + +JSON_HEADERS = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", +} + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +def _pipeline(era: str) -> StatelessIngressPipeline: + wire_mode = "reject" if era == "modern" else "stateless" + return StatelessIngressPipeline( + IngressContext( + "srv", + "1.0.0", + MODERN_PROTOCOL_VERSION, + wire_mode=wire_mode, + protocol_era=era, + ) + ) + + +def _middleware(era: str) -> StatelessTransportMiddleware: + return StatelessTransportMiddleware( + lambda scope, receive, send: None, + pipeline=_pipeline(era), + ) + + +def _echo_http_app(era: str): + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name=f"Deprecated{era.title()}", controllers=[EchoController]) + class EchoModule: + pass + + @mcp_app(module=EchoModule, server=ServerConfig(name=f"deprecated-{era}", protocol_era=era)) + class EchoApp: + pass + + app = asyncio.run(McpApplicationFactory.create(EchoApp)) + return app, app.get_combined_app(json_response=True) + + +def test_initialize_is_not_in_retired_method_table(): + assert "initialize" not in DEPRECATED_MODERN_METHODS + assert "notifications/initialized" not in DEPRECATED_MODERN_METHODS + assert deprecated_method_message("initialize") is None + + +@pytest.mark.parametrize("method", sorted(DEPRECATED_MODERN_METHODS)) +def test_retired_methods_error_only_on_modern(method: str): + assert rejects_deprecated_method(method, "modern") is True + assert rejects_deprecated_method(method, "auto") is False + assert rejects_deprecated_method(method, "legacy") is False + + +@pytest.mark.parametrize("entry", ["pipeline", "replay", "http_post", "http_get"]) +def test_modern_initialize_is_method_not_found_on_every_route(entry: str, monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + body = json.dumps(INITIALIZE_BODY).encode() + expected = "Method not found: initialize" + + if entry == "pipeline": + + async def _run(): + status, resp = await _pipeline("modern").handle_post(body, {}) + assert status == 200 + assert resp["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + assert resp["error"]["message"] == expected + + asyncio.run(_run()) + return + + if entry == "replay": + rejected = _middleware("modern")._reject_replay_headers(body, {}) + assert rejected is not None + status, resp = rejected + assert status == 200 + assert resp["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + assert resp["error"]["message"] == expected + return + + _, http_app = _echo_http_app("modern") + try: + if entry == "http_post": + with TestClient(http_app) as client: + response = client.post("/mcp", headers=JSON_HEADERS, json=INITIALIZE_BODY) + assert response.status_code == 200 + assert response.json()["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + assert response.json()["error"]["message"] == expected + return + status, payload = asyncio.run( + _asgi_json_get(http_app, {**JSON_HEADERS, "Mcp-Method": "initialize"}) + ) + assert status == 200 + assert payload["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + assert payload["error"]["message"] == expected + finally: + DIContainer.reset() + + +async def _asgi_json_get(app, headers: dict[str, str]) -> tuple[int, dict]: + status = 0 + chunks: list[bytes] = [] + sent_request = False + finished = asyncio.Event() + + async def receive(): + nonlocal sent_request + if not sent_request: + sent_request = True + return {"type": "http.request", "body": b"", "more_body": False} + await finished.wait() + return {"type": "http.disconnect"} + + async def send(message): + nonlocal status + if message["type"] == "http.response.start": + status = message["status"] + elif message["type"] == "http.response.body": + chunks.append(message.get("body") or b"") + if not message.get("more_body", False): + finished.set() + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": "/mcp", + "raw_path": b"/mcp", + "query_string": b"", + "headers": [ + (key.lower().encode("latin-1"), value.encode("latin-1")) + for key, value in headers.items() + ], + "client": ("testclient", 50000), + "server": ("test", 80), + } + + async def _http(): + await app(scope, receive, send) + await asyncio.wait_for(finished.wait(), timeout=2) + return status, json.loads(b"".join(chunks) or b"{}") + + started = asyncio.Event() + messages: asyncio.Queue = asyncio.Queue() + await messages.put({"type": "lifespan.startup"}) + + async def life_receive(): + return await messages.get() + + async def life_send(message): + if message["type"] == "lifespan.startup.complete": + started.set() + + task = asyncio.create_task(app({"type": "lifespan"}, life_receive, life_send)) + await asyncio.wait_for(started.wait(), timeout=2) + try: + return await _http() + finally: + await messages.put({"type": "lifespan.shutdown"}) + await asyncio.wait_for(task, timeout=2) + + +@pytest.mark.parametrize("entry", ["pipeline", "replay", "http_post"]) +def test_auto_initialize_is_not_the_deprecated_error(entry: str, monkeypatch): + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + body = json.dumps(INITIALIZE_BODY).encode() + + if entry == "pipeline": + + async def _run(): + status, resp = await _pipeline("auto").handle_post(body, {}) + assert status == 200 + assert "error" not in resp + assert resp["result"]["protocolVersion"] == "2025-06-18" + + asyncio.run(_run()) + return + + if entry == "replay": + rejected = _middleware("auto")._reject_replay_headers(body, {}) + assert rejected is None + return + + _, http_app = _echo_http_app("auto") + try: + with TestClient(http_app) as client: + response = client.post("/mcp", headers=JSON_HEADERS, json=INITIALIZE_BODY) + assert response.status_code == 200 + payload = response.json() + assert "error" not in payload + assert payload["result"]["protocolVersion"] == "2025-06-18" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +@pytest.mark.parametrize("method", sorted(DEPRECATED_MODERN_METHODS)) +@pytest.mark.parametrize("entry", ["pipeline", "replay"]) +def test_modern_retired_methods_match_on_post_and_replay(method: str, entry: str): + body = json.dumps({"jsonrpc": "2.0", "id": 7, "method": method, "params": {}}).encode() + headers = {"Mcp-Method": method} + expected = deprecated_method_message(method) + + async def _pipeline_reject(): + status, resp = await _pipeline("modern").handle_post(body, headers) + assert status == 200 + assert resp["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + assert resp["error"]["message"] == expected + + if entry == "pipeline": + asyncio.run(_pipeline_reject()) + return + + rejected = _middleware("modern")._reject_replay_headers(body, headers) + assert rejected is not None + status, resp = rejected + assert status == 200 + assert resp["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + assert resp["error"]["message"] == expected + + +def test_auto_official_handler_still_lists_tasks(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app, _ = _echo_http_app("auto") + try: + handler = app.mcp_server.request_handlers[types.ListTasksRequest] + result = asyncio.run(handler(types.ListTasksRequest(method="tasks/list", params={}))) + assert result.tasks == [] + finally: + DIContainer.reset() + + +def test_modern_official_handler_rejects_tasks_list(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app, _ = _echo_http_app("modern") + try: + handler = app.mcp_server.request_handlers[types.ListTasksRequest] + with pytest.raises(McpError) as exc: + asyncio.run(handler(types.ListTasksRequest(method="tasks/list", params={}))) + assert exc.value.error.code == types.METHOD_NOT_FOUND + assert exc.value.error.message == deprecated_method_message("tasks/list") + finally: + DIContainer.reset() diff --git a/tests/test_mcp20_extensions_cache_observability.py b/tests/test_mcp20_extensions_cache_observability.py new file mode 100644 index 0000000..0443606 --- /dev/null +++ b/tests/test_mcp20_extensions_cache_observability.py @@ -0,0 +1,256 @@ +"""Tests for MCP 2.0 extensions, cache hints, and observability.""" + +import asyncio +import os +import sys + +import mcp.types as types +from pydantic import BaseModel, Field + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, resource, tool +from nitrostack.core.additional_decorators import cache +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.protocol.cache_hints import ( + build_list_endpoint_cache_hint_meta, + resolve_resource_cache_hint_meta, + resolve_tool_cache_hint_meta, +) +from nitrostack.protocol.contracts import MCP_CACHE_HINT_KEY, build_cache_hint_meta +from nitrostack.protocol.discovery import build_discover_result +from nitrostack.protocol.extensions import MCPExtensionId +from nitrostack.protocol.meta import RequestMeta +from nitrostack.protocol.observability import extract_trace_context, trace_context_from_request_meta +from nitrostack.testing import NitroTestingModule + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class TestTraceContext: + def test_extract_bare_keys(self): + meta = { + "traceparent": "00-abc-def-01", + "tracestate": "vendor=1", + "baggage": "userId=alice", + } + trace = extract_trace_context(meta) + assert trace is not None + assert trace.traceparent == "00-abc-def-01" + assert trace.tracestate == "vendor=1" + assert trace.baggage == "userId=alice" + + def test_extract_prefixed_keys(self): + meta = { + "io.modelcontextprotocol/traceparent": "00-prefixed-01", + } + trace = extract_trace_context(meta) + assert trace is not None + assert trace.traceparent == "00-prefixed-01" + + def test_returns_none_when_empty(self): + assert extract_trace_context(None) is None + assert extract_trace_context({}) is None + + def test_from_request_meta(self): + meta = RequestMeta(traceparent="00-from-meta-01", baggage="tenant=acme") + trace = trace_context_from_request_meta(meta) + assert trace is not None + assert trace.traceparent == "00-from-meta-01" + assert trace.baggage == "tenant=acme" + + +class TestCacheHintResolution: + def test_resolve_tool_from_cache_decorator(self): + class _Cfg: + metadata = {} + + @cache(ttl=120) + async def handler(_input, _ctx): + return {} + + meta = resolve_tool_cache_hint_meta(_Cfg(), handler) + assert meta is not None + assert meta[MCP_CACHE_HINT_KEY]["ttlMs"] == 120_000 + assert meta[MCP_CACHE_HINT_KEY]["cacheScope"] == "private" + + def test_resolve_tool_from_explicit_metadata(self): + class _Cfg: + metadata = {"cacheHint": {"ttlMs": 30_000, "cacheScope": "public"}} + + meta = resolve_tool_cache_hint_meta(_Cfg(), None) + assert meta[MCP_CACHE_HINT_KEY]["ttlMs"] == 30_000 + assert meta[MCP_CACHE_HINT_KEY]["cacheScope"] == "public" + + def test_resolve_resource_from_cache_max_age(self): + class _Cfg: + metadata = {"cacheMaxAge": 45} + + meta = resolve_resource_cache_hint_meta(_Cfg()) + assert meta[MCP_CACHE_HINT_KEY]["ttlMs"] == 45_000 + + def test_list_endpoint_cache_hint(self): + meta = build_list_endpoint_cache_hint_meta() + assert meta[MCP_CACHE_HINT_KEY]["ttlMs"] == 60_000 + + +class TestDiscoveryExtensions: + def test_custom_extensions_merge(self): + result = build_discover_result( + server_name="srv", + server_version="1.0.0", + advertise_tasks=False, + custom_extensions={"acme.corp/enterprise_rbac": "2026-07-28"}, + ) + extensions = result["capabilities"]["extensions"] + assert "acme.corp/enterprise_rbac" in extensions + assert MCPExtensionId.TASKS.value not in extensions + + def test_tasks_extension_only_when_advertised(self): + result = build_discover_result( + server_name="srv", + server_version="1.0.0", + advertise_tasks=True, + ) + assert MCPExtensionId.TASKS.value in result["capabilities"]["extensions"] + assert result["resultType"] == "complete" + assert result["ttlMs"] == 60_000 + assert result["cacheScope"] == "private" + + +class TestAppIntegration: + def test_tool_list_includes_cache_hint_from_decorator(self): + @injectable() + class CacheController: + @tool(name="cached_echo", description="cached", input_schema=EchoInput) + @cache(ttl=90) + async def cached_echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="CacheHints", controllers=[CacheController]) + class CacheModule: + pass + + async def _run(): + harness = await NitroTestingModule.create(CacheModule) + handler = harness.app.mcp_server.request_handlers[types.ListToolsRequest] + response = await handler(types.ListToolsRequest(method="tools/list", params={})) + result = response.root + tool = next(t for t in result.tools if t.name == "cached_echo") + hint = (tool.meta or {}).get(MCP_CACHE_HINT_KEY) + assert hint["ttlMs"] == 90_000 + assert result.meta is not None + assert result.meta[MCP_CACHE_HINT_KEY]["ttlMs"] == 60_000 + + asyncio.run(_run()) + + def test_resource_list_includes_cache_hint(self): + @injectable() + class ResourceController: + @resource( + uri="mcp://cache/metrics", + name="metrics", + description="metrics", + metadata={"cacheMaxAge": 10}, + ) + async def metrics(self, context: ExecutionContext) -> str: + return "{}" + + @module(name="CacheResource", controllers=[ResourceController]) + class ResourceModule: + pass + + async def _run(): + harness = await NitroTestingModule.create(ResourceModule) + handler = harness.app.mcp_server.request_handlers[types.ListResourcesRequest] + response = await handler(types.ListResourcesRequest(method="resources/list", params={})) + result = response.root + resource = next(r for r in result.resources if r.name == "metrics") + hint = (resource.meta or {}).get(MCP_CACHE_HINT_KEY) + assert hint["ttlMs"] == 10_000 + + asyncio.run(_run()) + + def test_advertise_tasks_only_when_tool_supports_tasks(self): + @injectable() + class SyncController: + @tool(name="sync_only", description="sync", input_schema=EchoInput, task_support="forbidden") + async def sync_only(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="NoTasksExt", controllers=[SyncController]) + class NoTasksModule: + pass + + @mcp_app(module=NoTasksModule, server=ServerConfig(name="no-tasks")) + class NoTasksApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(NoTasksApp) + assert app._advertise_tasks_extension() is False + + asyncio.run(_run()) + + def test_trace_context_attached_on_tool_call(self): + captured: dict[str, ExecutionContext | None] = {"ctx": None} + + @injectable() + class TraceController: + @tool(name="trace_echo", description="trace", input_schema=EchoInput) + async def trace_echo(self, input: EchoInput, context: ExecutionContext) -> str: + captured["ctx"] = context + return input.value + + @module(name="TraceExt", controllers=[TraceController]) + class TraceModule: + pass + + async def _run(): + from nitrostack.runtime.request_ctx import RequestContext, RequestParamsMeta, request_ctx + + harness = await NitroTestingModule.create(TraceModule) + handler = harness.app.mcp_server.request_handlers[types.CallToolRequest] + meta = { + "traceparent": "00-trace-test-01", + "baggage": "session=abc", + } + request = types.CallToolRequest( + method="tools/call", + params=types.CallToolRequestParams( + name="trace_echo", + arguments={"value": "hi"}, + meta=meta, + ), + ) + token = request_ctx.set( + RequestContext( + request_id="req-trace", + meta=meta, + session=None, + lifespan_context={}, + ) + ) + try: + await handler(request) + finally: + request_ctx.reset(token) + + ctx = captured["ctx"] + assert ctx is not None + assert ctx.trace is not None + assert ctx.trace.traceparent == "00-trace-test-01" + assert ctx.trace.baggage == "session=abc" + + asyncio.run(_run()) diff --git a/tests/test_mcp20_foundation.py b/tests/test_mcp20_foundation.py new file mode 100644 index 0000000..23ee0c3 --- /dev/null +++ b/tests/test_mcp20_foundation.py @@ -0,0 +1,461 @@ +"""Tests for MCP 2.0 architectural foundation.""" + +import pytest + +from nitrostack.core.app import ServerConfig, resolve_http_host +from nitrostack.protocol.constants import ( + LEGACY_SESSION_HEADER, + MAX_CIMD_BYTES, + MAX_SCHEMA_DEPTH, +) +from nitrostack.protocol.extensions import MCPExtensionId +from nitrostack.protocol.layers import RuntimeLayer +from nitrostack.protocol.version import ( + LEGACY_PROTOCOL_VERSION, + MODERN_PROTOCOL_VERSION, + SUPPORTED_PROTOCOL_VERSIONS, + accepts_sessionless_initialize, + http_engine_for_era, + resolve_http_engine, + rejects_legacy_initialize, + needs_modern_engine, + needs_sessionful_engine, + protocol_era_for_wire_mode, + protocol_version_for_era, + resolve_protocol_era, + resolve_protocol_era_resolution, + stateless_for_era, + supported_protocol_versions_for_era, + wire_mode_for_era, +) +from nitrostack.runtime.stateless import ( + DEFAULT_STATELESS_INVARIANTS, + assert_stateless_headers, + has_incoming_session_id, + is_unsupported_protocol_version, + request_protocol_version, +) +from nitrostack.tasks.store import TaskStore +from nitrostack.tasks.types import TaskAccessContext, TERMINAL_TASK_STATUSES + + +class TestProtocolVersion: + def test_modern_protocol_version(self): + assert MODERN_PROTOCOL_VERSION == "2026-07-28" + + def test_legacy_protocol_version(self): + assert LEGACY_PROTOCOL_VERSION == "2025-06-18" + + def test_supported_versions(self): + assert MODERN_PROTOCOL_VERSION in SUPPORTED_PROTOCOL_VERSIONS + + def test_supported_versions_per_era(self): + assert supported_protocol_versions_for_era("modern") == frozenset( + {MODERN_PROTOCOL_VERSION} + ) + assert supported_protocol_versions_for_era("legacy") == frozenset( + {LEGACY_PROTOCOL_VERSION} + ) + assert supported_protocol_versions_for_era("auto") == frozenset( + {MODERN_PROTOCOL_VERSION, LEGACY_PROTOCOL_VERSION} + ) + assert protocol_era_for_wire_mode("reject") == "modern" + assert protocol_era_for_wire_mode("stateless") == "auto" + assert protocol_era_for_wire_mode("sessionful") == "legacy" + + def test_unsupported_protocol_version_contract(self): + assert request_protocol_version("2026-07-28", "2025-06-18") == "2026-07-28" + assert request_protocol_version(None, "2025-06-18") == "2025-06-18" + assert request_protocol_version(None, None) is None + assert is_unsupported_protocol_version("1999-01-01", "auto") is True + assert is_unsupported_protocol_version("2026-07-28", "auto") is False + assert is_unsupported_protocol_version("2025-06-18", "auto") is False + assert is_unsupported_protocol_version("2025-06-18", "modern") is True + assert is_unsupported_protocol_version(None, "modern") is False + + +class TestDefenseInDepthConstants: + def test_cimd_size_limit(self): + assert MAX_CIMD_BYTES == 5120 + + def test_schema_depth_limit(self): + assert MAX_SCHEMA_DEPTH == 64 + + def test_legacy_session_header_name(self): + assert LEGACY_SESSION_HEADER == "Mcp-Session-Id" + + +class TestRuntimeLayers: + def test_all_architecture_layers_defined(self): + layers = {layer.value for layer in RuntimeLayer} + assert layers == { + "transport_security", + "protocol_dispatcher", + "registries", + "task_management", + } + + +class TestExtensions: + def test_canonical_extension_ids(self): + assert MCPExtensionId.APP.value == "io.modelcontextprotocol/app" + assert MCPExtensionId.TASKS.value == "io.modelcontextprotocol/tasks" + + +class TestStatelessInvariants: + def test_default_invariants(self): + inv = DEFAULT_STATELESS_INVARIANTS + assert inv.emit_session_headers is False + assert inv.require_initialize_handshake is False + assert inv.allow_session_stickiness is False + + def test_rejects_session_header_on_response(self): + with pytest.raises(ValueError, match="Mcp-Session-Id"): + assert_stateless_headers({"Mcp-Session-Id": "abc"}) + + def test_allows_responses_without_session_header(self): + assert_stateless_headers({"Content-Type": "application/json"}) + + def test_detects_incoming_session_id(self): + assert has_incoming_session_id({"Mcp-Session-Id": "abc"}) is True + assert has_incoming_session_id({"mcp-session-id": "abc"}) is True + assert has_incoming_session_id({"Content-Type": "application/json"}) is False + assert has_incoming_session_id({"Mcp-Session-Id": " "}) is False + + def test_sessionless_engines_strip_incoming_session_id(self): + from nitrostack.runtime.stateless import sessionless_strips_incoming_session_id + + assert sessionless_strips_incoming_session_id("reject") is True + assert sessionless_strips_incoming_session_id("stateless") is True + assert sessionless_strips_incoming_session_id("sessionful") is False + + +class TestTaskStoreContract: + def test_task_store_is_abstract(self): + with pytest.raises(TypeError): + TaskStore() # type: ignore[abstract] + + def test_task_store_defines_required_methods(self): + required = {"get", "set", "delete", "has", "list", "cleanup_expired"} + assert required.issubset(set(TaskStore.__abstractmethods__)) + + +class TestTaskAccessContext: + def test_wire_alias_population(self): + ctx = TaskAccessContext.model_validate( + {"userId": "u1", "tenantId": "t1", "sessionId": "s1"} + ) + assert ctx.user_id == "u1" + assert ctx.tenant_id == "t1" + assert ctx.session_id == "s1" + + +class TestTerminalTaskStatuses: + def test_terminal_states(self): + assert TERMINAL_TASK_STATUSES == frozenset({"completed", "failed", "cancelled"}) + + +class TestServerConfigMcp20: + def test_default_stateless_config(self): + cfg = ServerConfig(name="test-server") + assert cfg.protocol_version == MODERN_PROTOCOL_VERSION + assert cfg.stateless is False + + +class TestTypescriptCompatibleProtocolEra: + @pytest.mark.parametrize( + "raw,era", + [ + ("auto", "auto"), + ("both", "auto"), + ("dual", "auto"), + ("dual-spec", "auto"), + ("AUTO", "auto"), + ("2026-07-28", "modern"), + ("2026", "modern"), + ("modern", "modern"), + ("latest", "modern"), + ("Modern", "modern"), + ("legacy", "legacy"), + ("2025-06-18", "legacy"), + ("2025-11-25", "legacy"), + ("2025", "legacy"), + ("", "auto"), + (None, "auto"), + ("unknown-era", "auto"), + ("mcp-2026", "auto"), + ("2026-06-18", "auto"), + ], + ) + def test_resolve_protocol_era(self, raw, era, monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + assert resolve_protocol_era(raw) == era + + def test_reads_nitro_mcp_protocol_version_env(self, monkeypatch): + monkeypatch.delenv("MCP_STATELESS", raising=False) + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "2026-07-28") + assert resolve_protocol_era() == "modern" + assert stateless_for_era(resolve_protocol_era()) is True + assert protocol_version_for_era("modern") == MODERN_PROTOCOL_VERSION + + def test_unknown_env_alias_is_auto_not_modern(self, monkeypatch): + monkeypatch.delenv("MCP_STATELESS", raising=False) + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "unknown-era") + assert resolve_protocol_era() == "auto" + + def test_legacy_era_is_sessionful(self): + assert stateless_for_era("legacy") is False + assert protocol_version_for_era("legacy") == LEGACY_PROTOCOL_VERSION + + def test_unset_era_defaults_to_auto_without_forcing_stateless(self, monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + resolution = resolve_protocol_era_resolution() + assert resolution.era == "auto" + assert resolution.source == "default" + assert resolution.log_line() == "protocol era=auto (source=default)" + assert resolve_protocol_era() == "auto" + assert stateless_for_era("auto") is None + assert protocol_version_for_era("auto") == MODERN_PROTOCOL_VERSION + + def test_mcp_stateless_true_forces_modern(self, monkeypatch): + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "legacy") + resolution = resolve_protocol_era_resolution(stateless_override="true") + assert resolution.era == "modern" + assert resolution.source == "mcp_stateless" + assert resolve_protocol_era(stateless_override="true") == "modern" + + def test_mcp_stateless_false_forces_legacy(self, monkeypatch): + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "modern") + monkeypatch.setenv("MCP_STATELESS", "false") + resolution = resolve_protocol_era_resolution() + assert resolution.era == "legacy" + assert resolution.source == "mcp_stateless" + assert resolution.log_line() == "protocol era=legacy (source=mcp_stateless)" + assert resolve_protocol_era() == "legacy" + assert stateless_for_era("legacy") is False + + def test_auto_is_not_modern_and_does_not_force_stateless(self): + assert resolve_protocol_era("auto") != resolve_protocol_era("modern") + assert stateless_for_era("auto") is None + assert stateless_for_era("modern") is True + assert wire_mode_for_era("auto") == "stateless" + assert wire_mode_for_era("modern") == "reject" + assert wire_mode_for_era("legacy") == "sessionful" + assert needs_modern_engine("auto") is True + assert needs_modern_engine("modern") is True + assert needs_modern_engine("legacy") is False + assert needs_sessionful_engine("auto") is False + assert needs_sessionful_engine("legacy") is True + assert http_engine_for_era("auto") == "sessionless" + assert http_engine_for_era("modern") == "sessionless" + assert http_engine_for_era("legacy") == "sessionful" + assert resolve_http_engine("auto", stateless=False) == "sessionless" + assert resolve_http_engine("modern", http_engine="sessionful") == "sessionless" + assert resolve_http_engine("legacy") == "sessionful" + assert resolve_http_engine("legacy", stateless=True) == "sessionless" + assert accepts_sessionless_initialize("auto") is True + assert accepts_sessionless_initialize("modern") is False + assert accepts_sessionless_initialize("legacy") is False + assert rejects_legacy_initialize("modern") is True + assert rejects_legacy_initialize("auto") is False + assert rejects_legacy_initialize("legacy") is False + + def test_config_era_used_when_env_unset(self, monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + resolution = resolve_protocol_era_resolution(config_value="legacy") + assert resolution.era == "legacy" + assert resolution.source == "config" + assert resolve_protocol_era(config_value="legacy") == "legacy" + assert resolve_protocol_era(config_value="modern") == "modern" + assert resolve_protocol_era(config_value="2026-07-28") == "modern" + + def test_env_wins_over_config_era(self, monkeypatch): + monkeypatch.delenv("MCP_STATELESS", raising=False) + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + resolution = resolve_protocol_era_resolution(config_value="legacy") + assert resolution.era == "auto" + assert resolution.source == "env" + assert resolve_protocol_era(config_value="legacy") == "auto" + + def test_stateless_override_wins_over_config_era(self, monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.setenv("MCP_STATELESS", "true") + assert resolve_protocol_era(config_value="legacy") == "modern" + + def test_invalid_config_era_matches_invalid_env(self, monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + assert resolve_protocol_era("unknown-era") == "auto" + assert resolve_protocol_era(config_value="unknown-era") == "auto" + + def test_server_config_protocol_era_default_is_unset(self): + cfg = ServerConfig(name="test-server") + assert cfg.protocol_era is None + resolution = resolve_protocol_era_resolution(config_value=cfg.protocol_era) + assert resolution.era == "auto" + assert resolution.source == "default" + assert resolve_protocol_era(config_value=cfg.protocol_era) == "auto" + + +class TestTypescriptCompatibleHost: + def test_host_defaults_to_loopback(self, monkeypatch): + monkeypatch.delenv("HOST", raising=False) + assert resolve_http_host() == "127.0.0.1" + + def test_empty_host_defaults_to_loopback(self, monkeypatch): + monkeypatch.setenv("HOST", " ") + assert resolve_http_host() == "127.0.0.1" + + def test_host_all_interfaces_still_available(self, monkeypatch): + monkeypatch.setenv("HOST", "0.0.0.0") + assert resolve_http_host() == "0.0.0.0" + + def test_host_env_matches_typescript(self, monkeypatch): + monkeypatch.setenv("HOST", "localhost") + assert resolve_http_host() == "localhost" + + +class TestTrustedReverseProxy: + def test_default_ignores_forwarded_host(self, monkeypatch): + from nitrostack.transports.proxy import public_origin, request_host_for_cimd + + monkeypatch.delenv("TRUSTED_PROXIES", raising=False) + monkeypatch.delenv("MCP_TRUSTED_PROXIES", raising=False) + headers = { + "Host": "internal:3000", + "X-Forwarded-Host": "mcp.example.com", + "X-Forwarded-Proto": "https", + "X-Forwarded-For": "10.0.0.5", + } + assert public_origin(headers, peer="8.8.8.8") == "http://internal:3000" + assert request_host_for_cimd(headers, peer="8.8.8.8") == "internal" + + def test_trusted_proxy_honors_forwarded_host(self, monkeypatch): + from nitrostack.transports.proxy import public_url, request_host_for_cimd + + monkeypatch.setenv("TRUSTED_PROXIES", "10.0.0.5") + headers = { + "Host": "internal:3000", + "X-Forwarded-Host": "mcp.example.com", + "X-Forwarded-Proto": "https", + } + assert public_url(headers, peer="10.0.0.5", path="/mcp") == "https://mcp.example.com/mcp" + assert request_host_for_cimd(headers, peer="10.0.0.5") == "mcp.example.com" + + def test_cidr_allow_list(self, monkeypatch): + from nitrostack.transports.proxy import peer_is_trusted + + monkeypatch.setenv("TRUSTED_PROXIES", "10.0.0.0/8") + assert peer_is_trusted("10.1.2.3") is True + assert peer_is_trusted("11.0.0.1") is False + + def test_x_forwarded_for_does_not_grant_trust(self, monkeypatch): + from nitrostack.transports.proxy import public_origin + + monkeypatch.setenv("TRUSTED_PROXIES", "10.0.0.5") + headers = { + "Host": "internal:3000", + "X-Forwarded-For": "10.0.0.5", + "X-Forwarded-Host": "evil.example", + "X-Forwarded-Proto": "https", + } + assert public_origin(headers, peer="8.8.8.8") == "http://internal:3000" + + def test_untrusted_forwarded_host_cannot_rebind_cimd(self, monkeypatch): + from nitrostack.auth.cimd import cimd_host_matches_request, request_host_for_cimd + + monkeypatch.delenv("TRUSTED_PROXIES", raising=False) + headers = { + "Host": "mcp.nitrostack.io", + "X-Forwarded-Host": "evil.example", + } + url = "https://evil.example/oauth/client-metadata.json" + assert request_host_for_cimd(headers, peer="8.8.8.8") == "mcp.nitrostack.io" + assert cimd_host_matches_request(url, headers, peer="8.8.8.8") is False + + def test_trusted_proxy_cimd_host_matches_forwarded_host(self, monkeypatch): + from nitrostack.auth.cimd import cimd_host_matches_request, request_host_for_cimd + + monkeypatch.setenv("TRUSTED_PROXIES", "10.0.0.5") + headers = { + "Host": "internal:3000", + "X-Forwarded-Host": "app.nitrostack.io", + } + url = "https://app.nitrostack.io/oauth/client-metadata.json" + assert request_host_for_cimd(headers, peer="10.0.0.5") == "app.nitrostack.io" + assert cimd_host_matches_request(url, headers, peer="10.0.0.5") is True + + +class TestOfficialMcpV2: + def test_official_mcp_v2_server_import(self): + from mcp.server import MCPServer + + assert MCPServer is not None + + def test_tools_call_uses_v2_not_sidecar(self, monkeypatch): + import asyncio + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> dict: + return {"value": input.value} + + @module(name="V2Wire", controllers=[EchoController]) + class V2Module: + pass + + @mcp_app(module=V2Module, server=ServerConfig(name="v2-wire")) + class V2App: + pass + + app = asyncio.run(McpApplicationFactory.create(V2App)) + http_app = app.get_combined_app(json_response=True) + + with TestClient(http_app) as client: + resp = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + "MCP-Protocol-Version": "2026-07-28", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "echo", + "arguments": {"value": "v2"}, + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {}, + }, + }, + }, + ) + assert resp.status_code == 200, resp.text + result = resp.json()["result"] + payload = result.get("structuredContent") or {} + assert payload.get("value") == "v2" + finally: + DIContainer.reset() diff --git a/tests/test_mcp20_jsonrpc_wire.py b/tests/test_mcp20_jsonrpc_wire.py new file mode 100644 index 0000000..89e81aa --- /dev/null +++ b/tests/test_mcp20_jsonrpc_wire.py @@ -0,0 +1,490 @@ +"""Tests for MCP 2.0 JSON-RPC wire specification.""" + +import asyncio +import json + +import pytest + +from nitrostack.core.errors import ResourceNotFoundError, ToolExecutionError, ValidationError +from nitrostack.protocol.deprecated import deprecated_method_message +from nitrostack.protocol.errors import ERROR_CODE_MESSAGES, JsonRpcErrorCode +from nitrostack.protocol.jsonrpc import ( + HEADER_BODY_MISMATCH, + PARSE_ERROR, + HeaderBodyMismatchError, + InvalidParamsError, + InvalidRequestError, + JsonRpcParseError, + build_tool_error_result, + map_exception_to_jsonrpc, + parse_jsonrpc_request, + validate_header_body_name, + validate_header_body_method, + validate_required_mcp_name, +) +from nitrostack.core.context import AuthContext, ExecutionContext +from nitrostack.protocol.meta import ( + bind_request_envelope, + envelope_identity_is_ignored, + envelope_protocol_version, + extract_request_meta, + split_params_and_meta, +) +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.transports.dispatch import IngressContext, StatelessIngressPipeline + + +class TestErrorCodes: + def test_standard_error_codes(self): + assert int(JsonRpcErrorCode.PARSE_ERROR) == -32700 + assert int(JsonRpcErrorCode.INVALID_REQUEST) == -32600 + assert int(JsonRpcErrorCode.METHOD_NOT_FOUND) == -32601 + assert int(JsonRpcErrorCode.INVALID_PARAMS) == -32602 + assert int(JsonRpcErrorCode.INTERNAL_ERROR) == -32603 + assert int(JsonRpcErrorCode.HEADER_BODY_MISMATCH) == -32020 + assert int(JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION) == -32022 + assert ( + ERROR_CODE_MESSAGES[JsonRpcErrorCode.UNSUPPORTED_PROTOCOL_VERSION] + == "Unsupported protocol version" + ) + + +class TestMetaEnvelope: + def test_extract_reverse_dns_keys(self): + params = { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": {"name": "Claude Desktop", "version": "1.5.0"}, + "io.modelcontextprotocol/traceparent": "00-abc", + }, + "name": "demo", + } + business, meta = split_params_and_meta(params) + assert business == {"name": "demo"} + assert meta.protocol_version == "2026-07-28" + assert meta.client_info["name"] == "Claude Desktop" + assert meta.traceparent == "00-abc" + + def test_extract_bare_keys(self): + meta = extract_request_meta( + {"_meta": {"traceparent": "00-bare", "clientInfo": {"name": "x"}}} + ) + assert meta.traceparent == "00-bare" + assert meta.client_info == {"name": "x"} + + def test_extracts_nested_mcp_protocol_version(self): + meta = extract_request_meta( + {"_meta": {"mcp": {"protocolVersion": "2025-06-18"}, "trace": {"id": "t"}}} + ) + assert meta.protocol_version == "2025-06-18" + assert envelope_protocol_version(meta) == "2025-06-18" + + def test_parse_request_strips_meta(self): + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "calc", + "_meta": {"io.modelcontextprotocol/protocolVersion": "2026-07-28"}, + }, + } + ).encode() + req = parse_jsonrpc_request(body) + assert req.params == {"name": "calc"} + assert req.meta.protocol_version == "2026-07-28" + + def test_bind_maps_trace_and_header_protocol_version(self): + envelope = bind_request_envelope( + raw_meta={"trace": {"id": "span-1"}, "protocolVersion": "2025-06-18"}, + mcp_headers={"MCP-Protocol-Version": "2026-07-28", "Mcp-Method": "tools/call"}, + ) + assert envelope.meta.trace == {"id": "span-1"} + assert envelope.protocol_version == "2026-07-28" + assert envelope.mcp_headers["Mcp-Method"] == "tools/call" + + def test_unsigned_identity_stays_off_context_user(self): + envelope = bind_request_envelope( + raw_meta={"userId": "spoofed", "tenantId": "evil", "trace": {"id": "t"}}, + ) + assert envelope.meta.trace == {"id": "t"} + assert envelope.meta.raw["userId"] == "spoofed" + assert envelope_identity_is_ignored(envelope.meta.raw) is True + + ctx = ExecutionContext( + request_id="env-1", + protocol_version=envelope.protocol_version, + rpc_meta=envelope.meta, + mcp_headers=dict(envelope.mcp_headers), + ) + assert ctx.user is None + assert ctx.rpc_meta.raw["userId"] == "spoofed" + + ctx.auth = AuthContext(subject="alice") + assert ctx.user == "alice" + + def test_apply_request_envelope_ignores_spoofed_userid(self): + from types import SimpleNamespace + + from nitrostack.runtime.request_ctx import RequestContext, RequestParamsMeta + from mcp import types + + from nitrostack.core.app import _apply_request_envelope + + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta.model_validate( + { + "trace": {"id": "span-2"}, + "userId": "spoofed", + "protocolVersion": "2025-11-25", + } + ), + session=None, + lifespan_context=None, + request=SimpleNamespace( + headers={ + "MCP-Protocol-Version": "2026-07-28", + "Mcp-Method": "tools/call", + "Authorization": "Bearer ignore-me", + } + ), + ) + ctx = ExecutionContext(request_id="env-http") + _apply_request_envelope(ctx, rc) + assert ctx.protocol_version == "2026-07-28" + assert ctx.rpc_meta is not None + assert ctx.rpc_meta.trace == {"id": "span-2"} + assert ctx.rpc_meta.raw["userId"] == "spoofed" + assert ctx.user is None + assert ctx.mcp_headers["MCP-Protocol-Version"] == "2026-07-28" + assert ctx.mcp_headers["Mcp-Method"] == "tools/call" + assert "Authorization" not in ctx.mcp_headers + assert "authorization" not in {key.lower() for key in ctx.mcp_headers} + + def test_apply_request_envelope_sets_verified_jwt_user(self): + from types import SimpleNamespace + + from nitrostack.runtime.request_ctx import RequestContext, RequestParamsMeta + from mcp import types + + from nitrostack.auth.jwt import JWTService + from nitrostack.core.app import _apply_request_envelope + from nitrostack.core.di import DIContainer + + DIContainer.reset() + try: + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + token = jwt.create_token({"sub": "alice", "tenant_id": "acme"}) + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta.model_validate({"userId": "eve"}), + session=None, + lifespan_context=None, + request=SimpleNamespace(headers={"authorization": f"Bearer {token}"}), + ) + ctx = ExecutionContext(request_id="env-jwt") + _apply_request_envelope(ctx, rc) + assert ctx.user == "alice" + assert ctx.auth is not None + assert ctx.auth.claims["tenant_id"] == "acme" + assert ctx.rpc_meta.raw["userId"] == "eve" + finally: + DIContainer.reset() + + +class TestHeaderBodyMismatch: + def test_mcp_method_mismatch(self): + with pytest.raises(HeaderBodyMismatchError) as exc: + validate_header_body_method("tools/list", "tools/call") + assert int(exc.value.code) == HEADER_BODY_MISMATCH + + def test_mcp_name_mismatch(self): + with pytest.raises(HeaderBodyMismatchError): + validate_header_body_name("get_weather", "get_forecast") + + def test_optional_name_allows_missing_header(self): + validate_header_body_name(None, "echo") + + def test_required_mcp_name_rejects_missing_header(self): + with pytest.raises(HeaderBodyMismatchError) as exc: + validate_required_mcp_name(None, "echo") + assert int(exc.value.code) == HEADER_BODY_MISMATCH + assert "required" in exc.value.message + + def test_required_mcp_name_rejects_mismatch(self): + with pytest.raises(HeaderBodyMismatchError) as exc: + validate_required_mcp_name("foo", "bar") + assert int(exc.value.code) == HEADER_BODY_MISMATCH + + def test_required_mcp_name_accepts_exact_match(self): + validate_required_mcp_name("echo", "echo") + + def test_required_mcp_method_rejects_missing_header(self): + from nitrostack.protocol.jsonrpc import validate_required_mcp_method + + with pytest.raises(HeaderBodyMismatchError) as exc: + validate_required_mcp_method(None, "tools/call") + assert int(exc.value.code) == HEADER_BODY_MISMATCH + assert "required" in exc.value.message + + def test_required_mcp_method_rejects_mismatch(self): + from nitrostack.protocol.jsonrpc import validate_required_mcp_method + + with pytest.raises(HeaderBodyMismatchError): + validate_required_mcp_method("tools/list", "tools/call") + + def test_required_mcp_method_accepts_exact_match(self): + from nitrostack.protocol.jsonrpc import validate_required_mcp_method + + validate_required_mcp_method("tools/call", "tools/call") + + def test_protocol_version_header_meta_mismatch(self): + from nitrostack.protocol.jsonrpc import validate_protocol_version_header_meta + + with pytest.raises(HeaderBodyMismatchError) as exc: + validate_protocol_version_header_meta("2026-07-28", "2025-06-18") + assert int(exc.value.code) == HEADER_BODY_MISMATCH + + def test_protocol_version_header_only_is_allowed(self): + from nitrostack.protocol.jsonrpc import validate_protocol_version_header_meta + + validate_protocol_version_header_meta("2026-07-28", None) + validate_protocol_version_header_meta(None, "2025-06-18") + validate_protocol_version_header_meta("2026-07-28", "2026-07-28") + + def test_pipeline_rejects_header_and_meta_mismatch(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "ping", + "params": {"_meta": {"mcp": {"protocolVersion": "2025-06-18"}}}, + } + ).encode() + status, resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "2026-07-28"} + ) + assert status == 400 + assert resp["error"]["code"] == HEADER_BODY_MISMATCH + + asyncio.run(_run()) + + def test_pipeline_header_only_proceeds(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "2026-07-28"} + ) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + def test_pipeline_ignores_unrelated_meta_keys(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "ping", + "params": {"_meta": {"trace": {"id": "span"}, "userId": "eve"}}, + } + ).encode() + status, resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "2026-07-28"} + ) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + +class TestUnsupportedProtocolVersion: + def test_unknown_header_is_rejected(self): + from nitrostack.protocol.jsonrpc import ( + UNSUPPORTED_PROTOCOL_VERSION, + UnsupportedProtocolVersionError, + validate_supported_protocol_version, + ) + + with pytest.raises(UnsupportedProtocolVersionError) as exc: + validate_supported_protocol_version("1999-01-01", {"2026-07-28"}) + assert int(exc.value.code) == UNSUPPORTED_PROTOCOL_VERSION + assert str(exc.value) == "Unsupported protocol version" + + def test_pipeline_rejects_unknown_header(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "1999-01-01"} + ) + assert status == 400 + assert resp["error"]["code"] == -32022 + assert resp["error"]["message"] == "Unsupported protocol version" + + asyncio.run(_run()) + + def test_pipeline_accepts_supported_header(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping"}).encode() + modern, modern_resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "2026-07-28"} + ) + legacy, legacy_resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "2025-06-18"} + ) + absent, absent_resp = await pipeline.handle_post(body, {}) + assert modern == 200 + assert modern_resp["result"] == {} + assert legacy == 200 + assert legacy_resp["result"] == {} + assert absent == 200 + assert absent_resp["result"] == {} + + asyncio.run(_run()) + + def test_pipeline_rejects_envelope_only_unknown_version(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "ping", + "params": {"_meta": {"mcp": {"protocolVersion": "1999-01-01"}}}, + } + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 400 + assert resp["error"]["code"] == -32022 + + asyncio.run(_run()) + + def test_modern_rejects_legacy_dated_header(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "2025-06-18", "Mcp-Method": "ping"} + ) + assert status == 400 + assert resp["error"]["code"] == -32022 + + asyncio.run(_run()) + + +class TestToolVsProtocolErrors: + def test_tool_error_result_uses_is_error_flag(self): + result = build_tool_error_result("User not found") + assert result["isError"] is True + assert result["content"][0]["text"] == "User not found" + + def test_resource_not_found_maps_to_invalid_params(self): + resp = map_exception_to_jsonrpc(ResourceNotFoundError("missing uri"), "1") + assert resp["error"]["code"] == int(JsonRpcErrorCode.INVALID_PARAMS) + + def test_validation_error_maps_to_invalid_params(self): + resp = map_exception_to_jsonrpc(ValidationError("bad schema"), 2) + assert resp["error"]["code"] == -32602 + + def test_tool_execution_error_maps_to_result_not_rpc_error(self): + resp = map_exception_to_jsonrpc(ToolExecutionError("payment declined"), 3) + assert "error" not in resp + assert resp["result"]["isError"] is True + + +class TestDeprecatedMethods: + @pytest.mark.parametrize( + "method", + ["tasks/result", "tasks/list", "resources/subscribe", "logging/setLevel"], + ) + def test_deprecated_methods_have_messages(self, method): + assert deprecated_method_message(method) is not None + + def test_pipeline_rejects_tasks_list(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext( + "srv", + "1.0.0", + MODERN_PROTOCOL_VERSION, + wire_mode="reject", + protocol_era="modern", + ) + ) + body = json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "tasks/list", "params": {}} + ).encode() + status, resp = await pipeline.handle_post(body, {"Mcp-Method": "tasks/list"}) + assert status == 200 + assert resp["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + + asyncio.run(_run()) + + def test_pipeline_auto_forwards_tasks_list(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext( + "srv", + "1.0.0", + MODERN_PROTOCOL_VERSION, + wire_mode="stateless", + protocol_era="auto", + ) + ) + body = json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "tasks/list", "params": {}} + ).encode() + assert await pipeline.handle_post(body, {"Mcp-Method": "tasks/list"}) is None + + asyncio.run(_run()) + + +class TestJsonRpcParseErrors: + def test_invalid_json(self): + with pytest.raises(JsonRpcParseError) as exc: + parse_jsonrpc_request(b"{") + assert int(exc.value.code) == PARSE_ERROR + + def test_non_object_is_invalid_request(self): + with pytest.raises(InvalidRequestError) as exc: + parse_jsonrpc_request(b"[]") + assert int(exc.value.code) == int(JsonRpcErrorCode.INVALID_REQUEST) + + def test_wrong_jsonrpc_version_is_invalid_request(self): + with pytest.raises(InvalidRequestError): + parse_jsonrpc_request(b'{"jsonrpc":"1.0","id":1,"method":"ping"}') + + def test_missing_method_is_invalid_request(self): + with pytest.raises(InvalidRequestError): + parse_jsonrpc_request(b'{"jsonrpc":"2.0","id":1}') + + def test_wire_error_to_response(self): + err = InvalidParamsError("amount must be positive", data={"param": "amount"}) + resp = err.to_response("req-12345") + assert resp["id"] == "req-12345" + assert resp["error"]["code"] == -32602 + assert resp["error"]["data"]["param"] == "amount" diff --git a/tests/test_mcp20_mrtr.py b/tests/test_mcp20_mrtr.py new file mode 100644 index 0000000..505d541 --- /dev/null +++ b/tests/test_mcp20_mrtr.py @@ -0,0 +1,146 @@ +"""Tests for MCP 2.0 MRTR multi round-trip requests.""" + +import asyncio +import os +import sys + +from pydantic import BaseModel, Field + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import injectable +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app, parse_tool_input +from nitrostack.core.context import ExecutionContext +from nitrostack.core.decorators import tool +from nitrostack.core.di import DIContainer +from nitrostack.core.module import module +from nitrostack.protocol.mrtr import ( + InputRequest, + accepted_content, + build_input_required_jsonrpc_result, + input_required, + split_mrtr_tool_params, +) + + +class TransferInput(BaseModel): + amount: float = Field(description="Amount to transfer") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class TestMrtrHelpers: + def test_accepted_content_returns_none_when_missing(self): + assert accepted_content(None, "confirm") is None + assert accepted_content({}, "confirm") is None + + def test_accepted_content_returns_value(self): + assert accepted_content({"confirm": "DELETE"}, "confirm") == "DELETE" + + def test_input_required_wire_shape(self): + result = input_required( + requests=[ + InputRequest( + id="otp_code", + message="Enter 6-digit code", + schema={"type": "string", "pattern": "^[0-9]{6}$"}, + ) + ], + request_state={"transactionId": "tx_89712398", "stage": "awaiting_2fa"}, + ) + wire = result.to_wire_dict() + assert wire["resultType"] == "input_required" + assert wire["inputRequests"][0]["id"] == "otp_code" + assert wire["requestState"]["transactionId"] == "tx_89712398" + + def test_split_mrtr_tool_params(self): + params = { + "name": "transfer_funds", + "arguments": {"recipient": "alice@example.com", "amount": 500}, + "inputResponses": {"otp_code": "481920"}, + "requestState": {"transactionId": "tx_89712398", "stage": "awaiting_2fa"}, + } + arguments, responses, state = split_mrtr_tool_params(params) + assert arguments["amount"] == 500 + assert responses["otp_code"] == "481920" + assert state["stage"] == "awaiting_2fa" + + def test_jsonrpc_envelope(self): + result = input_required( + requests=[InputRequest(id="confirm_delete", message="Type DELETE")], + request_state={"step": 1, "db": "prod_users"}, + ) + envelope = build_input_required_jsonrpc_result("req-99", result) + assert envelope["id"] == "req-99" + assert envelope["result"]["resultType"] == "input_required" + + +class TestMrtrToolExecution: + def test_tool_returns_input_required_result(self): + @injectable() + class MrtrController: + @tool(name="delete_database", description="delete", input_schema=TransferInput) + async def delete_database(self, input: TransferInput, context: ExecutionContext): + confirmation = accepted_content(context.input_responses, "confirm_delete") + if confirmation is None: + return input_required( + requests=[ + InputRequest( + id="confirm_delete", + message="Type DELETE to confirm", + schema={"type": "string"}, + ) + ], + request_state={"step": 1, "db": "prod_users"}, + ) + if confirmation == "DELETE": + return {"status": "dropped", "db": "prod_users"} + return {"status": "cancelled"} + + @module(name="MrtrModule", controllers=[MrtrController]) + class MrtrModule: + pass + + @mcp_app(module=MrtrModule, server=ServerConfig(name="mrtr")) + class MrtrApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(MrtrApp) + + first = await app._call_tool("delete_database", {"amount": 0}) + assert first.structured_content["resultType"] == "input_required" + assert first.structured_content["inputRequests"][0]["id"] == "confirm_delete" + + second = await app._call_tool( + "delete_database", + { + "amount": 0, + "inputResponses": {"confirm_delete": "DELETE"}, + "requestState": {"step": 1, "db": "prod_users"}, + }, + ) + assert second.structured_content["status"] == "dropped" + + asyncio.run(_run()) + + def test_split_mrtr_preserves_tool_arguments_for_validation(self): + arguments, responses, state = split_mrtr_tool_params( + { + "arguments": { + "amount": 12.5, + "inputResponses": {"confirm": True}, + "requestState": {"stage": 2}, + } + } + ) + model = parse_tool_input(TransferInput, arguments) + assert model.amount == 12.5 + assert responses["confirm"] is True + assert state["stage"] == 2 diff --git a/tests/test_mcp20_oauth_cimd.py b/tests/test_mcp20_oauth_cimd.py new file mode 100644 index 0000000..b269faa --- /dev/null +++ b/tests/test_mcp20_oauth_cimd.py @@ -0,0 +1,268 @@ +"""Tests for MCP 2.0 OAuth 2.1, CIMD, and SSRF security.""" + +import asyncio +import json +import socket +from unittest.mock import patch + +import pytest + +from nitrostack.auth.cimd import ( + CimdFetchError, + CimdValidationError, + _fetch_cimd_bytes, + assert_safe_fetch_target, + cimd_host_matches_request, + cimd_peer_is_acceptable, + is_blocked_ip, + request_host_for_cimd, + resolve_cimd, + validate_client_identifier_url, +) +from nitrostack.auth.oauth_module import apply_cimd_to_registration_body, build_protected_resource_metadata +from nitrostack.auth.oauth_security import ( + AuthorizationIssuerMismatchError, + validate_authorization_iss, +) +from nitrostack.protocol.constants import MAX_CIMD_BYTES + + +class _StubOAuthService: + resource_uri = "https://mcp.nitrostack.io/mcp" + authorization_servers = ["https://auth.nitrostack.io"] + scopes_supported = ["mcp:tools", "mcp:resources", "mcp:prompts"] + + +class TestClientIdentifierUrlValidation: + def test_accepts_https_with_path(self): + url = "https://app.nitrostack.io/oauth/client-metadata.json" + assert validate_client_identifier_url(url) == url + + def test_rejects_bare_domain(self): + with pytest.raises(CimdValidationError, match="non-root path"): + validate_client_identifier_url("https://example.com/") + + def test_rejects_userinfo(self): + with pytest.raises(CimdValidationError, match="userinfo"): + validate_client_identifier_url("https://user:pass@example.com/oauth/client.json") + + def test_rejects_fragment(self): + with pytest.raises(CimdValidationError, match="fragment"): + validate_client_identifier_url("https://example.com/oauth/client.json#x") + + def test_rejects_path_traversal(self): + with pytest.raises(CimdValidationError, match="\\.\\."): + validate_client_identifier_url("https://example.com/oauth/../client.json") + + def test_allows_loopback_http_when_enabled(self): + url = "http://127.0.0.1/oauth/client-metadata.json" + assert validate_client_identifier_url(url, allow_loopback=True) == url + + +class TestBlockedIpRanges: + @pytest.mark.parametrize( + "ip", + [ + "127.0.0.1", + "10.0.0.1", + "169.254.169.254", + "192.168.1.10", + "::1", + "fc00::1", + ], + ) + def test_blocks_special_use_addresses(self, ip): + assert is_blocked_ip(ip) is True + + def test_allows_public_ipv4(self): + assert is_blocked_ip("8.8.8.8") is False + + +class TestCimdResolver: + def test_blocks_private_dns_resolution(self): + async def _run(): + url = "https://metadata.example.com/oauth/client.json" + with patch( + "socket.getaddrinfo", + return_value=[(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.5", 0))], + ): + with pytest.raises(CimdFetchError, match="blocked IP"): + await assert_safe_fetch_target(url) + + asyncio.run(_run()) + + def test_rejects_redirects(self): + async def _run(): + url = "https://app.nitrostack.io/oauth/client-metadata.json" + + def fake_fetch(_url, *, timeout_sec, pinned_ip=None): + raise CimdFetchError("HTTP redirects are not allowed for CIMD fetch") + + with patch("nitrostack.auth.cimd.assert_safe_fetch_target", return_value=None): + with patch("nitrostack.auth.cimd._fetch_cimd_bytes", side_effect=fake_fetch): + with pytest.raises(CimdFetchError, match="redirect"): + await resolve_cimd(url) + + asyncio.run(_run()) + + def test_rejects_oversized_payload(self): + async def _run(): + url = "https://app.nitrostack.io/oauth/client-metadata.json" + oversized = b"x" * (MAX_CIMD_BYTES + 1) + + with patch("nitrostack.auth.cimd.assert_safe_fetch_target", return_value=None): + with patch("nitrostack.auth.cimd._fetch_cimd_bytes", return_value=oversized): + with pytest.raises(CimdFetchError, match="maximum size"): + await resolve_cimd(url) + + asyncio.run(_run()) + + def test_rejects_client_id_mismatch(self): + async def _run(): + url = "https://app.nitrostack.io/oauth/client-metadata.json" + body = json.dumps( + { + "client_id": "https://evil.example/oauth/client-metadata.json", + "redirect_uris": ["https://app.nitrostack.io/cb"], + } + ).encode() + + with patch("nitrostack.auth.cimd.assert_safe_fetch_target", return_value=None): + with patch("nitrostack.auth.cimd._fetch_cimd_bytes", return_value=body): + with pytest.raises(CimdValidationError, match="does not match"): + await resolve_cimd(url) + + asyncio.run(_run()) + + def test_returns_document_when_valid(self): + async def _run(): + url = "https://app.nitrostack.io/oauth/client-metadata.json" + body = json.dumps( + { + "client_id": url, + "client_name": "NitroStudio", + "redirect_uris": ["https://app.nitrostack.io/auth/callback"], + } + ).encode() + + with patch("nitrostack.auth.cimd.assert_safe_fetch_target", return_value=None): + with patch("nitrostack.auth.cimd._fetch_cimd_bytes", return_value=body): + doc = await resolve_cimd(url) + assert doc["client_name"] == "NitroStudio" + + asyncio.run(_run()) + + + def test_fetch_connects_to_pinned_ip(self): + connected: dict[str, object] = {} + + def fake_create(address, timeout=None): + connected["addr"] = address + raise OSError("stop before handshake") + + with patch("nitrostack.auth.cimd.socket.create_connection", side_effect=fake_create): + with pytest.raises(CimdFetchError): + _fetch_cimd_bytes( + "https://app.nitrostack.io/oauth/client-metadata.json", + timeout_sec=1.0, + pinned_ip="8.8.8.8", + ) + assert connected["addr"][0] == "8.8.8.8" + + +class TestCimdTrustedProxyHost: + def test_direct_peer_must_match_pin(self, monkeypatch): + monkeypatch.delenv("TRUSTED_PROXIES", raising=False) + assert cimd_peer_is_acceptable("8.8.8.8", "8.8.8.8") is True + assert cimd_peer_is_acceptable("1.2.3.4", "8.8.8.8") is False + + def test_trusted_proxy_peer_is_acceptable(self, monkeypatch): + monkeypatch.setenv("TRUSTED_PROXIES", "10.0.0.5") + assert cimd_peer_is_acceptable("10.0.0.5", "8.8.8.8") is True + assert cimd_peer_is_acceptable("10.0.0.9", "8.8.8.8") is False + + def test_untrusted_forwarded_host_cannot_rebind_cimd(self, monkeypatch): + monkeypatch.delenv("TRUSTED_PROXIES", raising=False) + headers = { + "Host": "mcp.nitrostack.io", + "X-Forwarded-Host": "evil.example", + } + url = "https://evil.example/oauth/client-metadata.json" + assert request_host_for_cimd(headers, peer="8.8.8.8") == "mcp.nitrostack.io" + assert cimd_host_matches_request(url, headers, peer="8.8.8.8") is False + + def test_apply_cimd_records_untrusted_request_host(self, monkeypatch): + monkeypatch.delenv("TRUSTED_PROXIES", raising=False) + url = "https://app.nitrostack.io/oauth/client-metadata.json" + headers = { + "Host": "mcp.nitrostack.io", + "X-Forwarded-Host": "app.nitrostack.io", + } + with patch( + "nitrostack.auth.oauth_module.resolve_cimd_sync", + return_value={"client_id": url, "client_name": "Studio"}, + ): + body = apply_cimd_to_registration_body( + {"client_id": url}, + headers=headers, + peer="8.8.8.8", + ) + assert body["_cimd_request_host"] == "mcp.nitrostack.io" + assert body["_cimd_host_matches_request"] is False + + def test_apply_cimd_honors_trusted_forwarded_host(self, monkeypatch): + monkeypatch.setenv("TRUSTED_PROXIES", "10.0.0.5") + url = "https://app.nitrostack.io/oauth/client-metadata.json" + headers = { + "Host": "internal:3000", + "X-Forwarded-Host": "app.nitrostack.io", + } + with patch( + "nitrostack.auth.oauth_module.resolve_cimd_sync", + return_value={"client_id": url, "client_name": "Studio"}, + ): + body = apply_cimd_to_registration_body( + {"client_id": url}, + headers=headers, + peer="10.0.0.5", + ) + assert body["_cimd_request_host"] == "app.nitrostack.io" + assert body["_cimd_host_matches_request"] is True + + +class TestCimdRegistrationWiring: + def test_applies_cimd_document_to_registration_body(self): + url = "https://app.nitrostack.io/oauth/client-metadata.json" + with patch( + "nitrostack.auth.oauth_module.resolve_cimd_sync", + return_value={"client_id": url, "client_name": "Studio"}, + ): + body = apply_cimd_to_registration_body({"client_id": url, "redirect_uris": []}) + assert body["client_id"] == url + assert body["_cimd"]["client_name"] == "Studio" + + def test_leaves_non_url_client_id_unchanged(self): + body = apply_cimd_to_registration_body({"client_id": "static-client"}) + assert body["client_id"] == "static-client" + assert "_cimd" not in body + + +class TestProtectedResourceMetadata: + def test_includes_bearer_methods_supported(self): + metadata = build_protected_resource_metadata(_StubOAuthService()) + assert metadata["resource"] == "https://mcp.nitrostack.io/mcp" + assert metadata["authorization_servers"] == ["https://auth.nitrostack.io"] + assert metadata["bearer_methods_supported"] == ["header"] + + +class TestRfc9207IssValidation: + def test_accepts_matching_issuer(self): + validate_authorization_iss("https://auth.nitrostack.io/", "https://auth.nitrostack.io") + + def test_rejects_missing_iss(self): + with pytest.raises(AuthorizationIssuerMismatchError, match="missing"): + validate_authorization_iss(None, "https://auth.nitrostack.io") + + def test_rejects_mismatched_iss(self): + with pytest.raises(AuthorizationIssuerMismatchError, match="mismatch"): + validate_authorization_iss("https://evil.example", "https://auth.nitrostack.io") diff --git a/tests/test_mcp20_stateless_http.py b/tests/test_mcp20_stateless_http.py new file mode 100644 index 0000000..615dec1 --- /dev/null +++ b/tests/test_mcp20_stateless_http.py @@ -0,0 +1,3324 @@ +"""Tests for MCP 2.0 stateless HTTP transport.""" + +import asyncio +import json + +import pytest + +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER +from nitrostack.protocol.contracts import MCP_CACHE_HINT_KEY +from nitrostack.protocol.discovery import DISCOVER_RESULT_TYPE, SERVER_DISCOVER_METHOD, build_discover_result +from nitrostack.protocol.jsonrpc import ( + HEADER_BODY_MISMATCH, + PARSE_ERROR, + HeaderBodyMismatchError, + JsonRpcParseError, + build_ping_response, + parse_jsonrpc_request, + validate_header_body_method, +) +from nitrostack.protocol.method_contract import ( + MODERN_METHOD_CONTRACTS, + mcp_name_field, + mcp_name_is_required, +) +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.transports.cors import build_cors_headers, cors_preflight_response_headers, resolve_allowed_origin +from nitrostack.transports.dispatch import ( + IngressContext, + StatelessIngressPipeline, + is_header_only_ping, + is_task_wire_interception, +) +from nitrostack.transports.headers import ( + MAX_MCP_PARAM_VALUE_BYTES, + build_mcp_echo_headers, + build_mcp_response_headers, + build_sse_stream_headers, + extract_mcp_param_headers, + first_oversized_mcp_param, + handled_protocol_version, + merge_mcp_param_headers, + scope_with_header_snapshot, + scope_without_session_headers, + snapshot_validated_asgi_headers, + strip_legacy_session_headers, + strip_legacy_session_headers_asgi, +) +from nitrostack.transports.sse import format_sse_message, sse_notification + + +class TestRequestHeaders: + def test_strip_legacy_session_header(self): + headers = {"Mcp-Session-Id": "legacy", "Content-Type": "application/json"} + cleaned = strip_legacy_session_headers(headers) + assert LEGACY_SESSION_HEADER not in cleaned + assert cleaned["Content-Type"] == "application/json" + + def test_response_headers_include_protocol_version(self): + headers = build_mcp_response_headers() + assert headers["MCP-Protocol-Version"] == MODERN_PROTOCOL_VERSION + assert headers["Vary"] == "Origin" + assert "Mcp-Method" not in headers + echoed = build_mcp_response_headers(method="tools/call") + assert echoed["Mcp-Method"] == "tools/call" + + def test_strip_asgi_session_headers_from_inner_scope(self): + headers = [ + (b"content-type", b"application/json"), + (b"mcp-session-id", b"forged"), + (b"Mcp-Session-Id", b"also"), + ] + cleaned = strip_legacy_session_headers_asgi(headers) + assert cleaned == [(b"content-type", b"application/json")] + scope = scope_without_session_headers({"type": "http", "headers": headers}) + assert all(key.lower() != b"mcp-session-id" for key, _ in scope["headers"]) + assert scope["type"] == "http" + + def test_replay_snapshot_keeps_contract_headers_and_drops_session_id(self): + headers = [ + (b"content-type", b"application/json"), + (b"mcp-method", b"tools/call"), + (b"mcp-name", b"echo"), + (b"mcp-protocol-version", b"2026-07-28"), + (b"mcp-session-id", b"forged"), + ] + snapshot = snapshot_validated_asgi_headers(headers) + assert all(key.lower() != b"mcp-session-id" for key, _ in snapshot) + by_name = {key.lower(): value for key, value in snapshot} + assert by_name[b"mcp-method"] == b"tools/call" + assert by_name[b"mcp-name"] == b"echo" + assert by_name[b"mcp-protocol-version"] == b"2026-07-28" + live = {"type": "http", "headers": list(headers)} + live["headers"].append((b"mcp-session-id", b"later")) + replayed = scope_with_header_snapshot(live, snapshot) + assert all(key.lower() != b"mcp-session-id" for key, _ in replayed["headers"]) + replayed_by_name = {key.lower(): value for key, value in replayed["headers"]} + assert replayed_by_name[b"mcp-method"] == b"tools/call" + assert replayed_by_name[b"mcp-name"] == b"echo" + + +class TestHeaderCompatPreservesProtocolVersion: + def test_inner_app_sees_original_protocol_version(self): + from starlette.testclient import TestClient + + from nitrostack.transports.http import HeaderCompatMiddleware + + captured: dict[str, list] = {} + + async def inner(scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + captured["headers"] = list(scope.get("headers") or []) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send({"type": "http.response.body", "body": b"{}"}) + + wrapped = HeaderCompatMiddleware(inner, drop_session_headers=True) + with TestClient(wrapped) as client: + response = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "*/*", + "MCP-Protocol-Version": "2026-07-28", + "Mcp-Session-Id": "forged", + }, + json={"jsonrpc": "2.0", "id": 1, "method": "ping"}, + ) + assert response.status_code == 200 + headers = {key.decode("latin-1"): value.decode("latin-1") for key, value in captured["headers"]} + assert headers["mcp-protocol-version"] == "2026-07-28" + assert "mcp-session-id" not in headers + assert "application/json" in headers["accept"] + assert "text/event-stream" in headers["accept"] + + def test_unknown_protocol_version_is_not_dropped(self): + from starlette.testclient import TestClient + + from nitrostack.transports.http import HeaderCompatMiddleware + + captured: dict[str, list] = {} + + async def inner(scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + captured["headers"] = list(scope.get("headers") or []) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send({"type": "http.response.body", "body": b"{}"}) + + wrapped = HeaderCompatMiddleware(inner) + with TestClient(wrapped) as client: + client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "1999-01-01", + }, + json={"jsonrpc": "2.0", "id": 1, "method": "ping"}, + ) + headers = {key.decode("latin-1"): value.decode("latin-1") for key, value in captured["headers"]} + assert headers["mcp-protocol-version"] == "1999-01-01" + + +class TestMcpParamHeaderMirroring: + def test_extract_is_case_insensitive(self): + params = extract_mcp_param_headers( + {"Mcp-Param-City": "Boston", "mcp-param-limit": "10", "Mcp-Name": "echo"} + ) + assert params["City"] == "Boston" + assert params["limit"] == "10" + assert "Name" not in params + + def test_merge_fills_missing_and_keeps_body(self): + merged = merge_mcp_param_headers( + {"value": "body", "extra": ""}, + {"value": "header", "extra": "from-header", "unknown": "x"}, + allowed_fields={"value", "extra"}, + ) + assert merged["value"] == "body" + assert merged["extra"] == "from-header" + assert "unknown" not in merged + + def test_merge_does_not_override_name(self): + merged = merge_mcp_param_headers( + {"name": "echo", "value": ""}, + {"name": "other", "value": "ok"}, + allowed_fields={"name", "value"}, + ) + assert merged["name"] == "echo" + assert merged["value"] == "ok" + + def test_oversized_param_is_detected(self): + huge = "x" * (MAX_MCP_PARAM_VALUE_BYTES + 1) + assert first_oversized_mcp_param({"Mcp-Param-City": huge}) == "City" + assert first_oversized_mcp_param({"Mcp-Param-City": "Boston"}) is None + + def test_pipeline_name_check_uses_body_not_mirrored_params(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {}}, + } + ).encode() + result = await pipeline.handle_post( + body, + { + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + "Mcp-Param-Name": "other", + }, + ) + assert result is None + + asyncio.run(_run()) + + def test_pipeline_rejects_oversized_param(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, + {"Mcp-Param-City": "x" * (MAX_MCP_PARAM_VALUE_BYTES + 1)}, + ) + assert status == 400 + assert resp["error"]["code"] == -32600 + + asyncio.run(_run()) + + def test_http_header_fills_missing_tool_argument(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> dict: + return { + "value": input.value, + "mirrored": context.mcp_param_headers, + } + + @module(name="McpParamHttp", controllers=[EchoController]) + class McpParamModule: + pass + + @mcp_app(module=McpParamModule, server=ServerConfig(name="mcp-param-http")) + class McpParamApp: + pass + + app = asyncio.run(McpApplicationFactory.create(McpParamApp)) + http_app = app.get_combined_app(json_response=True) + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + "MCP-Protocol-Version": "2025-06-18", + } + with TestClient(http_app) as client: + filled = client.post( + "/mcp", + headers={**json_headers, "Mcp-Param-value": "from-header"}, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {}}, + }, + ) + kept = client.post( + "/mcp", + headers={**json_headers, "Mcp-Param-value": "from-header"}, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "from-body"}}, + }, + ) + assert filled.status_code == 200, filled.text + filled_body = filled.json()["result"] + filled_payload = filled_body.get("structuredContent") or json.loads( + filled_body["content"][0]["text"] + ) + assert filled_payload["value"] == "from-header" + assert filled_payload["mirrored"]["value"] == "from-header" + assert kept.status_code == 200, kept.text + kept_body = kept.json()["result"] + kept_payload = kept_body.get("structuredContent") or json.loads( + kept_body["content"][0]["text"] + ) + assert kept_payload["value"] == "from-body" + assert kept_payload["mirrored"]["value"] == "from-header" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestCors: + def test_cors_allow_methods(self): + headers = build_cors_headers() + assert "GET, POST, DELETE, OPTIONS" in headers["Access-Control-Allow-Methods"] + assert "Mcp-Method" in headers["Access-Control-Allow-Headers"] + expose = headers["Access-Control-Expose-Headers"] + assert "MCP-Protocol-Version" in expose + assert "Mcp-Method" in expose + assert "Mcp-Name" in expose + assert "Mcp-Session-Id" not in expose + + def test_preflight_headers(self): + headers = cors_preflight_response_headers( + { + "Origin": "https://app.example.com", + "Access-Control-Request-Headers": "Mcp-Name, Mcp-Method, MCP-Protocol-Version, Mcp-Param-City", + } + ) + assert "Access-Control-Allow-Origin" in headers + allow = headers["Access-Control-Allow-Headers"] + assert "Mcp-Name" in allow + assert "Mcp-Method" in allow + assert "MCP-Protocol-Version" in allow + assert "Mcp-Param-City" in allow + + def test_does_not_reflect_arbitrary_origin(self): + headers = build_cors_headers(origin="https://evil.example") + assert headers["Access-Control-Allow-Origin"] == "*" + + def test_allowlist_echoes_only_listed_origin(self): + assert resolve_allowed_origin( + "https://app.example.com", + allowed_origins=("https://app.example.com",), + ) == "https://app.example.com" + assert resolve_allowed_origin( + "https://evil.example", + allowed_origins=("https://app.example.com",), + ) == "https://app.example.com" + + +class TestJsonRpcParsing: + def test_parse_valid_request(self): + body = json.dumps( + {"jsonrpc": "2.0", "id": "1", "method": "tools/call", "params": {"name": "x"}} + ).encode() + req = parse_jsonrpc_request(body) + assert req.method == "tools/call" + assert req.id == "1" + + def test_parse_invalid_json(self): + with pytest.raises(JsonRpcParseError) as exc: + parse_jsonrpc_request(b"{bad") + assert exc.value.code == PARSE_ERROR + + def test_header_body_mismatch(self): + with pytest.raises(HeaderBodyMismatchError) as exc: + validate_header_body_method("tools/list", "tools/call") + assert exc.value.code == HEADER_BODY_MISMATCH + + +class TestPingFastPath: + def test_ping_response(self): + resp = build_ping_response("ping-1") + assert resp == {"jsonrpc": "2.0", "id": "ping-1", "result": {}} + + def test_header_only_ping_empty_or_non_jsonrpc_body(self): + headers = {"Mcp-Method": "ping"} + assert is_header_only_ping(b"", headers) is True + assert is_header_only_ping(b" \n", headers) is True + assert is_header_only_ping(b"not-json", headers) is True + assert is_header_only_ping(b"{}", headers) is True + + def test_header_only_ping_does_not_apply_to_other_methods(self): + assert is_header_only_ping(b"", {"Mcp-Method": "tools/call"}) is False + assert is_header_only_ping(b"", {}) is False + + def test_parsed_jsonrpc_body_is_not_header_only(self): + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping"}).encode() + assert is_header_only_ping(body, {"Mcp-Method": "ping"}) is False + call = json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "echo"}} + ).encode() + assert is_header_only_ping(call, {"Mcp-Method": "ping"}) is False + + +class TestDiscovery: + def test_discover_result_shape(self): + result = build_discover_result( + server_name="my-stateless-mcp-server", + server_version="1.0.0", + ) + assert result["protocolVersion"] == "2026-07-28" + assert result["supportedVersions"] == ["2026-07-28"] + assert result["serverInfo"]["name"] == "my-stateless-mcp-server" + assert result["capabilities"]["resources"]["subscribe"] is False + assert "io.modelcontextprotocol/tasks" in result["capabilities"]["extensions"] + assert result["resultType"] == DISCOVER_RESULT_TYPE + assert isinstance(result["ttlMs"], int) and result["ttlMs"] >= 0 + assert result["cacheScope"] in ("public", "private") + assert result["_meta"][MCP_CACHE_HINT_KEY]["ttlMs"] == result["ttlMs"] + assert result["_meta"][MCP_CACHE_HINT_KEY]["cacheScope"] == result["cacheScope"] + + +class TestDispatchPipeline: + def test_handles_ping(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + body = json.dumps({"jsonrpc": "2.0", "id": 9, "method": "ping"}).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + def test_header_only_ping_empty_body(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + status, resp = await pipeline.handle_post(b"", {"Mcp-Method": "ping"}) + assert status == 200 + assert resp == {"jsonrpc": "2.0", "id": None, "result": {}} + assert "error" not in resp + + asyncio.run(_run()) + + def test_header_only_ping_non_jsonrpc_body(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + status, resp = await pipeline.handle_post(b"not-json", {"Mcp-Method": "ping"}) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + def test_empty_body_without_ping_header_is_parse_error(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + status, resp = await pipeline.handle_post(b"", {"Mcp-Method": "tools/call"}) + assert status == 400 + assert resp["error"]["code"] == PARSE_ERROR + + asyncio.run(_run()) + + def test_tools_call_body_with_ping_header_is_mismatch(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo"}, + } + ).encode() + status, resp = await pipeline.handle_post(body, {"Mcp-Method": "ping"}) + assert status == 400 + assert resp["error"]["code"] == HEADER_BODY_MISMATCH + + asyncio.run(_run()) + + def test_forwards_server_discover_to_engine_handler(self): + async def _run(): + payload = build_discover_result( + server_name="srv", + server_version="1.0.0", + protocol_version=MODERN_PROTOCOL_VERSION, + ) + + def _discover(_request): + return payload + + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION), + discover_handler=_discover, + ) + body = json.dumps( + {"jsonrpc": "2.0", "id": "req-001", "method": SERVER_DISCOVER_METHOD, "params": {}} + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + result = resp["result"] + assert result == payload + assert result["protocolVersion"] == "2026-07-28" + assert result["resultType"] == DISCOVER_RESULT_TYPE + assert MODERN_PROTOCOL_VERSION in result["supportedVersions"] + + asyncio.run(_run()) + + def test_server_discover_without_engine_handler_is_not_answered(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + body = json.dumps( + {"jsonrpc": "2.0", "id": "req-001", "method": SERVER_DISCOVER_METHOD, "params": {}} + ).encode() + assert await pipeline.handle_post(body, {}) is None + + asyncio.run(_run()) + + def test_passes_through_tools_call(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "demo"}, + } + ).encode() + assert await pipeline.handle_post( + body, {"Mcp-Name": "demo", "Mcp-Method": "tools/call"} + ) is None + + asyncio.run(_run()) + + def test_tools_call_requires_mcp_name(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "demo"}, + } + ).encode() + status, resp = await pipeline.handle_post(body, {"Mcp-Method": "tools/call"}) + assert status == 400 + assert resp["error"]["code"] == -32020 + assert "Mcp-Name" in resp["error"]["message"] + + asyncio.run(_run()) + + def test_tools_call_rejects_mcp_name_mismatch(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "bar"}, + } + ).encode() + status, resp = await pipeline.handle_post( + body, {"Mcp-Name": "foo", "Mcp-Method": "tools/call"} + ) + assert status == 400 + assert resp["error"]["code"] == -32020 + assert "foo" in resp["error"]["message"] + assert "bar" in resp["error"]["message"] + + asyncio.run(_run()) + + def test_ping_does_not_require_mcp_name(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + body = json.dumps({"jsonrpc": "2.0", "id": 9, "method": "ping"}).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + def test_modern_reject_initialize(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2026-07-28", + "capabilities": {}, + "clientInfo": {"name": "t", "version": "1"}, + }, + } + ).encode() + status, resp = await pipeline.handle_post(body, {"Mcp-Method": "initialize"}) + assert status == 200 + assert resp["error"]["code"] == -32601 + assert resp["error"]["message"] == "Method not found: initialize" + + asyncio.run(_run()) + + def test_modern_reject_initialized(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext( + "srv", + "1.0.0", + MODERN_PROTOCOL_VERSION, + wire_mode="reject", + protocol_era="modern", + ) + ) + body = json.dumps( + {"jsonrpc": "2.0", "method": "notifications/initialized"} + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + assert resp["error"]["code"] == -32601 + assert resp["error"]["message"] == "Method not found: notifications/initialized" + + asyncio.run(_run()) + + def test_modern_reject_legacy_protocol_version(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 2, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, {"MCP-Protocol-Version": "2025-06-18", "Mcp-Method": "ping"} + ) + assert status == 400 + assert resp["error"]["code"] == -32022 + + asyncio.run(_run()) + + def test_modern_ignores_incoming_session_id(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 3, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, {LEGACY_SESSION_HEADER: "session-1", "Mcp-Method": "ping"} + ) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + def test_auto_ignores_incoming_session_id(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 4, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, + {LEGACY_SESSION_HEADER: "session-1", "Mcp-Method": "ping"}, + ) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + def test_legacy_sessionful_pipeline_keeps_session_id(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="sessionful") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 5, "method": "ping"}).encode() + status, resp = await pipeline.handle_post( + body, {LEGACY_SESSION_HEADER: "session-1"} + ) + assert status == 200 + assert resp["result"] == {} + + asyncio.run(_run()) + + def test_auto_accepts_initialize(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2025-06-18", "capabilities": {}}, + } + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + result = resp["result"] + assert result["protocolVersion"] == "2025-06-18" + assert result["serverInfo"]["name"] == "srv" + assert "capabilities" in result + + asyncio.run(_run()) + + def test_auto_acks_initialized_without_forwarding(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext( + "srv", + "1.0.0", + MODERN_PROTOCOL_VERSION, + wire_mode="stateless", + protocol_era="auto", + ) + ) + body = json.dumps( + {"jsonrpc": "2.0", "method": "notifications/initialized"} + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 202 + assert resp == {} + + asyncio.run(_run()) + + def test_legacy_sessionful_pipeline_forwards_initialized(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext( + "srv", + "1.0.0", + MODERN_PROTOCOL_VERSION, + wire_mode="sessionful", + protocol_era="legacy", + ) + ) + body = json.dumps( + {"jsonrpc": "2.0", "method": "notifications/initialized"} + ).encode() + assert await pipeline.handle_post(body, {}) is None + + asyncio.run(_run()) + + +class TestTaskInterceptionDetection: + def test_tasks_method_prefix(self): + assert is_task_wire_interception("tasks/get", {"taskId": "x"}) + + def test_tools_call_with_task_param(self): + assert is_task_wire_interception("tools/call", {"name": "x", "task": {"ttl": 1000}}) + + def test_sync_tools_call_not_intercepted(self): + assert not is_task_wire_interception("tools/call", {"name": "x"}) + + +class TestSse: + def test_sse_stream_headers_disable_buffering(self): + headers = build_sse_stream_headers() + assert headers["X-Accel-Buffering"] == "no" + assert headers["Cache-Control"] == "no-transform" + assert headers["Content-Type"] == "text/event-stream" + + def test_sse_notification_format(self): + frame = sse_notification("notifications/tools/list_changed") + text = frame.decode() + assert text.startswith("event: message") + assert "notifications/tools/list_changed" in text + + def test_sse_task_status_notification(self): + frame = format_sse_message( + {"method": "notifications/tasks/status", "params": {"taskId": "abc"}} + ) + assert b"notifications/tasks/status" in frame + + +class TestReplayDoesNotSynthesizeDisconnect: + def test_waits_for_real_disconnect(self): + from nitrostack.transports.middleware import StatelessTransportMiddleware + + calls = {"n": 0} + + async def original_receive(): + calls["n"] += 1 + return {"type": "http.disconnect"} + + replay = StatelessTransportMiddleware._replay_receive(b"{}", original_receive) + + async def _run(): + first = await replay() + assert first == {"type": "http.request", "body": b"{}", "more_body": False} + assert calls["n"] == 0 + second = await replay() + assert second == {"type": "http.disconnect"} + assert calls["n"] == 1 + + asyncio.run(_run()) + + +class TestReplayUsesHeaderSnapshot: + def test_tools_call_replay_keeps_name_and_method(self): + from starlette.testclient import TestClient + + from nitrostack.transports.middleware import wrap_stateless_transport + + captured: dict[str, list] = {} + + async def inner(scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + captured["headers"] = list(scope.get("headers") or []) + await receive() + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send( + { + "type": "http.response.body", + "body": b'{"jsonrpc":"2.0","id":1,"result":{}}', + } + ) + + wrapped = wrap_stateless_transport( + inner, + server_name="srv", + server_version="1.0.0", + protocol_version=MODERN_PROTOCOL_VERSION, + wire_mode="stateless", + ) + with TestClient(wrapped) as client: + response = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + "MCP-Protocol-Version": "2026-07-28", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo"}, + }, + ) + assert response.status_code == 200 + forwarded = {key.lower(): value for key, value in captured["headers"]} + assert forwarded.get(b"mcp-method") == b"tools/call" + assert forwarded.get(b"mcp-name") == b"echo" + assert forwarded.get(b"mcp-protocol-version") == b"2026-07-28" + assert b"mcp-session-id" not in forwarded + + def test_session_id_is_rejected_before_replay(self): + from starlette.testclient import TestClient + + from nitrostack.transports.middleware import wrap_stateless_transport + + called = {"inner": False} + + async def inner(scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + called["inner"] = True + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send({"type": "http.response.body", "body": b"{}"}) + + wrapped = wrap_stateless_transport( + inner, + server_name="srv", + server_version="1.0.0", + protocol_version=MODERN_PROTOCOL_VERSION, + wire_mode="stateless", + ) + with TestClient(wrapped) as client: + response = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + LEGACY_SESSION_HEADER: "forged", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo"}, + }, + ) + assert response.status_code == 400 + assert response.json()["error"]["code"] == -32600 + assert called["inner"] is False + + +class TestOptionsScopedToMcpPath: + def test_options_mcp_is_204_health_is_forwarded(self): + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + from pydantic import BaseModel, Field + + class EchoInput(BaseModel): + value: str = Field(default="") + + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="OptionsHttp", controllers=[EchoController]) + class OptionsModule: + pass + + @mcp_app(module=OptionsModule, server=ServerConfig(name="options-http", stateless=True)) + class OptionsApp: + pass + + app = asyncio.run(McpApplicationFactory.create(OptionsApp)) + http_app = app.get_combined_app(stateless=True, json_response=True) + with TestClient(http_app) as client: + mcp_opt = client.options("/mcp") + health_opt = client.options("/mcp/health") + bare_health = client.options("/health") + oauth_opt = client.options("/oauth/v2/register") + health_preflight = client.options( + "/mcp/health", + headers={ + "Origin": "https://app.example.com", + "Access-Control-Request-Method": "GET", + }, + ) + oauth_preflight = client.options( + "/oauth/v2/register", + headers={ + "Origin": "https://app.example.com", + "Access-Control-Request-Method": "POST", + }, + ) + call = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + assert mcp_opt.status_code == 204 + assert "Access-Control-Allow-Origin" in mcp_opt.headers + assert health_opt.status_code != 204 + assert bare_health.status_code != 204 + assert oauth_opt.status_code != 204 + assert health_preflight.status_code != 204 + assert oauth_preflight.status_code != 204 + assert call.status_code == 200 + assert call.json()["result"]["content"][0]["text"] == "ok" + finally: + DIContainer.reset() + + def test_options_mcp_is_not_sidecar_204_when_cors_off(self): + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + from pydantic import BaseModel, Field + + class EchoInput(BaseModel): + value: str = Field(default="") + + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="OptionsCorsOffHttp", controllers=[EchoController]) + class OptionsCorsOffModule: + pass + + @mcp_app(module=OptionsCorsOffModule, server=ServerConfig(name="options-cors-off")) + class OptionsCorsOffApp: + pass + + app = asyncio.run(McpApplicationFactory.create(OptionsCorsOffApp)) + http_app = app.get_combined_app(stateless=True, json_response=True, enable_cors=False) + with TestClient(http_app) as client: + mcp_opt = client.options("/mcp") + assert mcp_opt.status_code != 204 + finally: + DIContainer.reset() + + def test_options_allow_headers_include_2026_mcp_names(self, monkeypatch): + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + from pydantic import BaseModel, Field + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + monkeypatch.setenv("MCP_CORS_ALLOWED_ORIGINS", "https://app.example.com") + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="CorsAllowHeadersHttp", controllers=[EchoController]) + class CorsAllowHeadersModule: + pass + + @mcp_app(module=CorsAllowHeadersModule, server=ServerConfig(name="cors-allow-headers")) + class CorsAllowHeadersApp: + pass + + app = asyncio.run(McpApplicationFactory.create(CorsAllowHeadersApp)) + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + allowed = client.options( + "/mcp", + headers={ + "Origin": "https://app.example.com", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": ( + "Mcp-Name, Mcp-Method, MCP-Protocol-Version, Mcp-Param-City" + ), + }, + ) + unknown = client.options( + "/mcp", + headers={ + "Origin": "https://evil.example", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": "Mcp-Name", + }, + ) + assert allowed.status_code == 204 + allow = allowed.headers.get("access-control-allow-headers", "") + assert "Mcp-Name" in allow + assert "Mcp-Method" in allow + assert "MCP-Protocol-Version" in allow + assert "Mcp-Param-City" in allow + assert allowed.headers.get("access-control-allow-origin") == "https://app.example.com" + assert unknown.headers.get("access-control-allow-origin") != "https://evil.example" + finally: + DIContainer.reset() + monkeypatch.delenv("MCP_CORS_ALLOWED_ORIGINS", raising=False) + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + + +class TestProviderToolDiscovery: + def test_provider_tools_are_registered(self): + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.decorators import ToolConfig + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + assert ToolConfig(name="x", description="d", input_schema={}).task_support == "forbidden" + + DIContainer.reset() + try: + @injectable() + class ToolProvider: + @tool(name="from_provider", description="provider", input_schema=EchoInput) + async def from_provider(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="ProviderTools", providers=[ToolProvider], controllers=[]) + class ProviderModule: + pass + + @mcp_app(module=ProviderModule, server=ServerConfig(name="provider-tools")) + class ProviderApp: + pass + + app = asyncio.run(McpApplicationFactory.create(ProviderApp)) + assert "from_provider" in app._tools + assert app._tools["from_provider"].config.task_support == "forbidden" + finally: + DIContainer.reset() + + +class TestNitroMcpProtocolVersionEnv: + def test_modern_era_enables_stateless_without_mcp_stateless(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "2026-07-28") + monkeypatch.delenv("MCP_STATELESS", raising=False) + + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="EraHttp", controllers=[EchoController]) + class EraModule: + pass + + @mcp_app(module=EraModule, server=ServerConfig(name="era-http")) + class EraApp: + pass + + app = asyncio.run(McpApplicationFactory.create(EraApp)) + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + mcp_opt = client.options("/mcp") + call = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "ping", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "ping", + }, + ) + assert mcp_opt.status_code == 204 + assert call.status_code == 200 + assert call.json()["result"] == {} + assert app.protocol_era == "modern" + inner = getattr(http_app, "app", http_app) + assert inner.state.protocol_era == "modern" + assert inner.state.wire_mode == "reject" + assert inner.state.stateless is True + assert inner.state.http_engine == "sessionless" + assert app.mcp_server.http_engine == "sessionless" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_auto_era_keeps_era_and_does_not_force_stateless(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="AutoEraHttp", controllers=[EchoController]) + class AutoEraModule: + pass + + @mcp_app(module=AutoEraModule, server=ServerConfig(name="auto-era-http")) + class AutoEraApp: + pass + + app = asyncio.run(McpApplicationFactory.create(AutoEraApp)) + http_app = app.get_combined_app(json_response=True) + state = _http_app_state(http_app) + assert app.protocol_era == "auto" + assert state.protocol_era == "auto" + assert state.wire_mode == "stateless" + assert state.stateless is True + assert state.http_engine == "sessionless" + assert state.streamable_http_manager_count == 1 + assert app.mcp_server.http_engine == "sessionless" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_legacy_era_uses_sessionful_engine(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "legacy") + monkeypatch.delenv("MCP_STATELESS", raising=False) + + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="LegacyEngineHttp", controllers=[EchoController]) + class LegacyEngineModule: + pass + + @mcp_app(module=LegacyEngineModule, server=ServerConfig(name="legacy-engine-http")) + class LegacyEngineApp: + pass + + app = asyncio.run(McpApplicationFactory.create(LegacyEngineApp)) + http_app = app.get_combined_app(json_response=True) + state = _http_app_state(http_app) + assert app.protocol_era == "legacy" + assert state.http_engine == "sessionful" + assert state.sessionful is True + assert state.stateless is False + assert state.session_manager.stateless is False + assert app.mcp_server.http_engine == "sessionful" + assert app.mcp_server.sessionful is True + assert "echo" in app._tools + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +def _http_app_state(asgi): + current = asgi + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + state = getattr(current, "state", None) + if state is not None and getattr(state, "protocol_era", None) is not None: + return state + current = getattr(current, "app", None) + raise AssertionError("HTTP app is missing protocol era state") + + +class TestEraSessionfulCoexistence: + def _echo_app(self, monkeypatch, era_value: str): + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", era_value) + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name=f"EraCoexist{era_value.title()}", controllers=[EchoController]) + class EraModule: + pass + + @mcp_app(module=EraModule, server=ServerConfig(name=f"era-coexist-{era_value}")) + class EraApp: + pass + + return asyncio.run(McpApplicationFactory.create(EraApp)) + + def test_auto_stateless_false_does_not_start_sessionful_manager(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + app = self._echo_app(monkeypatch, "auto") + http_app = app.get_combined_app(json_response=True, stateless=False) + state = _http_app_state(http_app) + assert state.http_engine == "sessionless" + assert state.sessionful is False + assert state.session_manager.stateless is True + assert state.streamable_http_manager_count == 1 + assert app.mcp_server.sessionful is False + assert "echo" in app._tools + + with TestClient(http_app) as client: + init = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + }, + ) + assert init.status_code == 200, init.text + assert "result" in init.json() + init_headers = {key.lower(): value for key, value in init.headers.items()} + assert LEGACY_SESSION_HEADER.lower() not in init_headers + assert not state.session_manager._server_instances + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_modern_stateless_false_still_rejects_initialize(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + app = self._echo_app(monkeypatch, "modern") + http_app = app.get_combined_app(json_response=True, stateless=False) + state = _http_app_state(http_app) + assert state.http_engine == "sessionless" + assert state.sessionful is False + assert "echo" in app._tools + + with TestClient(http_app) as client: + init = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + }, + ) + assert init.status_code == 200 + assert init.json()["error"]["code"] == -32601 + init_headers = {key.lower(): value for key, value in init.headers.items()} + assert LEGACY_SESSION_HEADER.lower() not in init_headers + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_legacy_initialize_creates_session_without_modern_headers(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + app = self._echo_app(monkeypatch, "legacy") + http_app = app.get_combined_app(json_response=True) + state = _http_app_state(http_app) + assert state.http_engine == "sessionful" + assert state.sessionful is True + assert state.streamable_http_manager_count == 1 + assert "echo" in app._tools + + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + with TestClient(http_app) as client: + init = client.post( + "/mcp", + headers=json_headers, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + }, + ) + session_id = init.headers.get("mcp-session-id") or init.headers.get( + "Mcp-Session-Id" + ) + assert init.status_code == 200, init.text + assert session_id + client.post( + "/mcp", + headers={**json_headers, LEGACY_SESSION_HEADER: session_id}, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}, + ) + call = client.post( + "/mcp", + headers={**json_headers, LEGACY_SESSION_HEADER: session_id}, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + assert call.status_code == 200, call.text + assert call.json()["result"]["content"][0]["text"] == "ok" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestUnsetEnvDefault: + def _echo_app(self): + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + + class EchoInput(BaseModel): + value: str = Field(default="") + + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="UnsetEnvDefault", controllers=[EchoController]) + class UnsetEnvModule: + pass + + @mcp_app(module=UnsetEnvModule, server=ServerConfig(name="unset-env-default")) + class UnsetEnvApp: + pass + + return asyncio.run(McpApplicationFactory.create(UnsetEnvApp)) + + def test_unset_env_http_is_auto_sessionless(self, monkeypatch, caplog): + import logging + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + with caplog.at_level(logging.INFO, logger="nitrostack.core.app"): + app = self._echo_app() + http_app = app.get_combined_app(json_response=True) + state = _http_app_state(http_app) + assert app.protocol_era == "auto" + assert app.protocol_era_source == "default" + assert state.protocol_era == "auto" + assert state.http_engine == "sessionless" + assert "protocol era=auto (source=default)" in caplog.text + + with TestClient(http_app) as client: + init = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + }, + ) + assert init.status_code == 200, init.text + assert "result" in init.json() + assert init.json()["result"]["protocolVersion"] == "2025-06-18" + init_headers = {key.lower(): value for key, value in init.headers.items()} + assert LEGACY_SESSION_HEADER.lower() not in init_headers + assert not state.session_manager._server_instances + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + os.environ.pop("MCP_STATELESS", None) + + def test_mcp_stateless_false_still_forces_legacy(self, monkeypatch, caplog): + import logging + import os + + from nitrostack.core.di import DIContainer + + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.setenv("MCP_STATELESS", "false") + DIContainer.reset() + try: + with caplog.at_level(logging.INFO, logger="nitrostack.core.app"): + app = self._echo_app() + http_app = app.get_combined_app(json_response=True) + state = _http_app_state(http_app) + assert app.protocol_era == "legacy" + assert app.protocol_era_source == "mcp_stateless" + assert state.protocol_era == "legacy" + assert state.http_engine == "sessionful" + assert "protocol era=legacy (source=mcp_stateless)" in caplog.text + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + os.environ.pop("MCP_STATELESS", None) + + +class TestAutoEraOneMcpDualClients: + def test_auto_serves_initialize_and_discover_on_one_mcp(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="DualClientHttp", controllers=[EchoController]) + class DualClientModule: + pass + + @mcp_app(module=DualClientModule, server=ServerConfig(name="dual-client-http")) + class DualClientApp: + pass + + app = asyncio.run(McpApplicationFactory.create(DualClientApp)) + http_app = app.get_combined_app(json_response=True) + state = _http_app_state(http_app) + assert state.protocol_era == "auto" + assert state.streamable_http_manager_count == 1 + assert state.session_manager.stateless is True + + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + modern_headers = { + **json_headers, + "MCP-Protocol-Version": MODERN_PROTOCOL_VERSION, + } + + with TestClient(http_app) as client: + init = client.post( + "/mcp", + headers=json_headers, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + }, + ) + initialized = client.post( + "/mcp", + headers=json_headers, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}, + ) + discover = client.post( + "/mcp", + headers={**modern_headers, "Mcp-Method": "server/discover"}, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "server/discover", + "params": {}, + }, + ) + call = client.post( + "/mcp", + headers={ + **json_headers, + "MCP-Protocol-Version": "2025-06-18", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + }, + json={ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + + assert init.status_code == 200, init.text + init_body = init.json()["result"] + expected_init = app.handle_sessionless_initialize("2025-06-18") + assert init_body == expected_init + assert app.mcp_server.handle_sessionless_initialize("2025-06-18") == expected_init + assert init_body["protocolVersion"] == "2025-06-18" + assert "serverInfo" in init_body + assert "capabilities" in init_body + init_headers = {key.lower(): value for key, value in init.headers.items()} + assert LEGACY_SESSION_HEADER.lower() not in init_headers + assert not state.session_manager._server_instances + + assert initialized.status_code == 202, initialized.text + assert initialized.json() == {} + initialized_headers = { + key.lower(): value for key, value in initialized.headers.items() + } + assert LEGACY_SESSION_HEADER.lower() not in initialized_headers + assert not state.session_manager._server_instances + + assert discover.status_code == 200, discover.text + discover_body = discover.json()["result"] + expected_discover = app.handle_server_discover() + assert discover_body == expected_discover + assert app.mcp_server.handle_server_discover() == expected_discover + assert discover_body["protocolVersion"] == MODERN_PROTOCOL_VERSION + assert discover.headers.get("MCP-Protocol-Version") == MODERN_PROTOCOL_VERSION + + assert call.status_code == 200, call.text + assert call.json()["result"]["content"][0]["text"] == "ok" + assert call.headers.get("MCP-Protocol-Version") == "2025-06-18" + assert call.headers.get("Mcp-Method") == "tools/call" + call_headers = {key.lower(): value for key, value in call.headers.items()} + assert LEGACY_SESSION_HEADER.lower() not in call_headers + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestModernEraRejectsLegacyWire: + def test_modern_initialize_is_jsonrpc_error_without_session(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "modern") + monkeypatch.delenv("MCP_STATELESS", raising=False) + + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="ModernRejectHttp", controllers=[EchoController]) + class ModernRejectModule: + pass + + @mcp_app(module=ModernRejectModule, server=ServerConfig(name="modern-reject-http")) + class ModernRejectApp: + pass + + app = asyncio.run(McpApplicationFactory.create(ModernRejectApp)) + http_app = app.get_combined_app(json_response=True) + state = _http_app_state(http_app) + assert state.wire_mode == "reject" + + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + initialize_body = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + } + initialized_body = {"jsonrpc": "2.0", "method": "notifications/initialized"} + + with TestClient(http_app) as client: + init = client.post("/mcp", headers=headers, json=initialize_body) + initialized = client.post( + "/mcp", headers=headers, json=initialized_body + ) + init_with_session = client.post( + "/mcp", + headers={**headers, LEGACY_SESSION_HEADER: "forged"}, + json=initialize_body, + ) + initialized_with_session = client.post( + "/mcp", + headers={**headers, LEGACY_SESSION_HEADER: "forged"}, + json=initialized_body, + ) + + assert init.status_code == 200 + assert init.json()["error"]["code"] == -32601 + assert init.json()["error"]["message"] == "Method not found: initialize" + init_headers = {key.lower(): value for key, value in init.headers.items()} + assert LEGACY_SESSION_HEADER.lower() not in init_headers + assert not state.session_manager._server_instances + + assert initialized.status_code == 200 + assert initialized.json()["error"]["code"] == -32601 + assert ( + initialized.json()["error"]["message"] + == "Method not found: notifications/initialized" + ) + initialized_headers = { + key.lower(): value for key, value in initialized.headers.items() + } + assert LEGACY_SESSION_HEADER.lower() not in initialized_headers + + assert init_with_session.status_code == 400 + assert init_with_session.json()["error"]["code"] == -32600 + session_init_headers = { + key.lower(): value for key, value in init_with_session.headers.items() + } + assert LEGACY_SESSION_HEADER.lower() not in session_init_headers + + assert initialized_with_session.status_code == 400 + assert initialized_with_session.json()["error"]["code"] == -32600 + assert not state.session_manager._server_instances + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestHealthAdvertisesEra: + def _health_body(self, monkeypatch, era_value: str) -> dict: + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", era_value) + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="HealthEraHttp", controllers=[EchoController]) + class HealthEraModule: + pass + + @mcp_app(module=HealthEraModule, server=ServerConfig(name="health-era-http")) + class HealthEraApp: + pass + + app = asyncio.run(McpApplicationFactory.create(HealthEraApp)) + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + response = client.get("/mcp/health") + assert response.status_code == 200 + return response.json() + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_modern_health_reports_modern_era(self, monkeypatch): + body = self._health_body(monkeypatch, "modern") + assert body["status"] == "ok" + assert body["transport"] == "streamable-http" + assert body["protocolEra"] == "modern" + assert body["protocolVersion"] == "2026-07-28" + assert body["statelessCapable"] is True + assert "stateless" in body + assert "uptimeSeconds" in body + + def test_auto_health_reports_auto_not_modern(self, monkeypatch): + body = self._health_body(monkeypatch, "auto") + assert body["protocolEra"] == "auto" + assert body["protocolEra"] != "modern" + assert body["statelessCapable"] is True + + def test_legacy_health_reports_legacy_era(self, monkeypatch): + body = self._health_body(monkeypatch, "legacy") + assert body["protocolEra"] == "legacy" + assert body["protocolVersion"] == "2025-06-18" + assert body["statelessCapable"] is False + + +class TestTrustedProxyPublicUrl: + def _http_app(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.delenv("MCP_STATELESS", raising=False) + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + DIContainer.reset() + + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="ProxyHttp", controllers=[EchoController]) + class ProxyModule: + pass + + @mcp_app(module=ProxyModule, server=ServerConfig(name="proxy-http")) + class ProxyApp: + pass + + app = asyncio.run(McpApplicationFactory.create(ProxyApp)) + return app.get_combined_app(json_response=True) + + def test_health_ignores_untrusted_forwarded_host(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + http_app = self._http_app(monkeypatch) + with TestClient(http_app, client=("8.8.8.8", 4321)) as client: + response = client.get( + "/mcp/health", + headers={ + "Host": "internal:3000", + "X-Forwarded-Host": "mcp.example.com", + "X-Forwarded-Proto": "https", + }, + ) + docs = client.get( + "/", + headers={ + "Host": "internal:3000", + "X-Forwarded-Host": "mcp.example.com", + "X-Forwarded-Proto": "https", + }, + ) + oauth = client.get( + "/.well-known/oauth-authorization-server", + headers={ + "Host": "internal:3000", + "X-Forwarded-Host": "mcp.example.com", + "X-Forwarded-Proto": "https", + }, + ) + assert response.status_code == 200 + assert "mcp.example.com" not in response.json()["publicUrl"] + assert "mcp.example.com" not in docs.text + assert "mcp.example.com" not in oauth.json()["error_description"] + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_health_honors_trusted_forwarded_host(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + monkeypatch.setenv("TRUSTED_PROXIES", "10.0.0.5") + try: + http_app = self._http_app(monkeypatch) + with TestClient(http_app, client=("10.0.0.5", 4321)) as client: + headers = { + "Host": "internal:3000", + "X-Forwarded-Host": "mcp.example.com", + "X-Forwarded-Proto": "https", + } + response = client.get("/mcp/health", headers=headers) + docs = client.get("/", headers=headers) + oauth = client.get("/.well-known/oauth-authorization-server", headers=headers) + version = client.get("/json/version", headers=headers) + assert response.status_code == 200 + assert response.json()["publicUrl"] == "https://mcp.example.com/mcp" + assert "https://mcp.example.com/mcp" in docs.text + assert "https://mcp.example.com/mcp" in oauth.json()["error_description"] + assert version.json()["publicUrl"] == "https://mcp.example.com/mcp" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + os.environ.pop("TRUSTED_PROXIES", None) + + +class TestHeaderTriggeredPing: + def _http_app(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.delenv("MCP_STATELESS", raising=False) + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + DIContainer.reset() + + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="HeaderPingHttp", controllers=[EchoController]) + class HeaderPingModule: + pass + + @mcp_app(module=HeaderPingModule, server=ServerConfig(name="header-ping-http")) + class HeaderPingApp: + pass + + app = asyncio.run(McpApplicationFactory.create(HeaderPingApp)) + return app.get_combined_app(json_response=True) + + def test_empty_body_ping_header_succeeds(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + http_app = self._http_app(monkeypatch) + with TestClient(http_app) as client: + response = client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Mcp-Method": "ping", + }, + content=b"", + ) + assert response.status_code == 200 + body = response.json() + assert body["result"] == {} + assert "error" not in body + assert body.get("error", {}).get("code") != PARSE_ERROR + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_tools_call_body_with_ping_header_is_mismatch(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + http_app = self._http_app(monkeypatch) + with TestClient(http_app) as client: + response = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "ping", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "x"}}, + }, + ) + assert response.status_code == 400 + assert response.json()["error"]["code"] == HEADER_BODY_MISMATCH + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestEnvelopeOnHandlerContext: + def test_tool_reads_trace_and_protocol_version_not_spoofed_user(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EmptyInput(BaseModel): + pass + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EnvelopeController: + @tool(name="envelope", description="read envelope", input_schema=EmptyInput) + async def envelope(self, input: EmptyInput, context: ExecutionContext) -> dict: + return { + "protocolVersion": context.protocol_version, + "trace": context.rpc_meta.trace if context.rpc_meta else None, + "user": context.user, + } + + @module(name="EnvelopeHttp", controllers=[EnvelopeController]) + class EnvelopeModule: + pass + + @mcp_app(module=EnvelopeModule, server=ServerConfig(name="envelope-http")) + class EnvelopeApp: + pass + + app = asyncio.run(McpApplicationFactory.create(EnvelopeApp)) + http_app = app.get_combined_app(json_response=True) + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "2025-06-18", + "Mcp-Method": "tools/call", + "Mcp-Name": "envelope", + } + with TestClient(http_app) as client: + response = client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "envelope", + "arguments": {}, + "_meta": { + "trace": {"id": "http-span"}, + "userId": "spoofed", + }, + }, + }, + ) + + assert response.status_code == 200, response.text + result = response.json()["result"] + body = result.get("structuredContent") or json.loads(result["content"][0]["text"]) + assert body["protocolVersion"] == "2025-06-18" + assert body["trace"] == {"id": "http-span"} + assert body["user"] is None + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestAuthFromEnvelopeAndHeaders: + def test_tool_user_comes_from_header_jwt_not_unsigned_meta(self, monkeypatch): + import os + + from pydantic import BaseModel + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.auth.jwt import JWTService + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EmptyInput(BaseModel): + pass + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + token = jwt.create_token({"sub": "alice", "tenant_id": "acme"}) + try: + @injectable() + class AuthController: + @tool(name="whoami", description="identity", input_schema=EmptyInput) + async def whoami(self, input: EmptyInput, context: ExecutionContext) -> dict: + return {"user": context.user} + + @module(name="AuthHttp", controllers=[AuthController]) + class AuthModule: + pass + + @mcp_app(module=AuthModule, server=ServerConfig(name="auth-http")) + class AuthApp: + pass + + app = asyncio.run(McpApplicationFactory.create(AuthApp)) + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + response = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "2025-06-18", + "Mcp-Method": "tools/call", + "Mcp-Name": "whoami", + "Authorization": f"Bearer {token}", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "whoami", + "arguments": {}, + "_meta": { + "userId": "eve", + "tenantId": "evil", + "io.modelcontextprotocol/auth": {"userId": "mallory"}, + }, + }, + }, + ) + + assert response.status_code == 200, response.text + result = response.json()["result"] + body = result.get("structuredContent") or json.loads(result["content"][0]["text"]) + assert body["user"] == "alice" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestIncomingSessionIdRejection: + def _echo_app(self, monkeypatch, era_value: str): + import os + + from pydantic import BaseModel, Field + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", era_value) + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="SessionRejectHttp", controllers=[EchoController]) + class SessionRejectModule: + pass + + @mcp_app(module=SessionRejectModule, server=ServerConfig(name="session-reject-http")) + class SessionRejectApp: + pass + + return asyncio.run(McpApplicationFactory.create(SessionRejectApp)) + + def test_modern_post_with_session_id_is_rejected(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + app = self._echo_app(monkeypatch, "modern") + http_app = app.get_combined_app(json_response=True) + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + LEGACY_SESSION_HEADER: "forged", + } + with TestClient(http_app) as client: + rejected = client.post( + "/mcp", + headers=headers, + json={"jsonrpc": "2.0", "id": 1, "method": "ping"}, + ) + allowed = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "ping", + }, + json={"jsonrpc": "2.0", "id": 2, "method": "ping"}, + ) + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == -32600 + assert allowed.status_code == 200 + assert allowed.json()["result"] == {} + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_auto_post_with_session_id_is_rejected(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + app = self._echo_app(monkeypatch, "auto") + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + rejected = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + LEGACY_SESSION_HEADER: "forged", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + allowed = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + }, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == -32600 + assert allowed.status_code == 200 + assert allowed.json()["result"]["content"][0]["text"] == "ok" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_legacy_post_accepts_session_id_after_initialize(self, monkeypatch): + import os + + from starlette.testclient import TestClient + + from nitrostack.core.di import DIContainer + + try: + app = self._echo_app(monkeypatch, "legacy") + http_app = app.get_combined_app(json_response=True) + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + with TestClient(http_app) as client: + init = client.post( + "/mcp", + headers=json_headers, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + }, + ) + session_id = init.headers.get("mcp-session-id") or init.headers.get( + "Mcp-Session-Id" + ) + assert init.status_code == 200, init.text + assert session_id + client.post( + "/mcp", + headers={**json_headers, LEGACY_SESSION_HEADER: session_id}, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}, + ) + call = client.post( + "/mcp", + headers={ + **json_headers, + LEGACY_SESSION_HEADER: session_id, + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + }, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + assert call.status_code == 200, call.text + assert call.json()["result"]["content"][0]["text"] == "ok" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestSessionIdNotForwarded: + def test_forwarded_scope_omits_session_header(self): + from starlette.testclient import TestClient + + from nitrostack.transports.middleware import wrap_stateless_transport + + captured: dict[str, list] = {} + + async def inner(scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + captured["headers"] = list(scope.get("headers") or []) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send( + { + "type": "http.response.body", + "body": b'{"jsonrpc":"2.0","id":1,"result":{}}', + } + ) + + wrapped = wrap_stateless_transport( + inner, + server_name="srv", + server_version="1.0.0", + protocol_version=MODERN_PROTOCOL_VERSION, + wire_mode="sessionful", + ) + with TestClient(wrapped) as client: + response = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + LEGACY_SESSION_HEADER: "forged", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo"}, + }, + ) + assert response.status_code == 200 + assert "headers" in captured + assert all(key.lower() != b"mcp-session-id" for key, _ in captured["headers"]) + forwarded = {key.lower(): value for key, value in captured["headers"]} + assert forwarded.get(b"mcp-method") == b"tools/call" + assert forwarded.get(b"mcp-name") == b"echo" + + def test_forged_session_id_cannot_associate_two_modern_calls(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "modern") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="NoSharedSessionHttp", controllers=[EchoController]) + class NoSharedSessionModule: + pass + + @mcp_app(module=NoSharedSessionModule, server=ServerConfig(name="no-shared-session")) + class NoSharedSessionApp: + pass + + app = asyncio.run(McpApplicationFactory.create(NoSharedSessionApp)) + http_app = app.get_combined_app(json_response=True) + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + LEGACY_SESSION_HEADER: "forged-shared", + } + body = { + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + } + with TestClient(http_app) as client: + first = client.post("/mcp", headers=headers, json={**body, "id": 1}) + second = client.post("/mcp", headers=headers, json={**body, "id": 2}) + assert first.status_code == 400 + assert second.status_code == 400 + assert first.json()["error"]["code"] == -32600 + assert second.json()["error"]["code"] == -32600 + state = _http_app_state(http_app) + assert not getattr(state.session_manager, "_server_instances", True) + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + def test_modern_get_mcp_rejects_session_id(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "modern") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="ModernGetSessionHttp", controllers=[EchoController]) + class ModernGetSessionModule: + pass + + @mcp_app(module=ModernGetSessionModule, server=ServerConfig(name="modern-get-session")) + class ModernGetSessionApp: + pass + + app = asyncio.run(McpApplicationFactory.create(ModernGetSessionApp)) + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + response = client.get( + "/mcp", + headers={ + "Accept": "text/event-stream", + LEGACY_SESSION_HEADER: "forged", + }, + ) + assert response.status_code == 400 + assert response.json()["error"]["code"] == -32600 + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestRequiredMcpName: + def test_http_tools_call_missing_and_mismatch_and_match(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="RequiredNameHttp", controllers=[EchoController]) + class RequiredNameModule: + pass + + @mcp_app(module=RequiredNameModule, server=ServerConfig(name="required-name-http")) + class RequiredNameApp: + pass + + app = asyncio.run(McpApplicationFactory.create(RequiredNameApp)) + http_app = app.get_combined_app(json_response=True) + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + } + call_body = { + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + } + with TestClient(http_app) as client: + missing = client.post( + "/mcp", + headers=json_headers, + json={**call_body, "id": 1}, + ) + mismatch = client.post( + "/mcp", + headers={**json_headers, "Mcp-Name": "foo"}, + json={**call_body, "id": 2}, + ) + matched = client.post( + "/mcp", + headers={**json_headers, "Mcp-Name": "echo"}, + json={**call_body, "id": 3}, + ) + assert missing.status_code == 200, missing.text + assert missing.json()["result"]["content"][0]["text"] == "ok" + assert mismatch.status_code == 400 + assert mismatch.json()["error"]["code"] == -32020 + assert matched.status_code == 200, matched.text + assert matched.json()["result"]["content"][0]["text"] == "ok" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestRequiredMcpMethod: + def test_pipeline_tools_call_requires_mcp_method_on_modern(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext( + "srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject" + ) + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "demo"}, + } + ).encode() + missing = await pipeline.handle_post(body, {"Mcp-Name": "demo"}) + mismatch = await pipeline.handle_post( + body, {"Mcp-Name": "demo", "Mcp-Method": "tools/list"} + ) + matched = await pipeline.handle_post( + body, {"Mcp-Name": "demo", "Mcp-Method": "tools/call"} + ) + assert missing is not None + assert missing[0] == 400 + assert missing[1]["error"]["code"] == -32020 + assert mismatch is not None + assert mismatch[0] == 400 + assert mismatch[1]["error"]["code"] == -32020 + assert matched is None + + asyncio.run(_run()) + + def test_auto_tools_call_accepts_body_only_name_and_method(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext( + "srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless" + ) + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "demo", "arguments": {}}, + } + ).encode() + legacy = await pipeline.handle_post(body, {}) + mismatch = await pipeline.handle_post( + body, {"Mcp-Method": "tools/list", "Mcp-Name": "demo"} + ) + assert legacy is None + assert mismatch is not None + assert mismatch[0] == 400 + assert mismatch[1]["error"]["code"] == -32020 + + asyncio.run(_run()) + + def test_auto_initialize_does_not_require_mcp_method(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2025-06-18", "capabilities": {}}, + } + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + assert resp["result"]["protocolVersion"] == "2025-06-18" + + asyncio.run(_run()) + + def test_auto_list_methods_do_not_require_mcp_method(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + for method in ("tools/list", "resources/list", "prompts/list"): + body = json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": method, "params": {}} + ).encode() + result = await pipeline.handle_post(body, {}) + assert result is None, method + + asyncio.run(_run()) + + def test_modern_ping_requires_mcp_method(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject") + ) + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping"}).encode() + missing = await pipeline.handle_post(body, {}) + matched = await pipeline.handle_post(body, {"Mcp-Method": "ping"}) + assert missing is not None + assert missing[0] == 400 + assert missing[1]["error"]["code"] == -32020 + assert matched is not None + assert matched[0] == 200 + assert matched[1]["result"] == {} + + asyncio.run(_run()) + + def test_http_tools_call_missing_and_mismatch_and_match(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="RequiredMethodHttp", controllers=[EchoController]) + class RequiredMethodModule: + pass + + @mcp_app(module=RequiredMethodModule, server=ServerConfig(name="required-method-http")) + class RequiredMethodApp: + pass + + app = asyncio.run(McpApplicationFactory.create(RequiredMethodApp)) + http_app = app.get_combined_app(json_response=True) + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Name": "echo", + } + call_body = { + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + } + with TestClient(http_app) as client: + missing = client.post( + "/mcp", + headers=json_headers, + json={**call_body, "id": 1}, + ) + mismatch = client.post( + "/mcp", + headers={**json_headers, "Mcp-Method": "tools/list"}, + json={**call_body, "id": 2}, + ) + matched = client.post( + "/mcp", + headers={**json_headers, "Mcp-Method": "tools/call"}, + json={**call_body, "id": 3}, + ) + assert missing.status_code == 200, missing.text + assert missing.json()["result"]["content"][0]["text"] == "ok" + assert mismatch.status_code == 400 + assert mismatch.json()["error"]["code"] == -32020 + assert matched.status_code == 200, matched.text + assert matched.json()["result"]["content"][0]["text"] == "ok" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +def _contract_params(method: str) -> dict: + if mcp_name_is_required(method): + field = mcp_name_field(method) + if field == "uri": + return {"uri": "mcp://demo/item"} + return {"name": "demo"} + if method.startswith("tasks/"): + return {"taskId": "task-1"} + return {} + + +class TestSep2243AllModernMethods: + def test_table_covers_modern_surface(self): + methods = {row.method for row in MODERN_METHOD_CONTRACTS} + assert "server/discover" in methods + assert "resources/read" in methods + assert "prompts/get" in methods + assert "completion/complete" in methods + assert "tasks/get" in methods + assert "tools/call" in methods + + @pytest.mark.parametrize( + "contract", + [ + row + for row in MODERN_METHOD_CONTRACTS + if row.method not in {"initialize", "notifications/initialized"} + ], + ids=lambda row: row.method, + ) + def test_modern_missing_mcp_method_is_header_mismatch(self, contract): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject") + ) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": contract.method, + "params": _contract_params(contract.method), + } + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 400 + assert resp["error"]["code"] == HEADER_BODY_MISMATCH + + asyncio.run(_run()) + + @pytest.mark.parametrize( + "contract", + [ + row + for row in MODERN_METHOD_CONTRACTS + if row.method not in {"initialize", "notifications/initialized"} + ], + ids=lambda row: row.method, + ) + def test_modern_unsupported_protocol_version(self, contract): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="reject") + ) + params = _contract_params(contract.method) + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": contract.method, + "params": params, + } + ).encode() + headers = { + "Mcp-Method": contract.method, + "MCP-Protocol-Version": "1999-01-01", + } + if mcp_name_is_required(contract.method): + field = mcp_name_field(contract.method) + headers["Mcp-Name"] = params[field] + status, resp = await pipeline.handle_post(body, headers) + assert status == 400 + assert resp["error"]["code"] == -32022 + + asyncio.run(_run()) + + def test_resources_read_and_prompts_get_require_mcp_name(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION) + ) + read_body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "resources/read", + "params": {"uri": "mcp://demo/item"}, + } + ).encode() + prompt_body = json.dumps( + { + "jsonrpc": "2.0", + "id": 2, + "method": "prompts/get", + "params": {"name": "demo"}, + } + ).encode() + read_missing = await pipeline.handle_post( + read_body, {"Mcp-Method": "resources/read"} + ) + prompt_missing = await pipeline.handle_post( + prompt_body, {"Mcp-Method": "prompts/get"} + ) + assert read_missing is not None + assert read_missing[0] == 400 + assert read_missing[1]["error"]["code"] == HEADER_BODY_MISMATCH + assert prompt_missing is not None + assert prompt_missing[0] == 400 + assert prompt_missing[1]["error"]["code"] == HEADER_BODY_MISMATCH + assert read_missing[1]["error"]["code"] == prompt_missing[1]["error"]["code"] + + asyncio.run(_run()) + + def test_auto_discover_still_allows_missing_mcp_method(self): + async def _run(): + pipeline = StatelessIngressPipeline( + IngressContext("srv", "1.0.0", MODERN_PROTOCOL_VERSION, wire_mode="stateless") + ) + body = json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "server/discover", "params": {}} + ).encode() + assert await pipeline.handle_post(body, {}) is None + + asyncio.run(_run()) + + +class TestProtocolVersionCrossCheck: + def test_http_mismatch_is_rejected_and_header_only_succeeds(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="VersionCrossCheckHttp", controllers=[EchoController]) + class VersionCrossCheckModule: + pass + + @mcp_app(module=VersionCrossCheckModule, server=ServerConfig(name="version-cross-check")) + class VersionCrossCheckApp: + pass + + app = asyncio.run(McpApplicationFactory.create(VersionCrossCheckApp)) + http_app = app.get_combined_app(json_response=True) + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + } + with TestClient(http_app) as client: + mismatch = client.post( + "/mcp", + headers={**json_headers, "MCP-Protocol-Version": "2026-07-28"}, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "echo", + "arguments": {"value": "ok"}, + "_meta": {"mcp": {"protocolVersion": "2025-06-18"}}, + }, + }, + ) + header_only = client.post( + "/mcp", + headers={**json_headers, "MCP-Protocol-Version": "2025-06-18"}, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + assert mismatch.status_code == 400 + assert mismatch.json()["error"]["code"] == -32020 + assert header_only.status_code == 200, header_only.text + assert header_only.json()["result"]["content"][0]["text"] == "ok" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestUnsupportedProtocolVersionHttp: + def test_http_unknown_version_is_rejected_and_supported_proceeds(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="UnsupportedVersionHttp", controllers=[EchoController]) + class UnsupportedVersionModule: + pass + + @mcp_app(module=UnsupportedVersionModule, server=ServerConfig(name="unsupported-version")) + class UnsupportedVersionApp: + pass + + app = asyncio.run(McpApplicationFactory.create(UnsupportedVersionApp)) + http_app = app.get_combined_app(json_response=True) + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + ping = {"jsonrpc": "2.0", "id": 1, "method": "ping"} + with TestClient(http_app) as client: + unknown = client.post( + "/mcp", + headers={**json_headers, "MCP-Protocol-Version": "1999-01-01"}, + json=ping, + ) + supported = client.post( + "/mcp", + headers={**json_headers, "MCP-Protocol-Version": "2026-07-28"}, + json={**ping, "id": 2}, + ) + legacy = client.post( + "/mcp", + headers={**json_headers, "MCP-Protocol-Version": "2025-06-18"}, + json={**ping, "id": 3}, + ) + assert unknown.status_code == 400 + assert unknown.json()["error"]["code"] == -32022 + assert unknown.json()["error"]["message"] == "Unsupported protocol version" + assert supported.status_code == 200, supported.text + assert supported.json()["result"] == {} + assert legacy.status_code == 200, legacy.text + assert legacy.json()["result"] == {} + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestResponseEchoHeaders: + def test_echo_uses_supported_request_version(self): + headers = build_mcp_echo_headers( + {"MCP-Protocol-Version": "2025-06-18", "Mcp-Method": "tools/call"}, + protocol_version=MODERN_PROTOCOL_VERSION, + supported_versions={"2026-07-28", "2025-06-18"}, + ) + assert headers["MCP-Protocol-Version"] == "2025-06-18" + assert headers["Mcp-Method"] == "tools/call" + assert handled_protocol_version( + {"MCP-Protocol-Version": "1999-01-01"}, + fallback=MODERN_PROTOCOL_VERSION, + supported={"2026-07-28", "2025-06-18"}, + ) == MODERN_PROTOCOL_VERSION + + def test_http_success_and_error_echo_headers(self, monkeypatch): + import os + + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="EchoHeadersHttp", controllers=[EchoController]) + class EchoHeadersModule: + pass + + @mcp_app(module=EchoHeadersModule, server=ServerConfig(name="echo-headers")) + class EchoHeadersApp: + pass + + app = asyncio.run(McpApplicationFactory.create(EchoHeadersApp)) + http_app = app.get_combined_app(json_response=True) + json_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "2025-06-18", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + } + with TestClient(http_app) as client: + success = client.post( + "/mcp", + headers=json_headers, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + mismatch = client.post( + "/mcp", + headers={**json_headers, "Mcp-Name": "other"}, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + unsupported = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "1999-01-01", + }, + json={"jsonrpc": "2.0", "id": 3, "method": "ping"}, + ) + assert success.status_code == 200, success.text + assert success.headers.get("MCP-Protocol-Version") == "2025-06-18" + assert success.headers.get("Mcp-Method") == "tools/call" + assert mismatch.status_code == 400 + assert mismatch.json()["error"]["code"] == -32020 + assert mismatch.headers.get("MCP-Protocol-Version") == "2025-06-18" + assert mismatch.headers.get("Mcp-Method") == "tools/call" + assert unsupported.status_code == 400 + assert unsupported.json()["error"]["code"] == -32022 + assert unsupported.headers.get("MCP-Protocol-Version") == MODERN_PROTOCOL_VERSION + assert unsupported.headers.get("Mcp-Method") == "ping" + finally: + DIContainer.reset() + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +class TestCorsExposeHeaders: + def test_http_mcp_response_exposes_echo_headers_not_session_id(self, monkeypatch): + from pydantic import BaseModel, Field + from starlette.testclient import TestClient + + from nitrostack import ExecutionContext, injectable, module, tool + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.core.di import DIContainer + + class EchoInput(BaseModel): + value: str = Field(default="") + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "auto") + monkeypatch.delenv("MCP_STATELESS", raising=False) + DIContainer.reset() + try: + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="CorsExposeHttp", controllers=[EchoController]) + class CorsExposeModule: + pass + + @mcp_app(module=CorsExposeModule, server=ServerConfig(name="cors-expose")) + class CorsExposeApp: + pass + + app = asyncio.run(McpApplicationFactory.create(CorsExposeApp)) + http_app = app.get_combined_app(json_response=True) + origin = "https://app.example.com" + with TestClient(http_app) as client: + response = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Origin": origin, + "MCP-Protocol-Version": "2025-06-18", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"value": "ok"}}, + }, + ) + preflight = client.options( + "/mcp", + headers={ + "Origin": origin, + "Access-Control-Request-Method": "POST", + }, + ) + assert response.status_code == 200, response.text + expose = response.headers.get("access-control-expose-headers", "") + assert "MCP-Protocol-Version" in expose + assert "Mcp-Method" in expose + assert "Mcp-Session-Id" not in expose + assert preflight.status_code == 204 + preflight_expose = preflight.headers.get("access-control-expose-headers", "") + assert "MCP-Protocol-Version" in preflight_expose + assert "Mcp-Method" in preflight_expose + assert "Mcp-Session-Id" not in preflight_expose + finally: + DIContainer.reset() + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) diff --git a/tests/test_mcp20_stdio.py b/tests/test_mcp20_stdio.py new file mode 100644 index 0000000..94a4b23 --- /dev/null +++ b/tests/test_mcp20_stdio.py @@ -0,0 +1,189 @@ +"""Official mcp 2.x stdio follows the active protocol era.""" + +from __future__ import annotations + +import asyncio +import os +import sys + +import anyio +from mcp.shared.memory import create_client_server_memory_streams +from mcp.shared.message import SessionMessage +from mcp_types import PROTOCOL_VERSION_META_KEY, jsonrpc_message_adapter +from pydantic import BaseModel, Field +from starlette.testclient import TestClient + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION +from nitrostack.transports.stdio import serve_stdio_streams + +MODERN_META = { + PROTOCOL_VERSION_META_KEY: MODERN_PROTOCOL_VERSION, + "io.modelcontextprotocol/clientInfo": {"name": "stdio-test", "version": "1.0"}, + "io.modelcontextprotocol/clientCapabilities": {}, +} + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +def _session_message(method: str, params: dict, req_id: int = 1) -> SessionMessage: + return SessionMessage( + jsonrpc_message_adapter.validate_python( + {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params} + ) + ) + + +def _message_payload(session_message: SessionMessage) -> dict: + return session_message.message.model_dump(by_alias=True, mode="json") + + +def _echo_app(era: str): + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name=f"StdioEra{era.title()}", controllers=[EchoController]) + class EchoModule: + pass + + @mcp_app(module=EchoModule, server=ServerConfig(name=f"stdio-{era}", protocol_era=era)) + class EchoApp: + pass + + return asyncio.run(McpApplicationFactory.create(EchoApp)) + + +async def _roundtrip(app, era: str, method: str, params: dict) -> dict: + async with create_client_server_memory_streams() as (client, server): + client_read, client_write = client + server_read, server_write = server + + async def run_server(): + await serve_stdio_streams(app.mcp_server, server_read, server_write, era) + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + await client_write.send(_session_message(method, params)) + response = await asyncio.wait_for(client_read.receive(), timeout=2) + tg.cancel_scope.cancel() + return _message_payload(response) + + +def test_auto_stdio_calls_tool_without_initialize(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app = _echo_app("auto") + payload = asyncio.run( + _roundtrip( + app, + "auto", + "tools/call", + { + "name": "echo", + "arguments": {"value": "ok"}, + "_meta": MODERN_META, + }, + ) + ) + assert "result" in payload + assert payload["result"]["content"][0]["text"] == "ok" + + +def test_auto_stdio_discovers_without_initialize(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app = _echo_app("auto") + payload = asyncio.run( + _roundtrip(app, "auto", "server/discover", {"_meta": MODERN_META}) + ) + assert "result" in payload + result = payload["result"] + assert result.get("protocolVersion") == MODERN_PROTOCOL_VERSION or "supportedVersions" in result + assert "echo" in app._tools + + +def test_modern_stdio_rejects_initialize(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app = _echo_app("modern") + payload = asyncio.run( + _roundtrip( + app, + "modern", + "initialize", + { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy", "version": "1"}, + }, + ) + ) + assert "error" in payload + assert "initialize" in payload["error"]["message"] + + +def test_legacy_stdio_still_answers_initialize(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app = _echo_app("legacy") + payload = asyncio.run( + _roundtrip( + app, + "legacy", + "initialize", + { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "legacy", "version": "1"}, + }, + ) + ) + assert "result" in payload + assert payload["result"]["protocolVersion"] == "2025-06-18" + assert payload["result"]["serverInfo"]["name"] == "stdio-legacy" + + +def test_dual_http_and_stdio_share_the_same_tools(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app = _echo_app("auto") + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + listed = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Method": "tools/list", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/list", + "params": {"_meta": MODERN_META}, + }, + ) + assert listed.status_code == 200, listed.text + http_names = {tool["name"] for tool in listed.json()["result"]["tools"]} + stdio_payload = asyncio.run( + _roundtrip(app, "auto", "tools/list", {"_meta": MODERN_META}) + ) + stdio_names = {tool["name"] for tool in stdio_payload["result"]["tools"]} + assert http_names == stdio_names == {"echo"} diff --git a/tests/test_mcp20_subscriptions.py b/tests/test_mcp20_subscriptions.py new file mode 100644 index 0000000..de0ecc0 --- /dev/null +++ b/tests/test_mcp20_subscriptions.py @@ -0,0 +1,235 @@ +"""MCP 2026 ``/subscriptions/listen`` attach path.""" + +from __future__ import annotations + +import asyncio +import os +import sys + +from pydantic import BaseModel, Field + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.auth.jwt import JWTModule, JWTService +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.transports.headers import SSE_SUBSCRIPTIONS_PATH + +SSE_HEADERS = {"Accept": "text/event-stream"} + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def _make_app(name: str = "listen-app"): + @injectable() + class EchoController: + @tool(name="echo", description="echo", input_schema=EchoInput) + async def echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name=name, controllers=[EchoController]) + class EchoModule: + pass + + @mcp_app(module=EchoModule, server=ServerConfig(name=name)) + class App: + pass + + return asyncio.run(McpApplicationFactory.create(App)) + + +def _header_list(headers: dict[str, str]) -> list[tuple[bytes, bytes]]: + return [(key.lower().encode("latin-1"), value.encode("latin-1")) for key, value in headers.items()] + + +async def _with_lifespan(app, http_call): + started = asyncio.Event() + messages: asyncio.Queue = asyncio.Queue() + await messages.put({"type": "lifespan.startup"}) + + async def receive(): + return await messages.get() + + async def send(message): + if message["type"] == "lifespan.startup.complete": + started.set() + + task = asyncio.create_task(app({"type": "lifespan"}, receive, send)) + await asyncio.wait_for(started.wait(), timeout=2) + try: + return await http_call() + finally: + await messages.put({"type": "lifespan.shutdown"}) + await asyncio.wait_for(task, timeout=2) + + +async def _asgi_listen( + app, + *, + method: str = "GET", + headers: dict[str, str] | None = None, + body: bytes = b"", + after_start=None, +) -> tuple[int, dict[str, str], bytes]: + async def http_call(): + status = 0 + response_headers: dict[str, str] = {} + chunks: list[bytes] = [] + sent_request = False + allow_disconnect = asyncio.Event() + + async def receive(): + nonlocal sent_request + if not sent_request: + sent_request = True + return {"type": "http.request", "body": body, "more_body": False} + await allow_disconnect.wait() + return {"type": "http.disconnect"} + + async def send(message): + nonlocal status + if message["type"] == "http.response.start": + status = message["status"] + response_headers.update( + { + key.decode("latin-1"): value.decode("latin-1") + for key, value in message.get("headers") or [] + } + ) + elif message["type"] == "http.response.body": + chunks.append(message.get("body") or b"") + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": method, + "scheme": "http", + "path": SSE_SUBSCRIPTIONS_PATH, + "raw_path": SSE_SUBSCRIPTIONS_PATH.encode("ascii"), + "query_string": b"", + "headers": _header_list(headers or SSE_HEADERS), + "client": ("testclient", 50000), + "server": ("test", 80), + } + request_task = asyncio.create_task(app(scope, receive, send)) + for _ in range(200): + if status and chunks: + break + if request_task.done(): + break + await asyncio.sleep(0.01) + if after_start is not None: + await after_start() + for _ in range(100): + if b"list_changed" in b"".join(chunks): + break + await asyncio.sleep(0.01) + allow_disconnect.set() + if not request_task.done(): + request_task.cancel() + try: + await request_task + except asyncio.CancelledError: + pass + else: + await request_task + return status, response_headers, b"".join(chunks) + + return await _with_lifespan(app, http_call) + + +def test_modern_and_auto_listen_attach_and_ack(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + app = _make_app("listen-auto") + http_app = app.get_combined_app(json_response=True) + status, headers, body = asyncio.run(_asgi_listen(http_app)) + assert status == 200 + assert "text/event-stream" in headers.get("content-type", "") + assert b"notifications/subscriptions/acknowledged" in body + assert b"toolsListChanged" in body + + +def test_modern_listen_route(monkeypatch): + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "modern") + try: + app = _make_app("listen-modern") + http_app = app.get_combined_app(json_response=True) + status, headers, body = asyncio.run(_asgi_listen(http_app)) + assert status == 200 + assert "text/event-stream" in headers.get("content-type", "") + assert b"subscriptions/acknowledged" in body + finally: + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +def test_legacy_listen_is_absent(monkeypatch): + from starlette.testclient import TestClient + + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "legacy") + try: + app = _make_app("listen-legacy") + http_app = app.get_combined_app(json_response=True) + with TestClient(http_app) as client: + response = client.get(SSE_SUBSCRIPTIONS_PATH, headers=SSE_HEADERS) + assert response.status_code == 404 + finally: + os.environ.pop("NITRO_MCP_PROTOCOL_VERSION", None) + + +def test_listen_forwards_bus_events(monkeypatch): + from mcp.shared.subscriptions import ToolsListChanged + + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + app = _make_app("listen-bus") + bus = app.mcp_server.subscription_bus + http_app = app.get_combined_app(json_response=True) + + async def publish(): + await bus.publish(ToolsListChanged()) + + status, _headers, body = asyncio.run(_asgi_listen(http_app, after_start=publish)) + assert status == 200 + assert b"notifications/subscriptions/acknowledged" in body + assert b"notifications/tools/list_changed" in body + + +def test_unauthenticated_listen_denied_when_jwt_configured(monkeypatch): + monkeypatch.setenv("JWT_SECRET", "listen-jwt") + JWTModule.for_root(secret_env_var="JWT_SECRET", audience="mcp", issuer="nitro") + token = DIContainer.get_instance().resolve(JWTService).create_token({"sub": "ada"}) + app = _make_app("listen-jwt") + + denied = asyncio.run(_asgi_listen(app.get_combined_app(json_response=True))) + assert denied[0] == 401 + assert b"unauthorized" in denied[2] + + status, headers, body = asyncio.run( + _asgi_listen( + app.get_combined_app(json_response=True), + headers={**SSE_HEADERS, "Authorization": f"Bearer {token}"}, + ) + ) + assert status == 200 + assert "text/event-stream" in headers.get("content-type", "") + assert b"subscriptions/acknowledged" in body + + +def test_official_listen_handler_is_registered(): + app = _make_app("listen-handler") + handlers = getattr(app.mcp_server, "_request_handlers", {}) + assert "subscriptions/listen" in handlers + assert app.mcp_server.listen_handler is not None + assert app.mcp_server.subscription_bus is not None diff --git a/tests/test_mcp20_task_authorization.py b/tests/test_mcp20_task_authorization.py new file mode 100644 index 0000000..3c6eb35 --- /dev/null +++ b/tests/test_mcp20_task_authorization.py @@ -0,0 +1,377 @@ +"""Tests for MCP 2.0 task authorization and tenant isolation.""" + +import asyncio +import os +import sys + +import mcp.types as types +import pytest +from mcp import MCPError as McpError +from nitrostack.runtime.request_ctx import Experimental, RequestContext, RequestParamsMeta, request_ctx +from pydantic import BaseModel, Field + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.core.errors import TaskNotFoundError +from nitrostack.core.task import TaskManager +from nitrostack.tasks.authorization import ( + check_task_access, + entry_matches_access_context, + extract_task_access_context, + list_task_wire_data_for_context, +) +from nitrostack.tasks.types import TaskAccessContext, TaskEntry, TaskWireData, utc_now + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class TestCheckTaskAccess: + def test_allows_matching_tenant_and_user(self): + entry = TaskEntry( + task_id="t1", + data=TaskWireData(task_id="t1"), + owner_id="user-a", + tenant_id="tenant-a", + ) + ctx = TaskAccessContext(user_id="user-a", tenant_id="tenant-a") + check_task_access(entry, ctx) + + def test_denies_cross_tenant_with_not_found(self): + entry = TaskEntry( + task_id="t1", + data=TaskWireData(task_id="t1"), + owner_id="user-a", + tenant_id="tenant-a", + ) + ctx = TaskAccessContext(user_id="user-a", tenant_id="tenant-b") + with pytest.raises(TaskNotFoundError) as exc: + check_task_access(entry, ctx) + assert exc.value.task_id == "t1" + + def test_denies_cross_user_with_not_found(self): + entry = TaskEntry( + task_id="t1", + data=TaskWireData(task_id="t1"), + owner_id="user-a", + tenant_id="tenant-a", + ) + ctx = TaskAccessContext(user_id="user-b", tenant_id="tenant-a") + with pytest.raises(TaskNotFoundError): + check_task_access(entry, ctx) + + def test_none_context_allows_internal_access(self): + entry = TaskEntry(task_id="t1", data=TaskWireData(task_id="t1"), owner_id="user-a") + check_task_access(entry, None) + + def test_empty_http_context_denies_owned_task(self): + entry = TaskEntry( + task_id="t1", + data=TaskWireData(task_id="t1"), + owner_id="user-a", + tenant_id="tenant-a", + ) + with pytest.raises(TaskNotFoundError): + check_task_access(entry, TaskAccessContext()) + + def test_partial_context_denies_when_task_has_tenant(self): + entry = TaskEntry( + task_id="t1", + data=TaskWireData(task_id="t1"), + owner_id="alice", + tenant_id="acme", + ) + with pytest.raises(TaskNotFoundError): + check_task_access(entry, TaskAccessContext(user_id="alice")) + + +class TestListFiltering: + def test_filters_by_tenant_and_sorts_desc(self): + import datetime + + now = utc_now() + older = now - datetime.timedelta(seconds=10) + entries = [ + TaskEntry( + task_id="older", + data=TaskWireData(task_id="older", created_at=older, last_updated_at=older), + tenant_id="tenant-a", + ), + TaskEntry( + task_id="newer", + data=TaskWireData(task_id="newer", created_at=now, last_updated_at=now), + tenant_id="tenant-a", + ), + TaskEntry( + task_id="other-tenant", + data=TaskWireData(task_id="other-tenant", created_at=now, last_updated_at=now), + tenant_id="tenant-b", + ), + ] + ctx = TaskAccessContext(tenant_id="tenant-a") + page, next_cursor = list_task_wire_data_for_context(entries, ctx, limit=10) + assert [item.task_id for item in page] == ["newer", "older"] + assert next_cursor is None + assert entry_matches_access_context(entries[2], ctx) is False + + +class TestTaskManagerAuthorization: + def test_get_task_enforces_access_context(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task( + ttl_ms=60_000, + owner_id="alice", + tenant_id="acme", + ) + await manager.get_task(task.id, access_context=TaskAccessContext(user_id="alice", tenant_id="acme")) + with pytest.raises(TaskNotFoundError): + await manager.get_task(task.id, access_context=TaskAccessContext(user_id="bob", tenant_id="acme")) + + asyncio.run(_run()) + + def test_cancel_cross_tenant_returns_not_found(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task(owner_id="alice", tenant_id="acme") + with pytest.raises(TaskNotFoundError): + await manager.cancel_task( + task.id, + access_context=TaskAccessContext(user_id="alice", tenant_id="evil"), + ) + + asyncio.run(_run()) + + +class TestWireHandlersAntiEnumeration: + def test_tasks_get_returns_identical_error_for_missing_and_forbidden(self): + @injectable() + class OwnerController: + @tool(name="owner_tool", description="owner", input_schema=EchoInput, task_support="optional") + async def owner_tool(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="TaskAuth", controllers=[OwnerController]) + class AuthModule: + pass + + @mcp_app(module=AuthModule, server=ServerConfig(name="task-auth")) + class AuthApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(AuthApp) + task = await app.task_manager.create_task( + owner_id="alice", + tenant_id="acme", + ttl_ms=60_000, + ) + handler = app.mcp_server.request_handlers[types.GetTaskRequest] + + missing_token = request_ctx.set( + RequestContext( + request_id="1", + meta=RequestParamsMeta(__pydantic_extra__={"tenantId": "acme", "userId": "alice"}), + session=None, + lifespan_context=None, + ) + ) + try: + with pytest.raises(McpError) as missing_exc: + await handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId="missing-task"), + ) + ) + finally: + request_ctx.reset(missing_token) + + forbidden_token = request_ctx.set( + RequestContext( + request_id="2", + meta=RequestParamsMeta(__pydantic_extra__={"tenantId": "evil", "userId": "bob"}), + session=None, + lifespan_context=None, + ) + ) + try: + with pytest.raises(McpError) as forbidden_exc: + await handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId=task.id), + ) + ) + finally: + request_ctx.reset(forbidden_token) + + assert missing_exc.value.error.code == forbidden_exc.value.error.code + assert missing_exc.value.error.code == types.INVALID_PARAMS + assert "not found" in missing_exc.value.error.message.lower() + assert "not found" in forbidden_exc.value.error.message.lower() + + asyncio.run(_run()) + + +class TestExtractTaskAccessContext: + def test_ignores_spoofed_meta_identity(self): + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta( + __pydantic_extra__={"userId": "u1", "tenantId": "t1", "sessionId": "s1"} + ), + session=None, + lifespan_context=None, + ) + ctx = extract_task_access_context(rc) + assert ctx is not None + assert ctx.user_id is None + assert ctx.tenant_id is None + assert ctx.session_id is None + + def test_uses_verified_jwt_not_meta(self): + from types import SimpleNamespace + + from nitrostack.auth.jwt import JWTService + + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + token = jwt.create_token({"sub": "alice", "tenant_id": "acme"}) + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta( + __pydantic_extra__={"userId": "eve", "tenantId": "evil"} + ), + session=None, + lifespan_context=None, + request=SimpleNamespace(headers={"authorization": f"Bearer {token}"}), + ) + ctx = extract_task_access_context(rc) + assert ctx is not None + assert ctx.user_id == "alice" + assert ctx.tenant_id == "acme" + + def test_failed_jwt_returns_empty_context(self): + from types import SimpleNamespace + + from nitrostack.auth.jwt import JWTService + + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + rc = RequestContext( + request_id="1", + meta=None, + session=None, + lifespan_context=None, + request=SimpleNamespace(headers={"authorization": "Bearer not-a-jwt"}), + ) + ctx = extract_task_access_context(rc) + assert ctx is not None + assert ctx.user_id is None + assert ctx.tenant_id is None + + def test_header_jwt_wins_over_envelope_auth_and_unsigned_identity(self): + from types import SimpleNamespace + + from nitrostack.auth.jwt import JWTService + + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + header_token = jwt.create_token({"sub": "alice", "tenant_id": "acme"}) + envelope_token = jwt.create_token({"sub": "mallory", "tenant_id": "evil"}) + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta.model_validate( + { + "userId": "eve", + "tenantId": "evil", + "io.modelcontextprotocol/auth": { + "authorization": f"Bearer {envelope_token}", + "userId": "mallory", + }, + } + ), + session=None, + lifespan_context=None, + request=SimpleNamespace(headers={"authorization": f"Bearer {header_token}"}), + ) + ctx = extract_task_access_context(rc) + assert ctx is not None + assert ctx.user_id == "alice" + assert ctx.tenant_id == "acme" + + def test_envelope_auth_token_used_when_header_absent(self): + from nitrostack.auth.jwt import JWTService + + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + token = jwt.create_token({"sub": "alice", "tenant_id": "acme"}) + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta.model_validate( + { + "userId": "eve", + "io.modelcontextprotocol/auth": {"authorization": f"Bearer {token}"}, + } + ), + session=None, + lifespan_context=None, + ) + ctx = extract_task_access_context(rc) + assert ctx is not None + assert ctx.user_id == "alice" + assert ctx.tenant_id == "acme" + + def test_unsigned_envelope_auth_identity_is_ignored(self): + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta.model_validate( + { + "io.modelcontextprotocol/auth": { + "userId": "eve", + "tenantId": "evil", + } + } + ), + session=None, + lifespan_context=None, + ) + ctx = extract_task_access_context(rc) + assert ctx is not None + assert ctx.user_id is None + assert ctx.tenant_id is None + + def test_failed_header_jwt_does_not_use_envelope_token(self): + from types import SimpleNamespace + + from nitrostack.auth.jwt import JWTService + + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + envelope_token = jwt.create_token({"sub": "alice", "tenant_id": "acme"}) + rc = RequestContext( + request_id="1", + meta=RequestParamsMeta.model_validate( + {"io.modelcontextprotocol/auth": {"authorization": f"Bearer {envelope_token}"}} + ), + session=None, + lifespan_context=None, + request=SimpleNamespace(headers={"authorization": "Bearer not-a-jwt"}), + ) + ctx = extract_task_access_context(rc) + assert ctx is not None + assert ctx.user_id is None + assert ctx.tenant_id is None diff --git a/tests/test_mcp20_task_store.py b/tests/test_mcp20_task_store.py new file mode 100644 index 0000000..5aa3b5f --- /dev/null +++ b/tests/test_mcp20_task_store.py @@ -0,0 +1,118 @@ +"""Tests for MCP 2.0 task store and distributed persistence.""" + +import asyncio +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack.core.task import TaskManager +from nitrostack.core.errors import TaskNotFoundError +from nitrostack.tasks.eviction import is_terminal_eviction_eligible, should_evict_terminal_task +from nitrostack.tasks.memory import InMemoryTaskStore +from nitrostack.tasks.store import TaskStore +from nitrostack.tasks.types import TaskEntry, TaskWireData, datetime_to_ms, utc_now + + +class TestTaskStoreContract: + def test_in_memory_store_implements_interface(self): + store = InMemoryTaskStore() + assert isinstance(store, TaskStore) + + def test_store_round_trip(self): + async def _run(): + store = InMemoryTaskStore() + now = utc_now() + wire = TaskWireData(task_id="t1", status="working", created_at=now, last_updated_at=now) + entry = TaskEntry(task_id="t1", data=wire, status="working") + await store.set("t1", entry) + assert await store.has("t1") is True + loaded = await store.get("t1") + assert loaded is not None + assert loaded.task_id == "t1" + assert await store.delete("t1") is True + assert await store.get("t1") is None + + asyncio.run(_run()) + + +class TestTerminalEvictionRules: + def test_active_tasks_not_eviction_eligible(self): + assert is_terminal_eviction_eligible("working") is False + assert is_terminal_eviction_eligible("input_required") is False + assert is_terminal_eviction_eligible("completed") is True + + def test_terminal_task_evicted_after_ttl_from_last_updated(self): + now = utc_now() + wire = TaskWireData( + task_id="t2", + status="completed", + created_at=now, + last_updated_at=now, + ttl_ms=1_000, + ) + entry = TaskEntry(task_id="t2", data=wire, status="completed") + now_ms = datetime_to_ms(now) + 2_000 + assert should_evict_terminal_task(entry, now_ms) is True + + def test_active_task_never_evicted_even_when_stale(self): + now = utc_now() + wire = TaskWireData( + task_id="t3", + status="working", + created_at=now, + last_updated_at=now, + ttl_ms=1, + ) + entry = TaskEntry(task_id="t3", data=wire, status="working") + now_ms = datetime_to_ms(now) + 10_000 + assert should_evict_terminal_task(entry, now_ms) is False + + +class TestTaskManagerStoreIntegration: + def test_manager_uses_injected_store(self): + async def _run(): + store = InMemoryTaskStore() + manager = TaskManager(store=store) + task = await manager.create_task(ttl_ms=500) + assert await store.has(task.id) + await manager.complete_task(task.id, {"done": True}) + entries = await store.list() + assert len(entries) == 1 + assert entries[0].result == {"done": True} + + asyncio.run(_run()) + + def test_cleanup_expired_removes_terminal_tasks_from_store(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task(ttl_ms=100) + await manager.complete_task(task.id, {"ok": True}) + entry = await manager._require_entry(task.id) + entry.data.last_updated_at = utc_now() + await manager._store.set(task.id, entry) + import datetime + + entry.data.last_updated_at = utc_now() - datetime.timedelta(milliseconds=200) + await manager._store.set(task.id, entry) + evicted = await manager.cleanup_expired() + assert evicted == 1 + with pytest.raises(TaskNotFoundError): + await manager.get_task(task.id) + + asyncio.run(_run()) + + def test_shared_store_across_managers(self): + """Simulates two replicas reading the same in-memory store.""" + async def _run(): + store = InMemoryTaskStore() + replica_a = TaskManager(store=store) + replica_b = TaskManager(store=store) + task = await replica_a.create_task(ttl_ms=60_000) + await replica_a.complete_task(task.id, {"from": "replica-a"}) + fetched = await replica_b.get_task(task.id) + assert fetched.result == {"from": "replica-a"} + + asyncio.run(_run()) diff --git a/tests/test_mcp20_tasks.py b/tests/test_mcp20_tasks.py new file mode 100644 index 0000000..da6bf1d --- /dev/null +++ b/tests/test_mcp20_tasks.py @@ -0,0 +1,282 @@ +"""Tests for MCP 2.0 Tasks protocol and lifecycle.""" + +import asyncio +import os +import sys + +import mcp.types as types +from mcp import MCPError as McpError +import pytest +from pydantic import BaseModel, Field + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.core.errors import TaskAlreadyTerminalError +from nitrostack.core.task import TaskManager, TaskStatus +from nitrostack.protocol.tasks import ( + DEFAULT_POLL_INTERVAL_MS, + DEFAULT_TASK_TTL_MS, + RESULT_TYPE_TASK, + build_task_create_jsonrpc_result, + ttl_ms_to_seconds, +) +from nitrostack.runtime.request_ctx import Experimental, RequestContext, RequestParamsMeta, request_ctx + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class TestTaskProtocolHelpers: + def test_ttl_ms_to_seconds(self): + assert ttl_ms_to_seconds(600_000) == 600 + assert ttl_ms_to_seconds(None) is None + + def test_task_create_envelope(self): + envelope = build_task_create_jsonrpc_result( + "req-1", + {"taskId": "t1", "status": "working"}, + ) + assert envelope["result"]["resultType"] == RESULT_TYPE_TASK + assert envelope["result"]["task"]["taskId"] == "t1" + + +class TestTaskManagerLifecycle: + def test_input_required_transition(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task() + await manager.require_input(task.id, {"resultType": "input_required", "inputRequests": []}) + snapshot = await manager.get_task(task.id) + assert snapshot.status == TaskStatus.INPUT_REQUIRED + assert snapshot.result["resultType"] == "input_required" + + asyncio.run(_run()) + + def test_resume_from_input_required(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task() + await manager.require_input(task.id, {"pause": True}) + await manager.resume_task(task.id) + assert (await manager.get_task(task.id)).status == TaskStatus.WORKING + + asyncio.run(_run()) + + def test_create_task_stores_ttl_ms(self): + async def _run(): + task = await TaskManager().create_task(ttl_ms=600_000) + assert task.ttl_ms == 600_000 + assert task.ttl == 600_000 + assert task.poll_interval == DEFAULT_POLL_INTERVAL_MS + + asyncio.run(_run()) + + +class TestTaskSupportNegotiation: + def test_forbidden_tool_rejects_task_augmentation(self): + @injectable() + class SyncController: + @tool(name="sync_only", description="sync", input_schema=EchoInput, task_support="forbidden") + async def sync_only(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="TasksForbidden", controllers=[SyncController]) + class ForbiddenModule: + pass + + @mcp_app(module=ForbiddenModule, server=ServerConfig(name="tasks-forbidden")) + class ForbiddenApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(ForbiddenApp) + token = request_ctx.set( + RequestContext( + request_id="1", + meta=None, + session=None, + lifespan_context=None, + experimental=Experimental(task_metadata=types.TaskMetadata(ttl=60_000)), + ) + ) + try: + with pytest.raises(McpError) as exc: + await app._call_tool("sync_only", {"value": "x"}) + assert exc.value.error.code == types.METHOD_NOT_FOUND + finally: + request_ctx.reset(token) + + asyncio.run(_run()) + + def test_required_tool_rejects_sync_call(self): + @injectable() + class HeavyController: + @tool(name="heavy_job", description="heavy", input_schema=EchoInput, task_support="required") + async def heavy_job(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="TasksRequired", controllers=[HeavyController]) + class RequiredModule: + pass + + @mcp_app(module=RequiredModule, server=ServerConfig(name="tasks-required")) + class RequiredApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(RequiredApp) + with pytest.raises(McpError) as exc: + await app._call_tool("heavy_job", {"value": "x"}) + assert exc.value.error.code == types.INVALID_REQUEST + + asyncio.run(_run()) + + +class TestTaskWireHandlers: + def test_tasks_get_embeds_completed_result(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task(ttl_ms=DEFAULT_TASK_TTL_MS) + await manager.complete_task( + task.id, + types.CallToolResult(content=[types.TextContent(type="text", text="done")]), + ) + + @injectable() + class DummyController: + @tool(name="noop", description="noop", input_schema=EchoInput) + async def noop(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="TasksGet", controllers=[DummyController]) + class GetModule: + pass + + @mcp_app(module=GetModule, server=ServerConfig(name="tasks-get")) + class GetApp: + pass + + app = await McpApplicationFactory.create(GetApp) + app.task_manager = manager + handler = app.mcp_server.request_handlers[types.GetTaskRequest] + response = await handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId=task.id), + ) + ) + payload = response.model_dump(by_alias=True) + assert payload["status"] == "completed" + assert payload["result"]["content"][0]["text"] == "done" + + asyncio.run(_run()) + + def test_tasks_list_deprecated_on_modern_protocol(self): + @injectable() + class DummyController: + @tool(name="noop2", description="noop", input_schema=EchoInput) + async def noop2(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="TasksList", controllers=[DummyController]) + class ListModule: + pass + + @mcp_app( + module=ListModule, + server=ServerConfig(name="tasks-list", protocol_era="modern"), + ) + class ListApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(ListApp) + handler = app.mcp_server.request_handlers[types.ListTasksRequest] + with pytest.raises(McpError) as exc: + await handler(types.ListTasksRequest(method="tasks/list", params={})) + assert exc.value.error.code == types.METHOD_NOT_FOUND + + asyncio.run(_run()) + + def test_cancel_terminal_task_raises_invalid_params(self): + async def _run(): + manager = TaskManager() + task = await manager.create_task() + await manager.complete_task(task.id, {"ok": True}) + + @injectable() + class DummyController: + @tool(name="noop3", description="noop", input_schema=EchoInput) + async def noop3(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="TasksCancel", controllers=[DummyController]) + class CancelModule: + pass + + @mcp_app(module=CancelModule, server=ServerConfig(name="tasks-cancel")) + class CancelApp: + pass + + app = await McpApplicationFactory.create(CancelApp) + app.task_manager = manager + handler = app.mcp_server.request_handlers[types.CancelTaskRequest] + with pytest.raises(McpError) as exc: + await handler( + types.CancelTaskRequest( + method="tasks/cancel", + params=types.CancelTaskRequestParams(taskId=task.id), + ) + ) + assert exc.value.error.code == types.INVALID_PARAMS + assert "terminal" in exc.value.error.message.lower() + + asyncio.run(_run()) + + def test_task_augmented_call_returns_task_handle(self): + @injectable() + class AsyncController: + @tool(name="slow_echo", description="slow", input_schema=EchoInput, task_support="optional") + async def slow_echo(self, input: EchoInput, context: ExecutionContext) -> str: + return input.value + + @module(name="TasksCreate", controllers=[AsyncController]) + class CreateModule: + pass + + @mcp_app(module=CreateModule, server=ServerConfig(name="tasks-create")) + class CreateApp: + pass + + async def _run(): + app = await McpApplicationFactory.create(CreateApp) + token = request_ctx.set( + RequestContext( + request_id="1", + meta=None, + session=None, + lifespan_context=None, + experimental=Experimental(task_metadata=types.TaskMetadata(ttl=120_000)), + ) + ) + try: + result = await app._call_tool("slow_echo", {"value": "async"}) + assert isinstance(result, types.CreateTaskResult) + assert result.task.status == "working" + assert result.task.ttl == 120_000 + finally: + request_ctx.reset(token) + + asyncio.run(_run()) diff --git a/tests/test_mcp20_tool_args.py b/tests/test_mcp20_tool_args.py new file mode 100644 index 0000000..f619cfe --- /dev/null +++ b/tests/test_mcp20_tool_args.py @@ -0,0 +1,223 @@ +"""Envelope keys are stripped from tool arguments before user handlers.""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys + +import mcp.types as types +from pydantic import BaseModel, ConfigDict, Field +from starlette.testclient import TestClient + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from nitrostack import ExecutionContext, injectable, module, tool +from nitrostack.auth.jwt import JWTService +from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app +from nitrostack.core.di import DIContainer +from nitrostack.core.task import TaskStatus +from nitrostack.protocol.meta import MCP_META_PREFIX, strip_tool_arguments + +CALL_HEADERS = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "2025-06-18", + "Mcp-Method": "tools/call", + "Mcp-Name": "echo", +} + + +class LooseInput(BaseModel): + model_config = ConfigDict(extra="allow") + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +def test_strip_drops_meta_and_namespaced_keys(): + cleaned = strip_tool_arguments( + { + "value": "ok", + "_meta": {"userId": "eve", "auth": {"token": "secret"}}, + f"{MCP_META_PREFIX}auth": {"userId": "mallory"}, + f"{MCP_META_PREFIX}protocolVersion": "2026-07-28", + } + ) + assert cleaned == {"value": "ok"} + assert "_meta" not in cleaned + assert all(not str(key).startswith(MCP_META_PREFIX) for key in cleaned) + + +def test_strip_drops_envelope_keys_inside_legacy_input_wrap(): + cleaned = strip_tool_arguments( + { + "input": { + "value": "ok", + "_meta": {"userId": "eve"}, + f"{MCP_META_PREFIX}trace": {"id": "nested"}, + }, + "_meta": {"trace": {"id": "outer"}}, + } + ) + assert cleaned == {"input": {"value": "ok"}} + + +def test_strip_leaves_nested_user_meta_values(): + cleaned = strip_tool_arguments({"payload": {"_meta": "user-owned"}}) + assert cleaned == {"payload": {"_meta": "user-owned"}} + + +def test_strip_empty_and_none(): + assert strip_tool_arguments(None) == {} + assert strip_tool_arguments({}) == {} + + +def _echo_app(seen: dict, *, task_support: str = "optional"): + @injectable() + class EchoController: + @tool( + name="echo", + description="echo", + input_schema=LooseInput, + task_support=task_support, + ) + async def echo(self, input: LooseInput, context: ExecutionContext) -> dict: + seen["keys"] = set(input.model_dump().keys()) + seen["extra"] = dict(getattr(input, "model_extra", None) or {}) + seen["value"] = input.value + seen["user"] = context.user + seen["trace"] = context.rpc_meta.trace if context.rpc_meta else None + return {"value": input.value, "keys": sorted(seen["keys"])} + + @module(name="StripToolArgs", controllers=[EchoController]) + class EchoModule: + pass + + @mcp_app(module=EchoModule, server=ServerConfig(name="strip-tool-args")) + class EchoApp: + pass + + return asyncio.run(McpApplicationFactory.create(EchoApp)) + + +def test_handler_input_never_contains_meta(): + seen: dict = {} + app = _echo_app(seen) + result = asyncio.run( + app._call_tool( + "echo", + { + "value": "ok", + "_meta": {"userId": "eve", "trace": {"id": "from-args"}}, + f"{MCP_META_PREFIX}auth": {"userId": "mallory"}, + }, + ) + ) + assert result.is_error is not True + assert seen["value"] == "ok" + assert "_meta" not in seen["keys"] + assert "_meta" not in seen["extra"] + assert all(not str(key).startswith(MCP_META_PREFIX) for key in seen["keys"]) + assert all(not str(key).startswith(MCP_META_PREFIX) for key in seen["extra"]) + + +def test_task_path_strips_arguments_before_handler(): + seen: dict = {} + app = _echo_app(seen, task_support="optional") + created = asyncio.run( + app._call_tool( + "echo", + { + "value": "task-ok", + "_meta": {"userId": "eve"}, + f"{MCP_META_PREFIX}protocolVersion": "2026-07-28", + }, + task=types.TaskMetadata(ttl=60_000), + ) + ) + task_id = created.task.task_id + finished = asyncio.run(app.task_manager.wait_until_done(task_id)) + assert finished.status == TaskStatus.COMPLETED + assert seen["value"] == "task-ok" + assert "_meta" not in seen["keys"] + assert "_meta" not in seen["extra"] + + +def test_http_strips_argument_meta_and_keeps_envelope_trace(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + seen: dict = {} + app = _echo_app(seen) + http_app = app.get_combined_app(json_response=True) + try: + with TestClient(http_app) as client: + response = client.post( + "/mcp", + headers=CALL_HEADERS, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "echo", + "arguments": { + "value": "ok", + "_meta": {"userId": "eve", "trace": {"id": "from-args"}}, + f"{MCP_META_PREFIX}auth": {"token": "secret"}, + }, + "_meta": {"trace": {"id": "from-envelope"}}, + }, + }, + ) + assert response.status_code == 200, response.text + result = response.json()["result"] + body = result.get("structuredContent") or json.loads(result["content"][0]["text"]) + assert body["value"] == "ok" + assert "_meta" not in body["keys"] + assert seen["trace"] == {"id": "from-envelope"} + assert seen["user"] is None + assert "_meta" not in seen["extra"] + finally: + DIContainer.reset() + + +def test_http_jwt_still_comes_from_header_not_argument_meta(monkeypatch): + monkeypatch.delenv("NITRO_MCP_PROTOCOL_VERSION", raising=False) + monkeypatch.delenv("MCP_STATELESS", raising=False) + seen: dict = {} + jwt = JWTService() + DIContainer.get_instance().register_value(JWTService, jwt) + token = jwt.create_token({"sub": "alice", "tenant_id": "acme"}) + app = _echo_app(seen) + http_app = app.get_combined_app(json_response=True) + try: + with TestClient(http_app) as client: + response = client.post( + "/mcp", + headers={**CALL_HEADERS, "Authorization": f"Bearer {token}"}, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "echo", + "arguments": { + "value": "ok", + "_meta": {"userId": "eve", "authorization": "Bearer spoofed"}, + }, + }, + }, + ) + assert response.status_code == 200, response.text + assert seen["user"] == "alice" + assert "_meta" not in seen["keys"] + finally: + DIContainer.reset() diff --git a/tests/test_oauth.py b/tests/test_oauth.py index b1e140e..df66193 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -319,6 +319,12 @@ def test_discovery_document_includes_registration_endpoint_when_given(): ) metadata = build_authorization_server_metadata(service, registration_endpoint="/oauth/v2/register") assert metadata["registration_endpoint"] == "/oauth/v2/register" + absolute = build_authorization_server_metadata( + service, + registration_endpoint="/oauth/v2/register", + public_origin="https://mcp.example.com", + ) + assert absolute["registration_endpoint"] == "https://mcp.example.com/oauth/v2/register" print("Success! registration_endpoint included when supplied.") @@ -334,6 +340,7 @@ def test_protected_resource_metadata_shape(): "resource": "https://api.example.com", "authorization_servers": ["https://idp.example.com"], "scopes_supported": ["read", "write"], + "bearer_methods_supported": ["header"], } print("Success! RFC 9728 document matches expected shape.") diff --git a/tests/test_pizzaz_widgets.py b/tests/test_pizzaz_widgets.py index 67b97fa..124e449 100644 --- a/tests/test_pizzaz_widgets.py +++ b/tests/test_pizzaz_widgets.py @@ -145,9 +145,9 @@ async def run(): ) ) result = resp.root - assert result.structuredContent is not None - assert len(result.structuredContent["shops"]) == 2 - assert result.structuredContent["totalShops"] == 2 + assert result.structured_content is not None + assert len(result.structured_content["shops"]) == 2 + assert result.structured_content["totalShops"] == 2 assert result.meta["ui"]["resourceUri"] == "ui://widget/pizza-list.html" assert result.meta["openai/outputTemplate"] == "ui://widget/pizza-list.html" embedded = next(b for b in result.content if getattr(b, "type", None) == "resource") diff --git a/tests/test_tasks.py b/tests/test_tasks.py index a258276..b1dad3c 100644 --- a/tests/test_tasks.py +++ b/tests/test_tasks.py @@ -32,25 +32,75 @@ TaskStatus, is_terminal_status, ) +from nitrostack.tasks.types import utc_now import mcp.types as types -from mcp.server.lowlevel.server import request_ctx, RequestContext -from mcp.server.experimental.request_context import Experimental +from nitrostack.runtime.request_ctx import Experimental, RequestContext, RequestParamsMeta, request_ctx # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -def _manager() -> TaskManager: - return TaskManager() +def run(coro): + return asyncio.run(coro) -def _force_expire(manager: TaskManager, task_id: str) -> None: - """Set expires_at in the past so the next access lazily expires the task.""" - entry = manager._tasks[task_id] - entry.data.expires_at = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( - seconds=1 - ) +class SyncManager: + """Sync facade over async ``TaskManager`` for unit tests.""" + + def __init__(self, manager: TaskManager | None = None) -> None: + self.m = manager or TaskManager() + + def create_task(self, *args, **kwargs): + return run(self.m.create_task(*args, **kwargs)) + + def get_task(self, task_id: str): + return run(self.m.get_task(task_id)) + + def update_progress(self, task_id: str, progress): + return run(self.m.update_progress(task_id, progress)) + + def complete_task(self, task_id: str, result): + return run(self.m.complete_task(task_id, result)) + + def fail_task(self, task_id: str, error): + return run(self.m.fail_task(task_id, error)) + + def cancel_task(self, task_id: str): + return run(self.m.cancel_task(task_id)) + + def list_tasks(self): + return run(self.m.list_tasks()) + + def has_task(self, task_id: str): + return run(self.m.has_task(task_id)) + + def get_result(self, task_id: str): + return run(self.m.get_result(task_id)) + + def cleanup_expired(self, now_ms: int | None = None): + return run(self.m.cleanup_expired(now_ms)) + + @property + def raw(self) -> TaskManager: + return self.m + + +def _manager() -> SyncManager: + return SyncManager() + + +async def _backdate_task(manager: TaskManager, task_id: str, *, age_ms: int) -> None: + entry = await manager._require_entry(task_id) + entry.data.last_updated_at = utc_now() - datetime.timedelta(milliseconds=age_ms) + await manager._store.set(task_id, entry) + + +async def _age_terminal_task(manager: TaskManager, task_id: str, *, age_ms: int) -> None: + """Backdate a terminal task so ``cleanup_expired`` can evict it.""" + entry = await manager._require_entry(task_id) + entry.data.last_updated_at = utc_now() - datetime.timedelta(milliseconds=age_ms) + await manager._store.set(task_id, entry) # =========================================================================== @@ -84,20 +134,19 @@ def test_creates_task_with_working_status(self): assert task.status == TaskStatus.WORKING assert task.id assert task.created_at is not None - assert task.progress == "Task started" + assert task.progress == "Task created" def test_default_ttl_is_none_never_expires(self): task = _manager().create_task() assert task.ttl_seconds is None + assert task.ttl_ms is None assert task.expires_at is None - def test_custom_ttl_sets_expires_at(self): - before = datetime.datetime.now(datetime.timezone.utc) + def test_custom_ttl_records_ttl_ms(self): task = _manager().create_task(ttl_seconds=60) - after = datetime.datetime.now(datetime.timezone.utc) assert task.ttl_seconds == 60 - assert task.expires_at is not None - assert before <= task.expires_at - datetime.timedelta(seconds=60) <= after + assert task.ttl_ms == 60_000 + assert task.expires_at is None def test_generates_unique_task_ids(self): manager = _manager() @@ -239,12 +288,12 @@ def test_list_tasks_empty(self): def test_wait_until_done_returns_completed(self): async def _run(): - manager = _manager() - task = manager.create_task() + manager = TaskManager() + task = await manager.create_task() async def finish(): await asyncio.sleep(0.05) - manager.complete_task(task.id, "async-ok") + await manager.complete_task(task.id, "async-ok") asyncio.create_task(finish()) done = await manager.wait_until_done(task.id) @@ -259,50 +308,36 @@ async def finish(): # =========================================================================== class TestTTLExpiration: - def test_task_without_ttl_never_expires(self): + def test_active_task_is_never_evicted_during_execution(self): manager = _manager() - task = manager.create_task() - # Even if we wait a bit, no expires_at means still WORKING - time.sleep(0.01) + task = manager.create_task(ttl_ms=1) + time.sleep(0.02) data = manager.get_task(task.id) assert data.status == TaskStatus.WORKING - assert data.expires_at is None - def test_expired_task_transitions_to_expired_on_read(self): - manager = _manager() - task = manager.create_task(ttl_seconds=1) - _force_expire(manager, task.id) - data = manager.get_task(task.id) - assert data.status == TaskStatus.EXPIRED - assert is_terminal_status(data.status) + def test_terminal_task_evicted_after_post_completion_ttl(self): + async def _run_flow(): + manager = TaskManager() + task = await manager.create_task(ttl_ms=100) + await manager.complete_task(task.id, {"ok": True}) + await _age_terminal_task(manager, task.id, age_ms=200) + evicted = await manager.cleanup_expired() + assert evicted == 1 + with pytest.raises(TaskNotFoundError): + await manager.get_task(task.id) - def test_update_progress_on_expired_raises_task_expired(self): - manager = _manager() - task = manager.create_task(ttl_seconds=1) - _force_expire(manager, task.id) - with pytest.raises(TaskExpiredError): - manager.update_progress(task.id, "nope") + run(_run_flow()) - def test_complete_on_expired_raises_task_expired(self): - manager = _manager() - task = manager.create_task(ttl_seconds=1) - _force_expire(manager, task.id) - with pytest.raises(TaskExpiredError): - manager.complete_task(task.id, "late") + def test_cleanup_skips_active_tasks(self): + async def _run_flow(): + manager = TaskManager() + task = await manager.create_task(ttl_ms=50) + await _backdate_task(manager, task.id, age_ms=200) + evicted = await manager.cleanup_expired() + assert evicted == 0 + assert (await manager.get_task(task.id)).status == TaskStatus.WORKING - def test_cancel_on_expired_raises_task_expired(self): - manager = _manager() - task = manager.create_task(ttl_seconds=1) - _force_expire(manager, task.id) - with pytest.raises(TaskExpiredError): - manager.cancel_task(task.id) - - def test_get_result_on_expired_raises(self): - manager = _manager() - task = manager.create_task(ttl_seconds=1) - _force_expire(manager, task.id) - with pytest.raises(TaskExpiredError): - manager.get_result(task.id) + run(_run_flow()) # =========================================================================== @@ -311,25 +346,33 @@ def test_get_result_on_expired_raises(self): class TestTaskContext: def test_update_progress_delegates_to_manager(self): - manager = _manager() - task = manager.create_task() - ctx = TaskContext(task.id, manager) - ctx.update_progress("via context") - assert ctx.progress_message == "via context" - assert manager.get_task(task.id).progress == "via context" + async def _run_flow(): + manager = TaskManager() + task = await manager.create_task() + ctx = TaskContext(task.id, manager) + ctx.update_progress("via context") + await asyncio.sleep(0.01) + assert ctx.progress_message == "via context" + assert (await manager.get_task(task.id)).progress == "via context" + + run(_run_flow()) def test_cancel_sets_cancelled(self): - manager = _manager() - task = manager.create_task() - ctx = TaskContext(task.id, manager) - ctx.cancel() - assert ctx.is_cancelled is True - assert manager.get_task(task.id).status == TaskStatus.CANCELLED + async def _run_flow(): + manager = TaskManager() + task = await manager.create_task() + ctx = TaskContext(task.id, manager) + ctx.cancel() + await asyncio.sleep(0.01) + assert ctx.is_cancelled is True + assert (await manager.get_task(task.id)).status == TaskStatus.CANCELLED + + run(_run_flow()) def test_throw_if_cancelled_raises(self): manager = _manager() task = manager.create_task() - ctx = TaskContext(task.id, manager) + ctx = TaskContext(task.id, manager.raw) manager.cancel_task(task.id) with pytest.raises(TaskCancelledError): ctx.throw_if_cancelled() @@ -337,15 +380,16 @@ def test_throw_if_cancelled_raises(self): def test_throw_if_cancelled_noop_when_working(self): manager = _manager() task = manager.create_task() - ctx = TaskContext(task.id, manager) + ctx = TaskContext(task.id, manager.raw) ctx.throw_if_cancelled() # should not raise - def test_update_progress_after_cancel_does_not_raise(self): + def test_update_progress_after_cancel_raises(self): manager = _manager() task = manager.create_task() - ctx = TaskContext(task.id, manager) + ctx = TaskContext(task.id, manager.raw) manager.cancel_task(task.id) - ctx.update_progress("ignored") # swallowed + with pytest.raises(TaskAlreadyTerminalError): + ctx.update_progress("ignored") # =========================================================================== @@ -434,7 +478,7 @@ async def _mcp_task_flow(): params=types.CallToolRequestParams( name="delayed_tool", arguments={"input": {"duration": 0.2}}, - task=types.TaskMetadata(ttl=60), + task=types.TaskMetadata(ttl=60_000), ), ) @@ -455,15 +499,7 @@ async def _mcp_task_flow(): request_ctx.reset(token) assert isinstance(response.root, types.CreateTaskResult) - task_id = response.root.task.taskId - - list_req = types.ListTasksRequest( - method="tasks/list", - params=types.PaginatedRequestParams(), - ) - list_handler = harness.app.mcp_server.request_handlers[types.ListTasksRequest] - list_res = await list_handler(list_req) - assert any(t.taskId == task_id for t in list_res.tasks) + task_id = response.root.task.task_id get_req = types.GetTaskRequest( method="tasks/get", @@ -473,17 +509,15 @@ async def _mcp_task_flow(): get_res = await get_handler(get_req) assert get_res.status == "working" - result_req = types.GetTaskPayloadRequest( - method="tasks/result", - params=types.GetTaskPayloadRequestParams(taskId=task_id), - ) - result_handler = harness.app.mcp_server.request_handlers[types.GetTaskPayloadRequest] - result_res = await result_handler(result_req) - assert result_res.isError is False - assert "Success payload!" in result_res.content[0].text + deadline = time.time() + 5 + get_res2 = get_res + while get_res2.status not in ("completed", "failed", "cancelled") and time.time() < deadline: + await asyncio.sleep(0.05) + get_res2 = await get_handler(get_req) - get_res2 = await get_handler(get_req) assert get_res2.status == "completed" + assert get_res2.result is not None + assert "Success payload!" in get_res2.result["content"][0]["text"] # Cancellation path req_cancel = types.CallToolRequest( @@ -491,7 +525,7 @@ async def _mcp_task_flow(): params=types.CallToolRequestParams( name="delayed_tool", arguments={"input": {"duration": 1.0}}, - task=types.TaskMetadata(ttl=60), + task=types.TaskMetadata(ttl=60_000), ), ) token = request_ctx.set( @@ -509,7 +543,7 @@ async def _mcp_task_flow(): finally: request_ctx.reset(token) - task_id_cancel = resp_cancel.root.task.taskId + task_id_cancel = resp_cancel.root.task.task_id cancel_req = types.CancelTaskRequest( method="tasks/cancel", params=types.CancelTaskRequestParams(taskId=task_id_cancel), @@ -518,13 +552,13 @@ async def _mcp_task_flow(): cancel_res = await cancel_handler(cancel_req) assert cancel_res.status == "cancelled" - result_req_c = types.GetTaskPayloadRequest( - method="tasks/result", - params=types.GetTaskPayloadRequestParams(taskId=task_id_cancel), + get_res_cancel = await get_handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId=task_id_cancel), + ) ) - payload_c = await result_handler(result_req_c) - assert payload_c.isError is True - assert "cancelled" in payload_c.content[0].text.lower() + assert get_res_cancel.status == "cancelled" def test_mcp_task_integration(): diff --git a/tests/test_tool_input_schema.py b/tests/test_tool_input_schema.py index ed144b5..8a9af3e 100644 --- a/tests/test_tool_input_schema.py +++ b/tests/test_tool_input_schema.py @@ -236,7 +236,7 @@ def test_listed_tools_expose_top_level_fields_not_input_wrap(): } assert set(tools) >= set(expected) for name, fields in expected.items(): - schema = tools[name].inputSchema + schema = tools[name].input_schema _assert_inspector_schema(schema, fields) properties = schema.get("properties") or {} if fields: diff --git a/tests/test_transport_http.py b/tests/test_transport_http.py index 9c0b8e6..1ebf277 100644 --- a/tests/test_transport_http.py +++ b/tests/test_transport_http.py @@ -75,7 +75,13 @@ def test_cors_disabled_rejects_disallowed_origin(monkeypatch): def test_stateful_tool_call_without_session_is_rejected(): app = _app() - http_app = build_http_app(app, enable_cors=True, stateless=False, json_response=True) + http_app = build_http_app( + app, + enable_cors=True, + protocol_era="legacy", + wire_mode="sessionful", + json_response=True, + ) with TestClient(http_app) as client: resp = client.post( "/mcp", diff --git a/tests/test_transports.py b/tests/test_transports.py index 23bcabc..ecf006d 100644 --- a/tests/test_transports.py +++ b/tests/test_transports.py @@ -18,8 +18,7 @@ from starlette.testclient import TestClient import mcp.types as types -from mcp.server.lowlevel.server import request_ctx, RequestContext -from mcp.server.experimental.request_context import Experimental +from nitrostack.runtime.request_ctx import Experimental, RequestContext, RequestParamsMeta, request_ctx from nitrostack import injectable, module, tool, ExecutionContext, DIContainer from nitrostack.core.app import McpApplication, McpApplicationFactory, ServerConfig, mcp_app from nitrostack.transports.http import build_http_app @@ -97,6 +96,11 @@ class _TestApp: return await McpApplicationFactory.create(_TestApp) +def _sessionful_http(app, **kwargs): + kwargs.setdefault("enable_cors", True) + return build_http_app(app, protocol_era="legacy", wire_mode="sessionful", **kwargs) + + def _initialize(client: TestClient) -> str: resp = client.post("/mcp", headers=JSON_HEADERS, json=INITIALIZE_BODY) assert resp.status_code == 200, resp.text @@ -150,7 +154,7 @@ def _extract_json_rpc(resp) -> dict: def test_http_health_and_cors(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app) as client: health = client.get("/mcp/health") @@ -165,7 +169,11 @@ def test_http_health_and_cors(): ) assert preflight.status_code == 200 assert preflight.headers.get("access-control-allow-origin") == "*" - assert "Mcp-Session-Id" in preflight.headers.get("access-control-allow-headers", "") + allow_headers = preflight.headers.get("access-control-allow-headers", "") + assert "Mcp-Session-Id" in allow_headers + assert "Mcp-Name" in allow_headers + assert "Mcp-Method" in allow_headers + assert "MCP-Protocol-Version" in allow_headers root = client.get("/") assert root.status_code == 200 @@ -196,7 +204,7 @@ def test_http_health_and_cors(): def test_http_tool_call_parity(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app) as client: session_id = _initialize(client) @@ -231,7 +239,7 @@ def test_http_tool_call_parity(): def test_session_isolation_and_termination(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app) as client: session_a = _initialize(client) @@ -260,7 +268,7 @@ def test_session_isolation_and_termination(): def test_max_sessions_cap(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True, max_sessions=1) + http_app = _sessionful_http(app, max_sessions=1) with TestClient(http_app) as client: _initialize(client) # first session: at capacity now @@ -308,7 +316,7 @@ def test_stateless_mode_skips_handshake(): def test_di_singletons_shared_across_transports(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app) as client: session_id = _initialize(client) @@ -357,7 +365,7 @@ async def _test_progress_notifications_pushed(): params=types.CallToolRequestParams( name="progress_task", arguments={"input": {"value": ""}}, - task=types.TaskMetadata(ttl=60), + task=types.TaskMetadata(ttl=60_000), _meta={"progressToken": "tok-abc"}, ), ) @@ -377,14 +385,18 @@ async def _test_progress_notifications_pushed(): request_ctx.reset(token) assert isinstance(response.root, types.CreateTaskResult) - task_id = response.root.task.taskId + task_id = response.root.task.task_id # Wait for the background task to finish (it does 3 quick progress updates). - result_handler = app.mcp_server.request_handlers[types.GetTaskPayloadRequest] - result_req = types.GetTaskPayloadRequest( - method="tasks/result", params=types.GetTaskPayloadRequestParams(taskId=task_id) + get_handler = app.mcp_server.request_handlers[types.GetTaskRequest] + get_req = types.GetTaskRequest( + method="tasks/get", params=types.GetTaskRequestParams(taskId=task_id) ) - await result_handler(result_req) + for _ in range(40): + get_res = await get_handler(get_req) + if get_res.status in ("completed", "failed", "cancelled"): + break + await asyncio.sleep(0.05) # Give the fire-and-forget notification tasks a moment to actually run. for _ in range(20): @@ -409,7 +421,7 @@ def test_progress_notifications_pushed(): async def _test_dual_mode_coordinated_shutdown(): app = await _build_app() - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) port = _free_port() stdio_started = asyncio.Event() @@ -474,7 +486,7 @@ def test_dual_mode_coordinated_shutdown(): def test_mcp_path_does_not_redirect(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app, follow_redirects=False) as client: init = client.post("/mcp", headers=JSON_HEADERS, json=INITIALIZE_BODY) @@ -520,7 +532,7 @@ def test_mcp_path_does_not_redirect(): def test_legacy_sse_messages_not_swallowed_by_streamable_http(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app) as client: # Trailing-slash path reaches SseServerTransport. Unknown session → 404 @@ -551,15 +563,13 @@ def test_legacy_sse_messages_not_swallowed_by_streamable_http(): # --------------------------------------------------------------------------- # 11. Client-header tolerance: `StreamableHTTPServerTransport` matches Accept -# media types with `startswith` (so `*/*` is rejected with 406) and rejects -# any MCP-Protocol-Version it doesn't know with 400. Both happen before the -# JSON-RPC layer, so the client just sees a stream open and close with no -# response on it. +# media types with `startswith` (so `*/*` is rejected with 406). Protocol +# version is left on the request so later checks see the client value. # --------------------------------------------------------------------------- def test_wildcard_and_missing_accept_are_honoured(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app) as client: init = client.post( @@ -597,29 +607,47 @@ def test_wildcard_and_missing_accept_are_honoured(): print("Success! Wildcard and absent Accept headers no longer 406 on /mcp.") -def test_unsupported_protocol_version_header_does_not_fail_request(): - app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) - - with TestClient(http_app) as client: - session_id = _initialize(client) +def test_header_compat_preserves_protocol_version_for_inner_app(): + from nitrostack.transports.http import HeaderCompatMiddleware + + captured: dict[str, list] = {} + + async def inner(scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + captured["headers"] = list(scope.get("headers") or []) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send({"type": "http.response.body", "body": b"{}"}) - listed = client.post( + wrapped = HeaderCompatMiddleware(inner) + with TestClient(wrapped) as client: + response = client.post( "/mcp", - headers={**JSON_HEADERS, "mcp-session-id": session_id, "MCP-Protocol-Version": "2026-06-18"}, + headers={**JSON_HEADERS, "MCP-Protocol-Version": "2026-07-28"}, json={"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}, ) - assert listed.status_code == 200, ( - f"a newer-than-supported protocol version should not fail the request, got: {listed.text}" - ) - assert "echo" in [t["name"] for t in _extract_json_rpc(listed)["result"]["tools"]] + assert response.status_code == 200, response.text + headers = {key.decode("latin-1"): value.decode("latin-1") for key, value in captured["headers"]} + assert headers["mcp-protocol-version"] == "2026-07-28" - print("Success! An unknown MCP-Protocol-Version no longer turns into a 400.") + print("Success! HeaderCompatMiddleware keeps MCP-Protocol-Version for the inner app.") def test_delete_terminates_live_session_and_404s_unknown_one(): app = asyncio.run(_build_app()) - http_app = build_http_app(app, enable_cors=True) + http_app = _sessionful_http(app) with TestClient(http_app) as client: session_id = _initialize(client) @@ -695,7 +723,7 @@ def test_oauth_configured_skips_not_supported_stubs(): test_mcp_path_does_not_redirect() test_legacy_sse_messages_not_swallowed_by_streamable_http() test_wildcard_and_missing_accept_are_honoured() - test_unsupported_protocol_version_header_does_not_fail_request() + test_header_compat_preserves_protocol_version_for_inner_app() test_delete_terminates_live_session_and_404s_unknown_one() test_oauth_register_returns_json_not_html() test_oauth_configured_skips_not_supported_stubs() diff --git a/tests/test_widget_metadata.py b/tests/test_widget_metadata.py index d474ba9..907c7ba 100644 --- a/tests/test_widget_metadata.py +++ b/tests/test_widget_metadata.py @@ -76,8 +76,12 @@ def test_widget_metadata_openai_mode(): assert meta.get("ui/template") == uri assert meta.get("openai/outputTemplate") == uri assert "ui" not in meta - assert getattr(target, "outputTemplate", None) == uri - schema = getattr(target, "outputSchema", None) + assert ( + getattr(target, "output_template", None) + or getattr(target, "outputTemplate", None) + or meta.get("openai/outputTemplate") + ) == uri + schema = getattr(target, "output_schema", None) or getattr(target, "outputSchema", None) assert isinstance(schema, dict) assert "status" in (schema.get("properties") or {}) diff --git a/tests/test_widget_parity.py b/tests/test_widget_parity.py index 72ceac8..84a2aa1 100644 --- a/tests/test_widget_parity.py +++ b/tests/test_widget_parity.py @@ -141,8 +141,8 @@ def _assert_project(label: str, project_dir: Path, family: str) -> None: ) ) payload = raw.root - assert payload.isError is not True, f"{label}: {name} error {getattr(payload, 'content', None)}" - assert payload.structuredContent, f"{label}: {name} has no structuredContent" + assert payload.is_error is not True, f"{label}: {name} error {getattr(payload, 'content', None)}" + assert payload.structured_content, f"{label}: {name} has no structuredContent" embedded = [block for block in payload.content if getattr(block, "type", None) == "resource"] assert embedded, f"{label}: {name} tools/call missing EmbeddedResource widget" assert str(embedded[0].resource.uri) == uri diff --git a/tests/test_widgets.py b/tests/test_widgets.py index 33a8cdc..2625533 100644 --- a/tests/test_widgets.py +++ b/tests/test_widgets.py @@ -16,8 +16,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) import mcp.types as types -from mcp.server.lowlevel.server import request_ctx, RequestContext -from mcp.server.experimental.request_context import Experimental +from nitrostack.runtime.request_ctx import Experimental, RequestContext, RequestParamsMeta, request_ctx from pydantic import BaseModel from starlette.testclient import TestClient @@ -285,7 +284,7 @@ class SchemaModule: harness = await NitroTestingModule.create(SchemaModule) tools = await _list_tools(harness) target = next(t for t in tools if t.name == "schema_tool") - schema = getattr(target, "outputSchema", None) + schema = getattr(target, "output_schema", None) or getattr(target, "outputSchema", None) assert isinstance(schema, dict) props = schema.get("properties") or {} assert "value" in props @@ -346,7 +345,11 @@ async def run(): assert meta["ui/template"] == "ui://widget/sample.html" assert meta["openai/outputTemplate"] == "ui://widget/sample.html" assert "ui" not in meta or "resourceUri" not in (meta.get("ui") or {}) - assert getattr(widget_tool, "outputTemplate", None) == "ui://widget/sample.html" + assert ( + getattr(widget_tool, "output_template", None) + or getattr(widget_tool, "outputTemplate", None) + or meta.get("openai/outputTemplate") + ) == "ui://widget/sample.html" assert "openai/outputTemplate" not in plain_meta assert "ui" not in plain_meta @@ -438,7 +441,7 @@ async def run(): assert len(contents) == 1 text = contents[0].text or contents[0].blob assert "widget" in (text or "") - assert contents[0].mimeType == RESOURCE_MIME_TYPE_MCP_APP + assert getattr(contents[0], "mime_type", None) or getattr(contents[0], "mimeType", None) == RESOURCE_MIME_TYPE_MCP_APP meta = contents[0].meta or {} assert meta.get("openai/widgetPrefersBorder") is True assert meta["ui"]["prefersBorder"] is True @@ -458,7 +461,7 @@ async def run(): with app_mode("universal"): resp = await _call_tool_raw(harness, "widget_tool", {"value": "hello"}) result = resp.root - assert result.structuredContent == {"value": "hello", "rendered": True} + assert result.structured_content == {"value": "hello", "rendered": True} assert result.meta["ui"]["resourceUri"] == "ui://widget/sample.html" assert result.meta["openai/outputTemplate"] == "ui://widget/sample.html" types_found = {getattr(block, "type", None) for block in result.content} @@ -520,25 +523,25 @@ async def run(): request_ctx.reset(token) assert isinstance(task_resp.root, types.CreateTaskResult) - task_id = task_resp.root.task.taskId + task_id = task_resp.root.task.task_id - payload_handler = harness.app.mcp_server.request_handlers[types.GetTaskPayloadRequest] + get_handler = harness.app.mcp_server.request_handlers[types.GetTaskRequest] task_payload = None for _ in range(50): await asyncio.sleep(0.02) - raw = await payload_handler( - types.GetTaskPayloadRequest( - method="tasks/result", - params=types.GetTaskPayloadRequestParams(taskId=task_id), + raw = await get_handler( + types.GetTaskRequest( + method="tasks/get", + params=types.GetTaskRequestParams(taskId=task_id), ) ) - task_payload = getattr(raw, "root", raw) - if hasattr(task_payload, "structuredContent"): + if raw.status == "completed" and raw.result is not None: + task_payload = types.CallToolResult(**raw.result) break else: raise AssertionError("task did not complete in time") - assert task_payload.structuredContent == direct.structuredContent + assert task_payload.structured_content == direct.structured_content assert task_payload.meta == direct.meta asyncio.run(run())