diff --git a/docs/source/run-workflows/fastmcp-server.md b/docs/source/run-workflows/fastmcp-server.md index 9201da98e6..d8a6fb6960 100644 --- a/docs/source/run-workflows/fastmcp-server.md +++ b/docs/source/run-workflows/fastmcp-server.md @@ -150,6 +150,14 @@ general: With this configuration, the MCP server is accessible at `http://localhost:9902/api/v1/mcp`. +### Per-User Workflows + +A per-user workflow can be served over FastMCP. Each user gets their own workflow instance, built on first use. [MCP authentication](../components/auth/mcp-auth/index.md) prescribes pairing `per_user_mcp_client` with a per-user workflow, such as `per_user_react_agent`. + +The user is taken from the Bearer token on the MCP request when [`server_auth`](#authentication) is configured, or from runtime context when it is already set. Per-user workflows require authenticated streamable-http requests so the server can derive a user ID. Without that, tool calls fail with an error saying the user ID could not be determined. + +The `/debug/tools/list` route returns an empty tool list for a per-user workflow because no shared workflow instance exists at startup. Use `nat mcp client tool list` to inspect the registered MCP tool and its input schema. + ## Inspecting and Running MCP Tools Published by a FastMCP Server Use `nat mcp client` to inspect and run tools exposed by an MCP server using the FastMCP server runtime. diff --git a/docs/source/run-workflows/mcp-server.md b/docs/source/run-workflows/mcp-server.md index 75370bec66..80b2f4286e 100644 --- a/docs/source/run-workflows/mcp-server.md +++ b/docs/source/run-workflows/mcp-server.md @@ -93,6 +93,14 @@ The `base_path` must start with a forward slash (`/`) and must not end with a fo The `base_path` feature requires the `streamable-http` transport. SSE transport does not support custom base paths. ::: +### Per-User Workflows + +A per-user workflow can be served over MCP. Each user gets their own workflow instance, built on first use. [MCP authentication](../components/auth/mcp-auth/index.md) prescribes pairing `per_user_mcp_client` with a per-user workflow, such as `per_user_react_agent`. + +The user is taken from the Bearer token on the MCP request when [`server_auth`](#authentication) is configured, or from runtime context when it is already set. Per-user workflows require authenticated streamable-http requests so the server can derive a user ID. Without that, tool calls fail with an error saying the user ID could not be determined. + +The `/debug/tools/list` route returns an empty tool list for a per-user workflow because no shared workflow instance exists at startup. Use `nat mcp client tool list` to inspect the registered MCP tool and its input schema. + ## Displaying MCP Tools published by an MCP server To list the tools published by the MCP server you can use the `nat mcp client tool list` command. This command acts as an MCP client and connects to the MCP server running on the specified URL (defaults to `http://localhost:9901/mcp` for streamable-http, with backwards compatibility for `http://localhost:9901/sse`). @@ -393,9 +401,15 @@ This is useful for health checks and monitoring. ## Security Considerations -### Authentication Limitations -- The `nat mcp serve` command currently starts an MCP server without built-in authentication. Server-side authentication is planned for a future release. -- NeMo Agent Toolkit workflows can still connect to protected third-party MCP servers through the MCP client auth provider. Refer to [MCP Authentication](../components/auth/mcp-auth/index.md) for more information. +### Authentication + +MCP servers using the **streamable-http** transport can validate OAuth2 bearer tokens when `server_auth` is configured in the front-end config. Validation uses JWKS or OIDC discovery for JWT access tokens, or token introspection for opaque access tokens. Refer to [MCP Authentication](../components/auth/mcp-auth/index.md) and the protected examples under `examples/MCP/`. + +The **SSE** transport does not apply `server_auth`. Use streamable-http when the MCP server must authenticate callers. + +Per-user workflows require authenticated requests so the server can derive a user ID from the Bearer token. See [Per-User Workflows](#per-user-workflows). + +NeMo Agent Toolkit workflows can still connect to protected third-party MCP servers through the MCP client auth provider. ### Local Development For local development, you can use `localhost` or `127.0.0.1` as the host (default). This limits access to your local machine only. diff --git a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin.py b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin.py index 6ee1e13c87..eb8ef6e2bc 100644 --- a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin.py +++ b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin.py @@ -66,49 +66,52 @@ async def run(self) -> None: # Get the worker instance worker = self._get_worker_instance() - # Let the worker create the FastMCP server (allows plugins to customize) - mcp = await worker.create_mcp_server() - - # Add routes through the worker (includes health endpoint and function registration) - await worker.add_routes(mcp, builder) - try: - if self.front_end_config.base_path: - if self.front_end_config.transport == "sse": - logger.warning( - "base_path is configured but SSE transport does not support mounting at sub-paths. " - "Use streamable-http transport for base_path support.") + # Let the worker create the FastMCP server (allows plugins to customize) + mcp = await worker.create_mcp_server() + + # Add routes through the worker (includes health endpoint and function registration) + await worker.add_routes(mcp, builder) + + try: + if self.front_end_config.base_path: + if self.front_end_config.transport == "sse": + logger.warning( + "base_path is configured but SSE transport does not support mounting at sub-paths. " + "Use streamable-http transport for base_path support.") + logger.info("Starting FastMCP server with SSE endpoint at /sse") + await mcp.run_async(transport="sse", + host=self.front_end_config.host, + port=self.front_end_config.port, + log_level=self.front_end_config.log_level.lower()) + else: + full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}{self.front_end_config.base_path}/mcp" + logger.info( + "Mounting FastMCP server at %s/mcp on %s:%s", + self.front_end_config.base_path, + self.front_end_config.host, + self.front_end_config.port, + ) + logger.info("FastMCP server URL: %s", full_url) + await self._run_with_mount(mcp, worker) + elif self.front_end_config.transport == "sse": logger.info("Starting FastMCP server with SSE endpoint at /sse") await mcp.run_async(transport="sse", host=self.front_end_config.host, port=self.front_end_config.port, log_level=self.front_end_config.log_level.lower()) else: - full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}{self.front_end_config.base_path}/mcp" - logger.info( - "Mounting FastMCP server at %s/mcp on %s:%s", - self.front_end_config.base_path, - self.front_end_config.host, - self.front_end_config.port, - ) + full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}/mcp" logger.info("FastMCP server URL: %s", full_url) - await self._run_with_mount(mcp, worker) - elif self.front_end_config.transport == "sse": - logger.info("Starting FastMCP server with SSE endpoint at /sse") - await mcp.run_async(transport="sse", - host=self.front_end_config.host, - port=self.front_end_config.port, - log_level=self.front_end_config.log_level.lower()) - else: - full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}/mcp" - logger.info("FastMCP server URL: %s", full_url) - await mcp.run_async(transport="streamable-http", - host=self.front_end_config.host, - port=self.front_end_config.port, - path="/mcp", - log_level=self.front_end_config.log_level.lower()) - except KeyboardInterrupt: - logger.info("FastMCP server shutdown requested (Ctrl+C). Shutting down gracefully.") + await mcp.run_async(transport="streamable-http", + host=self.front_end_config.host, + port=self.front_end_config.port, + path="/mcp", + log_level=self.front_end_config.log_level.lower()) + except KeyboardInterrupt: + logger.info("FastMCP server shutdown requested (Ctrl+C). Shutting down gracefully.") + finally: + await worker.cleanup() async def _run_with_mount(self, mcp: "FastMCP", worker: FastMCPFrontEndPluginWorkerBase) -> None: """Run FastMCP server mounted at configured base_path using FastAPI wrapper. diff --git a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin_worker.py b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin_worker.py index f711cae82e..83fbed3544 100644 --- a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin_worker.py +++ b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/front_end_plugin_worker.py @@ -52,6 +52,44 @@ def __init__(self, config: Config): """ self.full_config = config self.front_end_config: FastMCPFrontEndConfig = config.general.front_end + self._session_managers: list[SessionManager] = [] + + def _track_session_manager(self, session_manager: SessionManager) -> SessionManager: + """Track a session manager so ``cleanup()`` can shut it down.""" + self._session_managers.append(session_manager) + return session_manager + + def _is_primary_workflow_per_user(self) -> bool: + """Return True when the configured workflow is registered as per-user.""" + from nat.cli.type_registry import GlobalTypeRegistry + + workflow_registration = GlobalTypeRegistry.get().get_function(type(self.full_config.workflow)) + return workflow_registration.is_per_user + + def _resolve_workflow_tool_name(self) -> str: + """Return the MCP tool name for the configured workflow entry point.""" + workflow_config = self.full_config.workflow + alias = getattr(workflow_config, "workflow_alias", None) + return alias if alias else workflow_config.type + + async def cleanup(self) -> None: + """Shut down all tracked session managers. + + Attempts shutdown for every manager even when an individual call fails. + Clears the tracked list before re-raising the first shutdown error. + """ + errors: list[BaseException] = [] + try: + for session_manager in self._session_managers: + try: + await session_manager.shutdown() + except Exception as exc: + logger.exception("Failed to shut down SessionManager during FastMCP worker cleanup") + errors.append(exc) + finally: + self._session_managers.clear() + if errors: + raise errors[0] def _setup_health_endpoint(self, mcp: FastMCP): """Set up the HTTP health endpoint that exercises FastMCP ping handler.""" @@ -106,14 +144,27 @@ async def add_routes(self, mcp: FastMCP, builder: WorkflowBuilder): ... async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder) -> None: - """Default implementation for adding routes to FastMCP.""" + """Default route registration for FastMCP. + + Creates session managers via SessionManager.create (shared or per-user), + registers workflow tools, and exposes debug endpoints for introspection. + """ from nat.plugins.fastmcp.server.tool_converter import register_function_with_mcp # Set up the health endpoint self._setup_health_endpoint(mcp) - # Build the default workflow - workflow = await builder.build() + if self._is_primary_workflow_per_user(): + session_manager = self._track_session_manager(await SessionManager.create(config=self.full_config, + shared_builder=builder)) + tool_name = self._resolve_workflow_tool_name() + register_function_with_mcp(mcp, tool_name, session_manager, function=None) + self._setup_debug_endpoints(mcp, {}) + return + + primary_session_manager = self._track_session_manager(await SessionManager.create(config=self.full_config, + shared_builder=builder)) + workflow = primary_session_manager.workflow # Get all functions from the workflow functions = await self._get_all_functions(workflow) @@ -136,14 +187,11 @@ async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder) -> N for function_name, function in functions.items(): if isinstance(function, Workflow): logger.info("Function %s is a Workflow, using directly", function_name) - session_managers[function_name] = await SessionManager.create(config=self.full_config, - shared_builder=builder, - entry_function=None) + session_managers[function_name] = primary_session_manager else: logger.info("Function %s is a regular function, building entry workflow", function_name) - session_managers[function_name] = await SessionManager.create(config=self.full_config, - shared_builder=builder, - entry_function=function_name) + session_managers[function_name] = self._track_session_manager(await SessionManager.create( + config=self.full_config, shared_builder=builder, entry_function=function_name)) # Register each function with FastMCP, passing SessionManager for observability for function_name, session_manager in session_managers.items(): @@ -152,8 +200,7 @@ async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder) -> N if not session_managers: raise RuntimeError("No functions found in workflow. Please check your configuration.") - # After registration, expose debug endpoints for tool/schema inspection - debug_functions = {name: sm.workflow for name, sm in session_managers.items()} + debug_functions = {name: sm.workflow for name, sm in session_managers.items() if not sm.is_workflow_per_user} self._setup_debug_endpoints(mcp, debug_functions) async def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]: diff --git a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/tool_converter.py b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/tool_converter.py index 75de40416b..2613587a2e 100644 --- a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/tool_converter.py +++ b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/tool_converter.py @@ -27,12 +27,16 @@ from pydantic_core import PydanticUndefined from fastmcp import FastMCP +from fastmcp.server.context import Context from nat.builder.function import Function # type: ignore[reportMissingImports] from nat.builder.function_base import FunctionBase # type: ignore[reportMissingImports] from nat.runtime.session import SessionManager # type: ignore[reportMissingImports] logger = logging.getLogger(__name__) +# Reserved wrapper parameter for FastMCP request Context injection. +INJECTED_CONTEXT_PARAM = "_nat_mcp_context" + # Sentinel: marks "optional; let Pydantic supply default/factory" _USE_PYDANTIC_DEFAULT = object() @@ -186,6 +190,73 @@ def _is_chat_request_schema(schema: Any) -> bool: return schema_name == "ChatRequest" or "ChatRequest" in schema_qualname +def _schema_wrapper_parameter_names(input_schema: Any) -> set[str]: + """Return collision-safe wrapper parameter names derived from the input schema.""" + if _is_chat_request_schema(input_schema): + return {"query"} + if not hasattr(input_schema, "model_fields"): + return set() + + name_map = _build_name_mapping(list(input_schema.model_fields.keys())) + return set(name_map.values()) + + +def _validate_input_schema_for_context_injection(input_schema: Any) -> None: + """Reject schemas whose fields collide with the injected `Context` parameter.""" + schema_params = _schema_wrapper_parameter_names(input_schema) + if INJECTED_CONTEXT_PARAM in schema_params: + raise ValueError( + f"Workflow input schema cannot declare a field that maps to reserved MCP parameter " + f"{INJECTED_CONTEXT_PARAM!r}.", ) + if "ctx" in getattr(input_schema, "model_fields", {}): + raise ValueError( + "Workflow input schema cannot declare a field named 'ctx'; that name is reserved for " + "MCP request context injection.", ) + + +def _append_context_parameter(signature: Signature) -> Signature: + """Append the FastMCP `Context` parameter used for request injection.""" + param_names = {param.name for param in signature.parameters.values()} + if INJECTED_CONTEXT_PARAM in param_names: + raise ValueError( + f"Workflow input schema cannot declare a field that maps to reserved MCP parameter " + f"{INJECTED_CONTEXT_PARAM!r}.", ) + + ctx_param = Parameter( + INJECTED_CONTEXT_PARAM, + Parameter.KEYWORD_ONLY, + default=None, + annotation=Context | None, + ) + return signature.replace(parameters=[*signature.parameters.values(), ctx_param]) + + +async def _run_through_session_manager(session_manager: "SessionManager", payload: Any, ctx: Any | None = None) -> Any: + """Execute a payload through SessionManager, including per-user workflows.""" + if session_manager.is_workflow_per_user: + from nat.builder.context import Context + from nat.runtime.user_manager import UserManager + + user_id = Context.get().user_id + http_connection = None + if ctx is not None: + try: + http_connection = ctx.request_context.request + if user_id is None and http_connection is not None: + user_info = UserManager.extract_user_from_connection(http_connection) + if user_info is not None: + user_id = user_info.get_user_id() + except ValueError: + pass + + async with session_manager.session(user_id=user_id, http_connection=http_connection) as session: + async with session.run(payload) as runner: + return await runner.result() + + async with session_manager.run(payload) as runner: + return await runner.result() + + def create_function_wrapper( function_name: str, session_manager: "SessionManager", @@ -199,8 +270,9 @@ def create_function_wrapper( input_schema: Input schema for the workflow/function. """ signature, alias_map = _build_signature_from_schema(input_schema) + _validate_input_schema_for_context_injection(input_schema) - async def wrapper_func(**kwargs: Any) -> Any: + async def wrapper_func(_nat_mcp_context: Context | None = None, **kwargs: Any) -> Any: if _is_chat_request_schema(input_schema): from nat.data_models.api_server import ChatRequest # type: ignore[reportMissingImports] @@ -214,8 +286,7 @@ async def wrapper_func(**kwargs: Any) -> Any: payload = input_schema.model_validate(cleaned_kwargs) if hasattr(input_schema, "model_validate") else cleaned_kwargs - async with session_manager.run(payload) as runner: - result = await runner.result() + result = await _run_through_session_manager(session_manager, payload, ctx=_nat_mcp_context) if isinstance(result, str): return result @@ -223,8 +294,11 @@ async def wrapper_func(**kwargs: Any) -> Any: return json.dumps(result, default=str) return str(result) - wrapper_func.__signature__ = signature # type: ignore[attr-defined] - wrapper_func.__annotations__ = _build_annotations_from_schema(input_schema) + wrapper_func.__signature__ = _append_context_parameter(signature) # type: ignore[attr-defined] + annotations = _build_annotations_from_schema(input_schema) + annotations[INJECTED_CONTEXT_PARAM] = Context | None + annotations["return"] = Any + wrapper_func.__annotations__ = annotations wrapper_func.__name__ = function_name wrapper_func.__doc__ = "Auto-generated wrapper for a NeMo Agent Toolkit workflow." return wrapper_func @@ -273,19 +347,18 @@ def register_function_with_mcp(mcp: FastMCP, """ logger.info("Registering function %s with FastMCP", function_name) - # Get the workflow from the session manager - workflow = session_manager.workflow + if session_manager.is_workflow_per_user: + input_schema = session_manager.get_workflow_input_schema() + workflow_config = session_manager.config.workflow + function_description = getattr(workflow_config, "description", None) or function_name + else: + workflow = session_manager.workflow + target_function = function or workflow + input_schema = getattr(target_function, "input_schema", workflow.input_schema) + function_description = get_function_description(target_function) - # Prefer the function's schema/description when available, fall back to workflow - target_function = function or workflow - - # Get the input schema from the most specific object available - input_schema = getattr(target_function, "input_schema", workflow.input_schema) logger.info("Function %s has input schema: %s", function_name, input_schema) - # Get function description - function_description = get_function_description(target_function) - # Create and register the wrapper function with FastMCP wrapper_func = create_function_wrapper(function_name, session_manager, input_schema) mcp.tool(name=function_name, description=function_description)(wrapper_func) diff --git a/packages/nvidia_nat_fastmcp/tests/server/test_per_user_workflow.py b/packages/nvidia_nat_fastmcp/tests/server/test_per_user_workflow.py new file mode 100644 index 0000000000..4fb70d2433 --- /dev/null +++ b/packages/nvidia_nat_fastmcp/tests/server/test_per_user_workflow.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +import pytest +from mcp.server.lowlevel.server import request_ctx +from pydantic import BaseModel +from starlette.requests import Request +from starlette.testclient import TestClient + +from nat.builder.builder import Builder +from nat.builder.function_info import FunctionInfo +from nat.builder.workflow_builder import WorkflowBuilder +from nat.cli.register_workflow import register_function +from nat.cli.register_workflow import register_per_user_function +from nat.data_models.config import Config +from nat.data_models.config import GeneralConfig +from nat.data_models.function import FunctionBaseConfig +from nat.plugins.fastmcp.server.front_end_config import FastMCPFrontEndConfig +from nat.plugins.fastmcp.server.front_end_plugin import FastMCPFrontEndPlugin +from nat.plugins.fastmcp.server.front_end_plugin_worker import FastMCPFrontEndPluginWorker +from nat.runtime.session import SessionManager + + +class _Input(BaseModel): + message: str + + +class _Output(BaseModel): + result: str + + +class PerUserFastMCPWorkflowConfig(FunctionBaseConfig, name="per_user_fastmcp_test_workflow"): + """Per-user workflow config for FastMCP front-end tests.""" + + +class SharedFastMCPWorkflowConfig(FunctionBaseConfig, name="shared_fastmcp_test_workflow"): + """Shared workflow config for FastMCP front-end tests.""" + + +@pytest.fixture(name="registered_workflows", scope="module") +def fixture_registered_workflows(): + """Register test workflows in a pushed registry so they do not leak.""" + + @register_per_user_function(config_type=PerUserFastMCPWorkflowConfig, input_type=_Input, single_output_type=_Output) + async def _build_per_user(_config: PerUserFastMCPWorkflowConfig, _builder: Builder): + + async def _impl(inp: _Input) -> _Output: + return _Output(result=f"per-user: {inp.message}") + + yield FunctionInfo.from_fn(_impl) + + @register_function(config_type=SharedFastMCPWorkflowConfig) + async def _build_shared(_config: SharedFastMCPWorkflowConfig, _builder: Builder): + + async def _impl(inp: _Input) -> _Output: + return _Output(result=f"shared: {inp.message}") + + yield FunctionInfo.from_fn(_impl) + + +def _config(workflow) -> Config: + return Config( + general=GeneralConfig(front_end=FastMCPFrontEndConfig( + name="Test FastMCP Server", + host="localhost", + port=9903, + debug=False, + log_level="INFO", + )), + workflow=workflow, + ) + + +@pytest.fixture(name="per_user_config") +def fixture_per_user_config(registered_workflows) -> Config: + return _config(PerUserFastMCPWorkflowConfig()) + + +@pytest.fixture(name="shared_config") +def fixture_shared_config(registered_workflows) -> Config: + return _config(SharedFastMCPWorkflowConfig()) + + +class TestPerUserWorkflowStartup: + """The FastMCP front end must serve per-user workflows, not die building them.""" + + async def test_per_user_workflow_server_starts(self, per_user_config, monkeypatch): + """Startup used to raise "Must set a workflow before building".""" + + async def _no_serve(*_args, **_kwargs): + return None + + monkeypatch.setattr("fastmcp.server.server.FastMCP.run_async", _no_serve) + + await FastMCPFrontEndPlugin(full_config=per_user_config).run() + + async def test_shared_workflow_still_builds(self, shared_config, monkeypatch): + captured = {} + + original_create = SessionManager.create + + async def _capture_create(*args, **kwargs): + session_manager = await original_create(*args, **kwargs) + if not session_manager.is_workflow_per_user: + captured["workflow"] = session_manager.workflow + return session_manager + + async def _no_serve(*_args, **_kwargs): + return None + + monkeypatch.setattr(SessionManager, "create", _capture_create) + monkeypatch.setattr("fastmcp.server.server.FastMCP.run_async", _no_serve) + + await FastMCPFrontEndPlugin(full_config=shared_config).run() + + assert captured["workflow"] is not None + + async def test_per_user_session_manager_reaps_and_shuts_down(self, per_user_config): + worker = FastMCPFrontEndPluginWorker(per_user_config) + + async with WorkflowBuilder.from_config(config=per_user_config) as builder: + mcp = await worker.create_mcp_server() + await worker._default_add_routes(mcp, builder) + + assert len(worker._session_managers) == 1 + session_manager = worker._session_managers[0] + assert session_manager.is_workflow_per_user + assert session_manager._per_user_builders_cleanup_task is not None + + cleanup_task = session_manager._per_user_builders_cleanup_task + await worker.cleanup() + assert cleanup_task.done() + + async def test_register_function_skips_shared_workflow_lookup(self, per_user_config, monkeypatch): + from nat.plugins.fastmcp.server import tool_converter + + worker = FastMCPFrontEndPluginWorker(per_user_config) + + async with WorkflowBuilder.from_config(config=per_user_config) as builder: + session_manager = await SessionManager.create(config=per_user_config, shared_builder=builder) + mcp = await worker.create_mcp_server() + + get_schema = MagicMock(return_value=_Input) + monkeypatch.setattr(session_manager, "get_workflow_input_schema", get_schema) + + tool_converter.register_function_with_mcp(mcp, "per_user_fastmcp_test_workflow", session_manager) + + get_schema.assert_called_once() + + +class TestWorkerCleanup: + """Worker cleanup must shut down every manager and clear tracking.""" + + async def test_cleanup_shuts_down_all_managers_after_shutdown_failure(self, per_user_config): + worker = FastMCPFrontEndPluginWorker(per_user_config) + failing = MagicMock() + failing.shutdown = AsyncMock(side_effect=RuntimeError("boom")) + succeeding = MagicMock() + succeeding.shutdown = AsyncMock() + worker._session_managers = [failing, succeeding] + + with pytest.raises(RuntimeError, match="boom"): + await worker.cleanup() + + failing.shutdown.assert_awaited_once() + succeeding.shutdown.assert_awaited_once() + assert worker._session_managers == [] + + +class TestPerUserRequestIdentity: + """Per-user tool calls must reach session() with a resolved user id.""" + + async def test_run_through_session_manager_uses_context_user_id(self, monkeypatch): + from nat.plugins.fastmcp.server.tool_converter import _run_through_session_manager + + session_manager = MagicMock() + session_manager.is_workflow_per_user = True + session_manager.session = MagicMock() + session = MagicMock() + runner = MagicMock() + runner.result = AsyncMock(return_value="ok") + session.run.return_value.__aenter__ = AsyncMock(return_value=runner) + session.run.return_value.__aexit__ = AsyncMock(return_value=False) + session_manager.session.return_value.__aenter__ = AsyncMock(return_value=session) + session_manager.session.return_value.__aexit__ = AsyncMock(return_value=False) + + context = SimpleNamespace(user_id="alice") + monkeypatch.setattr("nat.builder.context.Context.get", lambda: context) + + payload = _Input(message="hello") + result = await _run_through_session_manager(session_manager, payload) + + assert result == "ok" + session_manager.session.assert_called_once_with(user_id="alice", http_connection=None) + + async def test_run_through_session_manager_resolves_user_from_mcp_request(self, monkeypatch): + from nat.plugins.fastmcp.server.tool_converter import _run_through_session_manager + + session_manager = MagicMock() + session_manager.is_workflow_per_user = True + session_manager.session = MagicMock() + session = MagicMock() + runner = MagicMock() + runner.result = AsyncMock(return_value="ok") + session.run.return_value.__aenter__ = AsyncMock(return_value=runner) + session.run.return_value.__aexit__ = AsyncMock(return_value=False) + session_manager.session.return_value.__aenter__ = AsyncMock(return_value=session) + session_manager.session.return_value.__aexit__ = AsyncMock(return_value=False) + + context = SimpleNamespace(user_id=None) + monkeypatch.setattr("nat.builder.context.Context.get", lambda: context) + + request = MagicMock() + ctx = SimpleNamespace(request_context=SimpleNamespace(request=request)) + user_info = MagicMock() + user_info.get_user_id.return_value = "bob" + monkeypatch.setattr( + "nat.runtime.user_manager.UserManager.extract_user_from_connection", + MagicMock(return_value=user_info), + ) + + payload = _Input(message="hello") + result = await _run_through_session_manager(session_manager, payload, ctx=ctx) + + assert result == "ok" + session_manager.session.assert_called_once_with(user_id="bob", http_connection=request) + + +class TestPerUserToolStreamableHttp: + """Per-user tool calls over streamable-http must receive injected request context.""" + + async def test_call_tool_resolves_user_from_bearer_token(self, per_user_config, monkeypatch): + from nat.plugins.fastmcp.server import tool_converter + + worker = FastMCPFrontEndPluginWorker(per_user_config) + + async with WorkflowBuilder.from_config(config=per_user_config) as builder: + session_manager = await SessionManager.create(config=per_user_config, shared_builder=builder) + mcp = await worker.create_mcp_server() + tool_converter.register_function_with_mcp(mcp, "per_user_fastmcp_test_workflow", session_manager) + + session = MagicMock() + runner = MagicMock() + runner.result = AsyncMock(return_value="ok") + session.run.return_value.__aenter__ = AsyncMock(return_value=runner) + session.run.return_value.__aexit__ = AsyncMock(return_value=False) + + captured: dict[str, object] = {} + session_cm = MagicMock() + session_cm.__aenter__ = AsyncMock(return_value=session) + session_cm.__aexit__ = AsyncMock(return_value=False) + + def _capture_session(user_id=None, http_connection=None): + captured["user_id"] = user_id + captured["http_connection"] = http_connection + return session_cm + + monkeypatch.setattr(session_manager, "session", _capture_session) + + user_info = MagicMock() + user_info.get_user_id.return_value = "bob" + monkeypatch.setattr( + "nat.runtime.user_manager.UserManager.extract_user_from_connection", + MagicMock(return_value=user_info), + ) + monkeypatch.setattr("nat.builder.context.Context.get", lambda: SimpleNamespace(user_id=None)) + + request = Request({ + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/mcp", + "raw_path": b"/mcp", + "query_string": b"", + "headers": [(b"authorization", b"Bearer test-token")], + "client": ("127.0.0.1", 1234), + "server": ("testserver", 80), + "root_path": "", + }) + + with TestClient(mcp.http_app(transport="streamable-http")): + token = request_ctx.set(SimpleNamespace(request=request, session=MagicMock())) + try: + await mcp.call_tool("per_user_fastmcp_test_workflow", {"message": "hello"}) + finally: + request_ctx.reset(token) + + assert captured["user_id"] == "bob" + assert captured["http_connection"] is request diff --git a/packages/nvidia_nat_fastmcp/tests/server/test_tool_converter.py b/packages/nvidia_nat_fastmcp/tests/server/test_tool_converter.py index a66f33f463..38984fafc4 100644 --- a/packages/nvidia_nat_fastmcp/tests/server/test_tool_converter.py +++ b/packages/nvidia_nat_fastmcp/tests/server/test_tool_converter.py @@ -17,8 +17,12 @@ from unittest.mock import AsyncMock from unittest.mock import MagicMock +import pytest +from fastmcp import FastMCP +from fastmcp.server.context import Context from pydantic import create_model +from nat.plugins.fastmcp.server.tool_converter import INJECTED_CONTEXT_PARAM from nat.plugins.fastmcp.server.tool_converter import _build_name_mapping from nat.plugins.fastmcp.server.tool_converter import _sanitize_parameter_name from nat.plugins.fastmcp.server.tool_converter import create_function_wrapper @@ -28,6 +32,7 @@ def _mock_session_manager(result_value="result"): """Create a mock SessionManager for testing.""" mock_sm = MagicMock(spec=SessionManager) + mock_sm.is_workflow_per_user = False mock_runner = MagicMock() mock_runner.__aenter__ = AsyncMock(return_value=mock_runner) mock_runner.__aexit__ = AsyncMock(return_value=None) @@ -160,3 +165,40 @@ def test_annotations_use_sanitized_names(self): wrapper = create_function_wrapper("tool", _mock_session_manager(), schema) assert "from_" in wrapper.__annotations__ + + def test_wrapper_declares_context_for_injection(self): + """FastMCP injects request context only when a Context parameter is annotated.""" + from fastmcp.server.dependencies import find_kwarg_by_type + + schema = create_model("Schema", **{"query": (str, ...)}) # type: ignore[call-overload] + wrapper = create_function_wrapper("tool", _mock_session_manager(), schema) + + assert find_kwarg_by_type(wrapper, Context) == INJECTED_CONTEXT_PARAM + + def test_create_wrapper_rejects_schema_field_named_ctx(self): + """A workflow field named `ctx` collides with MCP context injection.""" + schema = create_model("CtxSchema", **{"ctx": (str, ...)}) # type: ignore[call-overload] + + with pytest.raises(ValueError, match="cannot declare a field named 'ctx'"): + create_function_wrapper("tool", _mock_session_manager(), schema) + + async def test_registered_per_user_tool_excludes_ctx_from_client_schema(self): + """Registered tools expose workflow args but not the injected Context param.""" + from fastmcp.server.dependencies import find_kwarg_by_type + + schema = create_model("ToolSchema", **{"message": (str, ...)}) # type: ignore[call-overload] + mock_sm = _mock_session_manager() + mock_sm.is_workflow_per_user = True + mock_sm.get_workflow_input_schema = MagicMock(return_value=schema) + mock_sm.config = MagicMock(workflow=MagicMock(description="Per-user workflow")) + + mcp = FastMCP("test-server") + from nat.plugins.fastmcp.server.tool_converter import register_function_with_mcp + + register_function_with_mcp(mcp, "per_user_tool", mock_sm) + tools = await mcp.list_tools() + + assert len(tools) == 1 + assert set(tools[0].parameters["properties"]) == {"message"} + tool = await mcp.get_tool("per_user_tool") + assert find_kwarg_by_type(tool.fn, Context) == INJECTED_CONTEXT_PARAM diff --git a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin.py b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin.py index 21e6db65f3..1eafe473cb 100644 --- a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin.py +++ b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin.py @@ -66,43 +66,46 @@ async def run(self) -> None: # Get the worker instance worker = self._get_worker_instance() - # Let the worker create the MCP server (allows plugins to customize) - mcp = await worker.create_mcp_server() + try: + # Let the worker create the MCP server (allows plugins to customize) + mcp = await worker.create_mcp_server() - # Add routes through the worker (includes health endpoint and function registration) - await worker.add_routes(mcp, builder) + # Add routes through the worker (includes health endpoint and function registration) + await worker.add_routes(mcp, builder) - # Start the MCP server with configurable transport - # streamable-http is the default, but users can choose sse if preferred - try: - # If base_path is configured, mount server at sub-path using FastAPI wrapper - if self.front_end_config.base_path: - if self.front_end_config.transport == "sse": - logger.warning( - "base_path is configured but SSE transport does not support mounting at sub-paths. " - "Use streamable-http transport for base_path support.") + # Start the MCP server with configurable transport + # streamable-http is the default, but users can choose sse if preferred + try: + # If base_path is configured, mount server at sub-path using FastAPI wrapper + if self.front_end_config.base_path: + if self.front_end_config.transport == "sse": + logger.warning( + "base_path is configured but SSE transport does not support mounting at sub-paths. " + "Use streamable-http transport for base_path support.") + logger.info("Starting MCP server with SSE endpoint at /sse") + await mcp.run_sse_async() + else: + full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}{self.front_end_config.base_path}/mcp" + logger.info( + "Mounting MCP server at %s/mcp on %s:%s", + self.front_end_config.base_path, + self.front_end_config.host, + self.front_end_config.port, + ) + logger.info("MCP server URL: %s", full_url) + await self._run_with_mount(mcp) + # Standard behavior - run at root path + elif self.front_end_config.transport == "sse": logger.info("Starting MCP server with SSE endpoint at /sse") await mcp.run_sse_async() - else: - full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}{self.front_end_config.base_path}/mcp" - logger.info( - "Mounting MCP server at %s/mcp on %s:%s", - self.front_end_config.base_path, - self.front_end_config.host, - self.front_end_config.port, - ) + else: # streamable-http + full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}/mcp" logger.info("MCP server URL: %s", full_url) - await self._run_with_mount(mcp) - # Standard behavior - run at root path - elif self.front_end_config.transport == "sse": - logger.info("Starting MCP server with SSE endpoint at /sse") - await mcp.run_sse_async() - else: # streamable-http - full_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}/mcp" - logger.info("MCP server URL: %s", full_url) - await mcp.run_streamable_http_async() - except KeyboardInterrupt: - logger.info("MCP server shutdown requested (Ctrl+C). Shutting down gracefully.") + await mcp.run_streamable_http_async() + except KeyboardInterrupt: + logger.info("MCP server shutdown requested (Ctrl+C). Shutting down gracefully.") + finally: + await worker.cleanup() async def _run_with_mount(self, mcp: "FastMCP") -> None: """Run MCP server mounted at configured base_path using FastAPI wrapper. diff --git a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin_worker.py b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin_worker.py index 6ad294d420..957cdd428b 100644 --- a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin_worker.py +++ b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/front_end_plugin_worker.py @@ -62,6 +62,44 @@ def __init__(self, config: Config): log_interval=self.front_end_config.memory_profile_interval, top_n=self.front_end_config.memory_profile_top_n, log_level=self.front_end_config.memory_profile_log_level) + self._session_managers: list[SessionManager] = [] + + def _track_session_manager(self, session_manager: SessionManager) -> SessionManager: + """Track a session manager so ``cleanup()`` can shut it down.""" + self._session_managers.append(session_manager) + return session_manager + + def _is_primary_workflow_per_user(self) -> bool: + """Return True when the configured workflow is registered as per-user.""" + from nat.cli.type_registry import GlobalTypeRegistry + + workflow_registration = GlobalTypeRegistry.get().get_function(type(self.full_config.workflow)) + return workflow_registration.is_per_user + + def _resolve_workflow_tool_name(self) -> str: + """Return the MCP tool name for the configured workflow entry point.""" + workflow_config = self.full_config.workflow + alias = getattr(workflow_config, "workflow_alias", None) + return alias if alias else workflow_config.type + + async def cleanup(self) -> None: + """Shut down all tracked session managers. + + Attempts shutdown for every manager even when an individual call fails. + Clears the tracked list before re-raising the first shutdown error. + """ + errors: list[BaseException] = [] + try: + for session_manager in self._session_managers: + try: + await session_manager.shutdown() + except Exception as exc: + logger.exception("Failed to shut down SessionManager during MCP worker cleanup") + errors.append(exc) + finally: + self._session_managers.clear() + if errors: + raise errors[0] def _setup_health_endpoint(self, mcp: FastMCP): """Set up the HTTP health endpoint that exercises MCP ping handler.""" @@ -129,7 +167,8 @@ async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder): This method: - Sets up the health endpoint - - Builds the workflow and extracts all functions + - Creates session managers via SessionManager.create (shared or per-user) + - Extracts functions from the shared workflow when applicable - Filters functions based on tool_names config - Registers each function as an MCP tool - Sets up debug endpoints for tool introspection @@ -143,8 +182,17 @@ async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder): # Set up the health endpoint self._setup_health_endpoint(mcp) - # Build the default workflow - workflow = await builder.build() + if self._is_primary_workflow_per_user(): + session_manager = self._track_session_manager(await SessionManager.create(config=self.full_config, + shared_builder=builder)) + tool_name = self._resolve_workflow_tool_name() + register_function_with_mcp(mcp, tool_name, session_manager, self.memory_profiler) + self._setup_debug_endpoints(mcp, {}) + return + + primary_session_manager = self._track_session_manager(await SessionManager.create(config=self.full_config, + shared_builder=builder)) + workflow = primary_session_manager.workflow # Get all functions from the workflow functions = await self._get_all_functions(workflow) @@ -170,17 +218,12 @@ async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder): session_managers: dict[str, SessionManager] = {} for function_name, function in functions.items(): if isinstance(function, Workflow): - # Already a workflow, use it directly logger.info("Function %s is a Workflow, using directly", function_name) - session_managers[function_name] = await SessionManager.create(config=self.full_config, - shared_builder=builder, - entry_function=None) + session_managers[function_name] = primary_session_manager else: - # Regular function - build a workflow with this function as entry point logger.info("Function %s is a regular function, building entry workflow", function_name) - session_managers[function_name] = await SessionManager.create(config=self.full_config, - shared_builder=builder, - entry_function=function_name) + session_managers[function_name] = self._track_session_manager(await SessionManager.create( + config=self.full_config, shared_builder=builder, entry_function=function_name)) # Register each function with MCP, passing SessionManager for observability for function_name, session_manager in session_managers.items(): @@ -194,9 +237,7 @@ async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder): if not session_managers: raise RuntimeError("No functions found in workflow. Please check your configuration.") - # After registration, expose debug endpoints for tool/schema inspection - # Extract the entry functions from session managers for debug endpoints - debug_functions = {name: sm.workflow for name, sm in session_managers.items()} + debug_functions = {name: sm.workflow for name, sm in session_managers.items() if not sm.is_workflow_per_user} self._setup_debug_endpoints(mcp, debug_functions) async def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]: diff --git a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/tool_converter.py b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/tool_converter.py index 6313eee346..ff19a15040 100644 --- a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/tool_converter.py +++ b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/server/tool_converter.py @@ -29,6 +29,7 @@ from pydantic_core import PydanticUndefined from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp.server import Context from nat.builder.function import Function from nat.builder.function_base import FunctionBase @@ -38,6 +39,9 @@ logger = logging.getLogger(__name__) +# Reserved wrapper parameter for MCP request Context injection. +INJECTED_CONTEXT_PARAM = "_nat_mcp_context" + _USE_PYDANTIC_DEFAULT = object() @@ -128,6 +132,18 @@ def _build_name_mapping(field_names: list[str]) -> dict[str, str]: return name_map +def _validate_input_schema_for_context_injection(schema: type[BaseModel], name_map: dict[str, str]) -> None: + """Reject schemas whose fields collide with the injected `Context` parameter.""" + if INJECTED_CONTEXT_PARAM in name_map.values(): + raise ValueError( + f"Workflow input schema cannot declare a field that maps to reserved MCP parameter " + f"{INJECTED_CONTEXT_PARAM!r}.", ) + if "ctx" in schema.model_fields: + raise ValueError( + "Workflow input schema cannot declare a field named 'ctx'; that name is reserved for " + "MCP request context injection.", ) + + def is_field_optional(field: FieldInfo) -> tuple[bool, Any]: """Determine if a Pydantic field is optional and extract its default value for MCP signatures. @@ -164,6 +180,39 @@ def is_field_optional(field: FieldInfo) -> tuple[bool, Any]: return True, Parameter.empty +async def _run_through_session_manager(session_manager: 'SessionManager', + payload: BaseModel, + ctx: Any | None = None, + memory_profiler: 'MemoryProfiler | None' = None) -> Any: + """Execute a payload through SessionManager, including per-user workflows.""" + try: + if session_manager.is_workflow_per_user: + from nat.builder.context import Context + from nat.runtime.user_manager import UserManager + + user_id = Context.get().user_id + http_connection = None + if ctx is not None: + try: + http_connection = ctx.request_context.request + if user_id is None and http_connection is not None: + user_info = UserManager.extract_user_from_connection(http_connection) + if user_info is not None: + user_id = user_info.get_user_id() + except ValueError: + pass + + async with session_manager.session(user_id=user_id, http_connection=http_connection) as session: + async with session.run(payload) as runner: + return await runner.result() + + async with session_manager.run(payload) as runner: + return await runner.result() + finally: + if memory_profiler: + memory_profiler.on_request_complete() + + def create_function_wrapper( function_name: str, session_manager: 'SessionManager', @@ -208,6 +257,7 @@ def create_function_wrapper( # emit for each declared schema field. name_map = _build_name_mapping(list(param_fields.keys())) argument_name_map = {safe: orig for orig, safe in name_map.items() if orig != safe} + _validate_input_schema_for_context_injection(schema, name_map) parameters = [] for name, field in param_fields.items(): @@ -237,34 +287,30 @@ def create_function_wrapper( annotation=field_type, )) - # Create the function signature WITHOUT the ctx parameter - # We'll handle this in the wrapper function internally + # Create the function signature WITHOUT the injected Context parameter + # Context is declared on the wrapper for FastMCP injection but omitted from __signature__ + # so it is excluded from the client-facing tool schema. sig = Signature(parameters=parameters, return_annotation=str) - # Define the actual wrapper function that accepts ctx but doesn't expose it + # Define the actual wrapper function that accepts injected context but doesn't expose it def create_wrapper(): - async def wrapper_with_ctx(**kwargs): + async def wrapper_with_ctx(_nat_mcp_context: Context | None = None, **kwargs): """Internal wrapper that will be called by MCP. Uses SessionManager.run() which creates a Runner that automatically handles observability. """ - # MCP will add a ctx parameter, extract it - ctx = kwargs.get("ctx") - - # Remove ctx if present - if "ctx" in kwargs: - del kwargs["ctx"] - # FastMCP applies the wrapper field's factory to omitted arguments. # Remove its marker so the declared schema can apply the original # default factory and preserve model_fields_set semantics. kwargs = {k: v for k, v in kwargs.items() if v is not _USE_PYDANTIC_DEFAULT} # Process the function call - if ctx: - ctx.info("Calling function %s with args: %s", function_name, json.dumps(kwargs, default=str)) - await ctx.report_progress(0, 100) + if _nat_mcp_context: + _nat_mcp_context.info("Calling function %s with args: %s", + function_name, + json.dumps(kwargs, default=str)) + await _nat_mcp_context.report_progress(0, 100) try: # Prepare input payload @@ -287,16 +333,14 @@ async def wrapper_with_ctx(**kwargs): # 3. Execute the function/workflow # 4. Emit WORKFLOW_END/FUNCTION_END events # 5. Stop the exporter manager - async with session_manager.run(payload) as runner: - result = await runner.result() + result = await _run_through_session_manager(session_manager, + payload, + ctx=_nat_mcp_context, + memory_profiler=memory_profiler) # Report completion - if ctx: - await ctx.report_progress(100, 100) - - # Track request completion for memory profiling - if memory_profiler: - memory_profiler.on_request_complete() + if _nat_mcp_context: + await _nat_mcp_context.report_progress(100, 100) # Handle different result types for proper formatting if isinstance(result, str): @@ -305,12 +349,8 @@ async def wrapper_with_ctx(**kwargs): return json.dumps(result, default=str) return str(result) except Exception as e: - if ctx: - ctx.error("Error calling function %s: %s", function_name, str(e)) - - # Track request completion even on error - if memory_profiler: - memory_profiler.on_request_complete() + if _nat_mcp_context: + _nat_mcp_context.error("Error calling function %s: %s", function_name, str(e)) raise @@ -319,7 +359,7 @@ async def wrapper_with_ctx(**kwargs): # Create the wrapper function wrapper = create_wrapper() - # Set the signature on the wrapper function (WITHOUT ctx) + # Set the signature on the wrapper function (WITHOUT injected context) wrapper.__signature__ = sig # type: ignore wrapper.__name__ = function_name @@ -388,19 +428,18 @@ def register_function_with_mcp(mcp: FastMCP, """ logger.info("Registering function %s with MCP", function_name) - # Get the workflow from the session manager - workflow = session_manager.workflow - - # Prefer the function's schema/description when available, fall back to workflow - target_function = function or workflow + if session_manager.is_workflow_per_user: + input_schema = session_manager.get_workflow_input_schema() + workflow_config = session_manager.config.workflow + function_description = getattr(workflow_config, "description", None) or function_name + else: + workflow = session_manager.workflow + target_function = function or workflow + input_schema = getattr(target_function, "input_schema", workflow.input_schema) + function_description = get_function_description(target_function) - # Get the input schema from the most specific object available - input_schema = getattr(target_function, "input_schema", workflow.input_schema) logger.info("Function %s has input schema: %s", function_name, input_schema) - # Get function description - function_description = get_function_description(target_function) - # Create and register the wrapper function with MCP wrapper_func = create_function_wrapper(function_name, session_manager, input_schema, memory_profiler) mcp.tool(name=function_name, description=function_description)(wrapper_func) diff --git a/packages/nvidia_nat_mcp/tests/server/test_mcp_front_end_plugin.py b/packages/nvidia_nat_mcp/tests/server/test_mcp_front_end_plugin.py index 24536e59a5..be07aaa4f9 100644 --- a/packages/nvidia_nat_mcp/tests/server/test_mcp_front_end_plugin.py +++ b/packages/nvidia_nat_mcp/tests/server/test_mcp_front_end_plugin.py @@ -230,7 +230,7 @@ async def test_workflow_alias_with_function_groups(): async def test_session_manager_creation_for_workflow_vs_function(): - """Test that SessionManager.create is called with correct entry_function for workflows vs regular functions.""" + """Test SessionManager.create usage for shared workflows vs entry functions.""" from unittest.mock import AsyncMock from unittest.mock import MagicMock from unittest.mock import patch @@ -266,6 +266,7 @@ async def test_session_manager_creation_for_workflow_vs_function(): # Configure the mock to return a mock SessionManager mock_session_manager = MagicMock() mock_session_manager.workflow = mock_workflow + mock_session_manager.is_workflow_per_user = False mock_session_create.return_value = mock_session_manager # Patch register_function_with_mcp to avoid actual registration @@ -273,28 +274,11 @@ async def test_session_manager_creation_for_workflow_vs_function(): # Call the method we're testing await worker._default_add_routes(mock_mcp, mock_builder) - # Verify SessionManager.create was called twice (once for each function) + # Primary manager builds the shared workflow; regular functions get their own entry manager. assert mock_session_create.call_count == 2 - # Extract the calls - calls = mock_session_create.call_args_list + primary_call = mock_session_create.call_args_list[0] + function_call = mock_session_create.call_args_list[1] - # Find the call for the workflow and the call for the regular function - workflow_call = None - function_call = None - - for call in calls: - # Check the entry_function parameter - entry_function = call.kwargs.get('entry_function') - if entry_function is None: - workflow_call = call - else: - function_call = call - - # Verify workflow call used entry_function=None - assert workflow_call is not None, "Workflow should use entry_function=None" - assert workflow_call.kwargs['entry_function'] is None - - # Verify regular function call used entry_function=function_name - assert function_call is not None, "Function should use entry_function=" - assert function_call.kwargs['entry_function'] == "echo_function" + assert primary_call.kwargs.get('entry_function') is None + assert function_call.kwargs.get('entry_function') == "echo_function" diff --git a/packages/nvidia_nat_mcp/tests/server/test_per_user_workflow.py b/packages/nvidia_nat_mcp/tests/server/test_per_user_workflow.py new file mode 100644 index 0000000000..9ea6273240 --- /dev/null +++ b/packages/nvidia_nat_mcp/tests/server/test_per_user_workflow.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +import pytest +from pydantic import BaseModel + +from nat.builder.builder import Builder +from nat.builder.function_info import FunctionInfo +from nat.builder.workflow_builder import WorkflowBuilder +from nat.cli.register_workflow import register_function +from nat.cli.register_workflow import register_per_user_function +from nat.data_models.config import Config +from nat.data_models.config import GeneralConfig +from nat.data_models.function import FunctionBaseConfig +from nat.plugins.mcp.server.front_end_config import MCPFrontEndConfig +from nat.plugins.mcp.server.front_end_plugin import MCPFrontEndPlugin +from nat.plugins.mcp.server.front_end_plugin_worker import MCPFrontEndPluginWorker +from nat.runtime.session import SessionManager + + +class _Input(BaseModel): + message: str + + +class _Output(BaseModel): + result: str + + +class PerUserMCPWorkflowConfig(FunctionBaseConfig, name="per_user_mcp_test_workflow"): + """Per-user workflow config for MCP front-end tests.""" + + +class SharedMCPWorkflowConfig(FunctionBaseConfig, name="shared_mcp_test_workflow"): + """Shared workflow config for MCP front-end tests.""" + + +@pytest.fixture(name="registered_workflows", scope="module") +def fixture_registered_workflows(): + """Register test workflows in a pushed registry so they do not leak.""" + + @register_per_user_function(config_type=PerUserMCPWorkflowConfig, input_type=_Input, single_output_type=_Output) + async def _build_per_user(_config: PerUserMCPWorkflowConfig, _builder: Builder): + + async def _impl(inp: _Input) -> _Output: + return _Output(result=f"per-user: {inp.message}") + + yield FunctionInfo.from_fn(_impl) + + @register_function(config_type=SharedMCPWorkflowConfig) + async def _build_shared(_config: SharedMCPWorkflowConfig, _builder: Builder): + + async def _impl(inp: _Input) -> _Output: + return _Output(result=f"shared: {inp.message}") + + yield FunctionInfo.from_fn(_impl) + + +def _config(workflow) -> Config: + return Config( + general=GeneralConfig(front_end=MCPFrontEndConfig( + name="Test MCP Server", + host="localhost", + port=9902, + debug=False, + log_level="INFO", + )), + workflow=workflow, + ) + + +@pytest.fixture(name="per_user_config") +def fixture_per_user_config(registered_workflows) -> Config: + return _config(PerUserMCPWorkflowConfig()) + + +@pytest.fixture(name="shared_config") +def fixture_shared_config(registered_workflows) -> Config: + return _config(SharedMCPWorkflowConfig()) + + +class TestPerUserWorkflowStartup: + """The MCP front end must serve per-user workflows, not die building them.""" + + async def test_per_user_workflow_server_starts(self, per_user_config, monkeypatch): + """Startup used to raise "Must set a workflow before building".""" + + async def _no_serve(_self): + return None + + monkeypatch.setattr("mcp.server.fastmcp.server.FastMCP.run_streamable_http_async", _no_serve) + + await MCPFrontEndPlugin(full_config=per_user_config).run() + + async def test_shared_workflow_still_builds(self, shared_config, monkeypatch): + captured = {} + + original_create = SessionManager.create + + async def _capture_create(*args, **kwargs): + session_manager = await original_create(*args, **kwargs) + if not session_manager.is_workflow_per_user: + captured["workflow"] = session_manager.workflow + return session_manager + + async def _no_serve(_self): + return None + + monkeypatch.setattr(SessionManager, "create", _capture_create) + monkeypatch.setattr("mcp.server.fastmcp.server.FastMCP.run_streamable_http_async", _no_serve) + + await MCPFrontEndPlugin(full_config=shared_config).run() + + assert captured["workflow"] is not None + + async def test_per_user_session_manager_reaps_and_shuts_down(self, per_user_config): + worker = MCPFrontEndPluginWorker(per_user_config) + + async with WorkflowBuilder.from_config(config=per_user_config) as builder: + mcp = await worker.create_mcp_server() + await worker._default_add_routes(mcp, builder) + + assert len(worker._session_managers) == 1 + session_manager = worker._session_managers[0] + assert session_manager.is_workflow_per_user + assert session_manager._per_user_builders_cleanup_task is not None + + cleanup_task = session_manager._per_user_builders_cleanup_task + await worker.cleanup() + assert cleanup_task.done() + + async def test_register_function_skips_shared_workflow_lookup(self, per_user_config, monkeypatch): + from nat.plugins.mcp.server import tool_converter + + worker = MCPFrontEndPluginWorker(per_user_config) + + async with WorkflowBuilder.from_config(config=per_user_config) as builder: + session_manager = await SessionManager.create(config=per_user_config, shared_builder=builder) + mcp = await worker.create_mcp_server() + + get_schema = MagicMock(return_value=_Input) + monkeypatch.setattr(session_manager, "get_workflow_input_schema", get_schema) + + tool_converter.register_function_with_mcp(mcp, "per_user_mcp_test_workflow", session_manager) + + get_schema.assert_called_once() + + +class TestWorkerCleanup: + """Worker cleanup must shut down every manager and clear tracking.""" + + async def test_cleanup_shuts_down_all_managers_after_shutdown_failure(self, per_user_config): + worker = MCPFrontEndPluginWorker(per_user_config) + failing = MagicMock() + failing.shutdown = AsyncMock(side_effect=RuntimeError("boom")) + succeeding = MagicMock() + succeeding.shutdown = AsyncMock() + worker._session_managers = [failing, succeeding] + + with pytest.raises(RuntimeError, match="boom"): + await worker.cleanup() + + failing.shutdown.assert_awaited_once() + succeeding.shutdown.assert_awaited_once() + assert worker._session_managers == [] + + +class TestPerUserRequestIdentity: + """Per-user tool calls must reach session() with a resolved user id.""" + + async def test_run_through_session_manager_uses_context_user_id(self, monkeypatch): + from nat.plugins.mcp.server.tool_converter import _run_through_session_manager + + session_manager = MagicMock() + session_manager.is_workflow_per_user = True + session_manager.session = MagicMock() + session = MagicMock() + runner = MagicMock() + runner.result = AsyncMock(return_value="ok") + session.run.return_value.__aenter__ = AsyncMock(return_value=runner) + session.run.return_value.__aexit__ = AsyncMock(return_value=False) + session_manager.session.return_value.__aenter__ = AsyncMock(return_value=session) + session_manager.session.return_value.__aexit__ = AsyncMock(return_value=False) + + context = SimpleNamespace(user_id="alice") + monkeypatch.setattr("nat.builder.context.Context.get", lambda: context) + + payload = _Input(message="hello") + result = await _run_through_session_manager(session_manager, payload) + + assert result == "ok" + session_manager.session.assert_called_once_with(user_id="alice", http_connection=None) + + async def test_run_through_session_manager_resolves_user_from_mcp_request(self, monkeypatch): + from nat.plugins.mcp.server.tool_converter import _run_through_session_manager + + session_manager = MagicMock() + session_manager.is_workflow_per_user = True + session_manager.session = MagicMock() + session = MagicMock() + runner = MagicMock() + runner.result = AsyncMock(return_value="ok") + session.run.return_value.__aenter__ = AsyncMock(return_value=runner) + session.run.return_value.__aexit__ = AsyncMock(return_value=False) + session_manager.session.return_value.__aenter__ = AsyncMock(return_value=session) + session_manager.session.return_value.__aexit__ = AsyncMock(return_value=False) + + context = SimpleNamespace(user_id=None) + monkeypatch.setattr("nat.builder.context.Context.get", lambda: context) + + request = MagicMock() + ctx = SimpleNamespace(request_context=SimpleNamespace(request=request)) + user_info = MagicMock() + user_info.get_user_id.return_value = "bob" + monkeypatch.setattr( + "nat.runtime.user_manager.UserManager.extract_user_from_connection", + MagicMock(return_value=user_info), + ) + + payload = _Input(message="hello") + result = await _run_through_session_manager(session_manager, payload, ctx=ctx) + + assert result == "ok" + session_manager.session.assert_called_once_with(user_id="bob", http_connection=request) diff --git a/packages/nvidia_nat_mcp/tests/server/test_tool_converter.py b/packages/nvidia_nat_mcp/tests/server/test_tool_converter.py index d0a2dff4b1..dfe2c5a6e8 100644 --- a/packages/nvidia_nat_mcp/tests/server/test_tool_converter.py +++ b/packages/nvidia_nat_mcp/tests/server/test_tool_converter.py @@ -30,6 +30,7 @@ from nat.builder.function import Function from nat.builder.workflow import Workflow +from nat.plugins.mcp.server.tool_converter import INJECTED_CONTEXT_PARAM from nat.plugins.mcp.server.tool_converter import _build_name_mapping from nat.plugins.mcp.server.tool_converter import _sanitize_parameter_name from nat.plugins.mcp.server.tool_converter import create_function_wrapper @@ -120,6 +121,7 @@ def create_mock_session_manager(workflow=None, result_value="result"): result_value: The value to return from runner.result() """ mock_session_manager = MagicMock(spec=SessionManager) + mock_session_manager.is_workflow_per_user = False if workflow is None: workflow = create_mock_workflow_with_observability() @@ -290,6 +292,27 @@ def test_create_wrapper_for_regular_function(self): assert "name" in sig.parameters assert "age" in sig.parameters + def test_create_wrapper_declares_context_for_injection(self): + """FastMCP injects request context only when a Context parameter is annotated.""" + from mcp.server.fastmcp.utilities.context_injection import find_context_parameter + + mock_session_manager = create_mock_session_manager() + wrapper = create_function_wrapper("regular_function", mock_session_manager, MockRegularSchema) + + assert find_context_parameter(wrapper) == INJECTED_CONTEXT_PARAM + sig = getattr(wrapper, "__signature__", None) + assert sig is not None + assert INJECTED_CONTEXT_PARAM not in sig.parameters + + def test_create_wrapper_rejects_schema_field_named_ctx(self): + """A workflow field named `ctx` collides with MCP context injection.""" + from pydantic import create_model + + schema = create_model("CtxSchema", **{"ctx": (str, ...)}) # type: ignore[call-overload] + + with pytest.raises(ValueError, match="cannot declare a field named 'ctx'"): + create_function_wrapper("tool", create_mock_session_manager(), schema) + def test_create_wrapper_for_workflow(self): """Test creating wrapper for workflow function.""" # Arrange @@ -407,6 +430,7 @@ def test_register_function_with_mcp_uses_function_metadata(self, mock_logger, mo mock_function = MagicMock(spec=Function) mock_function.input_schema = "function_schema" mock_session_manager = MagicMock(spec=SessionManager) + mock_session_manager.is_workflow_per_user = False mock_session_manager.workflow = mock_workflow function_name = "test_function" @@ -436,6 +460,7 @@ def test_register_workflow_with_mcp_falls_back_to_workflow(self, mock_logger, mo mock_workflow = MagicMock(spec=Workflow) mock_workflow.input_schema = "workflow_schema" mock_session_manager = MagicMock(spec=SessionManager) + mock_session_manager.is_workflow_per_user = False mock_session_manager.workflow = mock_workflow function_name = "test_workflow" @@ -699,6 +724,7 @@ async def test_error_handling_in_wrapper_execution(self): # Arrange mock_workflow = create_mock_workflow_with_observability() mock_session_manager = MagicMock(spec=SessionManager) + mock_session_manager.is_workflow_per_user = False mock_session_manager.workflow = mock_workflow # Create mock runner that raises an error