[transforms] Progressive tool discovery (#62)

nitrostack/core/app.py

@@ -10,7 +10,9 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Mapping, Optional, Pattern, Set, Tuple, Type
-
+from nitrostack.transforms import (
+    SyntheticTool, ToolArgumentError, TransformContext, validate_tool_arguments,
+)
 import mcp.types as types
 from mcp import MCPError
 from mcp.server.context import ServerRequestContext
@@ -25,6 +27,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 +214,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."""
@@ -599,8 +610,19 @@ def __init__(self, app_class: Type):
         self.task_manager = TaskManager()
         self._in_flight = InFlightRegistry()
 
+        # Pipeline first: handlers consult it lazily when requests arrive.
+        self._transforms = TransformPipeline(
+            getattr(self.app_class, "_mcp_transforms", None) or []
+        )
+
         self._bootstrap()
 
+        # Tools are registered now, so transforms can see the catalog.
+        self._transforms.register(TransformRegistry(self._base_tool_list))
+
+    def _base_tool_list(self) -> List[types.Tool]:
+        return [self._build_tool_definition(entry) for entry in self._tools.values()]
+
     def _bootstrap(self) -> None:
         # 1. Construct the low-level server directly.
         self.mcp_server = NitroStackMcpServer(
@@ -880,9 +902,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(),
             )
 
@@ -1166,6 +1209,12 @@ def _pipeline_stages(self, method: Callable) -> Tuple[List[Type], List[Type], Li
             getattr(method, "_mcp_pipes", []),
             getattr(method, "_mcp_filters", []),
         )
+    def _base_lookup(self, name: str):
+        return self._tools.get(name)
+
+    async def _invoke_for_transform(self, name: str, arguments: dict, tctx: TransformContext):
+        # Re-enter the normal call path so guards and validation run for the real tool.
+        return await self._call_tool(name, arguments, _tctx=tctx)
 
     def _to_call_tool_result(
         self,
@@ -1234,6 +1283,52 @@ def _to_call_tool_result(
             isError=False,
         )
 
+    async def _call_tool(
+        self,
+        name: str,
+        arguments: Dict[str, Any],
+        task: Any = None,
+        _tctx: Optional[TransformContext] = None,
+    ):
+        """Resolve through the transform pipeline, then run the real tool path."""
+        if not self._transforms.enabled:
+            return await self._call_real_tool(name, arguments, task=task)
+
+        tctx = _tctx or TransformContext(_invoker=self._invoke_for_transform)
+        try:
+            resolved = await self._transforms.resolve_tool(name, self._base_lookup, tctx)
+        except Exception as exc:
+            logger.exception("Transform failed resolving tool %s", name)
+            return types.CallToolResult(
+                content=[types.TextContent(type="text", text=str(exc))], isError=True
+            )
+
+        if isinstance(resolved, SyntheticTool):
+            payload = arguments or {}
+            inner = payload.get("input")
+            if (
+                isinstance(inner, dict)
+                and set(payload) <= {"input"}
+                and "input" not in resolved.input_model.model_fields
+            ):
+                payload = inner  # accept the older {"input": {...}} wrap
+            try:
+                validated = validate_tool_arguments(resolved.name, resolved.input_model, payload)
+                result = await resolved.handler(validated, tctx)
+            except ToolArgumentError as exc:
+                return types.CallToolResult(
+                    content=[types.TextContent(type="text", text=f"{exc}: {exc.errors}")],
+                    isError=True,
+                )
+            except Exception as exc:
+                logger.exception("Synthetic tool %s failed", name)
+                return types.CallToolResult(
+                    content=[types.TextContent(type="text", text=str(exc))], isError=True
+                )
+            return self._to_call_tool_result(result)
+
+        return await self._call_real_tool(name, arguments, task=task)
+
     @staticmethod
     def _call_tool_result_meta(result: types.CallToolResult) -> Optional[Dict[str, Any]]:
         """Read caller ``_meta`` from the field or pydantic extra (alias vs extra)."""
@@ -1281,7 +1376,7 @@ def _widget_result_content(self, structured: Dict[str, Any], component: Optional
         )
         return content
 
-    async def _call_tool(self, name: str, arguments: Dict[str, Any], task: Any = None):
+    async def _call_real_tool(self, name: str, arguments: Dict[str, Any], task: Any = None):
         entry = self._tools.get(name)
         if entry is None:
             return types.CallToolResult(

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,21 @@
+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
+from nitrostack.transforms.synthetic import SyntheticTool
+
+__all__ = [
+    "CatalogTransform", "McpTransform", "McpTransformBase", "TransformContext",
+    "TransformError", "TransformPipeline", "TransformRegistry",
+    "ToolArgumentError", "validate_tool_arguments", "active_bypass", "bypassing", "SyntheticTool",
+]

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/search/__init__.py

@@ -0,0 +1,12 @@
+from nitrostack.transforms.search.base import BaseSearchTransform, is_browse_query
+from nitrostack.transforms.search.models import (
+    BROWSE_MAX_RESULTS, DEFAULT_SEARCH_LIMIT, MAX_SEARCH_RESULTS,
+    CallToolInput, RankOutcome, SearchToolsInput,
+)
+from nitrostack.transforms.search.regex import RegexSearchTransform
+
+__all__ = [
+    "BaseSearchTransform", "RegexSearchTransform", "RankOutcome",
+    "SearchToolsInput", "CallToolInput", "is_browse_query",
+    "MAX_SEARCH_RESULTS", "DEFAULT_SEARCH_LIMIT", "BROWSE_MAX_RESULTS",
+]