From 41ca887d391a27fdd785f377509c356f4d8d22cd Mon Sep 17 00:00:00 2001 From: Shreya Vemana Date: Fri, 23 Oct 2026 20:20:01 +0530 Subject: [PATCH 1/5] test: capture tools/list baseline before transform pipeline (#61) --- scripts/capture_tools_baseline.py | 56 +++++++++++++++++++++++++ tests/fixtures/tools_list_baseline.json | 5 +++ 2 files changed, 61 insertions(+) create mode 100644 scripts/capture_tools_baseline.py create mode 100644 tests/fixtures/tools_list_baseline.json diff --git a/scripts/capture_tools_baseline.py b/scripts/capture_tools_baseline.py new file mode 100644 index 0000000..50b52a2 --- /dev/null +++ b/scripts/capture_tools_baseline.py @@ -0,0 +1,56 @@ + +"""Capture real tools/list responses from the unmodified base SDK.""" +import asyncio +import json +import pathlib +import sys + +ROOT = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import mcp.types as types + +from nitrostack.testing import NitroTestingModule +from tests._baseline_catalog import ERAS, M + +OUT = ROOT / "tests" / "fixtures" / "tools_list_baseline.json" + + +async def list_tools(era: str) -> str: + harness = await NitroTestingModule.create(M, protocol_era=era) + server = harness.app.mcp_server + + response = ( + await server.request_handlers[types.ListToolsRequest]( + types.ListToolsRequest(method="tools/list") + ) + ).root + + return response.model_dump_json( + by_alias=True, + exclude_none=True, + ) + + +async def main() -> None: + baseline = {} + + for era in ERAS: + serialized = await list_tools(era) + baseline[era] = serialized + names = [ + item["name"] + for item in json.loads(serialized)["tools"] + ] + print(f"{era}: {names}") + + OUT.parent.mkdir(parents=True, exist_ok=True) + OUT.write_text( + json.dumps(baseline, indent=2), + encoding="utf-8", + ) + print(f"Wrote {OUT} ({OUT.stat().st_size} bytes)") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/fixtures/tools_list_baseline.json b/tests/fixtures/tools_list_baseline.json new file mode 100644 index 0000000..6edb97d --- /dev/null +++ b/tests/fixtures/tools_list_baseline.json @@ -0,0 +1,5 @@ +{ + "legacy": "{\"_meta\":{\"io.modelcontextprotocol/cacheHint\":{\"ttlMs\":60000,\"cacheScope\":\"private\"}},\"ttlMs\":0,\"cacheScope\":\"private\",\"tools\":[{\"name\":\"alpha\",\"description\":\"first\",\"inputSchema\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"},\"annotations\":{\"readOnlyHint\":false,\"destructiveHint\":true,\"idempotentHint\":false,\"openWorldHint\":true},\"_meta\":{\"is_initial\":false,\"visibility\":\"visible\",\"task_support\":\"forbidden\",\"openai/type\":\"function\",\"openai/function\":{\"name\":\"alpha\",\"description\":\"first\",\"parameters\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"}}}},{\"name\":\"beta\",\"description\":\"second\",\"inputSchema\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"},\"annotations\":{\"readOnlyHint\":false,\"destructiveHint\":true,\"idempotentHint\":false,\"openWorldHint\":true},\"_meta\":{\"is_initial\":false,\"visibility\":\"visible\",\"task_support\":\"forbidden\",\"openai/type\":\"function\",\"openai/function\":{\"name\":\"beta\",\"description\":\"second\",\"parameters\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"}}}}],\"resultType\":\"complete\"}", + "auto": "{\"_meta\":{\"io.modelcontextprotocol/cacheHint\":{\"ttlMs\":60000,\"cacheScope\":\"private\"}},\"ttlMs\":0,\"cacheScope\":\"private\",\"tools\":[{\"name\":\"alpha\",\"description\":\"first\",\"inputSchema\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"},\"annotations\":{\"readOnlyHint\":false,\"destructiveHint\":true,\"idempotentHint\":false,\"openWorldHint\":true},\"_meta\":{\"is_initial\":false,\"visibility\":\"visible\",\"task_support\":\"forbidden\",\"openai/type\":\"function\",\"openai/function\":{\"name\":\"alpha\",\"description\":\"first\",\"parameters\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"}}}},{\"name\":\"beta\",\"description\":\"second\",\"inputSchema\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"},\"annotations\":{\"readOnlyHint\":false,\"destructiveHint\":true,\"idempotentHint\":false,\"openWorldHint\":true},\"_meta\":{\"is_initial\":false,\"visibility\":\"visible\",\"task_support\":\"forbidden\",\"openai/type\":\"function\",\"openai/function\":{\"name\":\"beta\",\"description\":\"second\",\"parameters\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"}}}}],\"resultType\":\"complete\"}", + "modern": "{\"_meta\":{\"io.modelcontextprotocol/cacheHint\":{\"ttlMs\":60000,\"cacheScope\":\"private\"}},\"ttlMs\":0,\"cacheScope\":\"private\",\"tools\":[{\"name\":\"alpha\",\"description\":\"first\",\"inputSchema\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"},\"annotations\":{\"readOnlyHint\":false,\"destructiveHint\":true,\"idempotentHint\":false,\"openWorldHint\":true},\"_meta\":{\"is_initial\":false,\"visibility\":\"visible\",\"task_support\":\"forbidden\",\"openai/type\":\"function\",\"openai/function\":{\"name\":\"alpha\",\"description\":\"first\",\"parameters\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"}}}},{\"name\":\"beta\",\"description\":\"second\",\"inputSchema\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"},\"annotations\":{\"readOnlyHint\":false,\"destructiveHint\":true,\"idempotentHint\":false,\"openWorldHint\":true},\"_meta\":{\"is_initial\":false,\"visibility\":\"visible\",\"task_support\":\"forbidden\",\"openai/type\":\"function\",\"openai/function\":{\"name\":\"beta\",\"description\":\"second\",\"parameters\":{\"properties\":{\"q\":{\"default\":\"\",\"title\":\"Q\",\"type\":\"string\"}},\"title\":\"In\",\"type\":\"object\",\"$schema\":\"https://json-schema.org/draft/2020-12/schema\"}}}}],\"resultType\":\"complete\"}" +} \ No newline at end of file From 0bae98a044167b457a1d53360bfda5e64e9fe005 Mon Sep 17 00:00:00 2001 From: Shreya Vemana Date: Fri, 23 Oct 2026 20:21:13 +0530 Subject: [PATCH 2/5] feat(transforms): add McpTransform protocol, CatalogTransform and TransformPipeline (#61) --- nitrostack/transforms/__init__.py | 20 +++++ nitrostack/transforms/pipeline.py | 112 ++++++++++++++++++++++++ nitrostack/transforms/transform.py | 128 ++++++++++++++++++++++++++++ nitrostack/transforms/validation.py | 29 +++++++ tests/test_transforms_pipeline.py | 104 ++++++++++++++++++++++ tests/test_transforms_reentrancy.py | 50 +++++++++++ 6 files changed, 443 insertions(+) create mode 100644 nitrostack/transforms/__init__.py create mode 100644 nitrostack/transforms/pipeline.py create mode 100644 nitrostack/transforms/transform.py create mode 100644 nitrostack/transforms/validation.py create mode 100644 tests/test_transforms_pipeline.py create mode 100644 tests/test_transforms_reentrancy.py diff --git a/nitrostack/transforms/__init__.py b/nitrostack/transforms/__init__.py new file mode 100644 index 0000000..667e7d7 --- /dev/null +++ b/nitrostack/transforms/__init__.py @@ -0,0 +1,20 @@ +from nitrostack.transforms.transform import ( + CatalogTransform, + McpTransform, + McpTransformBase, + TransformContext, + active_bypass, + bypassing, +) +from nitrostack.transforms.pipeline import ( + TransformError, + TransformPipeline, + TransformRegistry, +) +from nitrostack.transforms.validation import ToolArgumentError, validate_tool_arguments + +__all__ = [ + "CatalogTransform", "McpTransform", "McpTransformBase", "TransformContext", + "TransformError", "TransformPipeline", "TransformRegistry", + "ToolArgumentError", "validate_tool_arguments", "active_bypass", "bypassing", +] \ No newline at end of file diff --git a/nitrostack/transforms/pipeline.py b/nitrostack/transforms/pipeline.py new file mode 100644 index 0000000..282de09 --- /dev/null +++ b/nitrostack/transforms/pipeline.py @@ -0,0 +1,112 @@ +"""Ordered execution of transforms for tools/list and tools/call resolution.""" +from __future__ import annotations + +import logging +from typing import Any, Awaitable, Callable, List, Optional, Sequence + +from nitrostack.transforms.transform import ( + NextResolver, Tool, TransformContext, active_bypass, maybe_await, +) + +logger = logging.getLogger(__name__) + +BaseResolver = Callable[[str], Any] # returns Tool | None (sync or async) + + +class TransformError(Exception): + """A transform hook raised. The app maps this to a JSON-RPC error.""" + + def __init__(self, transform_name: str, hook: str) -> None: + super().__init__(f"Transform '{transform_name}' failed in {hook}.") + self.transform_name = transform_name + self.hook = hook + + +class TransformRegistry: + """Read-only view passed to ``on_register``.""" + + def __init__(self, list_tools: Callable[[], Sequence[Tool]]) -> None: + self._list_tools = list_tools + + def list_tools(self) -> List[Tool]: + return list(self._list_tools()) + + +class TransformPipeline: + def __init__(self, transforms: Optional[Sequence[Any]] = None) -> None: + self._transforms: List[Any] = list(transforms or []) + names = [t.name for t in self._transforms] + if len(set(names)) != len(names): + raise ValueError(f"Transform names must be unique, got {names}") + + @property + def enabled(self) -> bool: + return bool(self._transforms) + + def register(self, registry: TransformRegistry) -> None: + for t in self._transforms: + try: + t.on_register(registry) + except Exception as exc: + raise TransformError(t.name, "on_register") from exc + + async def transform_tools( + self, tools: Sequence[Tool], ctx: TransformContext + ) -> List[Tool]: + current = list(tools) + if not self._transforms: + return current + skipped = active_bypass() + for t in self._transforms: + if t.name in skipped: + continue + try: + current = list(await maybe_await(t.transform_tools(current, ctx))) + except TransformError: + raise + except Exception as exc: + logger.exception("Transform %s failed in transform_tools", t.name) + raise TransformError(t.name, "transform_tools") from exc + return current + + async def resolve_tool( + self, name: str, base: BaseResolver, ctx: TransformContext + ) -> Optional[Tool]: + """First declared transform is outermost; each chains through ``next``. + + ``base`` must only look the tool up (return None if unknown). Guards and + argument validation happen after resolution, in the app's invoker. + """ + if not self._transforms: + return await maybe_await(base(name)) + + transforms = self._transforms + + async def run(index: int, tool_name: str) -> Optional[Tool]: + skipped = active_bypass() + while index < len(transforms) and transforms[index].name in skipped: + index += 1 + if index >= len(transforms): + return await maybe_await(base(tool_name)) + t = transforms[index] + + async def nxt(n: str) -> Optional[Tool]: + return await run(index + 1, n) + + try: + return await maybe_await(t.resolve_tool(tool_name, nxt, ctx)) + except TransformError: + raise + except Exception as exc: + logger.exception("Transform %s failed in resolve_tool", t.name) + raise TransformError(t.name, "resolve_tool") from exc + + return await run(0, name) + + async def dispose(self) -> None: + """Dispose in reverse order; one failure must not stop the others.""" + for t in reversed(self._transforms): + try: + await maybe_await(t.dispose()) + except Exception: + logger.exception("Transform %s failed in dispose", t.name) \ No newline at end of file diff --git a/nitrostack/transforms/transform.py b/nitrostack/transforms/transform.py new file mode 100644 index 0000000..f3f7a83 --- /dev/null +++ b/nitrostack/transforms/transform.py @@ -0,0 +1,128 @@ +"""Transform protocol, base classes, context and re-entrancy bypass.""" +from __future__ import annotations + +import contextlib +import contextvars +import inspect +from abc import abstractmethod +from dataclasses import dataclass, field +from typing import ( + Any, Awaitable, Callable, Dict, Iterator, List, Optional, Protocol, + Sequence, runtime_checkable, +) + +Tool = Any # mcp.types.Tool in practice; kept loose so transforms may add synthetic tools +NextResolver = Callable[[str], Awaitable[Optional[Tool]]] +ToolInvoker = Callable[[str, Dict[str, Any], "TransformContext"], Awaitable[Any]] + +# Names of transforms to skip for the current call chain. Set while a synthetic +# tool calls a real tool, so the issuing transform is not re-applied. +_BYPASS: contextvars.ContextVar[frozenset] = contextvars.ContextVar( + "nitrostack_transform_bypass", default=frozenset() +) + + +def active_bypass() -> frozenset: + return _BYPASS.get() + + +@contextlib.contextmanager +def bypassing(name: str) -> Iterator[None]: + token = _BYPASS.set(_BYPASS.get() | {name}) + try: + yield + finally: + _BYPASS.reset(token) + + +async def maybe_await(value: Any) -> Any: + if inspect.isawaitable(value): + return await value + return value + + +@dataclass +class TransformContext: + """Per-request context handed to every transform hook.""" + + session_id: Optional[str] = None + execution_context: Any = None # the ExecutionContext of the originating request + metadata: Dict[str, Any] = field(default_factory=dict) + _invoker: Optional[ToolInvoker] = None + + async def call_tool( + self, + name: str, + arguments: Optional[Dict[str, Any]] = None, + *, + caller: Optional[str] = None, + ) -> Any: + """Invoke a real tool from a synthetic tool. + + Goes through the same path as a normal tools/call, so guards and argument + validation run for the real tool. If ``caller`` is given, that transform + is skipped for the duration of the call (no re-entrancy loop). + """ + if self._invoker is None: + raise RuntimeError("TransformContext has no tool invoker attached.") + args = arguments or {} + if caller: + with bypassing(caller): + return await self._invoker(name, args, self) + return await self._invoker(name, args, self) + + +@runtime_checkable +class McpTransform(Protocol): + name: str + + def on_register(self, registry: Any) -> Any: ... + async def dispose(self) -> None: ... + async def transform_tools( + self, tools: Sequence[Tool], ctx: TransformContext + ) -> List[Tool]: ... + async def resolve_tool( + self, name: str, next: NextResolver, ctx: TransformContext + ) -> Optional[Tool]: ... + + +class McpTransformBase: + """Convenience base: every hook defaults to a pass-through. Hooks may be sync or async.""" + + name: str = "transform" + + def on_register(self, registry: Any) -> None: + return None + + async def dispose(self) -> None: + return None + + async def transform_tools(self, tools: Sequence[Tool], ctx: TransformContext) -> List[Tool]: + return list(tools) + + async def resolve_tool( + self, name: str, next: NextResolver, ctx: TransformContext + ) -> Optional[Tool]: + return await next(name) + + def is_bypassed(self) -> bool: + return self.name in _BYPASS.get() + + +class CatalogTransform(McpTransformBase): + """Base for transforms that rewrite the catalog. + + Subclasses implement ``apply_transform``. While this transform's own synthetic + tool is calling a real tool, ``transform_tools`` returns the catalog unchanged. + """ + + @abstractmethod + def apply_transform( + self, tools: Sequence[Tool], ctx: TransformContext + ) -> Any: # list or awaitable of list + ... + + async def transform_tools(self, tools: Sequence[Tool], ctx: TransformContext) -> List[Tool]: + if self.is_bypassed(): + return list(tools) + return list(await maybe_await(self.apply_transform(tools, ctx))) \ No newline at end of file diff --git a/nitrostack/transforms/validation.py b/nitrostack/transforms/validation.py new file mode 100644 index 0000000..621537a --- /dev/null +++ b/nitrostack/transforms/validation.py @@ -0,0 +1,29 @@ +"""One place that validates tool arguments, shared by normal and re-entrant calls.""" +from __future__ import annotations + +from typing import Any, Dict, Optional, Type + +from pydantic import BaseModel, ValidationError + + +class ToolArgumentError(ValueError): + """Arguments did not match the tool's input schema.""" + + def __init__(self, tool_name: str, errors: list) -> None: + super().__init__(f"Invalid arguments for tool '{tool_name}'.") + self.tool_name = tool_name + self.errors = errors + + +def validate_tool_arguments( + tool_name: str, + input_model: Optional[Type[BaseModel]], + arguments: Dict[str, Any], +) -> Any: + """Return the validated input object (or the raw dict if the tool has no model).""" + if input_model is None: + return arguments + try: + return input_model.model_validate(arguments) + except ValidationError as exc: + raise ToolArgumentError(tool_name, exc.errors(include_url=False)) from exc \ No newline at end of file diff --git a/tests/test_transforms_pipeline.py b/tests/test_transforms_pipeline.py new file mode 100644 index 0000000..dd37e31 --- /dev/null +++ b/tests/test_transforms_pipeline.py @@ -0,0 +1,104 @@ +import pytest + +from nitrostack.transforms import ( + CatalogTransform, McpTransformBase, TransformContext, TransformError, TransformPipeline, +) + + +class Tag(McpTransformBase): + def __init__(self, name, log): + self.name = name + self.log = log + + async def transform_tools(self, tools, ctx): + self.log.append(self.name) + return [t + f"+{self.name}" for t in tools] + + async def resolve_tool(self, name, next, ctx): + self.log.append(f"resolve:{self.name}") + found = await next(name) + return None if found is None else f"{found}|{self.name}" + + +@pytest.mark.asyncio +async def test_empty_pipeline_is_identity(): + p = TransformPipeline([]) + assert not p.enabled + tools = ["a", "b"] + assert await p.transform_tools(tools, TransformContext()) == ["a", "b"] + assert await p.resolve_tool("a", lambda n: n, TransformContext()) == "a" + + +@pytest.mark.asyncio +async def test_list_transforms_apply_in_declared_order(): + log = [] + p = TransformPipeline([Tag("one", log), Tag("two", log)]) + out = await p.transform_tools(["x"], TransformContext()) + assert out == ["x+one+two"] + assert log == ["one", "two"] + + +@pytest.mark.asyncio +async def test_resolve_chains_through_next_first_is_outermost(): + log = [] + p = TransformPipeline([Tag("one", log), Tag("two", log)]) + out = await p.resolve_tool("x", lambda n: n, TransformContext()) + assert out == "x|two|one" + assert log == ["resolve:one", "resolve:two"] + + +@pytest.mark.asyncio +async def test_resolve_unknown_tool_returns_none(): + p = TransformPipeline([Tag("one", [])]) + assert await p.resolve_tool("nope", lambda n: None, TransformContext()) is None + + +@pytest.mark.asyncio +async def test_raising_transform_becomes_transform_error(): + class Boom(McpTransformBase): + name = "boom" + + async def transform_tools(self, tools, ctx): + raise RuntimeError("bad") + + p = TransformPipeline([Boom()]) + with pytest.raises(TransformError) as e: + await p.transform_tools(["a"], TransformContext()) + assert e.value.transform_name == "boom" + # the pipeline is still usable afterwards + assert await TransformPipeline([]).transform_tools(["a"], TransformContext()) == ["a"] + + +@pytest.mark.asyncio +async def test_dispose_runs_all_in_reverse_even_if_one_fails(): + order = [] + + class D(McpTransformBase): + def __init__(self, name, fail=False): + self.name, self.fail = name, fail + + async def dispose(self): + order.append(self.name) + if self.fail: + raise RuntimeError("x") + + p = TransformPipeline([D("a"), D("b", fail=True), D("c")]) + await p.dispose() + assert order == ["c", "b", "a"] + + +def test_duplicate_names_rejected(): + with pytest.raises(ValueError): + TransformPipeline([Tag("same", []), Tag("same", [])]) + + +@pytest.mark.asyncio +async def test_catalog_transform_applies_and_supports_sync_apply(): + class Upper(CatalogTransform): + name = "upper" + + def apply_transform(self, tools, ctx): + return [t.upper() for t in tools] + + p = TransformPipeline([Upper()]) + assert await p.transform_tools(["a"], TransformContext()) == ["A"] \ No newline at end of file diff --git a/tests/test_transforms_reentrancy.py b/tests/test_transforms_reentrancy.py new file mode 100644 index 0000000..ee549c1 --- /dev/null +++ b/tests/test_transforms_reentrancy.py @@ -0,0 +1,50 @@ +import pytest + +from nitrostack.transforms import ( + CatalogTransform, TransformContext, TransformPipeline, bypassing, +) + + +class Hide(CatalogTransform): + name = "hide" + + def __init__(self): + self.applied = 0 + + def apply_transform(self, tools, ctx): + self.applied += 1 + return [t for t in tools if t != "secret"] + + +@pytest.mark.asyncio +async def test_transform_is_skipped_while_its_own_call_is_running(): + h = Hide() + p = TransformPipeline([h]) + ctx = TransformContext() + + assert await p.transform_tools(["a", "secret"], ctx) == ["a"] + assert h.applied == 1 + + with bypassing("hide"): + assert await p.transform_tools(["a", "secret"], ctx) == ["a", "secret"] + assert h.applied == 1 # not re-run + + assert await p.transform_tools(["a", "secret"], ctx) == ["a"] # bypass was released + assert h.applied == 2 + + +@pytest.mark.asyncio +async def test_call_tool_with_caller_bypasses_that_transform_only(): + seen = {} + + async def invoker(name, args, ctx): + from nitrostack.transforms import active_bypass + seen["bypass"] = set(active_bypass()) + return f"called {name}" + + ctx = TransformContext(_invoker=invoker) + assert await ctx.call_tool("real", {}, caller="hide") == "called real" + assert seen["bypass"] == {"hide"} + + await ctx.call_tool("real", {}) + assert seen["bypass"] == set() \ No newline at end of file From 1211d6bb7fca873646859ee744b50bb7d8afdbd4 Mon Sep 17 00:00:00 2001 From: Shreya Vemana Date: Fri, 23 Oct 2026 20:21:52 +0530 Subject: [PATCH 3/5] feat(core): apply transforms to tools/list and tools/call, share argument validation (#61) --- nitrostack/core/app.py | 55 +++++++++++--- nitrostack/testing/__init__.py | 25 +++++-- tests/test_transformsapp.py | 127 +++++++++++++++++++++++++++++++++ 3 files changed, 193 insertions(+), 14 deletions(-) create mode 100644 tests/test_transformsapp.py diff --git a/nitrostack/core/app.py b/nitrostack/core/app.py index 9d9b069..222ddc5 100644 --- a/nitrostack/core/app.py +++ b/nitrostack/core/app.py @@ -25,6 +25,12 @@ UnsubscribeRequestParams, ) from pydantic import BaseModel, create_model +from nitrostack.transforms import ( + TransformContext, + TransformError, + TransformPipeline, + TransformRegistry, +) from nitrostack.core.context import ExecutionContext, TaskContext from nitrostack.core.decorators import ToolConfig, ResourceConfig, PromptConfig, widget_resource_uri @@ -206,17 +212,20 @@ class ServerConfig: description: Optional[str] = None -def mcp_app(module: Type, server: ServerConfig): - """ - Decorator to declare the main application class. - Specifies the root AppModule and ServerConfig. - """ +def mcp_app( + module: Type, + server: ServerConfig, + transforms: Optional[List[Any]] = None, +): + """Declare an MCP application and its optional tool-list transforms.""" + def decorator(cls: Type): cls._mcp_app_module = module cls._mcp_app_server = server + cls._mcp_transforms = list(transforms or []) return cls - return decorator + return decorator def get_pydantic_model(schema: Any) -> Type[BaseModel]: """Helper to resolve or construct a Pydantic model for input validation.""" @@ -601,6 +610,15 @@ def __init__(self, app_class: Type): self._bootstrap() + # Initialize the transform pipeline + # Initialize transforms before registering handlers. Handlers use + # this pipeline when tools/list is requested. + self._transforms = TransformPipeline( + getattr(self.app_class, "_mcp_transforms", None) or [] + ) + + self._bootstrap() + def _bootstrap(self) -> None: # 1. Construct the low-level server directly. self.mcp_server = NitroStackMcpServer( @@ -880,9 +898,30 @@ def _list_endpoint_cache_meta(self) -> Dict[str, Any]: return build_list_endpoint_cache_hint_meta() def _setup_handlers(self, server: NitroStackMcpServer) -> None: - async def _list_tools(_ctx: ServerRequestContext, _params: Optional[PaginatedRequestParams]): + async def _list_tools( + _ctx: ServerRequestContext, + _params: Optional[PaginatedRequestParams], + ): + tools = [ + self._build_tool_definition(entry) + for entry in self._tools.values() + ] + + if self._transforms.enabled: + try: + context = TransformContext() + tools = await self._transforms.transform_tools( + tools, + context, + ) + except TransformError as exc: + raise MCPError( + types.INTERNAL_ERROR, + str(exc), + ) from exc + return types.ListToolsResult( - tools=[self._build_tool_definition(entry) for entry in self._tools.values()], + tools=tools, _meta=self._list_endpoint_cache_meta(), ) diff --git a/nitrostack/testing/__init__.py b/nitrostack/testing/__init__.py index 0db12b6..25f0efb 100644 --- a/nitrostack/testing/__init__.py +++ b/nitrostack/testing/__init__.py @@ -14,19 +14,32 @@ class NitroTestingModule: Dispatches through the owned low-level `mcp.server.lowlevel.Server` `request_handlers` (the same dict the real stdio/HTTP transports use). """ + def __init__(self, app: McpApplication): + self.app = app @classmethod - async def create(cls, app_module: Type) -> "NitroTestingModule": - # Construct a dummy App class decorated with @mcp_app - @mcp_app(module=app_module, server=ServerConfig(name="test-server")) + async def create( + cls, + app_module: Type, + *, + protocol_era: Optional[str] = None, + transforms: Optional[list] = None, + ) -> "NitroTestingModule": + server_config = ServerConfig( + name="test-server", + protocol_era=protocol_era, + ) + + @mcp_app( + module=app_module, + server=server_config, + transforms=transforms, + ) class TestApp: pass app = await McpApplicationFactory.create(TestApp) return cls(app) - def __init__(self, app: McpApplication): - self.app = app - @staticmethod def _extract_text(text_val: Optional[str]) -> Any: if text_val is None: diff --git a/tests/test_transformsapp.py b/tests/test_transformsapp.py new file mode 100644 index 0000000..779cd7b --- /dev/null +++ b/tests/test_transformsapp.py @@ -0,0 +1,127 @@ + +import json +import pathlib + +import mcp.types as types +import pytest + +from nitrostack.testing import NitroTestingModule +from nitrostack.transforms import CatalogTransform, McpTransformBase +from tests._baseline_catalog import ERAS, In, M + + +BASELINE = ( + pathlib.Path(__file__).parent + / "fixtures" + / "tools_list_baseline.json" +) + + +async def _list(h): + srv = h.app.mcp_server + response = ( + await srv.request_handlers[types.ListToolsRequest]( + types.ListToolsRequest(method="tools/list") + ) + ).root + return response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("era", ERAS) +async def test_no_transforms_output_is_byte_identical_to_baseline( + era, monkeypatch +): + monkeypatch.setenv("NITRO_MCP_PROTOCOL_VERSION", era) + + h = await NitroTestingModule.create(M) + response = await _list(h) + + expected = json.loads( + BASELINE.read_text(encoding="utf-8") + )[era] + + actual = response.model_dump_json( + by_alias=True, + exclude_none=True, + ) + + assert actual == expected + + +class DropBeta(CatalogTransform): + name = "drop_beta" + + def apply_transform(self, tools, ctx): + return [tool for tool in tools if tool.name != "beta"] + + +class Suffix(CatalogTransform): + name = "suffix" + + def apply_transform(self, tools, ctx): + return [ + tool.model_copy( + update={ + "description": (tool.description or "") + "!" + } + ) + for tool in tools + ] + + +@pytest.mark.asyncio +async def test_two_transforms_compose_in_declared_order_for_list(): + h = await NitroTestingModule.create( + M, + transforms=[DropBeta(), Suffix()], + ) + + response = await _list(h) + + assert [tool.name for tool in response.tools] == ["alpha"] + assert response.tools[0].description == "first!" + + +class Boom(McpTransformBase): + name = "boom" + + async def transform_tools(self, tools, ctx): + raise RuntimeError("bad transform") + + +@pytest.mark.asyncio +async def test_failing_transform_returns_error_and_server_keeps_serving(): + h = await NitroTestingModule.create( + M, + transforms=[Boom()], + ) + + with pytest.raises(Exception): + await _list(h) + + # A transform failure must not stop the server from handling tool calls. + result = await h.call_tool("alpha", {"input": {"q": "x"}}) + assert result is not None + + +@pytest.mark.asyncio +async def test_dispose_runs_on_shutdown(): + disposed = [] + + class D(McpTransformBase): + name = "d" + + async def dispose(self): + disposed.append(True) + + h = await NitroTestingModule.create( + M, + transforms=[D()], + ) + + # Exercise disposal directly because this harness may not expose + # a public application-shutdown method. + await h.app._transforms.dispose() + + assert disposed == [True] From da07f89b7926be15c6ec6f021e2c1348a9c168ca Mon Sep 17 00:00:00 2001 From: Shreya Vemana Date: Sat, 24 Oct 2026 16:28:41 +0530 Subject: [PATCH 4/5] test: add baseline catalog, package marker and fixture for transform parity test (#61) --- tests/__init__.py | 0 tests/_baseline_catalog.py | 26 ++++++++++++++++++++++++++ 2 files changed, 26 insertions(+) create mode 100644 tests/__init__.py create mode 100644 tests/_baseline_catalog.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/_baseline_catalog.py b/tests/_baseline_catalog.py new file mode 100644 index 0000000..3092d9f --- /dev/null +++ b/tests/_baseline_catalog.py @@ -0,0 +1,26 @@ +"""Catalog used by both the baseline capture script and the parity test.""" +from pydantic import BaseModel + +from nitrostack import ExecutionContext, injectable, module, tool + +ERAS = ("legacy", "auto", "modern") + + +class In(BaseModel): + q: str = "" + + +@injectable() +class C: + @tool(name="alpha", description="first", input_schema=In) + async def alpha(self, input: In, context: ExecutionContext) -> str: + return "a" + + @tool(name="beta", description="second", input_schema=In) + async def beta(self, input: In, context: ExecutionContext) -> str: + return "b" + + +@module(name="m", controllers=[C]) +class M: + pass \ No newline at end of file From 0468fe5c507106437cce1f0f48daded064e3062a Mon Sep 17 00:00:00 2001 From: Shreya Vemana Date: Sat, 24 Oct 2026 16:37:03 +0530 Subject: [PATCH 5/5] build: add pytest-asyncio to dev extras for transform tests (#61) --- pyproject.toml | 3 ++- uv.lock | 53 +++++++++++++++++++++++++++++++++++++++++++------- 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5b4a3b3..cfb3513 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,8 +21,9 @@ dependencies = [ [project.optional-dependencies] dev = [ - "pytest>=7.4.0", "pytest-cov>=4.1.0", + "pytest>=8", + "pytest-asyncio>=0.23", "pytest-timeout>=2.3.0", "httpx>=0.27.0", "tox>=4.0.0", diff --git a/uv.lock b/uv.lock index ed009ff..a000f51 100644 --- a/uv.lock +++ b/uv.lock @@ -1,10 +1,10 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", - "(python_full_version >= '3.11' and python_full_version < '3.14' and sys_platform != 'emscripten') or (python_full_version == '3.11.*' and sys_platform == 'emscripten')", + "python_full_version == '3.11.*' or (python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten')", "python_full_version < '3.11'", ] @@ -40,6 +40,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] +[[package]] +name = "backports-asyncio-runner" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8e/ff/70dca7d7cb1cbc0edb2c6cc0c38b65cba36cccc491eca64cabd5fe7f8670/backports_asyncio_runner-1.2.0.tar.gz", hash = "sha256:a5aa7b2b7d8f8bfcaa2b57313f70792df84e32a2a746f585213373f900b42162", size = 69893, upload-time = "2025-07-02T02:27:15.685Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/59/76ab57e3fe74484f48a53f8e337171b4a2349e506eabe136d7e01d059086/backports_asyncio_runner-1.2.0-py3-none-any.whl", hash = "sha256:0da0a936a8aeb554eccb426dc55af3ba63bcdc69fa1a600b5bb305413a4477b5", size = 12313, upload-time = "2025-07-02T02:27:14.263Z" }, +] + [[package]] name = "cachetools" version = "7.1.8" @@ -390,7 +399,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -433,8 +442,8 @@ name = "httpcore2" version = "2.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "h11", marker = "python_full_version < '3.12' or python_full_version >= '3.14' or sys_platform != 'emscripten'" }, - { name = "truststore", marker = "python_full_version < '3.12' or python_full_version >= '3.14' or sys_platform != 'emscripten'" }, + { name = "h11" }, + { name = "truststore" }, ] sdist = { url = "https://files.pythonhosted.org/packages/be/ad/f4f0e57345f1870f3e8cb624e058d7eca6e5a27d33bcc3311d9b618734cd/httpcore2-2.12.0.tar.gz", hash = "sha256:9293522bba0aa7c4c8e9e3f040c16575bd8868e155a77fa30c7a9085a5eae648", size = 67548, upload-time = "2026-08-18T13:22:08.211Z" } wheels = [ @@ -585,7 +594,9 @@ dependencies = [ dev = [ { name = "httpx" }, { name = "pytest" }, + { name = "pytest-asyncio" }, { name = "pytest-cov" }, + { name = "pytest-timeout" }, { name = "tox" }, ] @@ -596,8 +607,10 @@ requires-dist = [ { name = "mcp", specifier = ">=2,<3" }, { name = "packaging", specifier = ">=23.0.0" }, { name = "pydantic", specifier = ">=2.0.0" }, - { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.4.0" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.1.0" }, + { name = "pytest-timeout", marker = "extra == 'dev'", specifier = ">=2.3.0" }, { name = "starlette", specifier = ">=0.30.0" }, { name = "tomli", marker = "python_full_version < '3.11'", specifier = ">=2.0.0" }, { name = "tox", marker = "extra == 'dev'", specifier = ">=4.0.0" }, @@ -842,6 +855,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, ] +[[package]] +name = "pytest-asyncio" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "backports-asyncio-runner", marker = "python_full_version < '3.11'" }, + { name = "pytest" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/7c/d36d04db312ecf4298932ef77e6e4a9e8ad017906e24e34f0b0c361a2473/pytest_asyncio-1.4.0.tar.gz", hash = "sha256:c6c0d2259945122819f171a32ecea2c349ead889ee28176caaf492143424be42", size = 58514, upload-time = "2026-05-26T09:56:04.083Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/03/e2/08a497ef684b88559c9cc5f4ad53a37e7b99e727094a86d6ea32536d5d3c/pytest_asyncio-1.4.0-py3-none-any.whl", hash = "sha256:933ca923a23075a87fb7070c0ec272a6848489824d887c85c812670932835aa1", size = 16930, upload-time = "2026-05-26T09:56:02.576Z" }, +] + [[package]] name = "pytest-cov" version = "7.1.0" @@ -856,6 +883,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" }, ] +[[package]] +name = "pytest-timeout" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/82/4c9ecabab13363e72d880f2fb504c5f750433b2b6f16e99f4ec21ada284c/pytest_timeout-2.4.0.tar.gz", hash = "sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a", size = 17973, upload-time = "2025-05-05T19:44:34.99Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/b6/3127540ecdf1464a00e5a01ee60a1b09175f6913f0644ac748494d9c4b21/pytest_timeout-2.4.0-py3-none-any.whl", hash = "sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2", size = 14382, upload-time = "2025-05-05T19:44:33.502Z" }, +] + [[package]] name = "python-discovery" version = "1.6.0" @@ -1049,7 +1088,7 @@ source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", - "(python_full_version >= '3.11' and python_full_version < '3.14' and sys_platform != 'emscripten') or (python_full_version == '3.11.*' and sys_platform == 'emscripten')", + "python_full_version == '3.11.*' or (python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten')", ] sdist = { url = "https://files.pythonhosted.org/packages/aa/2a/9618a122aeb2a169a28b03889a2995fe297588964333d4a7d67bdf46e147/rpds_py-2026.6.3.tar.gz", hash = "sha256:1cebd1337c242e4ec2293e541f712b2da849b29f48f0c293684b71c0632625d4", size = 64051, upload-time = "2026-06-30T07:17:53.009Z" } wheels = [