nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/front_ends/fastapi/step_adaptor.py

@@ -0,0 +1,319 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+
+import html
+import logging
+from functools import reduce
+from textwrap import dedent
+
+from nat.data_models.api_server import ResponseIntermediateStep
+from nat.data_models.api_server import ResponseSerializable
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepCategory
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.invocation_node import InvocationNode
+from nat.data_models.step_adaptor import StepAdaptorConfig
+from nat.data_models.step_adaptor import StepAdaptorMode
+from nat.utils.type_utils import is_valid_json
+
+logger = logging.getLogger(__name__)
+
+
+class StepAdaptor:
+
+    def __init__(self, config: StepAdaptorConfig):
+
+        self._history: list[IntermediateStep] = []
+        self.config = config
+
+    def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool:
+        """
+        Returns True if this intermediate step should be included (based on the config.mode).
+        """
+
+        if config.mode == StepAdaptorMode.OFF:
+            return False
+
+        if config.mode == StepAdaptorMode.DEFAULT:
+            # default existing behavior: show LLM events + TOOL_END + FUNCTION events
+            if step.event_category == IntermediateStepCategory.LLM:
+                return True
+            if step.event_category == IntermediateStepCategory.TOOL:
+                return True
+            if step.event_category == IntermediateStepCategory.FUNCTION:
+                return True
+            return False
+
+        if config.mode == StepAdaptorMode.CUSTOM:
+            # pass only what the user explicitly listed
+            return step.event_type in config.custom_event_types
+
+        return False
+
+    def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
+        input_str: str | None = None
+        output_str: str | None = None
+
+        # Find the start in the history with matching run_id
+        start_step = next(
+            (x for x in self._history if x.event_type == IntermediateStepType.LLM_START and x.UUID == step.UUID), None)
+
+        if not start_step:
+            # If we don't have a start step, we can't do anything
+            return None
+
+        input_str = str(start_step.data.input)
+
+        if step.event_type == IntermediateStepType.LLM_NEW_TOKEN:
+
+            # Find all of the previous LLM chunks and concatenate them
+            output_str = reduce(
+                lambda x, y: x + y,
+                (str(x.data.chunk)
+                 for x in self._history if x.event_type == IntermediateStepType.LLM_NEW_TOKEN and x.UUID == step.UUID),
+                "")
+
+        elif step.event_type == IntermediateStepType.LLM_END:
+            output_str = str(step.data.output)
+
+        if not input_str and not output_str:
+            return None
+
+        escaped_input = html.escape(input_str, quote=False)
+
+        # Dont use f-strings here because the payload is markdown and screws up the dedent
+        payload = dedent("""
+        **Input:**
+        ```python
+        {input_value}
+        ```
+        """).strip("\n").format(input_value=escaped_input)
+
+        if (output_str):
+            escaped_output = html.escape(output_str, quote=False) if output_str else ""
+
+            # Dont use f-strings here because the payload is markdown and screws up the dedent
+            payload = dedent("""
+            {payload}
+
+            **Output:**
+            {output_value}
+            """).strip("\n").format(payload=payload, output_value=escaped_output)
+
+        event = ResponseIntermediateStep(id=step.UUID,
+                                         name=step.name or "",
+                                         payload=payload,
+                                         parent_id=ancestry.function_id)
+
+        return event
+
+    def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
+        """
+        Handles both TOOL_START and TOOL_END events
+        """
+        input_str: str | None = None
+        output_str: str | None = None
+
+        # Find the start in the history with matching run_id
+        start_step = next(
+            (x for x in self._history if x.event_type == IntermediateStepType.TOOL_START and x.UUID == step.UUID), None)
+
+        if not start_step:
+            # If we don't have a start step, we can't do anything
+            return None
+
+        input_str = str(start_step.data.input)
+
+        if step.event_type == IntermediateStepType.TOOL_END:
+            output_str = str(step.data.output)
+
+        if not input_str and not output_str:
+            return None
+
+        escaped_input = html.escape(input_str, quote=False)
+        format_input_type = "json" if is_valid_json(escaped_input) else "python"
+
+        # Dont use f-strings here because the payload is markdown and screws up the dedent
+        payload = dedent("""
+        **Input:**
+        ```{format_input_type}
+        {input_value}
+        ```
+        """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
+
+        if output_str:
+            escaped_output = html.escape(output_str, quote=False)
+            format_output_type = "json" if is_valid_json(escaped_output) else "python"
+
+            # Dont use f-strings here because the payload is markdown and screws up the dedent
+            payload = dedent("""
+            {payload}
+
+            **Output:**
+            ```{format_output_type}
+            {output_value}
+            ```
+            """).strip("\n").format(payload=payload, output_value=escaped_output, format_output_type=format_output_type)
+
+        event = ResponseIntermediateStep(id=step.UUID,
+                                         name=f"Tool: {step.name}",
+                                         payload=payload,
+                                         parent_id=ancestry.function_id)
+
+        return event
+
+    def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
+        """
+        Handles the FUNCTION_START and FUNCTION_END events
+        """
+        input_str: str | None = None
+        output_str: str | None = None
+
+        if step.event_type == IntermediateStepType.FUNCTION_START:
+            # For function start events, display input data
+            if step.data and hasattr(step.data, 'input'):
+                input_str = str(step.data.input)
+            elif step.data:
+                input_str = str(step.data)
+
+            if not input_str:
+                return None
+
+            escaped_input = html.escape(input_str, quote=False)
+            format_input_type = "json" if is_valid_json(escaped_input) else "python"
+
+            # Create payload for function start
+            payload_str = dedent("""
+            **Function Input:**
+            ```{format_input_type}
+            {input_value}
+            ```
+            """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
+
+            event = ResponseIntermediateStep(id=step.UUID,
+                                             name=f"Function Start: {step.name}",
+                                             payload=payload_str,
+                                             parent_id=ancestry.parent_id)
+            return event
+
+        if step.event_type == IntermediateStepType.FUNCTION_END:
+            # Find the start event with matching UUID
+            start_step = next(
+                (x
+                 for x in self._history if x.event_type == IntermediateStepType.FUNCTION_START and x.UUID == step.UUID),
+                None)
+
+            # For function end events, display output data
+            if step.data and hasattr(step.data, 'output'):
+                output_str = str(step.data.output)
+            elif step.data:
+                output_str = str(step.data)
+
+            if not output_str:
+                return None
+
+            escaped_output = html.escape(output_str, quote=False)
+            format_output_type = "json" if is_valid_json(escaped_output) else "python"
+
+            # Get input from start step if available
+            input_payload = ""
+            if start_step and start_step.data:
+                if hasattr(start_step.data, 'input'):
+                    input_str = str(start_step.data.input)
+                else:
+                    input_str = str(start_step.data)
+
+                if input_str:
+                    escaped_input = html.escape(input_str, quote=False)
+                    format_input_type = "json" if is_valid_json(escaped_input) else "python"
+                    input_payload = dedent("""
+                    **Function Input:**
+                    ```{format_input_type}
+                    {input_value}
+                    ```
+                    """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
+
+            # Create payload for function end
+            payload_str = dedent("""
+            {input_payload}**Function Output:**
+            ```{format_output_type}
+            {output_value}
+            ```
+            """).strip("\n").format(input_payload=input_payload,
+                                    output_value=escaped_output,
+                                    format_output_type=format_output_type)
+
+            event = ResponseIntermediateStep(id=step.UUID,
+                                             name=f"Function Complete: {step.name}",
+                                             payload=payload_str,
+                                             parent_id=ancestry.parent_id)
+            return event
+
+        return None
+
+    def _handle_custom(self, payload: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
+        """
+        Handles the CUSTOM event
+        """
+        escaped_payload = html.escape(str(payload), quote=False)
+        escaped_payload = escaped_payload.replace("\n", "")
+
+        # Attempt to determine type
+        format_type = "json" if is_valid_json(escaped_payload) else "python"
+
+        # Don't use f-strings here because the payload is markdown and screws up the dedent
+        payload_str = dedent("""
+        ```{format_type}
+        {payload}
+        ```
+        """).strip("\n").format(payload=escaped_payload, format_type=format_type)
+
+        # Return the event
+        event = ResponseIntermediateStep(id=payload.UUID,
+                                         name=f"{payload.event_type}",
+                                         payload=payload_str,
+                                         parent_id=ancestry.function_id)
+
+        return event
+
+    def process(self, step: IntermediateStep) -> ResponseSerializable | None:
+
+        # Track the chunk
+        self._history.append(step)
+        payload = step.payload
+        ancestry = step.function_ancestry
+
+        if not self._step_matches_filter(step, self.config):
+            return None
+
+        try:
+
+            if step.event_category == IntermediateStepCategory.LLM:
+                return self._handle_llm(payload, ancestry)
+
+            if step.event_category == IntermediateStepCategory.TOOL:
+                return self._handle_tool(payload, ancestry)
+
+            if step.event_category == IntermediateStepCategory.FUNCTION:
+                return self._handle_function(payload, ancestry)
+
+            if step.event_category == IntermediateStepCategory.CUSTOM:
+                return self._handle_custom(payload, ancestry)
+
+        except Exception as e:
+            logger.exception("Error processing intermediate step: %s", e)
+
+        return None

nat/front_ends/fastapi/utils.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+
+import importlib
+import os
+
+
+def get_config_file_path() -> str:
+    """
+    Get the path to the NAT configuration file from the environment variable NAT_CONFIG_FILE.
+    Raises ValueError if the environment variable is not set.
+    """
+    config_file_path = os.getenv("NAT_CONFIG_FILE")
+    if (not config_file_path):
+        raise ValueError("Config file not found in environment variable NAT_CONFIG_FILE.")
+
+    return os.path.abspath(config_file_path)
+
+
+def import_class_from_string(class_full_name: str) -> type:
+    """
+    Import a class from a string in the format 'module.submodule.ClassName'.
+    Raises ImportError if the class cannot be imported.
+    """
+    try:
+        class_name_parts = class_full_name.split(".")
+
+        module_name = ".".join(class_name_parts[:-1])
+        class_name = class_name_parts[-1]
+
+        module = importlib.import_module(module_name)
+
+        if not hasattr(module, class_name):
+            raise ValueError(f"Class '{class_full_name}' not found.")
+
+        return getattr(module, class_name)
+    except (ImportError, AttributeError) as e:
+        raise ImportError(f"Could not import {class_full_name}.") from e
+
+
+def get_class_name(cls: type) -> str:
+    """
+    Get the full class name including the module.
+    """
+    return f"{cls.__module__}.{cls.__qualname__}"

nat/front_ends/mcp/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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.

nat/front_ends/mcp/introspection_token_verifier.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+"""OAuth 2.0 Token Introspection verifier implementation for MCP servers."""
+
+import logging
+
+from mcp.server.auth.provider import AccessToken
+from mcp.server.auth.provider import TokenVerifier
+
+from nat.authentication.credential_validator.bearer_token_validator import BearerTokenValidator
+from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class IntrospectionTokenVerifier(TokenVerifier):
+    """Token verifier that delegates token verification to BearerTokenValidator."""
+
+    def __init__(self, config: OAuth2ResourceServerConfig):
+        """Create IntrospectionTokenVerifier from OAuth2ResourceServerConfig.
+
+        Args:
+            config: OAuth2ResourceServerConfig
+        """
+        issuer = config.issuer_url
+        scopes = config.scopes or []
+        audience = config.audience
+        jwks_uri = config.jwks_uri
+        introspection_endpoint = config.introspection_endpoint
+        discovery_url = config.discovery_url
+        client_id = config.client_id
+        client_secret = config.client_secret
+
+        self._bearer_token_validator = BearerTokenValidator(
+            issuer=issuer,
+            audience=audience,
+            scopes=scopes,
+            jwks_uri=jwks_uri,
+            introspection_endpoint=introspection_endpoint,
+            discovery_url=discovery_url,
+            client_id=client_id,
+            client_secret=client_secret,
+        )
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        """Verify token by delegating to BearerTokenValidator.
+
+        Args:
+            token: The Bearer token to verify
+
+        Returns:
+            AccessToken | None: AccessToken if valid, None if invalid
+        """
+        validation_result = await self._bearer_token_validator.verify(token)
+
+        if validation_result.active:
+            return AccessToken(token=token,
+                               expires_at=validation_result.expires_at,
+                               scopes=validation_result.scopes or [],
+                               client_id=validation_result.client_id or "")
+        return None

nat/front_ends/mcp/mcp_front_end_config.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+
+import logging
+from typing import Literal
+
+from pydantic import Field
+from pydantic import model_validator
+
+from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
+from nat.data_models.front_end import FrontEndBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
+    """MCP front end configuration.
+
+    A simple MCP (Model Context Protocol) front end for NeMo Agent toolkit.
+    """
+
+    name: str = Field(default="NeMo Agent Toolkit MCP",
+                      description="Name of the MCP server (default: NeMo Agent Toolkit MCP)")
+    host: str = Field(default="localhost", description="Host to bind the server to (default: localhost)")
+    port: int = Field(default=9901, description="Port to bind the server to (default: 9901)", ge=0, le=65535)
+    debug: bool = Field(default=False, description="Enable debug mode (default: False)")
+    log_level: str = Field(default="INFO", description="Log level for the MCP server (default: INFO)")
+    tool_names: list[str] = Field(default_factory=list,
+                                  description="The list of tools MCP server will expose (default: all tools)")
+    transport: Literal["sse", "streamable-http"] = Field(
+        default="streamable-http",
+        description="Transport type for the MCP server (default: streamable-http, backwards compatible with sse)")
+    runner_class: str | None = Field(
+        default=None, description="Custom worker class for handling MCP routes (default: built-in worker)")
+
+    server_auth: OAuth2ResourceServerConfig | None = Field(
+        default=None, description=("OAuth 2.0 Resource Server configuration for token verification."))
+
+    # Memory profiling configuration
+    enable_memory_profiling: bool = Field(default=False,
+                                          description="Enable memory profiling and diagnostics (default: False)")
+    memory_profile_interval: int = Field(default=50,
+                                         description="Log memory stats every N requests (default: 50)",
+                                         ge=1)
+    memory_profile_top_n: int = Field(default=10,
+                                      description="Number of top memory allocations to log (default: 10)",
+                                      ge=1,
+                                      le=50)
+    memory_profile_log_level: str = Field(default="DEBUG",
+                                          description="Log level for memory profiling output (default: DEBUG)")
+
+    @model_validator(mode="after")
+    def validate_security_configuration(self):
+        """Validate security configuration to prevent accidental misconfigurations."""
+        # Check if server is bound to a non-localhost interface without authentication
+        localhost_hosts = {"localhost", "127.0.0.1", "::1"}
+        if self.host not in localhost_hosts and self.server_auth is None:
+            logger.warning(
+                "MCP server is configured to bind to '%s' without authentication. "
+                "This may expose your server to unauthorized access. "
+                "Consider either: (1) binding to localhost for local-only access, "
+                "or (2) configuring server_auth for production deployments on public interfaces.",
+                self.host)
+
+        # Check if SSE transport is used (which doesn't support authentication)
+        if self.transport == "sse":
+            if self.server_auth is not None:
+                logger.warning("SSE transport does not support authentication. "
+                               "The configured server_auth will be ignored. "
+                               "For production use with authentication, use 'streamable-http' transport instead.")
+            elif self.host not in localhost_hosts:
+                logger.warning(
+                    "SSE transport does not support authentication and is bound to '%s'. "
+                    "This configuration is not recommended for production use. "
+                    "For production deployments, use 'streamable-http' transport with server_auth configured.",
+                    self.host)
+
+        return self

nat/front_ends/mcp/mcp_front_end_plugin.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+
+import logging
+import typing
+
+from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
+from nat.builder.front_end import FrontEndBase
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorkerBase
+
+logger = logging.getLogger(__name__)
+
+
+class MCPFrontEndPlugin(FrontEndBase[MCPFrontEndConfig]):
+    """MCP front end plugin implementation."""
+
+    def get_worker_class(self) -> type[MCPFrontEndPluginWorkerBase]:
+        """Get the worker class for handling MCP routes."""
+        from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorker
+
+        return MCPFrontEndPluginWorker
+
+    @typing.final
+    def get_worker_class_name(self) -> str:
+        """Get the worker class name from configuration or default."""
+        if self.front_end_config.runner_class:
+            return self.front_end_config.runner_class
+
+        worker_class = self.get_worker_class()
+        return f"{worker_class.__module__}.{worker_class.__qualname__}"
+
+    def _get_worker_instance(self) -> MCPFrontEndPluginWorkerBase:
+        """Get an instance of the worker class."""
+        # Import the worker class dynamically if specified in config
+        if self.front_end_config.runner_class:
+            module_name, class_name = self.front_end_config.runner_class.rsplit(".", 1)
+            import importlib
+            module = importlib.import_module(module_name)
+            worker_class = getattr(module, class_name)
+        else:
+            worker_class = self.get_worker_class()
+
+        return worker_class(self.full_config)
+
+    async def _create_token_verifier(self, token_verifier_config: OAuth2ResourceServerConfig):
+        """Create a token verifier based on configuration."""
+        from nat.front_ends.mcp.introspection_token_verifier import IntrospectionTokenVerifier
+
+        if not self.front_end_config.server_auth:
+            return None
+
+        return IntrospectionTokenVerifier(token_verifier_config)
+
+    async def run(self) -> None:
+        """Run the MCP server."""
+        # Import FastMCP
+        from mcp.server.fastmcp import FastMCP
+
+        # Create auth settings and token verifier if auth is required
+        auth_settings = None
+        token_verifier = None
+
+        # Build the workflow and add routes using the worker
+        async with WorkflowBuilder.from_config(config=self.full_config) as builder:
+
+            if self.front_end_config.server_auth:
+                from mcp.server.auth.settings import AuthSettings
+                from pydantic import AnyHttpUrl
+
+                server_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}"
+
+                auth_settings = AuthSettings(issuer_url=AnyHttpUrl(self.front_end_config.server_auth.issuer_url),
+                                             required_scopes=self.front_end_config.server_auth.scopes,
+                                             resource_server_url=AnyHttpUrl(server_url))
+
+                token_verifier = await self._create_token_verifier(self.front_end_config.server_auth)
+
+            # Create an MCP server with the configured parameters
+            mcp = FastMCP(name=self.front_end_config.name,
+                          host=self.front_end_config.host,
+                          port=self.front_end_config.port,
+                          debug=self.front_end_config.debug,
+                          auth=auth_settings,
+                          token_verifier=token_verifier)
+
+            # Get the worker instance and set up routes
+            worker = self._get_worker_instance()
+
+            # 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
+            if self.front_end_config.transport == "sse":
+                logger.info("Starting MCP server with SSE endpoint at /sse")
+                await mcp.run_sse_async()
+            else:  # streamable-http
+                logger.info("Starting MCP server with streamable-http endpoint at /mcp/")
+                await mcp.run_streamable_http_async()