feat: add per-session tool visibility

README.md

@@ -289,6 +289,15 @@ server = ServerConfig(name="my-server", transport_type="http", max_sessions=100,
 
 `ServerConfig.protocol_era` is used only when `MCP_STATELESS` and `NITRO_MCP_PROTOCOL_VERSION` are both unset. Unknown tokens become `auto`, same as an unknown env value.
 
+### Per-session tool visibility
+
+Tools can declare `visibility="hidden"` and `default_visible=False`, then handlers can
+call `context.enable_tools(["admin_tool"])` or `context.disable_tools([...])`. Overrides
+are scoped by transport session ID and verified token subject, with a bounded TTL/LRU
+store. Dynamic visibility requires a sessionful transport. The default `auto` era is
+sessionless, so changes warn and do nothing there; the SDK does not fall back to
+subject-only state because it would share overrides across that subject's sessions.
+
 ---
 
 ## Developing & Testing

nitrostack/core/app.py

@@ -35,6 +35,7 @@
 )
 
 from nitrostack.core.context import ExecutionContext, TaskContext
+from nitrostack.core.visibility import VisibilityTransform
 from nitrostack.core.decorators import ToolConfig, ResourceConfig, PromptConfig, widget_resource_uri
 from nitrostack.core.app_mode import get_app_mode, get_widget_mime_type, is_mcp_app_mode, is_openai_mode
 from nitrostack.core.di import DIContainer
@@ -609,6 +610,7 @@ def __init__(self, app_class: Type):
         self._initial_tools: List[Tuple[Any, Callable, ToolConfig]] = []
         self.task_manager = TaskManager()
         self._in_flight = InFlightRegistry()
+        self.visibility_transform = VisibilityTransform()
 
         # Pipeline first: handlers consult it lazily when requests arrive.
         self._transforms = TransformPipeline(
@@ -906,9 +908,11 @@ async def _list_tools(
             _ctx: ServerRequestContext,
             _params: Optional[PaginatedRequestParams],
         ):
+            session_key = self._visibility_session_key(_ctx)
             tools = [
                 self._build_tool_definition(entry)
                 for entry in self._tools.values()
+                if self.visibility_transform.is_visible(entry, session_key)
             ]
 
             if self._transforms.enabled:
@@ -932,6 +936,8 @@ async def _list_tools(
         async def _call_tool(ctx: ServerRequestContext, params: CallToolRequestParams):
             token = request_ctx.set(ctx)
             try:
+                if not self._tool_visible(params.name, ctx):
+                    return self._unknown_tool_result(params.name)
                 result = await self._call_tool(
                     params.name, params.arguments or {}, task=params.task
                 )
@@ -1015,14 +1021,22 @@ async def _discover(ctx: ServerRequestContext, _params: Optional[RequestParams])
         server.add_request_handler("prompts/get", GetPromptRequestParams, _get_prompt)
 
         async def handle_list_tools(req=None):
-            return ServerResult(await _list_tools(None, None))
+            ctx = req if req is not None and hasattr(req, "session") else request_ctx.get(None)
+            return ServerResult(await _list_tools(ctx, None))
 
         async def handle_call_tool(req=None):
             params = getattr(req, "params", None) if req is not None else None
             name = getattr(params, "name", "") if params is not None else ""
             arguments = (getattr(params, "arguments", None) or {}) if params is not None else {}
             task = getattr(params, "task", None) if params is not None else None
-            return ServerResult(await self._call_tool(name, arguments, task=task))
+            ctx = req if req is not None and hasattr(req, "session") else request_ctx.get(None)
+            if not self._tool_visible(name, ctx):
+                return ServerResult(self._unknown_tool_result(name))
+            token = request_ctx.set(ctx)
+            try:
+                return ServerResult(await self._call_tool(name, arguments, task=task))
+            finally:
+                request_ctx.reset(token)
 
         async def handle_list_resources(req=None):
             return ServerResult(await _list_resources(None, None))
@@ -1216,6 +1230,62 @@ async def _invoke_for_transform(self, name: str, arguments: dict, tctx: Transfor
         # Re-enter the normal call path so guards and validation run for the real tool.
         return await self._call_tool(name, arguments, _tctx=tctx)
 
+    @staticmethod
+    def _visibility_session_key(request_context: Any) -> tuple[str, str] | None:
+        session = getattr(request_context, "session", None)
+        session_id = getattr(session, "id", None) or getattr(session, "session_id", None)
+        if session_id is None:
+            return None
+        auth = auth_context_from_request(request_context)
+        return VisibilityTransform.session_key(
+            str(session_id), auth.subject if auth is not None else None
+        )
+
+    def _tool_visible(self, name: str, request_context: Any) -> bool:
+        entry = self._tools.get(name)
+        if entry is None:
+            # Transform-defined synthetic tools are not in the app registry.
+            # Let the transform resolver decide whether such names exist.
+            return True
+        return self.visibility_transform.is_visible(
+            entry, self._visibility_session_key(request_context)
+        )
+
+    def _configure_visibility_context(
+        self, context: ExecutionContext, request_context: Any
+    ) -> None:
+        key = self._visibility_session_key(request_context)
+        context._tool_visibility_controller = self.visibility_transform
+        context._visibility_session_key = key
+        context._visibility_app = self
+        context._visibility_supported = key is not None
+        context._visibility_change_callback = self._publish_tools_changed
+
+    @staticmethod
+    def _unknown_tool_result(name: str) -> types.CallToolResult:
+        return types.CallToolResult(
+            content=[types.TextContent(type="text", text=f"Tool '{name}' not found")],
+            isError=True,
+        )
+
+    def _publish_tools_changed(self, context: ExecutionContext) -> None:
+        session = getattr(request_ctx.get(None), "session", None)
+        sender = getattr(session, "send_notification", None) if session is not None else None
+        if sender is None:
+            logger.debug("Transport does not support tools/list_changed notification")
+            return
+
+        async def send_notification() -> None:
+            try:
+                await sender(types.ToolListChangedNotification())
+            except Exception:
+                logger.debug("Unable to send session tools/list_changed notification", exc_info=True)
+
+        try:
+            asyncio.create_task(send_notification())
+        except RuntimeError:
+            logger.debug("No running event loop for tools/list_changed notification")
+
     def _to_call_tool_result(
         self,
         result: Any,
@@ -1379,10 +1449,9 @@ def _widget_result_content(self, structured: Dict[str, Any], component: Optional
     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(
-                content=[types.TextContent(type="text", text=f"Tool '{name}' not found")],
-                isError=True,
-            )
+            return self._unknown_tool_result(name)
+        if not self._tool_visible(name, request_ctx.get(None)):
+            return self._unknown_tool_result(name)
 
         cfg = entry.config
         tool_arguments, input_responses, request_state = split_mrtr_from_arguments(
@@ -1466,6 +1535,7 @@ async def background_execution():
                     trace=trace,
                 )
                 _apply_request_envelope(task_ctx, rc)
+                self._configure_visibility_context(task_ctx, rc)
                 task_ctx.task = TaskContext(
                     task_id,
                     self.task_manager,
@@ -1520,6 +1590,7 @@ async def background_execution():
             trace=trace,
         )
         _apply_request_envelope(ctx, rc)
+        self._configure_visibility_context(ctx, rc)
         ticket = self._in_flight.register(correlation_id, jsonrpc_id=jsonrpc_id)
         try:
             result = await run_pipeline(

nitrostack/core/context.py

@@ -197,6 +197,10 @@ class ExecutionContext:
     rpc_meta: Optional["RequestMeta"] = None
     mcp_headers: Dict[str, str] = field(default_factory=dict)
     mcp_param_headers: Dict[str, str] = field(default_factory=dict)
+    _tool_visibility_controller: Any = field(default=None, repr=False, compare=False)
+    _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)
 
     def __post_init__(self) -> None:
         if not self.correlation_id:
@@ -208,3 +212,19 @@ def user(self) -> Optional[str]:
         if self.auth is None:
             return None
         return self.auth.subject
+
+    def enable_tools(self, names: List[str]) -> None:
+        """Enable registered tools for this verified transport session."""
+        self._change_tool_visibility(names, visible=True)
+
+    def disable_tools(self, names: List[str]) -> None:
+        """Disable registered tools for this verified transport session."""
+        self._change_tool_visibility(names, visible=False)
+
+    def _change_tool_visibility(self, names: List[str], *, visible: bool) -> None:
+        if not self._visibility_supported or self._tool_visibility_controller is None:
+            self.logger.warn("Tool visibility changes require a sessionful transport session")
+            return
+        changed = self._tool_visibility_controller.change(self, names, visible=visible)
+        if changed and self._visibility_change_callback is not None:
+            self._visibility_change_callback(self)

nitrostack/core/decorators.py

@@ -49,6 +49,7 @@ class ToolConfig:
     annotations: ToolAnnotations = field(default_factory=ToolAnnotations)
     task_support: Literal["forbidden", "optional", "required"] = "forbidden"
     visibility: Literal["visible", "hidden"] = "visible"
+    default_visible: Optional[bool] = None
     examples: Optional[ToolExamples] = None
     invocation: Optional[ToolInvocation] = None
     metadata: Dict[str, Any] = field(default_factory=dict)
@@ -98,6 +99,7 @@ def tool(
     annotations: Optional[ToolAnnotations] = None,
     task_support: Literal["forbidden", "optional", "required"] = "forbidden",
     visibility: Literal["visible", "hidden"] = "visible",
+    default_visible: Optional[bool] = None,
     examples: Optional[ToolExamples] = None,
     invocation: Optional[ToolInvocation] = None,
     metadata: Optional[Dict[str, Any]] = None,
@@ -121,6 +123,7 @@ def decorator(func: Callable):
             annotations=annotations,
             task_support=task_support,
             visibility=visibility,
+            default_visible=(visibility == "visible") if default_visible is None else default_visible,
             examples=examples,
             invocation=invocation,
             metadata=metadata,

nitrostack/core/visibility.py

@@ -0,0 +1,111 @@
+"""Per-session tool visibility with bounded TTL/LRU state."""
+
+from __future__ import annotations
+
+import time
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from threading import RLock
+from typing import Callable, Iterable
+
+
+@dataclass
+class _SessionVisibility:
+    expires_at: float
+    overrides: dict[str, bool] = field(default_factory=dict)
+
+
+class SessionVisibilityStore:
+    """Thread-safe TTL/LRU storage for tool visibility overrides."""
+
+    def __init__(
+        self,
+        *,
+        ttl_seconds: float = 3600,
+        max_entries: int = 10_000,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if ttl_seconds <= 0 or max_entries <= 0:
+            raise ValueError("ttl_seconds and max_entries must be positive")
+        self.ttl_seconds = ttl_seconds
+        self.max_entries = max_entries
+        self._clock = clock
+        self._entries: OrderedDict[tuple[str, str], _SessionVisibility] = OrderedDict()
+        self._lock = RLock()
+
+    def _expire(self, now: float) -> None:
+        for key, entry in list(self._entries.items()):
+            if entry.expires_at <= now:
+                del self._entries[key]
+
+    def _get(self, key: tuple[str, str], now: float) -> _SessionVisibility:
+        self._expire(now)
+        entry = self._entries.get(key)
+        if entry is None:
+            entry = _SessionVisibility(now + self.ttl_seconds)
+            self._entries[key] = entry
+        else:
+            self._entries.move_to_end(key)
+            entry.expires_at = now + self.ttl_seconds
+        while len(self._entries) > self.max_entries:
+            self._entries.popitem(last=False)
+        return entry
+
+    def is_visible(self, key: tuple[str, str], name: str, default: bool) -> bool:
+        with self._lock:
+            return self._get(key, self._clock()).overrides.get(name, default)
+
+    def update(self, key: tuple[str, str], names: Iterable[str], visible: bool) -> bool:
+        changed = False
+        with self._lock:
+            entry = self._get(key, self._clock())
+            for name in names:
+                if entry.overrides.get(name) != visible:
+                    entry.overrides[name] = visible
+                    changed = True
+        return changed
+
+
+class VisibilityTransform:
+    """Apply decorator defaults and per-session tool visibility overrides."""
+
+    def __init__(self, store: SessionVisibilityStore | None = None) -> None:
+        self.store = store or SessionVisibilityStore()
+
+    @staticmethod
+    def session_key(
+        transport_session_id: str | None,
+        verified_subject: str | None,
+    ) -> tuple[str, str] | None:
+        if not transport_session_id:
+            return None
+        return transport_session_id, verified_subject or ""
+
+    def is_visible(self, entry: object, key: tuple[str, str] | None) -> bool:
+        config = getattr(entry, "config")
+        default = config.default_visible
+        if default is None:
+            default = config.visibility == "visible"
+        if key is None:
+            return default
+        return self.store.is_visible(key, config.name, default)
+
+    def update(self, key: tuple[str, str], names: Iterable[str], visible: bool) -> bool:
+        """Update known tool names for a specific visibility key."""
+        return self.store.update(key, names, visible)
+
+    def change(self, context: object, names: Iterable[str], *, visible: bool) -> bool:
+        key = getattr(context, "_visibility_session_key", None)
+        if key is None:
+            getattr(context, "logger").warn(
+                "Tool visibility changes require a sessionful transport session"
+            )
+            return False
+        app = getattr(context, "_visibility_app")
+        valid: list[str] = []
+        for name in names:
+            if name not in app._tools:
+                getattr(context, "logger").warn(f"Unknown tool name in visibility change: {name}")
+                continue
+            valid.append(name)
+        return self.store.update(key, valid, visible) if valid else False