Add automatic tool discovery via Python entrypoints

Zsailer · Zsailer · commit 1764537e3564 · 2025-10-17T14:53:06.000-07:00
Implements automatic discovery of MCP tools from installed Python packages
using the 'jupyter_ai.tools' entrypoint group. Packages can now expose
tools without requiring manual configuration.

Key changes:
- Add `use_tool_discovery` trait to enable/disable automatic discovery (default: True)
- Implement `_discover_entrypoint_tools()` to find and load tools from entrypoints
- Support both list and function-based entrypoint values
- Unified registration logic in `_register_tools()` method
- Add comprehensive tests for discovery and error handling
- Update documentation with entrypoint usage examples

Tools from entrypoints are registered first, followed by manually configured
tools, allowing configuration to override discovered tools.
diff --git a/.gitignore b/.gitignore
@@ -208,3 +208,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+# pixi environments
+.pixi/*
+!.pixi/config.toml
diff --git a/README.md b/README.md
@@ -11,7 +11,8 @@ This extension provides a simplified, trait-based approach to exposing Jupyter f
 ## Key Features
 
 - **Simplified Architecture**: Direct function registration without complex abstractions
-- **Configurable Tool Loading**: Register tools via string specifications (`module:function`)  
+- **Configurable Tool Loading**: Register tools via string specifications (`module:function`)
+- **Automatic Tool Discovery**: Python packages can expose tools via entrypoints
 - **Jupyter Integration**: Seamless integration with Jupyter Server extension system
 - **HTTP Transport**: FastMCP-based HTTP server with proper MCP protocol support
 - **Traitlets Configuration**: Full configuration support through Jupyter's traitlets system
@@ -131,21 +132,58 @@ Jupyter Server extension that manages the MCP server lifecycle:
 
 **Configuration Traits:**
 - `mcp_name` - Server name (default: "Jupyter MCP Server")
-- `mcp_port` - Server port (default: 3001)  
+- `mcp_port` - Server port (default: 3001)
 - `mcp_tools` - List of tools to register (format: "module:function")
+- `use_tool_discovery` - Enable automatic tool discovery via entrypoints (default: True)
 
-### Tool Loading System
+### Tool Registration
 
-Tools are loaded using string specifications in the format `module_path:function_name`:
+Tools can be registered in two ways:
+
+#### 1. Manual Configuration
+
+Specify tools directly in your Jupyter configuration using `module:function` format:
+
+```python
+c.MCPExtensionApp.mcp_tools = [
+    "os:getcwd",
+    "jupyter_ai_tools.toolkits.notebook:read_notebook",
+]
+```
+
+#### 2. Automatic Discovery via Entrypoints
+
+Python packages can expose tools automatically using the `jupyter_ai.tools` entrypoint group.
+
+**In your package's `pyproject.toml`:**
+
+```toml
+[project.entry-points."jupyter_ai.tools"]
+my_package_tools = "my_package.tools:TOOLS"
+```
+
+**In `my_package/tools.py`:**
 
 ```python
-# Examples
-"os:getcwd"                                           # Standard library
-"jupyter_ai_tools.toolkits.notebook:read_notebook"   # External package
-"math:sqrt"                                           # Built-in modules
+# Option 1: Define as a list
+TOOLS = [
+    "my_package.operations:create_file",
+    "my_package.operations:delete_file",
+]
+
+# Option 2: Define as a function
+def get_tools():
+    return [
+        "my_package.operations:create_file",
+        "my_package.operations:delete_file",
+    ]
 ```
 
-The extension dynamically imports the module and registers the function with FastMCP.
+Tools from entrypoints are discovered automatically when the extension starts. To disable automatic discovery:
+
+```python
+c.MCPExtensionApp.use_tool_discovery = False
+```
 
 ## Configuration Examples
 
diff --git a/jupyter_server_mcp/extension.py b/jupyter_server_mcp/extension.py
@@ -3,10 +3,11 @@
 import asyncio
 import contextlib
 import importlib
+import importlib.metadata
 import logging
 
 from jupyter_server.extension.application import ExtensionApp
-from traitlets import Int, List, Unicode
+from traitlets import Bool, Int, List, Unicode
 
 from .mcp_server import MCPServer
 
@@ -38,6 +39,14 @@ class MCPExtensionApp(ExtensionApp):
         ),
     ).tag(config=True)
 
+    use_tool_discovery = Bool(
+        default_value=True,
+        help=(
+            "Whether to automatically discover and register tools from "
+            "Python entrypoints in the 'jupyter_ai.tools' group"
+        ),
+    ).tag(config=True)
+
     mcp_server_instance: object | None = None
     mcp_server_task: asyncio.Task | None = None
 
@@ -75,22 +84,94 @@ def _load_function_from_string(self, tool_spec: str):
             msg = f"Function '{function_name}' not found in module '{module_path}': {e}"
             raise AttributeError(msg) from e
 
-    def _register_configured_tools(self):
-        """Register tools specified in the mcp_tools configuration."""
-        if not self.mcp_tools:
+    def _register_tools(self, tool_specs: list[str], source: str = "configuration"):
+        """Register tools from a list of tool specifications.
+
+        Args:
+            tool_specs: List of tool specifications in 'module:function' format
+            source: Description of where tools came from (for logging)
+        """
+        if not tool_specs:
             return
 
-        logger.info(f"Registering {len(self.mcp_tools)} configured tools")
+        logger.info(f"Registering {len(tool_specs)} tools from {source}")
 
-        for tool_spec in self.mcp_tools:
+        for tool_spec in tool_specs:
             try:
                 function = self._load_function_from_string(tool_spec)
                 self.mcp_server_instance.register_tool(function)
-                logger.info(f"✅ Registered tool: {tool_spec}")
+                logger.info(f"✅ Registered tool from {source}: {tool_spec}")
             except Exception as e:
-                logger.error(f"❌ Failed to register tool '{tool_spec}': {e}")
+                logger.error(f"❌ Failed to register tool '{tool_spec}' from {source}: {e}")
                 continue
 
+    def _discover_entrypoint_tools(self) -> list[str]:
+        """Discover tools from Python entrypoints in the 'jupyter_ai.tools' group.
+
+        Returns:
+            List of tool specifications in 'module:function' format
+        """
+        if not self.use_tool_discovery:
+            return []
+
+        discovered_tools = []
+
+        try:
+            # Use importlib.metadata to discover entrypoints
+            entrypoints = importlib.metadata.entry_points()
+
+            # Handle both Python 3.10+ and 3.9 style entrypoint APIs
+            if hasattr(entrypoints, "select"):
+                tools_group = entrypoints.select(group="jupyter_ai.tools")
+            else:
+                tools_group = entrypoints.get("jupyter_ai.tools", [])
+
+            for entry_point in tools_group:
+                try:
+                    # Load the entrypoint value (can be a list or a function that returns a list)
+                    loaded_value = entry_point.load()
+
+                    # Get tool specs from either a list or callable
+                    if isinstance(loaded_value, list):
+                        tool_specs = loaded_value
+                    elif callable(loaded_value):
+                        tool_specs = loaded_value()
+                        if not isinstance(tool_specs, list):
+                            logger.warning(
+                                f"Entrypoint '{entry_point.name}' function returned "
+                                f"{type(tool_specs).__name__} instead of list, skipping"
+                            )
+                            continue
+                    else:
+                        logger.warning(
+                            f"Entrypoint '{entry_point.name}' is neither a list nor callable, skipping"
+                        )
+                        continue
+
+                    # Validate and collect tool specs
+                    valid_specs = [spec for spec in tool_specs if isinstance(spec, str)]
+                    invalid_count = len(tool_specs) - len(valid_specs)
+
+                    if invalid_count > 0:
+                        logger.warning(
+                            f"Skipped {invalid_count} non-string tool specs from '{entry_point.name}'"
+                        )
+
+                    discovered_tools.extend(valid_specs)
+                    logger.info(f"Discovered {len(valid_specs)} tools from entrypoint '{entry_point.name}'")
+
+                except Exception as e:
+                    logger.error(f"Failed to load entrypoint '{entry_point.name}': {e}")
+                    continue
+
+        except Exception as e:
+            logger.error(f"Failed to discover entrypoints: {e}")
+
+        if not discovered_tools:
+            logger.info("No tools discovered from entrypoints")
+
+        return discovered_tools
+
     def initialize(self):
         """Initialize the extension."""
         super().initialize()
@@ -115,8 +196,10 @@ async def start_extension(self):
                 parent=self, name=self.mcp_name, port=self.mcp_port
             )
 
-            # Register configured tools
-            self._register_configured_tools()
+            # Register tools from entrypoints, then from configuration
+            entrypoint_tools = self._discover_entrypoint_tools()
+            self._register_tools(entrypoint_tools, source="entrypoints")
+            self._register_tools(self.mcp_tools, source="configuration")
 
             # Start the MCP server in a background task
             self.mcp_server_task = asyncio.create_task(
@@ -126,12 +209,9 @@ async def start_extension(self):
             # Give the server a moment to start
             await asyncio.sleep(0.5)
 
+            registered_count = len(self.mcp_server_instance._registered_tools)
             self.log.info(f"✅ MCP server started on port {self.mcp_port}")
-            if self.mcp_tools:
-                registered_count = len(self.mcp_server_instance._registered_tools)
-                self.log.info(f"Registered {registered_count} tools from configuration")
-            else:
-                self.log.info("Use mcp_server_instance.register_tool() to add tools")
+            self.log.info(f"Total registered tools: {registered_count}")
 
         except Exception as e:
             self.log.error(f"Failed to start MCP server: {e}")
diff --git a/tests/test_extension.py b/tests/test_extension.py
@@ -340,3 +340,94 @@ async def test_start_extension_no_tools(self):
 
             # Should not register any tools
             mock_server.register_tool.assert_not_called()
+
+
+class TestEntrypointDiscovery:
+    """Test entrypoint discovery functionality."""
+
+    def test_discover_entrypoint_tools_multiple_types(self):
+        """Test discovering tools from both list and function entrypoints."""
+        extension = MCPExtensionApp()
+
+        # Create mock entrypoints - one list, one function
+        mock_ep1 = Mock()
+        mock_ep1.name = "package1_tools"
+        mock_ep1.value = "package1.tools:TOOLS"
+        mock_ep1.load.return_value = ["os:getcwd", "math:sqrt"]
+
+        mock_ep2 = Mock()
+        mock_ep2.name = "package2_tools"
+        mock_ep2.value = "package2.tools:get_tools"
+        mock_function = Mock(return_value=["json:dumps", "time:time"])
+        mock_ep2.load.return_value = mock_function
+
+        with patch("importlib.metadata.entry_points") as mock_ep_func:
+            mock_ep_func.return_value.select = Mock(return_value=[mock_ep1, mock_ep2])
+
+            tools = extension._discover_entrypoint_tools()
+            assert len(tools) == 4
+            assert set(tools) == {"os:getcwd", "math:sqrt", "json:dumps", "time:time"}
+            mock_function.assert_called_once()  # Function was called
+
+    def test_discover_entrypoint_tools_error_handling(self):
+        """Test that discovery handles invalid entrypoints gracefully."""
+        extension = MCPExtensionApp()
+
+        # Mix of valid and invalid entrypoints
+        valid_ep = Mock()
+        valid_ep.name = "valid"
+        valid_ep.load.return_value = ["os:getcwd"]
+
+        invalid_type_ep = Mock()
+        invalid_type_ep.name = "invalid_type"
+        invalid_type_ep.load.return_value = "not_a_list"
+
+        function_bad_return_ep = Mock()
+        function_bad_return_ep.name = "bad_function"
+        function_bad_return_ep.load.return_value = Mock(return_value={"not": "list"})
+
+        load_error_ep = Mock()
+        load_error_ep.name = "load_error"
+        load_error_ep.load.side_effect = ImportError("Module not found")
+
+        with patch("importlib.metadata.entry_points") as mock_ep_func:
+            mock_ep_func.return_value.select = Mock(
+                return_value=[valid_ep, invalid_type_ep, function_bad_return_ep, load_error_ep]
+            )
+
+            with patch("jupyter_server_mcp.extension.logger"):
+                tools = extension._discover_entrypoint_tools()
+                # Should only get the valid one
+                assert tools == ["os:getcwd"]
+
+    def test_discover_entrypoint_tools_disabled(self):
+        """Test that discovery returns empty list when disabled."""
+        extension = MCPExtensionApp()
+        extension.use_tool_discovery = False
+
+        # Should return empty without trying to discover
+        tools = extension._discover_entrypoint_tools()
+        assert tools == []
+
+    @pytest.mark.asyncio
+    async def test_start_extension_with_entrypoints_and_config(self):
+        """Test extension startup with both entrypoint and configured tools."""
+        extension = MCPExtensionApp()
+        extension.mcp_port = 3086
+        extension.use_tool_discovery = True
+        extension.mcp_tools = ["json:dumps"]
+
+        discovered_tools = ["os:getcwd"]
+
+        with patch("jupyter_server_mcp.extension.MCPServer") as mock_mcp_class:
+            mock_server = Mock()
+            mock_server.start_server = AsyncMock()
+            mock_server._registered_tools = {"getcwd": {}, "dumps": {}}
+            mock_mcp_class.return_value = mock_server
+
+            with patch.object(extension, "_discover_entrypoint_tools", return_value=discovered_tools):
+                await extension.start_extension()
+
+                # Should register both entrypoint (1) and configured (1) tools = 2 total
+                assert mock_server.register_tool.call_count == 2
+
