From 95c38c6f4d3dda94c553c4eaaaacf9a874c0789b Mon Sep 17 00:00:00 2001 From: kiwigitops Date: Wed, 27 May 2026 22:01:16 -0400 Subject: [PATCH] fix: support function workflows in mcp front ends Signed-off-by: kiwigitops --- .../fastmcp/server/front_end_plugin_worker.py | 8 +++++--- .../nvidia_nat_fastmcp/tests/test_fastmcp.py | 13 +++++++++++++ .../mcp/server/front_end_plugin_worker.py | 8 +++++--- .../tests/server/test_mcp_front_end_plugin.py | 18 ++++++++++++++++++ 4 files changed, 41 insertions(+), 6 deletions(-) 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 df96b3798d..c18c285761 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 @@ -172,10 +172,12 @@ async def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]: for function_group in workflow.function_groups.values(): functions.update(await function_group.get_accessible_functions()) - if workflow.config.workflow.workflow_alias: - functions[workflow.config.workflow.workflow_alias] = workflow + workflow_config = workflow.config.workflow + workflow_alias = getattr(workflow_config, "workflow_alias", None) + if workflow_alias: + functions[workflow_alias] = workflow else: - functions[workflow.config.workflow.type] = workflow + functions[workflow_config.type] = workflow return functions diff --git a/packages/nvidia_nat_fastmcp/tests/test_fastmcp.py b/packages/nvidia_nat_fastmcp/tests/test_fastmcp.py index 386a7809d4..631777035b 100644 --- a/packages/nvidia_nat_fastmcp/tests/test_fastmcp.py +++ b/packages/nvidia_nat_fastmcp/tests/test_fastmcp.py @@ -15,6 +15,7 @@ """Tests for FastMCP CLI and server wiring.""" from pathlib import Path +from types import SimpleNamespace import pytest from fastmcp import FastMCP @@ -181,6 +182,18 @@ async def test_fastmcp_auth_introspection_exposes_metadata(): assert any(route.path.startswith("/.well-known/oauth-protected-resource") for route in routes) +async def test_fastmcp_function_workflow_config_without_alias_uses_type(): + config = Config(general=GeneralConfig(front_end=FastMCPFrontEndConfig())) + worker = FastMCPFrontEndPluginWorker(config) + workflow = SimpleNamespace(functions={}, + function_groups={}, + config=SimpleNamespace(workflow=SimpleNamespace(type="langgraph_wrapper"))) + + functions = await worker._get_all_functions(workflow) + + assert functions == {"langgraph_wrapper": workflow} + + def test_fastmcp_debug_route_lists_tools(): config = Config(general=GeneralConfig(front_end=FastMCPFrontEndConfig())) worker = FastMCPFrontEndPluginWorker(config) 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 922a015ce9..6ad294d420 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 @@ -215,10 +215,12 @@ async def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]: for function_group in workflow.function_groups.values(): functions.update(await function_group.get_accessible_functions()) - if workflow.config.workflow.workflow_alias: - functions[workflow.config.workflow.workflow_alias] = workflow + workflow_config = workflow.config.workflow + workflow_alias = getattr(workflow_config, "workflow_alias", None) + if workflow_alias: + functions[workflow_alias] = workflow else: - functions[workflow.config.workflow.type] = workflow + functions[workflow_config.type] = workflow return functions 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 a1d822f3c1..ee297eea0a 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 @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace from unittest.mock import MagicMock from unittest.mock import patch @@ -146,6 +147,23 @@ async def test_workflow_alias_usage_in_mcp_front_end(): assert "func1" in functions +async def test_function_workflow_config_without_alias_uses_type(): + """Test function-style workflows that do not define workflow_alias.""" + from nat.data_models.config import Config + from nat.plugins.mcp.server.front_end_plugin_worker import MCPFrontEndPluginWorker + + mock_workflow = SimpleNamespace(functions={}, + function_groups={}, + config=SimpleNamespace(workflow=SimpleNamespace(type="langgraph_wrapper"))) + + config = Config(general=GeneralConfig(front_end=MCPFrontEndConfig()), workflow=EchoFunctionConfig()) + worker = MCPFrontEndPluginWorker(config) + + functions = await worker._get_all_functions(mock_workflow) + + assert functions == {"langgraph_wrapper": mock_workflow} + + async def test_workflow_alias_priority_over_type(): """Test that workflow_alias takes priority over workflow type when both are present.""" from unittest.mock import MagicMock