Add automatic tool discovery via Python entrypoints (#10)

* Add automatic tool discovery via Python entrypoints

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.

* Update entrypoint name

* fix unit test

* lintin

Author: Zsailer

.gitignore

@@ -208,3 +208,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+# pixi environments
+.pixi/*
+!.pixi/config.toml

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_server_mcp.tools` entrypoint group.
+
+**In your package's `pyproject.toml`:**
+
+```toml
+[project.entry-points."jupyter_server_mcp.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
 

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_server_mcp.tools' group"
+        ),
+    ).tag(config=True)
+
     mcp_server_instance: object | None = None
     mcp_server_task: asyncio.Task | None = None
 
@@ -75,22 +84,98 @@ 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_server_mcp.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_server_mcp.tools")
+            else:
+                tools_group = entrypoints.get("jupyter_server_mcp.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 +200,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 +213,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}")