diff --git a/README.md b/README.md index adfd506..0b7a741 100644 --- a/README.md +++ b/README.md @@ -457,3 +457,100 @@ async def test_add(): if __name__ == "__main__": asyncio.run(test_add()) ``` + +### Large tool results: data spillover + +Attach `DataSpilloverInterceptor` to a tool to store large results separately and +return a preview, `resource_uri`, and `original_size_bytes`. Results at or below +the threshold pass through unchanged. Sizes count serialized UTF-8 bytes, not +characters; lists and dictionaries use compact JSON. + +```python +from pydantic import BaseModel +from nitrostack import DataSpilloverInterceptor, MemorySpilloverStore, tool, use_interceptors + +spillover = DataSpilloverInterceptor( + threshold_bytes=64 * 1024, + store=MemorySpilloverStore( + max_bytes=10 * 1024 * 1024, + max_total_bytes=100 * 1024 * 1024, + ttl_seconds=300, + ), + resource_uri_prefix="resource://data-spillover/", +) + +class ReportInput(BaseModel): + count: int = 10000 + +class ReportTools: + @tool(name="report", description="Generate a report", input_schema=ReportInput) + @use_interceptors(spillover) + async def report(self, input: ReportInput, context): + return [{"row": index} for index in range(input.count)] +``` + +Register `ReportTools` as a module provider or controller as usual. A configured +interceptor instance works directly in the decorator. For default configuration, +`@use_interceptors(DataSpilloverInterceptor)` also works through dependency injection. + +The MCP result includes a text preview, structured content, and a resource link. +Clients retrieve the full result with `resources/read`, for example: + +```json +{ + "jsonrpc": "2.0", + "id": 2, + "method": "resources/read", + "params": {"uri": "resource://data-spillover/"} +} +``` + +Lists and dictionaries preview their first five items or entries and total counts. +Strings preview up to 512 characters. The complete preview stays within 2 KiB; +large nested values are represented as truncated text. Reading the full resource +can consume context again, so clients should retrieve it only when needed. + +For disk storage, replace the store with: + +```python +from nitrostack import FileSpilloverStore + +store = FileSpilloverStore( + ".nitrostack/spillover", + max_bytes=10 * 1024 * 1024, + max_total_bytes=100 * 1024 * 1024, + ttl_seconds=300, +) +``` + +The memory store evicts least recently used results when its payload capacity is +full. The file store rejects writes when its capacity is full; its capacity counts +actual JSON file bytes, including base64 and metadata overhead. A single result +larger than `max_bytes` raises `PayloadTooLargeError`. Both stores remove expired +results during operations; file storage also cleans up on initialization. Expiry +is absolute and reads do not extend it. Each file-store root must be owned by one +store instance in one process; shared storage across processes is not supported. +Keep the configured directory private to the server account. + +Resource access requires the same verified principal, issuer, client identity, +and tenant. Anonymous access is limited to the same server-side MCP session. +Requests without either a verified identity or trusted session cannot spill +results; invalid credentials do not fall back to anonymous access. Anonymous +HTTP clients need a persistent MCP session; sessionless HTTP clients should +authenticate. Expired results and results owned by another caller return +resource-not-found. Stored results are not enumerated in `resources/list`. + +HTTP ownership uses the server-issued session ID after the official transport +has validated it, rather than the identity of an SDK request/session proxy. +Stdio assigns a random identity to each live stream connection. Requests on that +connection can retrieve its anonymous results; a later connection cannot inherit +them. Concurrent HTTP requests cannot borrow a stdio connection's identity. + +Successful explicit MCP tool results are stored as JSON; errors and input-required +responses pass through. Decorated tools with output schemas advertise either the +original output or the spillover envelope. Spilled widget results return the +resource reference without invoking the widget's data transformer. + +Independent HTTP/stdio tests and reproducible large-output checks live in +`tests/test_spillover_transport.py`. Run the spillover tests with +`python -m pytest tests/test_spillover.py tests/test_spillover_transport.py`. diff --git a/nitrostack/__init__.py b/nitrostack/__init__.py index 9322d9c..8a95809 100644 --- a/nitrostack/__init__.py +++ b/nitrostack/__init__.py @@ -21,6 +21,16 @@ except (ImportError, KeyError, OSError, ValueError, TypeError): __version__ = "unknown" +from nitrostack.core.interceptors import ( + DataSpilloverInterceptor, + FileSpilloverStore, + MemorySpilloverStore, + PayloadTooLargeError, + SpilloverEntry, + SpilloverStore, + generate_preview, +) + from nitrostack.core.decorators import ( tool, resource, @@ -168,6 +178,13 @@ __all__ = [ "__version__", + "DataSpilloverInterceptor", + "FileSpilloverStore", + "MemorySpilloverStore", + "PayloadTooLargeError", + "SpilloverEntry", + "SpilloverStore", + "generate_preview", "tool", "resource", "prompt", diff --git a/nitrostack/core/app.py b/nitrostack/core/app.py index 0a2992e..f6813b4 100644 --- a/nitrostack/core/app.py +++ b/nitrostack/core/app.py @@ -52,6 +52,10 @@ from nitrostack.core.mcp_server import NitroStackMcpServer from nitrostack.core.pipeline import run_pipeline from nitrostack.core.additional_decorators import HealthCheckRegistry, ServerHealthPayload +from nitrostack.core.interceptors.data_spillover import ( + DataSpilloverInterceptor, SPILLOVER_SCHEMA, SpilloverRegistry, SpilloverResult, + bind_spillover_context, caller_owner, +) from nitrostack.core.task import TaskManager, TaskStatus from nitrostack.events.event_emitter import EventEmitter from nitrostack.protocol.schema import ( @@ -617,6 +621,7 @@ def __init__(self, app_class: Type): self.task_manager = TaskManager() self._in_flight = InFlightRegistry() self.visibility_transform = VisibilityTransform() + self._spillover = SpilloverRegistry() # Pipeline first: handlers consult it lazily when requests arrive. self._transforms = TransformPipeline( @@ -702,6 +707,11 @@ def _bootstrap(self) -> None: # 5. Wire every protocol handler (tools/resources/prompts/tasks/notifications) # onto the owned low-level server. This is factored out so additional server # instances (e.g. one per HTTP session) can be configured identically. + for entry in self._tools.values(): + for candidate in getattr(entry.method, "_mcp_interceptors", []): + interceptor = container.resolve(candidate) if isinstance(candidate, type) else candidate + if isinstance(interceptor, DataSpilloverInterceptor): + self._spillover.register(interceptor) self._setup_handlers(self.mcp_server) def _assert_declared_dependencies( @@ -1276,6 +1286,11 @@ def _build_tool_definition(self, entry: _ToolEntry) -> types.Tool: output_schema = tool_json_schema(cfg.output_schema) if output_schema is not None: + if self._has_spillover(entry.method): + output_schema = { + "type": "object", "anyOf": [output_schema, SPILLOVER_SCHEMA], + **({"$defs": output_schema["$defs"]} if "$defs" in output_schema else {}), + } tool_kwargs["outputSchema"] = output_schema return types.Tool(**tool_kwargs) @@ -1320,7 +1335,33 @@ def _build_prompt_definition(self, entry: _PromptEntry) -> types.Prompt: # Request dispatch # ------------------------------------------------------------------ - def _pipeline_stages(self, method: Callable) -> Tuple[List[Type], List[Type], List[Type], List[Type], List[Type]]: + def _has_spillover(self, method: Callable) -> bool: + container = DIContainer.get_instance() + return any( + isinstance(container.resolve(item) if isinstance(item, type) else item, + DataSpilloverInterceptor) + for item in getattr(method, "_mcp_interceptors", []) + ) + + async def _bind_spillover(self, context: ExecutionContext, rc: Any) -> None: + if not self._spillover.routes: + return + # OAuth guards may verify opaque tokens as well as the shared JWT path. + from nitrostack.auth.request import authorization_token_from_request, auth_context_from_payload + from nitrostack.auth.oauth import OAuthService + + token = authorization_token_from_request(rc) + container = DIContainer.get_instance() + if context.auth is None and token and container.has_value(OAuthService): + try: + payload = await container.resolve(OAuthService).introspect_token(token) + if payload.get("active"): + context.auth = auth_context_from_payload(payload) + except Exception: + pass # Invalid credentials remain denied by caller_owner. + bind_spillover_context(context, rc, self._spillover, self.protocol_era) + + def _pipeline_stages(self, method: Callable) -> Tuple[List[Any], List[Any], List[Any], List[Any], List[Any]]: return ( getattr(method, "_mcp_guards", []), getattr(method, "_mcp_middleware", []), @@ -1390,6 +1431,8 @@ def _to_call_tool_result( component: Optional[Component] = None, context: Optional[ExecutionContext] = None, ) -> types.CallToolResult: + if isinstance(result, SpilloverResult): + return result.to_mcp() if isinstance(result, InputRequiredResult): wire = result.to_wire_dict() return types.CallToolResult( @@ -1641,6 +1684,7 @@ async def background_execution(): ) _apply_request_envelope(task_ctx, rc) self._configure_visibility_context(task_ctx, rc) + await self._bind_spillover(task_ctx, rc) task_ctx.task = TaskContext( task_id, self.task_manager, @@ -1696,6 +1740,7 @@ async def background_execution(): ) _apply_request_envelope(ctx, rc) self._configure_visibility_context(ctx, rc) + await self._bind_spillover(ctx, rc) ticket = self._in_flight.register(correlation_id, jsonrpc_id=jsonrpc_id) try: result = await run_pipeline( @@ -1728,6 +1773,21 @@ async def background_execution(): return self._to_call_tool_result(result, entry.component, ctx) async def _read_resource(self, uri: str) -> List[ReadResourceContents]: + for prefix, store in self._spillover.routes.items(): + if uri.startswith(prefix): + ctx = ExecutionContext(request_id=str(uuid.uuid4())) + rc = request_ctx.get(None) + _apply_request_envelope(ctx, rc) + await self._bind_spillover(ctx, rc) + try: + entry = await store.get(uri[len(prefix):], owner=caller_owner(ctx)) + except (PermissionError, ValueError): + raise ResourceNotFoundError(uri) from None + if entry is None: + raise ResourceNotFoundError(uri) + return [ReadResourceContents( + content=entry.payload.decode("utf-8"), mime_type=entry.mime_type, + )] entry = self._resources.get(uri) path_kwargs: Dict[str, str] = {} diff --git a/nitrostack/core/context.py b/nitrostack/core/context.py index 9abe686..7766389 100644 --- a/nitrostack/core/context.py +++ b/nitrostack/core/context.py @@ -207,6 +207,8 @@ class ExecutionContext: _visibility_session_key: Any = field(default=None, repr=False, compare=False) _visibility_change_callback: Any = field(default=None, repr=False, compare=False) _visibility_supported: bool = field(default=False, repr=False, compare=False) + spillover_session_key: Optional[str] = field(default=None, repr=False) + spillover_invalid_auth: bool = field(default=False, repr=False) def __post_init__(self) -> None: if not self.correlation_id: diff --git a/nitrostack/core/errors.py b/nitrostack/core/errors.py index b5f1c50..97824a1 100644 --- a/nitrostack/core/errors.py +++ b/nitrostack/core/errors.py @@ -1,6 +1,10 @@ from typing import Any +class PayloadTooLargeError(ValueError): + """A serialized tool result exceeds the spillover store's per-result limit.""" + + class ToolExecutionError(Exception): """Raised when tool execution fails.""" pass diff --git a/nitrostack/core/interceptors/__init__.py b/nitrostack/core/interceptors/__init__.py new file mode 100644 index 0000000..4344182 --- /dev/null +++ b/nitrostack/core/interceptors/__init__.py @@ -0,0 +1,15 @@ +"""Built-in tool interceptors and spillover storage adapters.""" + +from .data_spillover import DataSpilloverInterceptor, generate_preview +from .spillover import ( + FileSpilloverStore, + MemorySpilloverStore, + PayloadTooLargeError, + SpilloverEntry, + SpilloverStore, +) + +__all__ = [ + "DataSpilloverInterceptor", "FileSpilloverStore", "MemorySpilloverStore", + "PayloadTooLargeError", "SpilloverEntry", "SpilloverStore", "generate_preview", +] diff --git a/nitrostack/core/interceptors/data_spillover.py b/nitrostack/core/interceptors/data_spillover.py new file mode 100644 index 0000000..ee9ae97 --- /dev/null +++ b/nitrostack/core/interceptors/data_spillover.py @@ -0,0 +1,186 @@ +"""Tool-result spillover, bounded previews, and trusted caller ownership.""" + +from __future__ import annotations + +import hashlib +import json +import uuid +import weakref +from dataclasses import dataclass +from itertools import islice +from typing import Any +from urllib.parse import urlsplit + +import mcp.types as types +from pydantic import BaseModel + +from nitrostack.auth.request import authorization_token_from_request, tenant_from_claims +from nitrostack.core.context import ExecutionContext +from nitrostack.protocol.mrtr import InputRequiredResult +from nitrostack.core.interceptors.spillover import MemorySpilloverStore, SpilloverStore +from nitrostack.protocol.version import wire_mode_for_era +from nitrostack.runtime.caller import session_id_from_request_context +from nitrostack.runtime.request_ctx import RequestContext, stdio_connection_id + + +def generate_preview(value: Any, *, max_bytes: int = 2048) -> dict: + """Include head items and counts, shrinking even deeply nested values to fit.""" + if max_bytes < 128: + raise ValueError("Preview budget must be at least 128 bytes") + if isinstance(value, list): + preview = {"type": "list", "count": len(value), "head": value[:5], "truncated": len(value) > 5} + elif isinstance(value, dict): + preview = {"type": "dict", "count": len(value), "head": dict(islice(value.items(), 5)), + "truncated": len(value) > 5} + else: + text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str) + preview = {"type": "string", "count": len(text), "text": text[:512], "truncated": len(text) > 512} + encoded = json.dumps(preview, ensure_ascii=False, default=str).encode("utf-8") + if len(encoded) > max_bytes: + text = json.dumps(preview.get("head", preview.get("text")), ensure_ascii=False, default=str) + preview.pop("head", None) + preview["text"] = text[:512] + preview["truncated"] = True + while len(json.dumps(preview, ensure_ascii=False, default=str).encode("utf-8")) > max_bytes: + preview["text"] = preview["text"][:len(preview["text"]) // 2] + return preview + + +@dataclass(frozen=True) +class SpilloverResult: + preview: dict + resource_uri: str + original_size_bytes: int + + def to_mcp(self) -> types.CallToolResult: + data = dict(preview=self.preview, resource_uri=self.resource_uri, + original_size_bytes=self.original_size_bytes) + return types.CallToolResult( + content=[ + types.TextContent(type="text", text=json.dumps(data, ensure_ascii=False, default=str)), + types.ResourceLink(type="resource_link", uri=self.resource_uri, + name="Full tool result"), + ], + structuredContent=data, isError=False, + ) + + +SPILLOVER_SCHEMA = { + "type": "object", + "properties": { + "preview": {"type": "object"}, + "resource_uri": {"type": "string"}, + "original_size_bytes": {"type": "integer", "minimum": 0}, + }, + "required": ["preview", "resource_uri", "original_size_bytes"], +} + + +class SpilloverRegistry: + """Per-application routes and weakly held anonymous session identities.""" + + def __init__(self): + self.routes: dict[str, SpilloverStore] = {} + self._sessions: dict[int, tuple[weakref.ref, str]] = {} + + def register(self, interceptor): + prefix, store = interceptor.resource_uri_prefix, interceptor.store + for existing, registered in self.routes.items(): + if prefix.startswith(existing) or existing.startswith(prefix): + if prefix != existing or store is not registered: + raise ValueError("Conflicting spillover resource prefixes") + self.routes[prefix] = store + + def session_key(self, rc, era=None): + request = getattr(rc, "request", None) + if getattr(request, "scope", {}).get("type") == "http": + # Only routed requests reach here: the official manager has validated + # this ID. SDK session proxies themselves are per-request objects. + if wire_mode_for_era(era or "auto") != "sessionful": + return None + session_id = session_id_from_request_context(rc) + return "http:" + session_id if session_id else None + connection_id = stdio_connection_id.get() + if connection_id is not None: + return "stdio:" + connection_id + # Explicit in-process contexts retain their weakly-held test/session key. + # Never infer ownership from an official SDK request's proxy identity. + if not isinstance(rc, RequestContext): + return None + session = getattr(rc, "session", None) + if session is None: + return None + identity = id(session) + existing = self._sessions.get(identity) + if existing is not None and existing[0]() is session: + return existing[1] + key = "in-process:" + uuid.uuid4().hex + try: + reference = weakref.ref(session, lambda _: self._sessions.pop(identity, None)) + except TypeError: + return None + self._sessions[identity] = (reference, key) + return key + + +def caller_owner(context: ExecutionContext) -> str: + auth = context.auth + if auth is not None and (auth.subject or auth.client_id): + identity = ["principal", auth.iss, auth.subject, auth.client_id, + tenant_from_claims(auth.claims)] + elif context.spillover_invalid_auth: + raise PermissionError("Verified identity required for spillover") + elif context.spillover_session_key: + identity = ["session", context.spillover_session_key] + else: + raise PermissionError("Spillover requires a verified caller or trusted MCP session") + return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest() + + +def bind_spillover_context(context, rc, registry, era=None): + context.spillover_session_key = registry.session_key(rc, era) + context.spillover_invalid_auth = bool( + authorization_token_from_request(rc) and context.auth is None + ) + + +class DataSpilloverInterceptor: + def __init__( + self, threshold_bytes: int = 64 * 1024, store: SpilloverStore | None = None, + resource_uri_prefix: str = "resource://data-spillover/", + ): + if type(threshold_bytes) is not int or threshold_bytes < 0: + raise ValueError("Spillover threshold must be a nonnegative integer") + parts = urlsplit(resource_uri_prefix) + if (not parts.scheme or not parts.netloc or parts.query or parts.fragment + or not resource_uri_prefix.endswith("/")): + raise ValueError("Resource prefix must be an absolute URI ending in '/'") + self.threshold_bytes = threshold_bytes + self.store = store if store is not None else MemorySpilloverStore() + self.resource_uri_prefix = resource_uri_prefix + + async def intercept(self, context, next_fn): + result = await next_fn() + if isinstance(result, (InputRequiredResult, SpilloverResult)): + return result + if isinstance(result, dict) and result.get("resultType") == "input_required": + return result + value = result + if isinstance(value, types.CallToolResult): + if value.is_error or value.result_type == "input_required": + return result + value = value.model_dump(mode="json", by_alias=True, exclude_none=True) + elif isinstance(value, dict) and value.get("isError"): + return result + elif isinstance(value, BaseModel): + value = value.model_dump(mode="json") + if isinstance(value, str): + payload, mime = value.encode("utf-8"), "text/plain" + else: + payload = json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=str).encode("utf-8") + mime = "application/json" + if len(payload) <= self.threshold_bytes: + return result + owner = caller_owner(context) + key = await self.store.put(payload, owner=owner, mime_type=mime) + return SpilloverResult(generate_preview(value), self.resource_uri_prefix + key, len(payload)) diff --git a/nitrostack/core/interceptors/spillover.py b/nitrostack/core/interceptors/spillover.py new file mode 100644 index 0000000..7ab5482 --- /dev/null +++ b/nitrostack/core/interceptors/spillover.py @@ -0,0 +1,218 @@ +"""Bounded, caller-scoped storage for serialized tool results.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import math +import os +import re +import threading +import time +import uuid +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +from nitrostack.core.errors import PayloadTooLargeError + + +@dataclass(frozen=True) +class SpilloverEntry: + payload: bytes + owner: str + mime_type: str + expires_at: float + + +class SpilloverStore(Protocol): + max_bytes: int + ttl_seconds: float + + async def put( + self, payload: bytes, *, owner: str, mime_type: str, + ttl_seconds: float | None = None, + ) -> str: ... + + async def get(self, key: str, *, owner: str) -> SpilloverEntry | None: ... + + async def delete(self, key: str, *, owner: str) -> bool: ... + + async def cleanup_expired(self) -> int: ... + + +def _valid_key(key: str) -> bool: + return isinstance(key, str) and re.fullmatch(r"[0-9a-f]{32}", key) is not None + + +class MemorySpilloverStore: + """Thread-safe payload cache with absolute TTL and LRU eviction.""" + + def __init__( + self, *, max_bytes: int = 10 * 1024 * 1024, + max_total_bytes: int = 100 * 1024 * 1024, ttl_seconds: float = 300, + ): + if any(type(limit) is not int or limit <= 0 for limit in (max_bytes, max_total_bytes)): + raise ValueError("Storage limits must be positive integers") + if not math.isfinite(ttl_seconds) or ttl_seconds <= 0: + raise ValueError("TTL must be positive and finite") + self.max_bytes = max_bytes + self.max_total_bytes = max_total_bytes + self.ttl_seconds = ttl_seconds + self._entries: OrderedDict[str, SpilloverEntry] = OrderedDict() + self._lock = threading.RLock() + self._size = 0 + + def _entry(self, payload, owner, mime_type, ttl_seconds): + if not isinstance(payload, bytes): + raise TypeError("Store payload must be bytes") + if not owner: + raise ValueError("A caller owner is required") + if len(payload) > self.max_bytes: + raise PayloadTooLargeError(f"Payload exceeds max_bytes={self.max_bytes}") + ttl = self.ttl_seconds if ttl_seconds is None else ttl_seconds + if not math.isfinite(ttl) or ttl <= 0: + raise ValueError("TTL must be positive and finite") + return SpilloverEntry(payload, owner, mime_type, time.time() + ttl) + + def _remove(self, key): + entry = self._entries.pop(key) + self._size -= len(entry.payload) + + def _cleanup(self): + expired = [key for key, entry in self._entries.items() if entry.expires_at <= time.time()] + for key in expired: + self._remove(key) + return len(expired) + + async def put(self, payload, *, owner, mime_type, ttl_seconds=None): + entry = self._entry(payload, owner, mime_type, ttl_seconds) + if len(payload) > self.max_total_bytes: + raise ValueError("Payload exceeds total store capacity") + with self._lock: + self._cleanup() + while self._size + len(payload) > self.max_total_bytes: + self._remove(next(iter(self._entries))) + key = uuid.uuid4().hex + self._entries[key] = entry + self._size += len(payload) + return key + + async def get(self, key, *, owner): + with self._lock: + self._cleanup() + entry = self._entries.get(key) if _valid_key(key) else None + if entry is None or entry.owner != owner: + return None + self._entries.move_to_end(key) + return entry + + async def delete(self, key, *, owner): + with self._lock: + entry = self._entries.get(key) if _valid_key(key) else None + if entry is None or entry.owner != owner: + return False + self._remove(key) + return True + + async def cleanup_expired(self): + with self._lock: + return self._cleanup() + + +class FileSpilloverStore(MemorySpilloverStore): + """Atomic disk store. Each root must be owned by one store instance/process. + + Total capacity counts actual JSON file bytes, including base64 and metadata. + File operations run in worker threads; no background cleanup task is needed. + """ + + def __init__(self, root: str | Path, **kwargs): + super().__init__(**kwargs) + self.root = Path(root).resolve() + self.root.mkdir(mode=0o700, parents=True, exist_ok=True) + self._cleanup_files() + + def _path(self, key): + if not _valid_key(key): + raise ValueError("Invalid spillover ID") + path = self.root / f"{key}.json" + if path.is_symlink() or path.resolve().parent != self.root: + raise ValueError("Spillover path escapes storage root") + return path + + def _read(self, path): + raw = json.loads(path.read_text(encoding="utf-8")) + return SpilloverEntry( + base64.b64decode(raw["payload"], validate=True), raw["owner"], + raw["mime_type"], raw["expires_at"], + ) + + def _cleanup_files(self): + count = 0 + for candidate in self.root.glob("*.json"): + if not _valid_key(candidate.stem): + continue + path = self._path(candidate.stem) + if self._read(path).expires_at <= time.time(): + path.unlink() + count += 1 + return count + + def _put_file(self, payload, owner, mime_type, ttl_seconds): + entry = self._entry(payload, owner, mime_type, ttl_seconds) + encoded = json.dumps({ + "payload": base64.b64encode(payload).decode("ascii"), + "owner": entry.owner, "mime_type": entry.mime_type, + "expires_at": entry.expires_at, + }).encode("utf-8") + with self._lock: + self._cleanup_files() + used = sum(self._path(p.stem).stat().st_size for p in self.root.glob("*.json") + if _valid_key(p.stem)) + if used + len(encoded) > self.max_total_bytes: + raise ValueError("File store capacity exceeded") + key = uuid.uuid4().hex + target = self._path(key) + temporary = self.root / f"{key}.tmp" + try: + descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(descriptor, "wb") as stream: + stream.write(encoded) + temporary.replace(target) + finally: + temporary.unlink(missing_ok=True) + return key + + def _get_file(self, key, owner, delete=False): + with self._lock: + path = self._path(key) + self._cleanup_files() + if not path.exists(): + return None + entry = self._read(path) + if entry.expires_at <= time.time(): + path.unlink() + return None + if entry.owner != owner: + return None + if delete: + path.unlink() + return entry + + async def put(self, payload, *, owner, mime_type, ttl_seconds=None): + return await asyncio.to_thread(self._put_file, payload, owner, mime_type, ttl_seconds) + + async def get(self, key, *, owner): + return await asyncio.to_thread(self._get_file, key, owner) + + async def delete(self, key, *, owner): + return await asyncio.to_thread(self._get_file, key, owner, True) is not None + + async def cleanup_expired(self): + def cleanup(): + with self._lock: + return self._cleanup_files() + return await asyncio.to_thread(cleanup) diff --git a/nitrostack/core/pipeline.py b/nitrostack/core/pipeline.py index 6bc5091..1b0a765 100644 --- a/nitrostack/core/pipeline.py +++ b/nitrostack/core/pipeline.py @@ -40,7 +40,7 @@ def decorator(func: Callable): return func return decorator -def use_interceptors(*interceptors: Type[Interceptor]): +def use_interceptors(*interceptors: Type[Interceptor] | Interceptor): def decorator(func: Callable): func._mcp_interceptors = getattr(func, "_mcp_interceptors", []) + list(interceptors) return func @@ -204,7 +204,7 @@ async def run_pipeline( context: ExecutionContext, guards: List[Type[Guard]], middleware: List[Type[Middleware]], - interceptors: List[Type[Interceptor]], + interceptors: List[Type[Interceptor] | Interceptor], pipes: List[Type[Pipe]], filters: List[Type[ExceptionFilter]], param_name: str = "input", @@ -248,7 +248,10 @@ async def call_target(): # Middleware -> Interceptors -> Handler current_next = call_target for interceptor_cls in reversed(interceptors): - interceptor = container.resolve(interceptor_cls) + interceptor = ( + container.resolve(interceptor_cls) + if isinstance(interceptor_cls, type) else interceptor_cls + ) def make_interceptor_next(nxt, icpt=interceptor): return lambda: icpt.intercept(context, nxt) diff --git a/nitrostack/runtime/request_ctx.py b/nitrostack/runtime/request_ctx.py index 485e849..99f2207 100644 --- a/nitrostack/runtime/request_ctx.py +++ b/nitrostack/runtime/request_ctx.py @@ -7,6 +7,10 @@ from typing import Any, Mapping, Optional request_ctx: ContextVar[Any] = ContextVar("nitrostack_request_ctx", default=None) +# Set by the stdio adapter for one live stream pair, never by client metadata. +stdio_connection_id: ContextVar[Optional[str]] = ContextVar( + "nitrostack_stdio_connection_id", default=None +) @dataclass diff --git a/nitrostack/transports/stdio.py b/nitrostack/transports/stdio.py index d76e580..d3fc9bd 100644 --- a/nitrostack/transports/stdio.py +++ b/nitrostack/transports/stdio.py @@ -1,11 +1,15 @@ import sys import contextlib +import uuid +import contextvars from typing import TYPE_CHECKING, Any, Generator +import anyio from mcp.server.runner import serve_loop from mcp.server.stdio import stdio_server from nitrostack.protocol.version import ProtocolEra +from nitrostack.runtime.request_ctx import stdio_connection_id if TYPE_CHECKING: from mcp.server.lowlevel import Server @@ -45,11 +49,63 @@ def safe_stdio_transport() -> Generator[None, None, None]: sys.stdout = original_stdout +class _StdioConnectionStream: + """Preserve sender context while binding a server-owned connection identity.""" + + def __init__(self, stream: Any, connection_id: str) -> None: + self._stream = stream + self._connection_id = connection_id + self.last_context = None + + async def receive(self): + message = await self._stream.receive() + sender_context = getattr(self._stream, "last_context", None) + self.last_context = sender_context.copy() if sender_context is not None else contextvars.copy_context() + self.last_context.run(stdio_connection_id.set, self._connection_id) + return message + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return await self.receive() + except anyio.EndOfStream: + raise StopAsyncIteration from None + + async def __aenter__(self): + await self._stream.__aenter__() + return self + + async def __aexit__(self, *args): + return await self._stream.__aexit__(*args) + + async def aclose(self): + await self._stream.aclose() + + async def serve_stdio_streams( server: "Server[Any]", read_stream: "ReadStream[SessionMessage | Exception]", write_stream: "WriteStream[SessionMessage]", era: ProtocolEra, +) -> None: + """Serve one connection with a trusted identity shared by its requests.""" + connection_id = uuid.uuid4().hex + token = stdio_connection_id.set(connection_id) + try: + await _serve_stdio_connection( + server, _StdioConnectionStream(read_stream, connection_id), write_stream, era + ) + finally: + stdio_connection_id.reset(token) + + +async def _serve_stdio_connection( + server: "Server[Any]", + read_stream: "ReadStream[SessionMessage | Exception]", + write_stream: "WriteStream[SessionMessage]", + era: ProtocolEra, ) -> None: """Drive official mcp 2.x over an already-open stdio stream pair. diff --git a/tests/spillover_transport_support.py b/tests/spillover_transport_support.py new file mode 100644 index 0000000..86b5278 --- /dev/null +++ b/tests/spillover_transport_support.py @@ -0,0 +1,136 @@ +"""Real application and protocol clients for spillover transport regressions.""" + +import asyncio +from contextlib import contextmanager +import json +from pathlib import Path +import sys +from typing import Any + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from pydantic import BaseModel +from starlette.testclient import TestClient + +from nitrostack import ( + DataSpilloverInterceptor, FileSpilloverStore, MemorySpilloverStore, + module, tool, use_interceptors, +) +from nitrostack.core.app import McpApplication, ServerConfig, mcp_app + +MODERN_META = { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": {"name": "spillover-client", "version": "1"}, + "io.modelcontextprotocol/clientCapabilities": {}, +} + + +class PayloadInput(BaseModel): + value: Any = None + case: str | None = None + + +def large_value(case): + if case == "threshold": + return "x" * 65536 + if case == "threshold-plus-one": + return "x" * 65537 + if case.endswith("MiB"): + return "x" * (int(case[:-3]) * 1024 * 1024) + if case == "unicode": + return "é🚀" * 200000 + if case == "escaped": + return '\\"\n\t' * 100000 + if case == "rows": + return [{"row": i, "value": "row"} for i in range(100000)] + if case == "dictionary": + return {str(i): i for i in range(50000)} + if case == "nested": + return [{"chunks": ["x" * (1024 * 1024)] * 6}] + return "x" * (10 * 1024 * 1024 + 1) + + +def build_app(backend, root, era="legacy", threshold=64 * 1024, + prefix="resource://data-spillover/"): + options = dict(max_bytes=10 * 1024 * 1024, max_total_bytes=64 * 1024 * 1024) + store = MemorySpilloverStore(**options) if backend == "memory" else FileSpilloverStore(root, **options) + interceptor = DataSpilloverInterceptor(threshold, store, prefix) + + class Tools: + @tool("payload", "Return a payload", PayloadInput, task_support="optional") + @use_interceptors(interceptor) + async def payload(self, input, context): + return large_value(input.case) if input.case is not None else input.value + + @module(name="spillover-wire", providers=[Tools]) + class Root: + pass + + @mcp_app(module=Root, server=ServerConfig(name="spillover-wire", protocol_era=era)) + class App: + pass + + return McpApplication(App) + + +class HttpPeer: + def __init__(self, client, era="legacy", token=None): + self.client, self.era = client, era + self.headers = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"} + if token: + self.headers["Authorization"] = "Bearer " + token + self.next_id = 1 + if era == "legacy": + response = self.rpc("initialize", { + "protocolVersion": "2025-06-18", "capabilities": {}, + "clientInfo": {"name": "spillover-client", "version": "1"}, + }) + assert "result" in response + self.headers["mcp-session-id"] = self.last_response.headers["mcp-session-id"] + initialized = client.post("/mcp", headers=self.headers, json={ + "jsonrpc": "2.0", "method": "notifications/initialized", + }) + assert initialized.status_code == 202 + else: + self.headers["MCP-Protocol-Version"] = "2026-07-28" + + def rpc(self, method, params): + params = dict(params) + headers = {**self.headers, "Mcp-Method": method} + if method == "tools/call": + headers["Mcp-Name"] = params["name"] + elif method == "resources/read": + headers["Mcp-Name"] = params["uri"] + if self.era != "legacy": + params["_meta"] = {**params.get("_meta", {}), **MODERN_META} + self.last_response = self.client.post("/mcp", headers=headers, json={ + "jsonrpc": "2.0", "id": self.next_id, "method": method, "params": params, + }) + self.next_id += 1 + assert self.last_response.status_code == 200, self.last_response.text + return self.last_response.json() + + def call(self, value, **extra): + return self.rpc("tools/call", {"name": "payload", "arguments": {"value": value}, **extra}) + + def read(self, uri): + return self.rpc("resources/read", {"uri": uri}) + + +@contextmanager +def http_app(backend, root, **options): + app = build_app(backend, root, **options) + with TestClient(app.get_combined_app(json_response=True)) as client: + yield client + + +async def stdio_main(): + from nitrostack.transports.stdio import run_stdio, safe_stdio_transport + + with safe_stdio_transport(): + app = build_app(sys.argv[1], sys.argv[2], era="auto", threshold=16) + await run_stdio(app.mcp_server, "auto") + + +if __name__ == "__main__": + asyncio.run(stdio_main()) diff --git a/tests/test_spillover.py b/tests/test_spillover.py new file mode 100644 index 0000000..25dc634 --- /dev/null +++ b/tests/test_spillover.py @@ -0,0 +1,390 @@ +"""Issue #67: bounded results, caller isolation, and safe spillover stores.""" + +import asyncio +import gc +import json +from types import SimpleNamespace +from unittest.mock import MagicMock + +import mcp.types as types +import pytest +from pydantic import BaseModel + +from nitrostack import ( + AuthContext, DataSpilloverInterceptor, DIContainer, ExecutionContext, + FileSpilloverStore, MemorySpilloverStore, PayloadTooLargeError, + ResourceNotFoundError, generate_preview, module, tool, use_interceptors, +) +from nitrostack.core.app import McpApplication +from nitrostack.core.interceptors.data_spillover import ( + SpilloverRegistry, SpilloverResult, caller_owner, +) +from nitrostack.protocol.mrtr import InputRequiredResult +from nitrostack.runtime.request_ctx import RequestContext, request_ctx + + +@pytest.fixture(autouse=True) +def reset_container(): + DIContainer.reset() + yield + DIContainer.reset() + + +@pytest.fixture(params=["memory", "file"]) +def store(request, tmp_path): + cls = MemorySpilloverStore if request.param == "memory" else FileSpilloverStore + args = () if request.param == "memory" else (tmp_path / "results",) + return cls(*args, max_bytes=4096, max_total_bytes=16384, ttl_seconds=10) + + +def test_store_roundtrip_owner_delete(store): + async def scenario(): + key = await store.put(b'{"answer":42}', owner="alice", mime_type="application/json") + assert await store.get(key, owner="bob") is None + assert not await store.delete(key, owner="bob") + entry = await store.get(key, owner="alice") + assert entry.payload == b'{"answer":42}' + assert entry.mime_type == "application/json" + assert await store.delete(key, owner="alice") + assert await store.get(key, owner="alice") is None + asyncio.run(scenario()) + + +def test_store_expiry_and_payload_limit(store, monkeypatch): + clock = [1000.0] + monkeypatch.setattr("nitrostack.core.interceptors.spillover.time.time", lambda: clock[0]) + async def scenario(): + key = await store.put(b"hello", owner="alice", mime_type="text/plain") + with pytest.raises(PayloadTooLargeError): + await store.put(b"x" * 4097, owner="alice", mime_type="text/plain") + clock[0] += 10 + assert await store.cleanup_expired() == 1 + assert await store.get(key, owner="alice") is None + asyncio.run(scenario()) + + +def test_memory_lru_and_total_capacity(): + store = MemorySpilloverStore(max_bytes=8, max_total_bytes=8) + async def scenario(): + first = await store.put(b"1111", owner="a", mime_type="text/plain") + second = await store.put(b"2222", owner="a", mime_type="text/plain") + await store.get(first, owner="a") + await store.put(b"3333", owner="a", mime_type="text/plain") + assert await store.get(second, owner="a") is None + assert await store.get(first, owner="a") is not None + assert store._size == 8 + asyncio.run(scenario()) + + +def test_file_capacity_restart_cleanup(tmp_path, monkeypatch): + clock = [1000.0] + monkeypatch.setattr("nitrostack.core.interceptors.spillover.time.time", lambda: clock[0]) + store = FileSpilloverStore(tmp_path, max_total_bytes=300, ttl_seconds=5) + async def scenario(): + key = await store.put(b"x" * 80, owner="a", mime_type="text/plain") + with pytest.raises(ValueError, match="capacity"): + await store.put(b"x" * 80, owner="a", mime_type="text/plain") + reopened = FileSpilloverStore(tmp_path, max_total_bytes=300) + assert (await reopened.get(key, owner="a")).payload == b"x" * 80 + assert list(tmp_path.glob("*.tmp")) == [] + clock[0] += 5 + FileSpilloverStore(tmp_path) + assert list(tmp_path.glob("*.json")) == [] + asyncio.run(scenario()) + + +@pytest.mark.parametrize("key", ["../escape", "..\\escape", "%2e%2e%2fescape", "/etc/passwd", "C:\\escape", "x?y"]) +def test_file_rejects_unsafe_ids(tmp_path, key): + store = FileSpilloverStore(tmp_path) + async def scenario(): + with pytest.raises(ValueError): + await store.get(key, owner="a") + with pytest.raises(ValueError): + await store.delete(key, owner="a") + asyncio.run(scenario()) + + +def test_file_rejects_symlink_escape(tmp_path): + root = tmp_path / "results" + store = FileSpilloverStore(root) + outside = tmp_path / "private.json" + outside.write_text("private") + try: + (root / ("a" * 32 + ".json")).symlink_to(outside) + except OSError: + pytest.skip("Creating symlinks requires OS privileges") + with pytest.raises(ValueError, match="escapes"): + asyncio.run(store.get("a" * 32, owner="a")) + assert outside.read_text() == "private" + + +def test_concurrent_store_operations(store): + async def scenario(): + keys = await asyncio.gather(*[ + store.put(str(i).encode(), owner="a", mime_type="text/plain") for i in range(20) + ]) + assert len(set(keys)) == 20 + entries = await asyncio.gather(*[store.get(key, owner="a") for key in keys]) + assert [entry.payload for entry in entries] == [str(i).encode() for i in range(20)] + asyncio.run(scenario()) + + +@pytest.mark.parametrize("value", [list(range(20)), {str(i): i for i in range(20)}, "é" * 10000, + [{"huge": "x" * 10000}], {"huge": ["x" * 10000]}], + ids=["list", "dict", "unicode", "nested-list", "nested-dict"]) +def test_previews_are_bounded(value): + preview = generate_preview(value) + assert preview["count"] == len(value) + assert preview["truncated"] + assert len(json.dumps(preview, ensure_ascii=False).encode()) <= 2048 + + +def test_threshold_boundary_and_full_payload(store): + interceptor = DataSpilloverInterceptor(threshold_bytes=4, store=store) + context = ExecutionContext(request_id="1", auth=AuthContext(subject="alice"), logger=MagicMock()) + async def scenario(): + for value in ["éé", {"a": 1}, [1, 2]]: + async def handler(): + return value + result = await interceptor.intercept(context, handler) + if value == "éé": + assert result is value + else: + assert isinstance(result, SpilloverResult) + entry = await store.get(result.resource_uri.rsplit("/", 1)[1], owner=caller_owner(context)) + assert json.loads(entry.payload) == value + assert result.original_size_bytes == len(entry.payload) + async def unicode_handler(): + return "ééé" + assert (await interceptor.intercept(context, unicode_handler)).original_size_bytes == 6 + asyncio.run(scenario()) + + +def test_control_results_and_missing_identity(): + interceptor = DataSpilloverInterceptor(threshold_bytes=1) + context = ExecutionContext(request_id="1", logger=MagicMock()) + async def scenario(): + controls = [types.CallToolResult(content=[], isError=True), + {"resultType": "input_required", "message": "more"}, + InputRequiredResult(input_requests=[], request_state={}, message="more")] + for value in controls: + async def handler(): + return value + assert await interceptor.intercept(context, handler) is value + async def large(): + return "large" + with pytest.raises(PermissionError): + await interceptor.intercept(context, large) + asyncio.run(scenario()) + + +class Session: + pass + + +class ToolInput(BaseModel): + text: str = "large result" + + +def make_app(interceptor, *, task_support="forbidden", output_schema=None): + class Tools: + @tool(name="large", description="Test spillover", input_schema=ToolInput, + task_support=task_support, output_schema=output_schema) + @use_interceptors(interceptor) + async def large(self, input, context): + return {"text": input.text} + @module(name="spillover-test", providers=[Tools]) + class Root: + pass + return McpApplication(Root) + + +async def call(app, rc, text="large result", task=None): + token = request_ctx.set(rc) + try: + if task is not None: + return await app._call_tool("large", {"text": text}, task=task) + handler = app.mcp_server.request_handlers[types.CallToolRequest] + return (await handler(types.CallToolRequest( + method="tools/call", params=types.CallToolRequestParams(name="large", arguments={"text": text}), + ))).root + finally: + request_ctx.reset(token) + + +async def read(app, rc, uri): + token = request_ctx.set(rc) + try: + handler = app.mcp_server.request_handlers[types.ReadResourceRequest] + return (await handler(types.ReadResourceRequest( + method="resources/read", params=types.ReadResourceRequestParams(uri=uri), + ))).root.contents[0].text + finally: + request_ctx.reset(token) + + +def test_mcp_resource_roundtrip_and_session_isolation(store): + app = make_app(DataSpilloverInterceptor(threshold_bytes=1, store=store)) + alice = RequestContext(session=Session()) + bob = RequestContext(session=Session()) + async def scenario(): + result = await call(app, alice) + assert not result.is_error + data = result.structured_content + assert result.content[1].type == "resource_link" + assert json.loads(await read(app, alice, data["resource_uri"])) == {"text": "large result"} + for caller in [bob, None]: + with pytest.raises(ResourceNotFoundError): + await read(app, caller, data["resource_uri"]) + with pytest.raises(ResourceNotFoundError): + await read(app, alice, "resource://data-spillover/../escape") + assert data["original_size_bytes"] == len(b'{"text":"large result"}') + asyncio.run(scenario()) + + +def test_principal_tenant_issuer_isolation(monkeypatch): + payloads = { + "alice": {"sub": "alice", "iss": "issuer", "tenant_id": "one"}, + "bob": {"sub": "bob", "iss": "issuer", "tenant_id": "one"}, + "tenant": {"sub": "alice", "iss": "issuer", "tenant_id": "two"}, + "issuer": {"sub": "alice", "iss": "other", "tenant_id": "one"}, + } + monkeypatch.setattr("nitrostack.auth.request.verify_bearer_payload", payloads.get) + app = make_app(DataSpilloverInterceptor(threshold_bytes=1)) + def rc(token): + return RequestContext(request=SimpleNamespace(headers={"authorization": "Bearer " + token}), + session=Session(), meta={"userId": "alice", "tenantId": "one"}) + async def scenario(): + result = await call(app, rc("alice")) + uri = result.structured_content["resource_uri"] + assert json.loads(await read(app, rc("alice"), uri)) == {"text": "large result"} + for other in ["bob", "tenant", "issuer", "invalid"]: + with pytest.raises(ResourceNotFoundError): + await read(app, rc(other), uri) + assert (await call(app, rc("invalid"))).is_error + asyncio.run(scenario()) + + +def test_class_decorator_custom_prefix_and_schema(): + class Configured(DataSpilloverInterceptor): + def __init__(self): + super().__init__(threshold_bytes=1, resource_uri_prefix="resource://custom/") + schema = {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]} + app = make_app(Configured, output_schema=schema) + async def scenario(): + rc = RequestContext(session=Session()) + result = await call(app, rc) + assert result.structured_content["resource_uri"].startswith("resource://custom/") + definition = app._build_tool_definition(app._tools["large"]) + original = definition.output_schema["anyOf"][0] + assert all(original[key] == value for key, value in schema.items()) + assert definition.output_schema["anyOf"][1]["required"] == ["preview", "resource_uri", "original_size_bytes"] + asyncio.run(scenario()) + + +def test_registry_conflicts_and_session_cleanup(): + registry = SpilloverRegistry() + interceptor = DataSpilloverInterceptor() + registry.register(interceptor) + registry.register(interceptor) + with pytest.raises(ValueError, match="Conflicting"): + registry.register(DataSpilloverInterceptor()) + rc = RequestContext(session=Session()) + assert registry.session_key(rc) == registry.session_key(rc) + del rc + gc.collect() + assert registry._sessions == {} + + +def test_mcp_expiry_and_store_errors(store, monkeypatch): + clock = [1000.0] + monkeypatch.setattr("nitrostack.core.interceptors.spillover.time.time", lambda: clock[0]) + app = make_app(DataSpilloverInterceptor(threshold_bytes=1, store=store)) + rc = RequestContext(session=Session()) + async def scenario(): + result = await call(app, rc) + uri = result.structured_content["resource_uri"] + clock[0] += 10 + with pytest.raises(ResourceNotFoundError): + await read(app, rc, uri) + oversized = await call(app, rc, text="x" * 5000) + assert oversized.is_error + assert "max_bytes" in oversized.content[0].text + assert "x" * 5000 not in oversized.content[0].text + asyncio.run(scenario()) + + +def test_models_mcp_results_and_widget_bypass(store): + interceptor = DataSpilloverInterceptor(threshold_bytes=1, store=store) + context = ExecutionContext(request_id="1", auth=AuthContext(subject="alice"), logger=MagicMock()) + app = make_app(interceptor) + async def scenario(): + for value in [ToolInput(text="full"), types.CallToolResult( + content=[types.TextContent(type="text", text="full")], structuredContent={"text": "full"}, + )]: + async def handler(): + return value + spilled = await interceptor.intercept(context, handler) + payload = await store.get(spilled.resource_uri.rsplit("/", 1)[1], owner=caller_owner(context)) + assert json.loads(payload.payload) == value.model_dump(mode="json", by_alias=True, exclude_none=True) + component = MagicMock() + result = app._to_call_tool_result(spilled, component, context) + assert result.structured_content["resource_uri"] == spilled.resource_uri + component.transformer.assert_not_called() + asyncio.run(scenario()) + + +def test_background_task_spills_and_can_be_read(monkeypatch): + # Task augmentation is supported on the legacy wire, not MCP 2026. + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", "legacy") + app = make_app(DataSpilloverInterceptor(threshold_bytes=1), task_support="optional") + rc = RequestContext(session=Session()) + async def scenario(): + result = await call(app, rc, task=types.TaskMetadata(ttl=60000)) + for _ in range(100): + entry = await app.task_manager._store.get(result.task.task_id) + if entry.result is not None: + break + await asyncio.sleep(0.01) + assert entry.result is not None + stored = entry.result + data = stored.structured_content if isinstance(stored, types.CallToolResult) else stored["structuredContent"] + assert json.loads(await read(app, rc, data["resource_uri"])) == {"text": "large result"} + asyncio.run(scenario()) + + +def test_opaque_oauth_identity_reads_resources(monkeypatch): + from nitrostack.auth.oauth import OAuthService + from unittest.mock import AsyncMock + + monkeypatch.setattr("nitrostack.auth.request.verify_bearer_payload", lambda token: None) + service = MagicMock() + service.introspect_token = AsyncMock(return_value={"active": True, "sub": "alice", "iss": "oauth"}) + DIContainer.get_instance().register_value(OAuthService, service) + app = make_app(DataSpilloverInterceptor(threshold_bytes=1)) + rc = RequestContext(request=SimpleNamespace(headers={"authorization": "Bearer opaque"})) + async def scenario(): + result = await call(app, rc) + uri = result.structured_content["resource_uri"] + assert json.loads(await read(app, rc, uri)) == {"text": "large result"} + service.introspect_token.return_value = {"active": False} + with pytest.raises(ResourceNotFoundError): + await read(app, rc, uri) + asyncio.run(scenario()) + + +def test_file_symlink_check_without_os_privileges(tmp_path, monkeypatch): + from pathlib import Path + + store = FileSpilloverStore(tmp_path) + monkeypatch.setattr(Path, "is_symlink", lambda path: path.name == "a" * 32 + ".json") + with pytest.raises(ValueError, match="escapes"): + asyncio.run(store.get("a" * 32, owner="a")) + + +@pytest.mark.parametrize("kwargs", [{"threshold_bytes": -1}, {"threshold_bytes": 0.5}, + {"resource_uri_prefix": "../results/"}, + {"resource_uri_prefix": "resource://results/?x/"}]) +def test_invalid_interceptor_configuration(kwargs): + with pytest.raises(ValueError): + DataSpilloverInterceptor(**kwargs) diff --git a/tests/test_spillover_transport.py b/tests/test_spillover_transport.py new file mode 100644 index 0000000..f5670f8 --- /dev/null +++ b/tests/test_spillover_transport.py @@ -0,0 +1,265 @@ +"""Independent spillover verification through official HTTP and stdio transports.""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor +import json +import os +from pathlib import Path +import subprocess +import sys +import time +from types import SimpleNamespace + +import anyio +from mcp.shared.memory import create_client_server_memory_streams +from mcp.shared.message import SessionMessage +from mcp_types import jsonrpc_message_adapter +import pytest + +from nitrostack import DIContainer +from nitrostack.auth.jwt import JWTModule, JWTService +from nitrostack.runtime.request_ctx import stdio_connection_id +from nitrostack.transports.stdio import serve_stdio_streams +from tests.spillover_transport_support import MODERN_META, HttpPeer, build_app, http_app, large_value + +ROOT = Path(__file__).resolve().parents[1] + + +@pytest.fixture(autouse=True) +def clean_context(monkeypatch): + DIContainer.reset() + for key in ("NITRO_MCP_PROTOCOL_VERSION", "MCP_STATELESS"): + monkeypatch.delenv(key, raising=False) + yield + DIContainer.reset() + + +@pytest.fixture(params=["memory", "file"]) +def backend(request): + return request.param + + +def envelope(response): + result = response["result"] + assert result.get("isError", False) is False + data = result["structuredContent"] + assert len(json.dumps(data["preview"], ensure_ascii=False).encode()) <= 2048 + assert result["content"][1]["type"] == "resource_link" + assert result["content"][1]["uri"] == data["resource_uri"] + return data + + +def assert_full(peer, data, value): + expected = (value if isinstance(value, str) else json.dumps( + value, ensure_ascii=False, separators=(",", ":"), + )).encode() + assert data["original_size_bytes"] == len(expected) + contents = peer.read(data["resource_uri"])["result"]["contents"][0] + assert contents["text"].encode() == expected + assert contents["mimeType"] == ("text/plain" if isinstance(value, str) else "application/json") + + +def test_anonymous_http_roundtrip_and_session_isolation(backend, tmp_path): + with http_app(backend, tmp_path, threshold=16) as client: + alice, bob = HttpPeer(client), HttpPeer(client) + value = {"rows": ["é🚀" * 100]} + data = envelope(alice.call(value)) + assert_full(alice, data, value) + assert "error" in bob.read(data["resource_uri"]) + original = bob.headers["mcp-session-id"] + bob.headers["mcp-session-id"] = "not-server-issued" + response = client.post("/mcp", headers=bob.headers, json={ + "jsonrpc": "2.0", "id": 99, "method": "resources/read", + "params": {"uri": data["resource_uri"]}, + }) + assert response.status_code == 404 + bob.headers["mcp-session-id"] = original + # Expired/deleted transport sessions cannot recover anonymous payloads. + deleted = client.delete("/mcp", headers=alice.headers) + assert deleted.status_code == 200 + assert "error" in HttpPeer(client).read(data["resource_uri"]) + + +@pytest.mark.parametrize("era", ["legacy", "modern", "auto"]) +def test_verified_identity_reconnect_and_auth_isolation(backend, tmp_path, monkeypatch, era): + monkeypatch.setenv("JWT_SECRET", "spillover-wire-regression-secret") + JWTModule.for_root(secret_env_var="JWT_SECRET") + jwt = DIContainer.get_instance().resolve(JWTService) + claims = {"sub": "alice", "iss": "issuer-a", "client_id": "client-a", "tenant_id": "tenant-a"} + token = jwt.create_token(claims) + with http_app(backend, tmp_path, era=era, threshold=16) as client: + alice = HttpPeer(client, era, token) + data = envelope(alice.call("private-value" * 100)) + assert_full(HttpPeer(client, era, token), data, "private-value" * 100) + for field in ("sub", "iss", "client_id", "tenant_id"): + other = jwt.create_token({**claims, field: "another"}) + assert "error" in HttpPeer(client, era, other).read(data["resource_uri"]) + invalid = HttpPeer(client, era, "invalid-signature") + assert "error" in invalid.read(data["resource_uri"]) + assert invalid.call("private-value" * 100)["result"]["isError"] is True + assert "error" in HttpPeer(client, era).read(data["resource_uri"]) + + +@pytest.mark.parametrize("era", ["modern", "auto"]) +def test_anonymous_sessionless_cannot_create_owner(backend, tmp_path, era): + with http_app(backend, tmp_path, era=era, threshold=16) as client: + peer = HttpPeer(client, era) + denied = peer.call("private-value" * 100)["result"] + assert denied["isError"] is True + assert "structuredContent" not in denied + # Unsigned metadata is not a substitute for a trusted transport identity. + denied = peer.rpc("tools/call", { + "name": "payload", "arguments": {"value": "private-value" * 100}, + "_meta": {"sessionId": "forged", "userId": "alice"}, + }) + assert denied["result"]["isError"] is True + + +def test_expired_missing_and_malformed_uris(backend, tmp_path, monkeypatch): + clock = [1000.0] + monkeypatch.setattr("nitrostack.core.interceptors.spillover.time", SimpleNamespace(time=lambda: clock[0])) + with http_app(backend, tmp_path, threshold=16) as client: + peer = HttpPeer(client) + data = envelope(peer.call("private-value" * 100)) + clock[0] += 301 + for uri in (data["resource_uri"], "resource://data-spillover/" + "0" * 32, + "resource://data-spillover/../escape", "resource://data-spillover/%2e%2e%2fescape"): + response = peer.read(uri) + assert "error" in response + assert "private-value" not in json.dumps(response) + + +def test_custom_prefix_and_background_task(backend, tmp_path): + with http_app(backend, tmp_path, threshold=16, prefix="resource://reports/") as client: + peer = HttpPeer(client) + value = "task-result" * 100 + response = peer.call(value, task={"ttl": 60000})["result"] + task_id = response["structuredContent"]["task"]["taskId"] + for _ in range(100): + result = peer.rpc("tasks/get", {"taskId": task_id}) + assert "result" in result, result + if result["result"]["status"] == "completed": + break + time.sleep(0.01) + else: + pytest.fail("Spillover task did not complete") + completed = result["result"]["result"] + data = envelope({"result": completed}) + assert data["resource_uri"].startswith("resource://reports/") + assert_full(peer, data, value) + + +CASES = ["threshold", "threshold-plus-one", "1MiB", "5MiB", "10MiB", + "unicode", "escaped", "rows", "dictionary", "nested", "too-large", "20MiB"] + + +@pytest.mark.parametrize("case", CASES) +def test_large_payloads_over_http(backend, tmp_path, case): + value = large_value(case) + with http_app(backend, tmp_path) as client: + peer = HttpPeer(client) + response = peer.rpc("tools/call", {"name": "payload", "arguments": {"case": case}}) + if case == "threshold": + result = response["result"] + assert result["content"][0]["text"] == value + assert "resource_uri" not in result.get("structuredContent", {}) + elif case in ("too-large", "20MiB"): + result = response["result"] + assert result["isError"] is True + assert "max_bytes" in result["content"][0]["text"] + assert len(json.dumps(result)) < 2048 + else: + data = envelope(response) + assert_full(peer, data, value) + assert "error" in HttpPeer(client).read(data["resource_uri"]) + + +class StdioPeer: + def __init__(self, backend, root): + self.process = subprocess.Popen( + [sys.executable, str(ROOT / "tests" / "spillover_transport_support.py"), backend, str(root)], + cwd=ROOT, env={**os.environ, "PYTHONPATH": str(ROOT)}, + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + text=True, encoding="utf-8", + ) + self.reader = ThreadPoolExecutor(max_workers=1) + self.next_id = 1 + + def rpc(self, method, params): + self.process.stdin.write(json.dumps({ + "jsonrpc": "2.0", "id": self.next_id, "method": method, + "params": {**params, "_meta": MODERN_META}, + }) + "\n") + self.process.stdin.flush() + line = self.reader.submit(self.process.stdout.readline).result(timeout=10) + response = json.loads(line) + assert response["id"] == self.next_id + self.next_id += 1 + return response + + def close(self, *, clean=False): + try: + if clean: + self.process.stdin.close() + self.process.wait(timeout=10) + assert self.process.returncode == 0, self.process.stderr.read() + assert self.process.stdout.read() == "" + finally: + if self.process.poll() is None: + self.process.kill() + self.process.wait(timeout=5) + self.reader.shutdown() + for stream in (self.process.stdin, self.process.stdout, self.process.stderr): + stream.close() + + +def test_stdio_roundtrip_connection_isolation_and_shutdown(backend, tmp_path): + # Both processes use one disk root sequentially to prove persistence does not + # let a later stdio connection inherit anonymous ownership of an earlier one. + first = StdioPeer(backend, tmp_path) + try: + value = "stdio-result" * 100 + data = envelope(first.rpc("tools/call", {"name": "payload", "arguments": {"value": value}})) + result = first.rpc("resources/read", {"uri": data["resource_uri"]}) + assert result["result"]["contents"][0]["text"] == value + first.close(clean=True) + finally: + first.close() + second = StdioPeer(backend, tmp_path) + try: + assert "error" in second.rpc("resources/read", {"uri": data["resource_uri"]}) + second.close(clean=True) + finally: + second.close() + + +def test_stdio_identity_is_reset_after_shutdown_and_not_inherited_by_other_task(tmp_path): + async def scenario(): + app = build_app("memory", tmp_path, era="auto", threshold=16) + async with create_client_server_memory_streams() as (client, server): + client_read, client_write = client + seen = [] + + async def serve(): + await serve_stdio_streams(app.mcp_server, *server, "auto") + seen.append(stdio_connection_id.get()) + + async with anyio.create_task_group() as group: + group.start_soon(serve) + await client_write.send(SessionMessage(jsonrpc_message_adapter.validate_python({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": { + "name": "payload", "arguments": {"value": "x" * 100}, "_meta": MODERN_META, + }, + }))) + response = await asyncio.wait_for(client_read.receive(), 5) + assert "resource_uri" in response.message.model_dump_json() + assert stdio_connection_id.get() is None + # The same app can serve HTTP concurrently without borrowing the + # anonymous identity belonging to this stdio connection. + from starlette.testclient import TestClient + with TestClient(app.get_combined_app(json_response=True)) as http: + assert HttpPeer(http, "auto").call("x" * 100)["result"]["isError"] is True + await client_write.aclose() + assert seen == [None] + assert stdio_connection_id.get() is None + asyncio.run(scenario())