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..fbf6c07 --- /dev/null +++ b/nitrostack/auth/cimd.py @@ -0,0 +1,223 @@ +"""Client ID Metadata Document (CIMD) resolution with SSRF defenses (Doc 08).""" + +from __future__ import annotations + +import asyncio +import ipaddress +import json +import socket +import urllib.error +import urllib.parse +import urllib.request +from typing import Any + +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 (Doc 08 §3). + + 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 + + +async def assert_safe_fetch_target(url_str: str, *, allow_loopback: bool = False) -> None: + """DNS pre-resolution and IP range filtering (Doc 08 §4.1).""" + 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: + return + + 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}") + + +def _fetch_cimd_bytes(url_str: str, *, timeout_sec: float) -> bytes: + class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): + raise urllib.error.HTTPError( + url_str, + code, + "HTTP redirects are not allowed for CIMD fetch", + headers, + fp, + ) + + opener = urllib.request.build_opener(_NoRedirectHandler) + 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 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 + + +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 (Doc 08 §4). + + 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) + await assert_safe_fetch_target(normalized, allow_loopback=allow_loopback) + + body = await asyncio.to_thread(_fetch_cimd_bytes, normalized, timeout_sec=timeout_sec) + 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) diff --git a/nitrostack/auth/oauth_module.py b/nitrostack/auth/oauth_module.py index 31f8bab..bf8a9df 100644 --- a/nitrostack/auth/oauth_module.py +++ b/nitrostack/auth/oauth_module.py @@ -62,6 +62,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,6 +73,8 @@ 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) diff --git a/nitrostack/auth/oauth_security.py b/nitrostack/auth/oauth_security.py new file mode 100644 index 0000000..b6586d6 --- /dev/null +++ b/nitrostack/auth/oauth_security.py @@ -0,0 +1,28 @@ +"""OAuth 2.1 security helpers (Doc 08).""" + +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 (Doc 08 §5.1). + + 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/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..c364207 100644 --- a/nitrostack/core/app.py +++ b/nitrostack/core/app.py @@ -12,6 +12,7 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Pattern, Set, Tuple, Type import mcp.types as types +from mcp.shared.exceptions import McpError from mcp.server.lowlevel.server import request_ctx from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.server.stdio import stdio_server @@ -34,6 +35,23 @@ 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 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 +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.protocol.observability import TraceContext, extract_trace_context +from nitrostack.protocol.deprecated import deprecated_method_message +from nitrostack.protocol.tasks import ( + DEFAULT_TASK_TTL_MS, + task_support_forbidden_message, + task_support_required_message, +) +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 @@ -53,14 +71,16 @@ class ServerConfig: name: str version: str = "1.0.0" transport_type: Optional[Literal["stdio", "http", "dual"]] = None + protocol_version: str = MODERN_PROTOCOL_VERSION # 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. - stateless: bool = False + stateless: bool = True 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 +179,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 +262,40 @@ 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) + data: Dict[str, Any] = {} + if raw_meta is None: + return data + + 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 + return data + + +def _trace_context_from_request_ctx(rc: Any) -> TraceContext | None: + return extract_trace_context(_request_meta_from_ctx(rc)) + + def _auth_metadata_from_request_ctx(rc: Any) -> Dict[str, Any]: """Copy host-sent auth slots from MCP request ``_meta`` into ExecutionContext. @@ -255,26 +309,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(): @@ -356,6 +391,7 @@ def _bootstrap(self) -> None: self._assert_declared_dependencies(resolved_modules, container) # Instantiate all providers and controllers to populate container + module_instances: List[Any] = [] for mod in resolved_modules: mod_config = getattr(mod, "_mcp_module_config", None) if mod_config: @@ -364,10 +400,10 @@ def _bootstrap(self) -> None: container.resolve(provider) # Register & Resolve all controllers for controller in mod_config.controllers: - container.resolve(controller) + module_instances.append(container.resolve(controller)) - # 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 instances only + for instance in module_instances: # Scan members of this instance for name, member in inspect.getmembers(instance): # Discover Tools @@ -509,16 +545,11 @@ 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) + 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 @@ -543,18 +574,39 @@ async def health_status_resource(context: ExecutionContext) -> str: # Protocol handler wiring (owned low-level `mcp.server.lowlevel.Server`) # ------------------------------------------------------------------ + def _advertise_tasks_extension(self) -> bool: + return any( + entry.config.task_support in ("optional", "required") + for entry in self._tools.values() + ) + + def _custom_extensions(self) -> Optional[Dict[str, str]]: + extensions = getattr(self.server_config, "extensions", None) + if not extensions: + return None + return dict(extensions) + + def _list_endpoint_cache_meta(self) -> Dict[str, Any]: + return build_list_endpoint_cache_hint_meta() + 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()] + async def _list_tools() -> types.ListToolsResult: + return types.ListToolsResult( + tools=[self._build_tool_definition(entry) for entry in self._tools.values()], + _meta=self._list_endpoint_cache_meta(), + ) @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 {}) @server.list_resources() - async def _list_resources() -> List[types.Resource]: - return [self._build_resource_definition(entry) for entry in self._resources.values()] + async def _list_resources() -> types.ListResourcesResult: + return types.ListResourcesResult( + resources=[self._build_resource_definition(entry) for entry in self._resources.values()], + _meta=self._list_endpoint_cache_meta(), + ) @server.list_resource_templates() async def _list_resource_templates() -> List[types.ResourceTemplate]: @@ -574,8 +626,11 @@ async def _unsubscribe_resource(uri: Any) -> None: return None @server.list_prompts() - async def _list_prompts() -> List[types.Prompt]: - return [self._build_prompt_definition(entry) for entry in self._prompts.values()] + async def _list_prompts() -> types.ListPromptsResult: + return types.ListPromptsResult( + prompts=[self._build_prompt_definition(entry) for entry in self._prompts.values()], + _meta=self._list_endpoint_cache_meta(), + ) @server.get_prompt() async def _get_prompt(name: str, arguments: Optional[Dict[str, str]]) -> types.GetPromptResult: @@ -612,7 +667,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 +698,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 +741,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 +794,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( @@ -826,11 +911,12 @@ async def _call_tool(self, name: str, arguments: Dict[str, Any]): ) cfg = entry.config + tool_arguments, input_responses, request_state = split_mrtr_from_arguments(arguments or {}) # 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 @@ -844,17 +930,41 @@ async def _call_tool(self, name: str, arguments: Dict[str, Any]): if 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 + progress_token = getattr(rc.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.ErrorData( + code=types.METHOD_NOT_FOUND, + message=task_support_forbidden_message(cfg.name), + ) + ) + if cfg.task_support == "required" and task_metadata is None: + raise McpError( + types.ErrorData( + code=types.INVALID_REQUEST, + message=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 async def background_execution(): @@ -862,6 +972,9 @@ async def background_execution(): request_id=str(uuid.uuid4()), tool_name=cfg.name, metadata={"input": input_instance, **auth_meta}, + input_responses=input_responses, + request_state=request_state, + trace=trace, ) task_ctx.task = TaskContext( task_id, @@ -884,12 +997,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 +1017,14 @@ 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}) + ctx = ExecutionContext( + request_id=str(uuid.uuid4()), + tool_name=cfg.name, + metadata={"input": input_instance, **auth_meta}, + input_responses=input_responses, + request_state=request_state, + trace=trace, + ) try: result = await run_pipeline( handler=entry.method, @@ -1020,7 +1147,7 @@ async def _get_prompt(self, name: str, arguments: Dict[str, str]) -> types.GetPr 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( + raise McpError( types.ErrorData( code=types.INVALID_PARAMS, message=f"Task {task.id} has expired", @@ -1032,53 +1159,85 @@ def _task_data_to_mcp_task(self, task) -> types.Task: 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, + 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.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, + } + 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 types.GetTaskResult(**payload) + def _register_task_handlers(self, server: NitroStackMcpServer) -> None: + modern_protocol = self.server_config.protocol_version == MODERN_PROTOCOL_VERSION + async def handle_list_tasks(req): + if modern_protocol: + message = deprecated_method_message("tasks/list") + raise McpError( + types.ErrorData(code=types.METHOD_NOT_FOUND, message=message or "Not supported") + ) tasks_list = [] - for t in self.task_manager.list_tasks(): + access = extract_task_access_context(request_ctx.get(None)) + cursor = getattr(req.params, "cursor", 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 + 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( + raise 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, - ) + return self._build_get_task_result(t) async def handle_cancel_task(req): task_id = req.params.taskId + 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( + raise McpError( types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") ) except TaskExpiredError: - raise types.McpError( + raise McpError( types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} has expired") ) except TaskAlreadyTerminalError as e: - raise types.McpError( + raise McpError( types.ErrorData(code=types.INVALID_PARAMS, message=str(e)) ) mcp_task = self._task_data_to_mcp_task(t) @@ -1093,11 +1252,17 @@ async def handle_cancel_task(req): ) async def handle_get_task_payload(req): + if modern_protocol: + message = deprecated_method_message("tasks/result") + raise McpError( + types.ErrorData(code=types.METHOD_NOT_FOUND, message=message or "Not supported") + ) task_id = req.params.taskId + 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( + raise McpError( types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") ) if t.status == TaskStatus.COMPLETED: @@ -1183,15 +1348,37 @@ def get_combined_app( """ from nitrostack.transports.http import build_http_app - return build_http_app( + effective_stateless = ( + self.server_config.stateless if stateless is None else stateless + ) + 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, ) + if effective_stateless: + from nitrostack.transports.middleware import wrap_stateless_transport + + has_widgets = any( + getattr(entry, "component", None) is not None + for entry in getattr(self, "_tools", {}).values() + ) + http_app = wrap_stateless_transport( + http_app, + server_name=self.server_config.name, + server_version=self.server_config.version, + protocol_version=self.server_config.protocol_version, + advertise_tasks=self._advertise_tasks_extension(), + advertise_app=has_widgets, + custom_extensions=self._custom_extensions(), + ) + + 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()) diff --git a/nitrostack/core/context.py b/nitrostack/core/context.py index 817b431..e316a58 100644 --- a/nitrostack/core/context.py +++ b/nitrostack/core/context.py @@ -2,10 +2,13 @@ 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.observability import TraceContext + # Logger protocol used by ExecutionContext (Section 13) class Logger(Protocol): def debug(self, message: str, meta: dict | None = None) -> None: ... @@ -124,9 +127,10 @@ def update_progress(self, message: str) -> None: manager = self._task_manager if manager is not None: try: - manager.update_progress(self.task_id, message) + import asyncio + + asyncio.create_task(manager.update_progress(self.task_id, message)) except Exception: - # Task may already be terminal/expired — ignore for handler ergonomics. pass self._push_progress_notification(message) @@ -155,7 +159,9 @@ def cancel(self) -> None: if manager is None: return try: - manager.cancel_task(self.task_id) + import asyncio + + asyncio.create_task(manager.cancel_task(self.task_id)) except Exception: pass @@ -163,7 +169,7 @@ 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 @@ -178,3 +184,6 @@ class ExecutionContext: 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 diff --git a/nitrostack/core/decorators.py b/nitrostack/core/decorators.py index 60f4323..960a7ba 100644 --- a/nitrostack/core/decorators.py +++ b/nitrostack/core/decorators.py @@ -47,7 +47,7 @@ class ToolConfig: title: Optional[str] = None output_schema: Optional[Any] = None annotations: ToolAnnotations = field(default_factory=ToolAnnotations) - task_support: Literal["forbidden", "optional", "required"] = "forbidden" + task_support: Literal["forbidden", "optional", "required"] = "optional" visibility: Literal["visible", "hidden"] = "visible" examples: Optional[ToolExamples] = None invocation: Optional[ToolInvocation] = None @@ -96,7 +96,7 @@ def tool( title: Optional[str] = None, output_schema: Optional[Any] = None, annotations: Optional[ToolAnnotations] = None, - task_support: Literal["forbidden", "optional", "required"] = "forbidden", + task_support: Literal["forbidden", "optional", "required"] = "optional", visibility: Literal["visible", "hidden"] = "visible", examples: Optional[ToolExamples] = None, invocation: Optional[ToolInvocation] = None, diff --git a/nitrostack/core/task.py b/nitrostack/core/task.py index 57c292a..f84b805 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 (Phase 1 / Doc 05–06). -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``. +Doc 06 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 Phase 1 / Doc 05.""" 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,333 @@ 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. - - Typical flow:: + Task lifecycle manager backed by a pluggable ``TaskStore``. - 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, + ) + 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, ) - 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: + 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) + + async def update_progress(self, task_id: str, progress: Any) -> None: + """Update progress for an active task.""" + 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 = 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 + await self._store.set(task_id, entry) + + 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)) + 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) + + async def cancel_task( + self, + task_id: str, + *, + access_context: Optional[TaskAccessContext] = None, + ) -> None: + """Transition an active task to ``cancelled``.""" + entry = await self._require_entry(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() + await self._store.set(task_id, entry) + handle = self._runtime.setdefault(task_id, _RuntimeTaskHandle()) + handle.cancelled = True + self._signal_done(task_id) + + 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 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 - - 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) + 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 _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 + 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..716a386 --- /dev/null +++ b/nitrostack/protocol/__init__.py @@ -0,0 +1,120 @@ +"""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 +from nitrostack.protocol.discovery import DISCOVER_RESULT_TYPE, build_discover_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 RequestMeta, extract_request_meta, split_params_and_meta +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, + bound_schema_depth, + 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, + SUPPORTED_PROTOCOL_VERSIONS, +) + +__all__ = [ + "LEGACY_PROTOCOL_VERSION", + "MODERN_PROTOCOL_VERSION", + "SUPPORTED_PROTOCOL_VERSIONS", + "MCPExtensionId", + "RuntimeLayer", + "LEGACY_SESSION_HEADER", + "MAX_CIMD_BYTES", + "MAX_SCHEMA_DEPTH", + "JsonRpcErrorCode", + "ERROR_CODE_MESSAGES", + "RequestMeta", + "extract_request_meta", + "split_params_and_meta", + "deprecated_method_message", + "DISCOVER_RESULT_TYPE", + "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", + "bound_schema_depth", + "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..1467550 --- /dev/null +++ b/nitrostack/protocol/cache_hints.py @@ -0,0 +1,78 @@ +"""Cache hint resolution for tools, resources, and list endpoints (Doc 09 §3).""" + +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 (Doc 09 §3.2). + + 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..d510aaf --- /dev/null +++ b/nitrostack/protocol/constants.py @@ -0,0 +1,10 @@ +"""Cross-cutting security and validation limits (Doc 00 — defense in depth).""" + +# RFC 6890 CIMD resolver payload cap (Doc 08; principle 4) +MAX_CIMD_BYTES = 5120 + +# JSON Schema depth bounding for DoS protection (SEP-2106; Doc 03) +MAX_SCHEMA_DEPTH = 64 + +# Legacy session header that MUST NOT be emitted in stateless mode (Doc 01) +LEGACY_SESSION_HEADER = "Mcp-Session-Id" diff --git a/nitrostack/protocol/contracts.py b/nitrostack/protocol/contracts.py new file mode 100644 index 0000000..f087c9f --- /dev/null +++ b/nitrostack/protocol/contracts.py @@ -0,0 +1,80 @@ +"""Wire contract builders for tools, resources, and prompts (Doc 03).""" + +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; Doc 03 §1.1).""" + 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 (Doc 03 §2.3).""" + 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 (Doc 03 §2.3).""" + 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 (Doc 03 §3.2).""" + 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 (Doc 03 §3.2).""" + return { + "description": description, + "messages": messages, + } diff --git a/nitrostack/protocol/deprecated.py b/nitrostack/protocol/deprecated.py new file mode 100644 index 0000000..707e15d --- /dev/null +++ b/nitrostack/protocol/deprecated.py @@ -0,0 +1,23 @@ +"""Deprecated MCP methods rejected on the 2026-07-28 wire (Doc 02 §4).""" + +from __future__ import annotations + +from typing import Optional + +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 rejection message if method is deprecated on modern wire.""" + return DEPRECATED_MODERN_METHODS.get(method) diff --git a/nitrostack/protocol/discovery.py b/nitrostack/protocol/discovery.py new file mode 100644 index 0000000..4b3c2c3 --- /dev/null +++ b/nitrostack/protocol/discovery.py @@ -0,0 +1,57 @@ +"""server/discover capability negotiation (Doc 01 §4).""" + +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 MODERN_PROTOCOL_VERSION, SUPPORTED_PROTOCOL_VERSIONS + +DISCOVER_RESULT_TYPE = "complete" + + +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 payload (Doc 01 §4 / Doc 09 §1).""" + 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), + } diff --git a/nitrostack/protocol/errors.py b/nitrostack/protocol/errors.py new file mode 100644 index 0000000..fe7ec20 --- /dev/null +++ b/nitrostack/protocol/errors.py @@ -0,0 +1,24 @@ +"""Standard JSON-RPC error codes for MCP 2026-07-28 (Doc 02 §2).""" + +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 + + +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", +} 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..991c0c6 --- /dev/null +++ b/nitrostack/protocol/jsonrpc.py @@ -0,0 +1,186 @@ +"""JSON-RPC 2.0 wire protocol (Doc 02).""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +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 (Doc 01) +PARSE_ERROR = int(JsonRpcErrorCode.PARSE_ERROR) +HEADER_BODY_MISMATCH = int(JsonRpcErrorCode.HEADER_BODY_MISMATCH) + + +@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) + + +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 JsonRpcParseError("Request must be a JSON object") + if payload.get("jsonrpc") != JSONRPC_VERSION: + raise JsonRpcParseError("jsonrpc must be '2.0'") + if "method" not in payload or not isinstance(payload["method"], str): + raise JsonRpcParseError("method is required and must be a string") + + params = payload.get("params") or {} + if not isinstance(params, dict): + raise JsonRpcParseError("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 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}'" + ) + + +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 (Doc 01 §3 step 3).""" + 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 (Doc 02 §3). + 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..5766eb7 --- /dev/null +++ b/nitrostack/protocol/layers.py @@ -0,0 +1,12 @@ +"""Runtime layer identifiers matching the Doc 00 architecture diagram.""" + +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..767e9eb --- /dev/null +++ b/nitrostack/protocol/meta.py @@ -0,0 +1,60 @@ +"""Request _meta envelope parsing (Doc 02 §5).""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Optional + +MCP_META_PREFIX = "io.modelcontextprotocol/" + + +@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 + 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") + 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") + + 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, + 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 diff --git a/nitrostack/protocol/mrtr.py b/nitrostack/protocol/mrtr.py new file mode 100644 index 0000000..5f6caaf --- /dev/null +++ b/nitrostack/protocol/mrtr.py @@ -0,0 +1,153 @@ +"""Multi Round-Trip Request (MRTR) helpers — SEP-2322 (Doc 04).""" + +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 (Doc 04 §3.1).""" + + 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 (Doc 04 §3.2).""" + + 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 from Doc 04 §4. + """ + 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 (Doc 04 §4).""" + 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..6fcc8c3 --- /dev/null +++ b/nitrostack/protocol/observability.py @@ -0,0 +1,50 @@ +"""W3C Trace Context extraction for MCP _meta envelopes (Doc 09 §2).""" + +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 (Doc 09 §2.1).""" + 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..8de9aa2 --- /dev/null +++ b/nitrostack/protocol/resources.py @@ -0,0 +1,63 @@ +"""Resource URI resolution for static resources and templates (Doc 03 §2).""" + +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 (Doc 03 §2.4): + 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..ce2ef86 --- /dev/null +++ b/nitrostack/protocol/schema.py @@ -0,0 +1,111 @@ +"""JSON Schema 2020-12 normalization and depth bounding (Doc 03 §1.2, SEP-2106).""" + +from __future__ import annotations + +from typing import Any + +from nitrostack.protocol.constants import MAX_SCHEMA_DEPTH + +JSON_SCHEMA_2020_12_URI = "https://json-schema.org/draft/2020-12/schema" + +_COMPOSITION_KEYS = frozenset({"allOf", "anyOf", "oneOf"}) +_NESTED_SCHEMA_KEYS = frozenset( + { + "properties", + "patternProperties", + "additionalProperties", + "items", + "prefixItems", + "contains", + "propertyNames", + "if", + "then", + "else", + "not", + } +) + + +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 _NESTED_SCHEMA_KEYS: + if 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 isinstance(value, list): + result[key] = [ + bound_schema_depth(item, max_depth=max_depth, depth=depth + 1) for item in value + ] + 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 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. + """ + 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.""" + 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..48a368a --- /dev/null +++ b/nitrostack/protocol/tasks.py @@ -0,0 +1,43 @@ +"""MCP Tasks protocol helpers — Doc 05 (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 Doc 05 §4.2 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..bd70188 --- /dev/null +++ b/nitrostack/protocol/version.py @@ -0,0 +1,6 @@ +"""MCP protocol version identifiers (Doc 00 — paradigm evolution).""" + +MODERN_PROTOCOL_VERSION = "2026-07-28" +LEGACY_PROTOCOL_VERSION = "2025-06-18" + +SUPPORTED_PROTOCOL_VERSIONS: tuple[str, ...] = (MODERN_PROTOCOL_VERSION,) diff --git a/nitrostack/runtime/__init__.py b/nitrostack/runtime/__init__.py new file mode 100644 index 0000000..e8cc7bc --- /dev/null +++ b/nitrostack/runtime/__init__.py @@ -0,0 +1,34 @@ +"""Stateless runtime policies and invariants (Doc 00).""" + +from nitrostack.runtime.conformance import ( + BLUEPRINT_CONFORMANCE_AREAS, + ConformanceArea, + assert_blueprint_layout, + protocol_version_matches_blueprint, + verify_package_layout, +) +from nitrostack.runtime.epic_acceptance import ( + ACCEPTANCE_CRITERIA, + EPIC_DELIVERABLES, + ImplementationEpic, + acceptance_criteria_registered, + epic_coverage_complete, + modern_protocol_target, +) +from nitrostack.runtime.stateless import StatelessInvariants, assert_stateless_headers + +__all__ = [ + "StatelessInvariants", + "assert_stateless_headers", + "ConformanceArea", + "BLUEPRINT_CONFORMANCE_AREAS", + "assert_blueprint_layout", + "protocol_version_matches_blueprint", + "verify_package_layout", + "ImplementationEpic", + "EPIC_DELIVERABLES", + "ACCEPTANCE_CRITERIA", + "epic_coverage_complete", + "acceptance_criteria_registered", + "modern_protocol_target", +] diff --git a/nitrostack/runtime/conformance.py b/nitrostack/runtime/conformance.py new file mode 100644 index 0000000..9580441 --- /dev/null +++ b/nitrostack/runtime/conformance.py @@ -0,0 +1,102 @@ +"""MCP 2.0 implementation blueprint conformance registry (Doc 10).""" + +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.version import MODERN_PROTOCOL_VERSION + +_PACKAGE_ROOT = Path(__file__).resolve().parent.parent + + +class ConformanceArea(str, Enum): + """Five verification areas from Doc 10 §4.""" + + 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 Doc 10 recommended layout (actual names). +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( + { + "tasks/result", + "tasks/list", + } +) + + +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/epic_acceptance.py b/nitrostack/runtime/epic_acceptance.py new file mode 100644 index 0000000..56afd3f --- /dev/null +++ b/nitrostack/runtime/epic_acceptance.py @@ -0,0 +1,209 @@ +"""PYTHONSDK-11 epic and acceptance-criteria registry (Doc 11 / Plane ticket).""" + +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 ImplementationEpic(IntEnum): + """Seven implementation epics from TICKET_PLANE_PROJECT_MANAGEMENT.md.""" + + 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 EpicDeliverable: + epic: ImplementationEpic + item: str + spec_doc: str + test_module: str + + +@dataclass(frozen=True) +class AcceptanceCriterion: + key: str + description: str + test_module: str + + +EPIC_DELIVERABLES: tuple[EpicDeliverable, ...] = ( + EpicDeliverable( + ImplementationEpic.CORE_JSONRPC, + "JSON-RPC error code hierarchy (-32700..-32603, -32020)", + "02_JSONRPC_AND_WIRE_SPECIFICATION.md", + "tests/test_mcp20_doc02_jsonrpc_wire.py", + ), + EpicDeliverable( + ImplementationEpic.CORE_JSONRPC, + "SEP-2164 missing resource returns -32602", + "02_JSONRPC_AND_WIRE_SPECIFICATION.md", + "tests/test_mcp20_doc03_contracts.py", + ), + EpicDeliverable( + ImplementationEpic.CORE_JSONRPC, + "Reject tasks/result and tasks/list on modern wire", + "02_JSONRPC_AND_WIRE_SPECIFICATION.md", + "tests/test_mcp20_doc05_tasks.py", + ), + EpicDeliverable( + ImplementationEpic.STATELESS_HTTP, + "POST /mcp stateless ingress without Mcp-Session-Id", + "01_STATELESS_HTTP_AND_LIFECYCLE.md", + "tests/test_mcp20_doc01_stateless_http.py", + ), + EpicDeliverable( + ImplementationEpic.STATELESS_HTTP, + "server/discover and ping fast path", + "01_STATELESS_HTTP_AND_LIFECYCLE.md", + "tests/test_mcp20_doc01_stateless_http.py", + ), + EpicDeliverable( + ImplementationEpic.STATELESS_HTTP, + "CORS headers for MCP preflight", + "01_STATELESS_HTTP_AND_LIFECYCLE.md", + "tests/test_mcp20_doc01_stateless_http.py", + ), + EpicDeliverable( + ImplementationEpic.REGISTRIES_SCHEMA, + "JSON Schema 2020-12 with depth bounding (max 64)", + "03_TOOL_RESOURCE_PROMPT_CONTRACTS.md", + "tests/test_mcp20_doc03_contracts.py", + ), + EpicDeliverable( + ImplementationEpic.REGISTRIES_SCHEMA, + "Resource URI templates and resolution", + "03_TOOL_RESOURCE_PROMPT_CONTRACTS.md", + "tests/test_mcp20_doc03_contracts.py", + ), + EpicDeliverable( + ImplementationEpic.ASYNC_TASKS, + "Task state machine and tasks/get embedded result", + "05_MCP_TASKS_PROTOCOL_AND_LIFECYCLE.md", + "tests/test_mcp20_doc05_tasks.py", + ), + EpicDeliverable( + ImplementationEpic.ASYNC_TASKS, + "Pluggable TaskStore and terminal-only TTL eviction", + "06_TASK_STORE_AND_DISTRIBUTED_PERSISTENCE.md", + "tests/test_mcp20_doc06_task_store.py", + ), + EpicDeliverable( + ImplementationEpic.MULTI_TENANT, + "TaskAccessContext isolation and anti-enumeration", + "07_TASK_AUTHORIZATION_AND_TENANT_ISOLATION.md", + "tests/test_mcp20_doc07_task_authorization.py", + ), + EpicDeliverable( + ImplementationEpic.OAUTH_CIMD, + "CIMD resolver with SSRF defenses and RFC 9207 iss", + "08_OAUTH21_CIMD_AND_SECURITY.md", + "tests/test_mcp20_doc08_oauth_cimd.py", + ), + EpicDeliverable( + ImplementationEpic.MRTR_OBSERVABILITY, + "MRTR input_required elicitation helpers", + "04_MRTR_MULTI_ROUND_TRIP_SPEC.md", + "tests/test_mcp20_doc04_mrtr.py", + ), + EpicDeliverable( + ImplementationEpic.MRTR_OBSERVABILITY, + "Extensions map, trace context, and cache hints", + "09_EXTENSIONS_CACHE_AND_OBSERVABILITY.md", + "tests/test_mcp20_doc09_extensions_cache_observability.py", + ), +) + +ACCEPTANCE_CRITERIA: tuple[AcceptanceCriterion, ...] = ( + AcceptanceCriterion( + "stateless", + "No Mcp-Session-Id emitted; stateless POST /mcp works", + "tests/test_mcp20_doc10_blueprint.py", + ), + AcceptanceCriterion( + "task_conformance", + "Task-augmented tools/call returns immediate TaskData", + "tests/test_mcp20_doc05_tasks.py", + ), + AcceptanceCriterion( + "task_cancellation", + "tasks/cancel marks cancelled; terminal cancel returns -32602", + "tests/test_mcp20_doc05_tasks.py", + ), + AcceptanceCriterion( + "ttl_safety", + "Active tasks never evicted; terminal TTL from lastUpdatedAt", + "tests/test_mcp20_doc06_task_store.py", + ), + AcceptanceCriterion( + "multi_tenant", + "Cross-tenant access returns TaskNotFoundError / -32602", + "tests/test_mcp20_doc07_task_authorization.py", + ), + AcceptanceCriterion( + "ssrf_security", + "CIMD blocks special-use IPs and oversized payloads", + "tests/test_mcp20_doc08_oauth_cimd.py", + ), + AcceptanceCriterion( + "error_parity", + "JSON-RPC codes match specification including SEP-2164", + "tests/test_mcp20_doc02_jsonrpc_wire.py", + ), + AcceptanceCriterion( + "automated_tests", + "pytest suite covers all seven epics via doc00–doc11 modules", + "tests/test_mcp20_doc11_epic_acceptance.py", + ), +) + +MCP20_SPEC_DOCS: tuple[str, ...] = ( + "00_OVERVIEW_AND_ARCHITECTURE.md", + "01_STATELESS_HTTP_AND_LIFECYCLE.md", + "02_JSONRPC_AND_WIRE_SPECIFICATION.md", + "03_TOOL_RESOURCE_PROMPT_CONTRACTS.md", + "04_MRTR_MULTI_ROUND_TRIP_SPEC.md", + "05_MCP_TASKS_PROTOCOL_AND_LIFECYCLE.md", + "06_TASK_STORE_AND_DISTRIBUTED_PERSISTENCE.md", + "07_TASK_AUTHORIZATION_AND_TENANT_ISOLATION.md", + "08_OAUTH21_CIMD_AND_SECURITY.md", + "09_EXTENSIONS_CACHE_AND_OBSERVABILITY.md", + "10_PYTHON_SDK_IMPLEMENTATION_BLUEPRINT.md", + "TICKET_PLANE_PROJECT_MANAGEMENT.md", +) + +MCP20_TEST_MODULES: tuple[str, ...] = tuple( + sorted({d.test_module for d in EPIC_DELIVERABLES} | {c.test_module for c in ACCEPTANCE_CRITERIA}) +) + + +def deliverables_for_epic(epic: ImplementationEpic) -> tuple[EpicDeliverable, ...]: + return tuple(item for item in EPIC_DELIVERABLES if item.epic == epic) + + +def epic_coverage_complete() -> bool: + """True when every epic has at least one mapped deliverable.""" + return all(deliverables_for_epic(epic) for epic in ImplementationEpic) + + +def acceptance_criteria_registered() -> bool: + return len(ACCEPTANCE_CRITERIA) >= 8 + + +def modern_protocol_target() -> str: + return MODERN_PROTOCOL_VERSION + + +def iter_epic_summary() -> Iterable[str]: + for epic in ImplementationEpic: + items = deliverables_for_epic(epic) + yield f"Epic {epic.value}: {len(items)} deliverable(s)" diff --git a/nitrostack/runtime/stateless.py b/nitrostack/runtime/stateless.py new file mode 100644 index 0000000..18866e3 --- /dev/null +++ b/nitrostack/runtime/stateless.py @@ -0,0 +1,26 @@ +"""Stateless HTTP invariants from Doc 00 §2 and Doc 01.""" + +from dataclasses import dataclass + +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER + + +@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 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..803d008 --- /dev/null +++ b/nitrostack/tasks/__init__.py @@ -0,0 +1,18 @@ +"""Async MCP task subsystem (Doc 00 architecture — 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..469bc2e --- /dev/null +++ b/nitrostack/tasks/authorization.py @@ -0,0 +1,161 @@ +"""Task authorization and multi-tenant isolation (Doc 07).""" + +from __future__ import annotations + +from typing import Any, List, Optional, Tuple + +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. + + Raises ``TaskNotFoundError`` (not Forbidden) on mismatch — anti-enumeration. + """ + if context is None: + return + + task_id = entry.data.task_id + + if entry.tenant_id and context.tenant_id and entry.tenant_id != context.tenant_id: + raise TaskNotFoundError(task_id) + + if entry.owner_id and context.user_id and entry.owner_id != context.user_id: + raise TaskNotFoundError(task_id) + + if entry.session_id and context.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 (Doc 07 §5). + """ + 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 _meta_dict(raw_meta: Any) -> dict[str, Any]: + 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) + else: + for key in ("userId", "user_id", "tenantId", "tenant_id", "sessionId", "session_id", "authorization", "headers"): + value = getattr(raw_meta, key, None) + if value is not None: + data[key] = value + return data + + +def _tenant_from_claims(claims: dict[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 extract_task_access_context(rc: Any) -> Optional[TaskAccessContext]: + """ + Build ``TaskAccessContext`` from an MCP request context. + + Reads explicit ``userId`` / ``tenantId`` / ``sessionId`` from ``_meta``, then + falls back to verified JWT claims when ``JWTService`` is registered. + """ + if rc is None: + return None + + meta = _meta_dict(getattr(rc, "meta", None)) + user_id = meta.get("userId") or meta.get("user_id") + tenant_id = meta.get("tenantId") or meta.get("tenant_id") + session_id = meta.get("sessionId") or meta.get("session_id") + + session = getattr(rc, "session", None) + if session is not None and not session_id: + session_id = getattr(session, "id", None) or getattr(session, "session_id", None) + if session_id is not None: + session_id = str(session_id) + + auth_header = meta.get("authorization") or meta.get("Authorization") + headers = meta.get("headers") + if isinstance(headers, dict): + auth_header = auth_header or headers.get("authorization") or headers.get("Authorization") + + request = getattr(rc, "request", None) + headers_obj = getattr(request, "headers", None) if request is not None else None + if headers_obj is not None and not auth_header: + try: + auth_header = headers_obj.get("authorization") or headers_obj.get("Authorization") + except Exception: + pass + + if isinstance(auth_header, str) and auth_header.startswith("Bearer "): + token = auth_header[len("Bearer ") :].strip() + if token: + try: + from nitrostack.core.di import DIContainer + from nitrostack.auth.jwt import JWTService + + payload = DIContainer.get_instance().resolve(JWTService).verify_token(token) + user_id = user_id or payload.get("sub") + tenant_id = tenant_id or _tenant_from_claims(payload) + except Exception: + pass + + if not any([user_id, tenant_id, session_id]): + return None + + 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..ca8d2db --- /dev/null +++ b/nitrostack/tasks/eviction.py @@ -0,0 +1,25 @@ +"""Terminal-only TTL eviction rules (Doc 06 §4).""" + +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..2365da5 --- /dev/null +++ b/nitrostack/tasks/memory.py @@ -0,0 +1,43 @@ +"""In-memory TaskStore for development and single-replica deployments (Doc 06).""" + +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] = {} + + async def get(self, task_id: str) -> Optional[TaskEntry]: + return self._entries.get(task_id) + + async def set(self, task_id: str, entry: TaskEntry) -> None: + self._entries[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/store.py b/nitrostack/tasks/store.py new file mode 100644 index 0000000..2d06f63 --- /dev/null +++ b/nitrostack/tasks/store.py @@ -0,0 +1,37 @@ +"""Pluggable task persistence interface (Doc 00 principle 3; Doc 06 contract).""" + +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..564783d --- /dev/null +++ b/nitrostack/tasks/types.py @@ -0,0 +1,79 @@ +"""Task subsystem shared types (Doc 05/06).""" + +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 (Doc 07).""" + + 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 (Doc 06 §3.1).""" + + 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`` (Doc 06 §3.2).""" + + 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/transports/__init__.py b/nitrostack/transports/__init__.py new file mode 100644 index 0000000..bbfd256 --- /dev/null +++ b/nitrostack/transports/__init__.py @@ -0,0 +1,22 @@ +"""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 + +__all__ = [ + "DispatchStage", + "IngressContext", + "StatelessIngressPipeline", + "StatelessTransportMiddleware", + "wrap_stateless_transport", + "is_task_wire_interception", + "format_sse_message", + "sse_connect_headers", + "sse_notification", +] diff --git a/nitrostack/transports/cors.py b/nitrostack/transports/cors.py new file mode 100644 index 0000000..40791d7 --- /dev/null +++ b/nitrostack/transports/cors.py @@ -0,0 +1,37 @@ +"""CORS configuration for stateless MCP HTTP (Doc 01 §2.3).""" + +from __future__ import annotations + +from typing import Mapping, Optional + +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, + get_header, +) + + +def build_cors_headers( + origin: Optional[str] = None, + *, + allow_origin: str = "*", +) -> dict[str, str]: + """Build CORS headers for MCP browser clients (SEP-2243 & SEP-2575).""" + resolved_origin = origin if origin else allow_origin + 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 cors_preflight_response_headers(request_headers: Mapping[str, str]) -> dict[str, str]: + """Headers for OPTIONS preflight — HTTP 204 No Content (Doc 01 §2.3).""" + origin = get_header(request_headers, "Origin") + return build_cors_headers(origin=origin or "*") diff --git a/nitrostack/transports/dispatch.py b/nitrostack/transports/dispatch.py new file mode 100644 index 0000000..9f287e9 --- /dev/null +++ b/nitrostack/transports/dispatch.py @@ -0,0 +1,150 @@ +"""Stateless HTTP ingress pipeline (Doc 01 §3).""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from enum import Enum +from typing import Any, Awaitable, Callable, Optional + +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.jsonrpc import ( + HeaderBodyMismatchError, + JsonRpcParseError, + JsonRpcRequest, + JsonRpcWireError, + build_ping_response, + jsonrpc_error, + jsonrpc_success, + parse_jsonrpc_request, + validate_header_body_method, + validate_header_body_name, +) +from nitrostack.transports.headers import HEADER_MCP_METHOD, HEADER_MCP_NAME, get_header + +TaskDispatchHandler = Callable[[JsonRpcRequest], Awaitable[Optional[dict[str, Any]]]] +RegistryDispatchHandler = Callable[[JsonRpcRequest], Awaitable[Optional[dict[str, Any]]]] + + +class DispatchStage(str, Enum): + """Six-step ingress lifecycle from Doc 01 §3.""" + + 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/" + + +@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 + + +def is_task_wire_interception(method: str, params: dict[str, Any]) -> bool: + """Doc 01 §3 step 4 — route to task subsystem when matched.""" + if method.startswith(TASK_METHOD_PREFIX): + return True + return method == "tools/call" and bool(params.get("task")) + + +class StatelessIngressPipeline: + """ + Deterministic JSON-RPC pre-dispatch for stateless POST /mcp. + + Handles ping and server/discover inline. Task and registry methods return + None so the underlying MCP server can handle them (Doc 05+ adds task handler). + """ + + def __init__( + self, + context: IngressContext, + *, + task_handler: Optional[TaskDispatchHandler] = None, + registry_handler: Optional[RegistryDispatchHandler] = None, + ) -> None: + self._context = context + self._task_handler = task_handler + self._registry_handler = registry_handler + + 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. + """ + 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) + + header_method = get_header(request_headers, HEADER_MCP_METHOD) + try: + validate_header_body_method(header_method, request.method) + except HeaderBodyMismatchError as exc: + return 400, exc.to_response(request.id) + + header_name = get_header(request_headers, HEADER_MCP_NAME) + body_name = request.params.get("name") or request.params.get("uri") + if isinstance(body_name, str): + try: + validate_header_body_name(header_name, body_name) + except HeaderBodyMismatchError as exc: + return 400, exc.to_response(request.id) + + deprecated_msg = deprecated_method_message(request.method) + if deprecated_msg is not None: + return 200, jsonrpc_error( + request.id, + int(JsonRpcErrorCode.METHOD_NOT_FOUND), + deprecated_msg, + ) + + if request.method == "ping": + return 200, build_ping_response(request.id) + + if request.method == "server/discover": + result = build_discover_result( + server_name=self._context.server_name, + server_version=self._context.server_version, + 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 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/headers.py b/nitrostack/transports/headers.py new file mode 100644 index 0000000..c6b21f6 --- /dev/null +++ b/nitrostack/transports/headers.py @@ -0,0 +1,108 @@ +"""MCP 2.0 HTTP header names and builders (Doc 01 §2).""" + +from __future__ import annotations + +from typing import Mapping, Optional + +from nitrostack.protocol.constants import LEGACY_SESSION_HEADER +from nitrostack.protocol.version import MODERN_PROTOCOL_VERSION + +# Request headers (Doc 01 §2.1) +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-" +HEADER_AUTHORIZATION = "Authorization" +HEADER_LAST_EVENT_ID = "Last-Event-ID" + +# Response headers (Doc 01 §2.2) +HEADER_VARY = "Vary" + +# CORS headers (Doc 01 §2.3) +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 (Doc 01 §5) +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_HEADERS = ( + "Content-Type, Authorization, MCP-Protocol-Version, Mcp-Method, Mcp-Name, " + "Mcp-Param-*, Last-Event-ID" +) +CORS_EXPOSE_HEADERS = "MCP-Protocol-Version, Mcp-Method, Mcp-Name" + +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 build_mcp_response_headers( + *, + content_type: str = MCP_JSON_CONTENT_TYPE, + protocol_version: str = MODERN_PROTOCOL_VERSION, + extra: Optional[Mapping[str, str]] = None, +) -> dict[str, str]: + """Standard MCP 2026-07-28 response headers.""" + headers = { + HEADER_CONTENT_TYPE: content_type, + HEADER_MCP_PROTOCOL_VERSION: protocol_version, + HEADER_VARY: "Origin", + } + if extra: + headers.update(extra) + return strip_legacy_session_headers(headers) + + +def build_sse_stream_headers( + protocol_version: str = MODERN_PROTOCOL_VERSION, +) -> dict[str, str]: + """SSE stream headers including proxy buffering guards (Doc 01 §5).""" + 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-* mirrored parameters from request headers.""" + params: dict[str, str] = {} + prefix = HEADER_MCP_PARAM_PREFIX.lower() + for key, value in headers.items(): + if key.lower().startswith(prefix): + param_name = key[len(HEADER_MCP_PARAM_PREFIX) :] + params[param_name] = value + return params diff --git a/nitrostack/transports/middleware.py b/nitrostack/transports/middleware.py new file mode 100644 index 0000000..fff0551 --- /dev/null +++ b/nitrostack/transports/middleware.py @@ -0,0 +1,196 @@ +"""Stateless HTTP ASGI middleware (Doc 01).""" + +from __future__ import annotations + +from typing import Any, Callable, Optional + +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 IngressContext, StatelessIngressPipeline +from nitrostack.transports.headers import ( + MCP_HTTP_PATH, + build_mcp_response_headers, + get_header, + strip_legacy_session_headers, +) + +ASGIApp = Callable[..., Any] + +MCP_POST_PATHS = (MCP_HTTP_PATH, f"{MCP_HTTP_PATH}/") + + +class StatelessTransportMiddleware: + """ + ASGI wrapper implementing Doc 01 transport invariants: + - OPTIONS 204 CORS preflight + - 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, + ) -> None: + self.app = app + self.pipeline = pipeline + self.mcp_paths = mcp_paths + + 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": + await self._send_options(scope, receive, send) + return + + if method == "POST" and path in self.mcp_paths and self.pipeline is not None: + body = await self._read_body(receive) + handled = await self._try_pre_dispatch(scope, body, send) + if handled: + return + receive = self._replay_receive(body) + + await self._forward_with_stateless_headers(scope, receive, send) + + 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 = build_mcp_response_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 [] + req_headers = strip_legacy_session_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, req_headers) + if result is None: + return False + + status, jsonrpc_response = result + origin = get_header(req_headers, "Origin") + cors = build_cors_headers(origin=origin or "*") + response_headers = build_mcp_response_headers(extra=cors) + assert_stateless_headers(response_headers) + + 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}) + return True + + async def _forward_with_stateless_headers( + self, + scope: dict[str, Any], + receive: Any, + send: Any, + ) -> 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 []) + } + merged = build_mcp_response_headers( + content_type=raw_headers.get("content-type", "application/json"), + extra={ + **build_cors_headers(origin=get_header(req_headers, "Origin") or "*"), + **strip_legacy_session_headers(raw_headers), + }, + ) + assert_stateless_headers(merged) + message = { + **message, + "headers": self._encode_headers(merged), + } + await send(message) + + await self.app(scope, receive, send_wrapper) + + @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) -> 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} + return {"type": "http.disconnect"} + + 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()] + + +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, +) -> ASGIApp: + """Wrap an ASGI app with Doc 01 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, + ) + ) + return StatelessTransportMiddleware(app, pipeline=pipeline) diff --git a/nitrostack/transports/sse.py b/nitrostack/transports/sse.py new file mode 100644 index 0000000..7fc84e3 --- /dev/null +++ b/nitrostack/transports/sse.py @@ -0,0 +1,27 @@ +"""SSE notification bus helpers (Doc 01 §5).""" + +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/tests/test_lifecycle_http.py b/tests/test_lifecycle_http.py index 86441b0..8c0e729 100644 --- a/tests/test_lifecycle_http.py +++ b/tests/test_lifecycle_http.py @@ -143,26 +143,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_doc00_foundation.py b/tests/test_mcp20_doc00_foundation.py new file mode 100644 index 0000000..c9468b2 --- /dev/null +++ b/tests/test_mcp20_doc00_foundation.py @@ -0,0 +1,109 @@ +"""Tests for MCP 2.0 architectural foundation (Doc 00).""" + +import pytest + +from nitrostack.core.app import ServerConfig +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, +) +from nitrostack.runtime.stateless import ( + DEFAULT_STATELESS_INVARIANTS, + assert_stateless_headers, +) +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 + + +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"}) + + +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 True diff --git a/tests/test_mcp20_doc01_stateless_http.py b/tests/test_mcp20_doc01_stateless_http.py new file mode 100644 index 0000000..d36429e --- /dev/null +++ b/tests/test_mcp20_doc01_stateless_http.py @@ -0,0 +1,181 @@ +"""Tests for MCP 2.0 stateless HTTP transport (Doc 01).""" + +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, 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.version import MODERN_PROTOCOL_VERSION +from nitrostack.transports.cors import build_cors_headers, cors_preflight_response_headers +from nitrostack.transports.dispatch import ( + IngressContext, + StatelessIngressPipeline, + is_task_wire_interception, +) +from nitrostack.transports.headers import ( + build_mcp_response_headers, + build_sse_stream_headers, + strip_legacy_session_headers, +) +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" + + +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"] + + def test_preflight_headers(self): + headers = cors_preflight_response_headers({"Origin": "https://app.example.com"}) + assert "Access-Control-Allow-Origin" in headers + + +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": {}} + + +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_handles_server_discover(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", "params": {}} + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + result = resp["result"] + assert result["protocolVersion"] == "2026-07-28" + assert result["resultType"] == DISCOVER_RESULT_TYPE + assert isinstance(result["ttlMs"], int) and result["ttlMs"] >= 0 + assert result["cacheScope"] in ("public", "private") + assert MODERN_PROTOCOL_VERSION in result["supportedVersions"] + + 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, {}) 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 diff --git a/tests/test_mcp20_doc02_jsonrpc_wire.py b/tests/test_mcp20_doc02_jsonrpc_wire.py new file mode 100644 index 0000000..4aed9aa --- /dev/null +++ b/tests/test_mcp20_doc02_jsonrpc_wire.py @@ -0,0 +1,143 @@ +"""Tests for MCP 2.0 JSON-RPC wire specification (Doc 02).""" + +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 JsonRpcErrorCode +from nitrostack.protocol.jsonrpc import ( + HEADER_BODY_MISMATCH, + PARSE_ERROR, + HeaderBodyMismatchError, + InvalidParamsError, + JsonRpcParseError, + build_tool_error_result, + map_exception_to_jsonrpc, + parse_jsonrpc_request, + validate_header_body_name, + validate_header_body_method, +) +from nitrostack.protocol.meta import 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 + + +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_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" + + +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") + + +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) + ) + body = json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "tasks/list", "params": {}} + ).encode() + status, resp = await pipeline.handle_post(body, {}) + assert status == 200 + assert resp["error"]["code"] == int(JsonRpcErrorCode.METHOD_NOT_FOUND) + + 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_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_doc03_contracts.py b/tests/test_mcp20_doc03_contracts.py new file mode 100644 index 0000000..6fe1568 --- /dev/null +++ b/tests/test_mcp20_doc03_contracts.py @@ -0,0 +1,124 @@ +"""Tests for MCP 2.0 tool/resource/prompt contracts (Doc 03).""" + +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, + 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"] == {} + + +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"] diff --git a/tests/test_mcp20_doc04_mrtr.py b/tests/test_mcp20_doc04_mrtr.py new file mode 100644 index 0000000..f95e16e --- /dev/null +++ b/tests/test_mcp20_doc04_mrtr.py @@ -0,0 +1,146 @@ +"""Tests for MCP 2.0 MRTR multi round-trip requests (Doc 04).""" + +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.structuredContent["resultType"] == "input_required" + assert first.structuredContent["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.structuredContent["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_doc05_tasks.py b/tests/test_mcp20_doc05_tasks.py new file mode 100644 index 0000000..3d90ce0 --- /dev/null +++ b/tests/test_mcp20_doc05_tasks.py @@ -0,0 +1,280 @@ +"""Tests for MCP 2.0 Tasks protocol and lifecycle (Doc 05).""" + +import asyncio +import os +import sys + +import mcp.types as types +from mcp.shared.exceptions import 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 mcp.server.lowlevel.server import request_ctx, RequestContext +from mcp.server.experimental.request_context import Experimental + + +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="TasksDoc05Forbidden", 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="TasksDoc05Required", 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="TasksDoc05Get", 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="TasksDoc05List", controllers=[DummyController]) + class ListModule: + pass + + @mcp_app(module=ListModule, server=ServerConfig(name="tasks-list")) + 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="TasksDoc05Cancel", 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="TasksDoc05Create", 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_doc06_task_store.py b/tests/test_mcp20_doc06_task_store.py new file mode 100644 index 0000000..621ede8 --- /dev/null +++ b/tests/test_mcp20_doc06_task_store.py @@ -0,0 +1,118 @@ +"""Tests for MCP 2.0 task store and distributed persistence (Doc 06).""" + +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_doc07_task_authorization.py b/tests/test_mcp20_doc07_task_authorization.py new file mode 100644 index 0000000..4e70cf6 --- /dev/null +++ b/tests/test_mcp20_doc07_task_authorization.py @@ -0,0 +1,224 @@ +"""Tests for MCP 2.0 task authorization and tenant isolation (Doc 07).""" + +import asyncio +import os +import sys + +import mcp.types as types +import pytest +from mcp.shared.exceptions import McpError +from mcp.server.experimental.request_context import Experimental +from mcp.server.lowlevel.server import request_ctx, RequestContext +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) + + +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="Doc07Auth", controllers=[OwnerController]) + class AuthModule: + pass + + @mcp_app(module=AuthModule, server=ServerConfig(name="doc07-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=types.RequestParams.Meta(__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=types.RequestParams.Meta(__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_reads_identity_from_request_meta(self): + rc = RequestContext( + request_id="1", + meta=types.RequestParams.Meta( + __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 == "u1" + assert ctx.tenant_id == "t1" + assert ctx.session_id == "s1" diff --git a/tests/test_mcp20_doc08_oauth_cimd.py b/tests/test_mcp20_doc08_oauth_cimd.py new file mode 100644 index 0000000..b3b715e --- /dev/null +++ b/tests/test_mcp20_doc08_oauth_cimd.py @@ -0,0 +1,170 @@ +"""Tests for MCP 2.0 OAuth 2.1, CIMD, and SSRF security (Doc 08).""" + +import asyncio +import json +import socket +from unittest.mock import patch + +import pytest + +from nitrostack.auth.cimd import ( + CimdFetchError, + CimdValidationError, + assert_safe_fetch_target, + is_blocked_ip, + resolve_cimd, + validate_client_identifier_url, +) +from nitrostack.auth.oauth_module import 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): + 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()) + + +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_doc09_extensions_cache_observability.py b/tests/test_mcp20_doc09_extensions_cache_observability.py new file mode 100644 index 0000000..9c7111f --- /dev/null +++ b/tests/test_mcp20_doc09_extensions_cache_observability.py @@ -0,0 +1,256 @@ +"""Tests for MCP 2.0 extensions, cache hints, and observability (Doc 09).""" + +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="Doc09Cache", 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://doc09/metrics", + name="metrics", + description="metrics", + metadata={"cacheMaxAge": 10}, + ) + async def metrics(self, context: ExecutionContext) -> str: + return "{}" + + @module(name="Doc09Resource", 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="Doc09NoTasks", 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="Doc09Trace", controllers=[TraceController]) + class TraceModule: + pass + + async def _run(): + from mcp.server.lowlevel.server import request_ctx, RequestContext + + 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_doc10_blueprint.py b/tests/test_mcp20_doc10_blueprint.py new file mode 100644 index 0000000..7b93295 --- /dev/null +++ b/tests/test_mcp20_doc10_blueprint.py @@ -0,0 +1,403 @@ +"""Blueprint conformance verification (Doc 10 §4).""" + +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 mcp.server.experimental.request_context import Experimental +from mcp.server.lowlevel.server import request_ctx, RequestContext +from mcp.shared.exceptions import 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) + ) + 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, + 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 + assert "Mcp-Method" in headers.get("Access-Control-Allow-Headers", "") + + +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")) + 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.taskId), + ) + ) + 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_doc11_epic_acceptance.py b/tests/test_mcp20_doc11_epic_acceptance.py new file mode 100644 index 0000000..1407fe8 --- /dev/null +++ b/tests/test_mcp20_doc11_epic_acceptance.py @@ -0,0 +1,251 @@ +"""Epic and acceptance-criteria verification for PYTHONSDK-11 (Doc 11).""" + +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 mcp.server.experimental.request_context import Experimental +from mcp.server.lowlevel.server import request_ctx, RequestContext +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.conformance import assert_blueprint_layout +from nitrostack.runtime.epic_acceptance import ( + ACCEPTANCE_CRITERIA, + EPIC_DELIVERABLES, + MCP20_SPEC_DOCS, + MCP20_TEST_MODULES, + ImplementationEpic, + acceptance_criteria_registered, + deliverables_for_epic, + epic_coverage_complete, + iter_epic_summary, + modern_protocol_target, +) + + +REPO_ROOT = Path(__file__).resolve().parents[1] +_SPEC_CANDIDATES = ( + REPO_ROOT.parent / "traker" / "stateless-doc-for-python", + REPO_ROOT.parent.parent / "traker" / "stateless-doc-for-python", +) + + +def _resolve_spec_root() -> Path | None: + for candidate in _SPEC_CANDIDATES: + if candidate.is_dir(): + return candidate + return None + + +SPEC_ROOT = _resolve_spec_root() + + +class EchoInput(BaseModel): + value: str = Field(default="") + + +def setup_function() -> None: + DIContainer.reset() + + +def teardown_function() -> None: + DIContainer.reset() + + +class TestEpicRegistry: + def test_all_seven_epics_have_deliverables(self): + assert epic_coverage_complete() + for line in iter_epic_summary(): + assert "0 deliverable" not in line + + def test_epic_deliverable_count(self): + assert len(EPIC_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("epic", list(ImplementationEpic)) + def test_each_epic_maps_to_spec_and_tests(self, epic: ImplementationEpic): + items = deliverables_for_epic(epic) + assert items, f"Epic {epic} has no deliverables" + for item in items: + assert item.spec_doc.endswith(".md") + assert item.test_module.startswith("tests/test_mcp20_doc") + + +class TestSpecAndTestTraceability: + 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 + + @pytest.mark.parametrize("doc_name", MCP20_SPEC_DOCS) + def test_spec_documents_exist(self, doc_name: str): + if SPEC_ROOT is None: + pytest.skip("Spec folder not available beside the SDK checkout") + path = SPEC_ROOT / doc_name + assert path.is_file(), path + assert path.stat().st_size > 0 + + 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="Doc11Tasks", controllers=[AsyncController]) + class TasksModule: + pass + + @mcp_app(module=TasksModule, server=ServerConfig(name="doc11-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_doc11", description="noop", input_schema=EchoInput) + async def noop_doc11(self, input: EchoInput, context: ExecutionContext) -> str: + return "ok" + + @module(name="Doc11Cancel", controllers=[DummyController]) + class CancelModule: + pass + + @mcp_app(module=CancelModule, server=ServerConfig(name="doc11-cancel")) + class CancelApp: + pass + + from mcp.shared.exceptions import 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_doc_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): + doc_tests = sorted(REPO_ROOT.glob("tests/test_mcp20_doc*.py")) + assert len(doc_tests) == 12 diff --git a/tests/test_oauth.py b/tests/test_oauth.py index b1e140e..3f4fb38 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -334,6 +334,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_tasks.py b/tests/test_tasks.py index a258276..02d7e9c 100644 --- a/tests/test_tasks.py +++ b/tests/test_tasks.py @@ -32,6 +32,7 @@ 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 @@ -41,16 +42,66 @@ # 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 +135,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 +289,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 +309,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 +347,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,13 +381,13 @@ 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): 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 @@ -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), ), ) @@ -457,14 +501,6 @@ async def _mcp_task_flow(): 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) - get_req = types.GetTaskRequest( method="tasks/get", params=types.GetTaskRequestParams(taskId=task_id), @@ -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( @@ -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_transports.py b/tests/test_transports.py index 23bcabc..3fc52b8 100644 --- a/tests/test_transports.py +++ b/tests/test_transports.py @@ -357,7 +357,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"}, ), ) @@ -380,11 +380,15 @@ async def _test_progress_notifications_pushed(): task_id = response.root.task.taskId # 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): diff --git a/tests/test_widgets.py b/tests/test_widgets.py index 33a8cdc..8f64593 100644 --- a/tests/test_widgets.py +++ b/tests/test_widgets.py @@ -522,18 +522,18 @@ async def run(): assert isinstance(task_resp.root, types.CreateTaskResult) task_id = task_resp.root.task.taskId - 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")