[transforms] Add a transform pipeline for the tool catalog : Issue #61

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(),
             )
 

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:

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",
+]

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)

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)))