diff --git a/README.md b/README.md index bf299b1..02fb724 100644 --- a/README.md +++ b/README.md @@ -315,6 +315,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 diff --git a/nitrostack/core/app.py b/nitrostack/core/app.py index 255d504..6a6afd3 100644 --- a/nitrostack/core/app.py +++ b/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 @@ -611,6 +612,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( @@ -908,9 +910,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: @@ -934,20 +938,26 @@ async def _list_tools( ) async def _call_tool(ctx: ServerRequestContext, params: CallToolRequestParams): - result = await self._call_tool( - params.name, params.arguments or {}, task=params.task - ) - if isinstance(result, types.CreateTaskResult): - task = result.task - return types.CallToolResult( - content=[types.TextContent(type="text", text=task.task_id)], - structuredContent={ - "resultType": "task", - "task": task.model_dump(by_alias=True), - }, - isError=False, + 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 ) - return result + if isinstance(result, types.CreateTaskResult): + task = result.task + return types.CallToolResult( + content=[types.TextContent(type="text", text=task.task_id)], + structuredContent={ + "resultType": "task", + "task": task.model_dump(by_alias=True), + }, + isError=False, + ) + return result + finally: + request_ctx.reset(token) async def _list_resources(_ctx: ServerRequestContext, _params: Optional[PaginatedRequestParams]): return types.ListResourcesResult( @@ -1007,14 +1017,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)) @@ -1056,6 +1074,7 @@ async def handle_get_prompt(req): def _build_transform_context(self, rc: Any, *, invoker: Any) -> TransformContext: """TransformContext for the in-flight request, identity from ``caller_identity``.""" session_id, subject = caller_identity(rc, self.protocol_era) + session_key = VisibilityTransform.session_key(session_id, subject) exec_ctx = None if rc is not None: correlation_id, jsonrpc_id = _bind_correlation(rc) @@ -1067,10 +1086,17 @@ def _build_transform_context(self, rc: Any, *, invoker: Any) -> TransformContext trace=_trace_context_from_request_ctx(rc), ) _apply_request_envelope(exec_ctx, rc) + self._configure_visibility_context(exec_ctx, rc) + visible_tools = [ + self._build_tool_definition(entry) + for entry in self._tools.values() + if self.visibility_transform.is_visible(entry, session_key) + ] return TransformContext( session_id=session_id, subject=subject, execution_context=exec_ctx, + visible_tools=visible_tools, _invoker=invoker, ) @@ -1239,6 +1265,55 @@ 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) + def _visibility_session_key(self, request_context: Any) -> tuple[str, str] | None: + session_id, subject = caller_identity(request_context, self.protocol_era) + return VisibilityTransform.session_key(session_id, subject) + + 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, @@ -1404,10 +1479,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( @@ -1491,6 +1565,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, @@ -1545,6 +1620,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( diff --git a/nitrostack/core/context.py b/nitrostack/core/context.py index ac76790..fef026e 100644 --- a/nitrostack/core/context.py +++ b/nitrostack/core/context.py @@ -202,6 +202,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: @@ -213,3 +217,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) diff --git a/nitrostack/core/decorators.py b/nitrostack/core/decorators.py index 60f4323..9a9b39f 100644 --- a/nitrostack/core/decorators.py +++ b/nitrostack/core/decorators.py @@ -53,6 +53,7 @@ class ToolConfig: invocation: Optional[ToolInvocation] = None metadata: Dict[str, Any] = field(default_factory=dict) is_initial: bool = False + default_visible: Optional[bool] = None @dataclass class ResourceConfig: @@ -101,6 +102,7 @@ def tool( examples: Optional[ToolExamples] = None, invocation: Optional[ToolInvocation] = None, metadata: Optional[Dict[str, Any]] = None, + default_visible: Optional[bool] = None, ): """ Decorator to register a method as an MCP Tool. @@ -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, diff --git a/nitrostack/core/visibility.py b/nitrostack/core/visibility.py new file mode 100644 index 0000000..529dfbe --- /dev/null +++ b/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 diff --git a/nitrostack/transforms/search/base.py b/nitrostack/transforms/search/base.py index 050ee28..482d3f7 100644 --- a/nitrostack/transforms/search/base.py +++ b/nitrostack/transforms/search/base.py @@ -108,14 +108,18 @@ async def rank(self, query: str, tools: Sequence[Any]) -> RankOutcome: # ---- handlers -------------------------------------------------------- - def _catalog(self) -> List[Any]: - if self._registry is None: + def _catalog(self, ctx: Optional[TransformContext] = None) -> List[Any]: + if ctx is not None and ctx.visible_tools is not None: + tools = list(ctx.visible_tools) + elif self._registry is not None: + tools = self._registry.list_tools() + else: return [] - tools = [t for t in self._registry.list_tools() if t.name not in self._synthetic] + tools = [t for t in tools if t.name not in self._synthetic] return sorted(tools, key=lambda t: t.name) async def _handle_search(self, args: SearchToolsInput, ctx: TransformContext) -> Dict[str, Any]: - catalog = self._catalog() + catalog = self._catalog(ctx) notes: List[str] = [] if is_browse_query(args.query): @@ -157,4 +161,4 @@ async def _handle_call(self, args: CallToolInput, ctx: TransformContext) -> Any: ) # Goes back through the pipeline: guards and argument validation run for the real tool, # and this transform is skipped for the nested call. - return await ctx.call_tool(target, args.arguments, caller=self.name) \ No newline at end of file + return await ctx.call_tool(target, args.arguments, caller=self.name) diff --git a/nitrostack/transforms/transform.py b/nitrostack/transforms/transform.py index d35109f..8706478 100644 --- a/nitrostack/transforms/transform.py +++ b/nitrostack/transforms/transform.py @@ -49,6 +49,7 @@ class TransformContext: subject: Optional[str] = None # verified bearer subject, never client-supplied execution_context: Any = None # the ExecutionContext of the originating request metadata: Dict[str, Any] = field(default_factory=dict) + visible_tools: Optional[Sequence[Tool]] = None _invoker: Optional[ToolInvoker] = None async def call_tool( @@ -126,4 +127,4 @@ def apply_transform( 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 + return list(await maybe_await(self.apply_transform(tools, ctx))) diff --git a/tests/test_tool_visibility.py b/tests/test_tool_visibility.py new file mode 100644 index 0000000..36b6955 --- /dev/null +++ b/tests/test_tool_visibility.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +from pydantic import BaseModel + +from nitrostack import ( + ExecutionContext, + ToolExamples, + ToolInvocation, + injectable, + module, + tool, +) +from nitrostack.core.app import McpApplication +from nitrostack.core.visibility import SessionVisibilityStore, VisibilityTransform +from nitrostack.runtime.request_ctx import RequestContext, request_ctx + + +class EmptyInput(BaseModel): + pass + + +class AdminInput(BaseModel): + secret_value: str + + +@injectable() +class VisibilityTools: + @tool("public_tool", "public", EmptyInput) + async def public(self, input: EmptyInput, context: ExecutionContext): + return "public" + + @tool("admin_tool", "admin", AdminInput, visibility="hidden", default_visible=False) + async def admin(self, input: AdminInput, context: ExecutionContext): + return input.secret_value + + @tool("unlock", "unlock", EmptyInput) + async def unlock(self, input: EmptyInput, context: ExecutionContext): + context.enable_tools(["admin_tool", "not_registered"]) + return "ok" + + @tool("lock", "lock", EmptyInput) + async def lock(self, input: EmptyInput, context: ExecutionContext): + context.disable_tools(["admin_tool"]) + return "ok" + + +@module(name="visibility", providers=[VisibilityTools]) +class VisibilityModule: + pass + + +def _app(): + app = McpApplication(VisibilityModule) + app.protocol_era = "legacy" + return app + + +def _request(session_id="session-a"): + return RequestContext(session=SimpleNamespace(id=session_id), request_id="1") + + +def test_session_store_lru_eviction(): + now = [0.0] + store = SessionVisibilityStore(ttl_seconds=10, max_entries=2, clock=lambda: now[0]) + store.update(("a", "u"), ["tool"], True) + store.update(("b", "u"), ["tool"], True) + assert store.is_visible(("a", "u"), "tool", False) + store.update(("c", "u"), ["tool"], True) + assert not store.is_visible(("b", "u"), "tool", False) + + +def test_session_store_ttl_expiry_uses_fake_clock(): + now = [0.0] + store = SessionVisibilityStore(ttl_seconds=10, max_entries=2, clock=lambda: now[0]) + store.update(("session", "subject"), ["tool"], True) + now[0] = 11 + assert not store.is_visible(("session", "subject"), "tool", False) + + +def test_session_key_requires_session_and_includes_verified_subject(): + assert VisibilityTransform.session_key(None, "user-a") is None + assert VisibilityTransform.session_key("client-session", "user-a") != ( + VisibilityTransform.session_key("client-session", "user-b") + ) + + +def test_tool_default_visible_preserves_existing_positional_arguments(): + examples = ToolExamples(input={"x": 1}, output={"y": 2}) + invocation = ToolInvocation(invoking="Working", invoked="Worked") + metadata = {"custom": "value"} + + @tool( + "positional_compat", + "compatibility", + EmptyInput, + "Compatibility", + None, + None, + "forbidden", + "hidden", + examples, + invocation, + metadata, + ) + async def positional_compat(self, input: EmptyInput, context: ExecutionContext): + return "ok" + + config = positional_compat._mcp_tool_config + assert config.examples is examples + assert config.invocation is invocation + assert config.metadata is metadata + assert config.visibility == "hidden" + assert config.default_visible is False + + +def test_tools_list_and_call_enforce_visibility_and_isolate_sessions(): + async def run(): + app = _app() + list_handler = app.mcp_server.request_handlers[ + __import__("mcp.types", fromlist=["ListToolsRequest"]).ListToolsRequest + ] + request = _request() + token = request_ctx.set(request) + try: + listed = await list_handler(SimpleNamespace(params=None)) + assert "admin_tool" not in [item.name for item in listed.root.tools] + denied = await app._call_tool("admin_tool", {}) + assert denied.is_error + + enabled = await app._call_tool("unlock", {}) + assert not enabled.is_error + listed = await list_handler(SimpleNamespace(params=None)) + assert "admin_tool" in [item.name for item in listed.root.tools] + + entry = app._tools["admin_tool"] + other_session = app._visibility_session_key(_request("session-b")) + assert not app.visibility_transform.is_visible(entry, other_session) + other_subject = ("session-a", "verified-user-b") + assert not app.visibility_transform.is_visible(entry, other_subject) + + disabled = await app._call_tool("lock", {}) + assert not disabled.is_error + listed = await list_handler(SimpleNamespace(params=None)) + assert "admin_tool" not in [item.name for item in listed.root.tools] + finally: + request_ctx.reset(token) + + asyncio.run(run()) + + +def test_sessionless_mutation_is_noop_but_static_default_applies(): + app = _app() + context = ExecutionContext(request_id="r") + app._configure_visibility_context(context, None) + context.enable_tools(["admin_tool"]) + assert not app.visibility_transform.is_visible(app._tools["admin_tool"], None) + assert app.visibility_transform.is_visible(app._tools["public_tool"], None) + + +def test_http_session_id_is_transport_owned_and_keyed_with_verified_subject(monkeypatch): + """Live HTTP requests use validated session ids and verified subjects for visibility.""" + from starlette.testclient import TestClient + + from nitrostack import DIContainer + from nitrostack.auth.jwt import JWTModule, JWTService + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.transports.http import build_http_app + + async def build(): + @mcp_app(module=VisibilityModule, server=ServerConfig(name="visibility-http", protocol_era="legacy")) + class HttpVisibilityApp: + pass + + return await McpApplicationFactory.create(HttpVisibilityApp) + + monkeypatch.setenv("JWT_SECRET", "visibility-test-secret") + DIContainer.reset() + JWTModule.for_root(secret_env_var="JWT_SECRET") + app = asyncio.run(build()) + http_app = build_http_app( + app, protocol_era="legacy", wire_mode="sessionful", enable_cors=True + ) + token = DIContainer.get_instance().resolve(JWTService).create_token({"sub": "owner"}) + auth_headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + init_body = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "visibility-test", "version": "1"}, + }, + } + + def decode_response(response): + if response.headers.get("content-type", "").startswith("text/event-stream"): + for line in response.text.replace("\r\n", "\n").splitlines(): + if line.startswith("data:"): + return __import__("json").loads(line[len("data:"):].strip()) + raise AssertionError(f"No JSON-RPC response in {response.text!r}") + return response.json() + + with TestClient(http_app) as client: + session_a = client.post("/mcp", headers=auth_headers, json=init_body) + session_b = client.post("/mcp", headers=auth_headers, json=init_body) + assert session_a.status_code == session_b.status_code == 200 + id_a = session_a.headers["mcp-session-id"] + id_b = session_b.headers["mcp-session-id"] + assert id_a != id_b + + for session_id in (id_a, id_b): + client.post( + "/mcp", + headers={**auth_headers, "mcp-session-id": session_id}, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}, + ) + + def rpc(session_id, method, params, request_id=2, headers=auth_headers): + response = client.post( + "/mcp", + headers={**headers, "mcp-session-id": session_id}, + json={"jsonrpc": "2.0", "id": request_id, "method": method, "params": params}, + ) + assert response.status_code == 200, response.text + return decode_response(response)["result"] + + def call(session_id, name, request_id=2, headers=auth_headers, arguments=None): + return rpc( + session_id, + "tools/call", + {"name": name, "arguments": arguments or {}}, + request_id, + headers, + ) + + denied_before_unlock = call(id_a, "admin_tool") + assert denied_before_unlock["isError"] is True + unlocked = call(id_a, "unlock", request_id=3) + assert unlocked["isError"] is False + + listed_a = rpc(id_a, "tools/list", {}, request_id=4) + listed_b = rpc(id_b, "tools/list", {}, request_id=5) + names_a = {entry["name"] for entry in listed_a["tools"]} + names_b = {entry["name"] for entry in listed_b["tools"]} + assert "admin_tool" in names_a + assert "admin_tool" not in names_b + + allowed_call = call( + id_a, "admin_tool", request_id=6, arguments={"secret_value": "private"} + ) + denied_other_session = call(id_b, "admin_tool", request_id=7) + assert allowed_call["isError"] is False + assert denied_other_session["isError"] is True + + # A different verified subject using the same validated session id gets + # an independent visibility key. + other_token = DIContainer.get_instance().resolve(JWTService).create_token({"sub": "intruder"}) + other_headers = {**auth_headers, "Authorization": f"Bearer {other_token}"} + listed_other_subject = rpc(id_a, "tools/list", {}, request_id=8, headers=other_headers) + assert "admin_tool" not in {entry["name"] for entry in listed_other_subject["tools"]} + + # A client-supplied random ID is rejected; it cannot select stored state. + arbitrary = client.post( + "/mcp", + headers={**auth_headers, "mcp-session-id": "client-chosen"}, + json={"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}, + ) + assert arbitrary.status_code == 404 + DIContainer.reset() + + +def test_search_tools_uses_current_request_visibility_for_names_and_schemas(monkeypatch): + from starlette.testclient import TestClient + + from nitrostack import DIContainer + from nitrostack.auth.jwt import JWTModule, JWTService + from nitrostack.core.app import McpApplicationFactory, ServerConfig, mcp_app + from nitrostack.transports.http import build_http_app + from nitrostack.transforms.search import RegexSearchTransform + + async def build(): + @mcp_app( + module=VisibilityModule, + server=ServerConfig(name="visibility-search-http", protocol_era="legacy"), + transforms=[RegexSearchTransform(always_visible=["unlock", "lock"])], + ) + class SearchVisibilityApp: + pass + + return await McpApplicationFactory.create(SearchVisibilityApp) + + monkeypatch.setenv("JWT_SECRET", "visibility-search-test-secret") + DIContainer.reset() + try: + JWTModule.for_root(secret_env_var="JWT_SECRET") + app = asyncio.run(build()) + http_app = build_http_app( + app, protocol_era="legacy", wire_mode="sessionful", enable_cors=True + ) + jwt = DIContainer.get_instance().resolve(JWTService) + headers = { + "Authorization": f"Bearer {jwt.create_token({'sub': 'search-owner'})}", + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + initialize = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "visibility-search-test", "version": "1"}, + }, + } + + def decode(response): + if response.headers.get("content-type", "").startswith("text/event-stream"): + import json + + data = next(line[5:].strip() for line in response.text.splitlines() if line.startswith("data:")) + return json.loads(data) + return response.json() + + with TestClient(http_app) as client: + initialized = [ + client.post("/mcp", headers=headers, json=initialize), + client.post("/mcp", headers=headers, json=initialize), + ] + assert all(response.status_code == 200 for response in initialized) + session_a, session_b = [response.headers["mcp-session-id"] for response in initialized] + + def call(session_id, name, request_id, arguments): + response = client.post( + "/mcp", + headers={**headers, "mcp-session-id": session_id}, + json={ + "jsonrpc": "2.0", + "id": request_id, + "method": "tools/call", + "params": {"name": name, "arguments": arguments}, + }, + ) + assert response.status_code == 200, response.text + return decode(response)["result"] + + enabled = call(session_a, "unlock", 2, {}) + assert enabled["isError"] is False + + visible = call( + session_a, + "search_tools", + 3, + {"query": "admin_tool", "detail": "full"}, + ) + hidden = call( + session_b, + "search_tools", + 4, + {"query": "admin_tool", "detail": "full"}, + ) + visible_results = visible["structuredContent"]["results"] + hidden_results = hidden["structuredContent"]["results"] + assert any(result["name"] == "admin_tool" for result in visible_results) + admin_result = next(result for result in visible_results if result["name"] == "admin_tool") + assert "secret_value" in str(admin_result) + assert all(result["name"] != "admin_tool" for result in hidden_results) + assert "secret_value" not in str(hidden_results) + finally: + DIContainer.reset()