Move away from global MCP instance (Initial work towards library pattern)

DEVELOPING.md

@@ -13,13 +13,11 @@ make setup
 ### Run using the MCP inspector
 
 ```bash
-uv run fastmcp dev src/mcp_server_datahub/__main__.py --with-editable .
+make mcp-inspector
 ```
 
 In the inspector UI, add environment variables for `DATAHUB_GMS_URL` and `DATAHUB_GMS_TOKEN`, then click Connect.
 
-> **Note:** Use `fastmcp dev` (not `mcp dev`), since this project uses the standalone FastMCP package.
-
 ### Run using an MCP client
 
 Use this configuration in your MCP client e.g. Claude Desktop, Cursor, etc.

Makefile

@@ -1,4 +1,4 @@
-.PHONY: setup clean format format-check lint lint-check test
+.PHONY: setup clean format format-check lint lint-check test mcp-inspector
 
 PY_FILES = src tests scripts
 
@@ -32,3 +32,8 @@ clean:
 	rm -rf .pytest_cache/
 	rm -rf .mypy_cache/
 	rm -rf .ruff_cache/ 
+
+# Launch MCP Inspector
+# Use `fastmcp dev` (not `mcp dev`), since this project uses the standalone FastMCP package.
+mcp-inspector:
+	uv run fastmcp dev src/mcp_server_datahub/__main__.py --with-editable .

scripts/smoke_check.py

@@ -648,13 +648,14 @@ async def run_smoke_check(
     # Now import and register
     from mcp_server_datahub.document_tools_middleware import DocumentToolsMiddleware
     from mcp_server_datahub.mcp_server import (
-        mcp,
+        create_mcp_server,
         register_all_tools,
         with_datahub_client,
     )
     from mcp_server_datahub.version_requirements import VersionFilterMiddleware
 
-    register_all_tools(is_oss=True)
+    mcp = create_mcp_server()
+    register_all_tools(mcp, is_oss=True)
 
     # Add middleware so list_tools reflects what a real client sees
     mcp.add_middleware(VersionFilterMiddleware())

scripts/test_main_tools.py

@@ -7,7 +7,14 @@
 from datahub.sdk.main_client import DataHubClient
 from fastmcp import Client
 
-from mcp_server_datahub.mcp_server import mcp, set_datahub_client
+from mcp_server_datahub.mcp_server import (
+    create_mcp_server,
+    register_all_tools,
+    set_datahub_client,
+)
+
+mcp = create_mcp_server()
+register_all_tools(mcp, is_oss=True)
 
 
 def _divider() -> None:

src/mcp_server_datahub/__main__.py

@@ -10,13 +10,19 @@
 from mcp_server_datahub._telemetry import TelemetryMiddleware
 from mcp_server_datahub._version import __version__
 from mcp_server_datahub.document_tools_middleware import DocumentToolsMiddleware
-from mcp_server_datahub.mcp_server import mcp, register_all_tools, with_datahub_client
+from mcp_server_datahub.mcp_server import (
+    create_mcp_server,
+    register_all_tools,
+    with_datahub_client,
+)
 from mcp_server_datahub.version_requirements import VersionFilterMiddleware
 
 logging.basicConfig(level=logging.INFO)
 
+mcp = create_mcp_server()
+
 # Register tools with OSS-compatible descriptions
-register_all_tools(is_oss=True)
+register_all_tools(mcp, is_oss=True)
 
 
 @click.command()

src/mcp_server_datahub/mcp_server.py

@@ -314,9 +314,9 @@ def _register_tool(
         TOOL_VERSION_REQUIREMENTS[name] = req
 
 
-mcp = FastMCP[None](
-    name="datahub",
-)
+def create_mcp_server(name: str = "datahub") -> FastMCP[None]:
+    """Create a new FastMCP server instance for DataHub tools."""
+    return FastMCP[None](name=name)
 
 
 _mcp_dh_client = contextvars.ContextVar[DataHubClient]("_mcp_dh_client")
@@ -2624,8 +2624,8 @@ def _find_upstream_lineage_path(
     }
 
 
-# Track if tools have been registered to prevent duplicate registration
-_tools_registered = False
+# Track MCP instances that have already had tools registered
+_registered_mcp_instances: set[int] = set()
 _tools_registration_lock = threading.Lock()
 
 
@@ -2796,8 +2796,8 @@ def register_search_tools(mcp_instance: FastMCP, is_oss: bool = False) -> None:
     _register_tool(mcp_instance, "grep_documents", grep_documents)
 
 
-def register_all_tools(is_oss: bool = False) -> None:
-    """Register all MCP tools on the global mcp instance.
+def register_all_tools(mcp_instance: FastMCP, is_oss: bool = False) -> None:
+    """Register all MCP tools on a specific MCP instance.
 
     Args:
         is_oss: If True, use OSS-compatible tool descriptions (limited sorting fields).
@@ -2806,26 +2806,30 @@ def register_all_tools(is_oss: bool = False) -> None:
     Note: Thread-safe. Can be called multiple times from different threads.
           Only the first call will register tools, subsequent calls are no-ops.
     """
-    global _tools_registered
+    mcp_instance_id = id(mcp_instance)
 
     # Thread-safe check-and-set using lock
     with _tools_registration_lock:
-        if _tools_registered:
-            logger.debug("Tools already registered, skipping duplicate registration")
+        if mcp_instance_id in _registered_mcp_instances:
+            logger.debug(
+                "Tools already registered for this MCP instance, "
+                "skipping duplicate registration"
+            )
             return
 
-        _tools_registered = True
+        _registered_mcp_instances.add(mcp_instance_id)
         logger.info(f"Registering MCP tools (is_oss={is_oss})")
 
-    # Call the core registration logic on the global mcp instance
-    register_search_tools(mcp, is_oss)
+    # Call the core registration logic on the provided mcp instance
+    register_search_tools(mcp_instance, is_oss)
 
-    register_mutation_tools(mcp, is_oss)
+    register_mutation_tools(mcp_instance, is_oss)
 
-    register_user_tools(mcp, is_oss)
+    register_user_tools(mcp_instance, is_oss)
 
 
 def get_valid_tools_from_mcp(
+    mcp_instance: FastMCP,
     filter_fn: Optional[Callable[[FastMCPTool], bool]] = None,
 ) -> List[FastMCPTool]:
     """Get valid tools from MCP, optionally filtered.
@@ -2847,7 +2851,7 @@ def get_valid_tools_from_mcp(
             filter_fn=lambda tool: bool((tool.tags or set()) & {"search", "user"})
         )
     """
-    tools = list(mcp._tool_manager._tools.values())
+    tools = list(mcp_instance._tool_manager._tools.values())
     if filter_fn:
         return [tool for tool in tools if filter_fn(tool)]
     return tools

tests/test_async_setup.py

@@ -5,7 +5,14 @@
 import fastmcp.tools.tool
 import pytest
 
-from mcp_server_datahub.mcp_server import async_background, mcp
+from mcp_server_datahub.mcp_server import (
+    async_background,
+    create_mcp_server,
+    register_all_tools,
+)
+
+mcp = create_mcp_server()
+register_all_tools(mcp, is_oss=True)
 
 
 @pytest.mark.anyio

tests/test_mcp/test_async_setup.py

@@ -5,7 +5,14 @@
 import fastmcp.tools.tool
 import pytest
 
-from datahub_integrations.mcp.mcp_server import async_background, mcp
+from datahub_integrations.mcp.mcp_server import (
+    async_background,
+    create_mcp_server,
+    register_all_tools,
+)
+
+mcp = create_mcp_server()
+register_all_tools(mcp, is_oss=True)
 
 
 @pytest.mark.anyio

tests/test_mcp_integration.py

@@ -13,10 +13,16 @@
 from mcp.types import TextContent
 from loguru import logger
 from mcp_server_datahub._telemetry import TelemetryMiddleware
-from mcp_server_datahub.mcp_server import mcp, register_all_tools, with_datahub_client
+from mcp_server_datahub.mcp_server import (
+    create_mcp_server,
+    register_all_tools,
+    with_datahub_client,
+)
+
+mcp = create_mcp_server()
 
 # Register tools with OSS-compatible descriptions for testing
-register_all_tools(is_oss=True)
+register_all_tools(mcp, is_oss=True)
 
 _test_urn = "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.pet_details,PROD)"
 _test_domain = "urn:li:domain:0da1ef03-8870-45db-9f47-ef4f592f095c"  # "urn:li:domain:7186eeff-a860-4b0a-989f-69473a0c9c67"