From 0d468b10fbd83ee7ee2e3e2c3fbe3ad91cc8e51b Mon Sep 17 00:00:00 2001 From: Jeremy Tuloup Date: Mon, 27 Oct 2025 09:47:11 +0100 Subject: [PATCH 1/4] Add support for prompts --- README.md | 230 ++++++++++++++++++++++++++++--- jupyter_server_mcp/extension.py | 119 +++++++++++++++- jupyter_server_mcp/mcp_server.py | 64 ++++++++- tests/test_extension.py | 205 +++++++++++++++++++++++++++ tests/test_mcp_server.py | 150 ++++++++++++++++++++ 5 files changed, 742 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 4a521f1..c7bb792 100644 --- a/README.md +++ b/README.md @@ -6,13 +6,13 @@ https://github.com/user-attachments/assets/aa779b1c-a443-48d7-b3eb-13f27a4333b3 ## Overview -This extension provides a simplified, trait-based approach to exposing Jupyter functionality through the MCP protocol. It can dynamically load and register tools from various Python packages, making them available to AI assistants and other MCP clients. +This extension provides a simplified, trait-based approach to exposing Jupyter functionality through the MCP protocol. It can dynamically load and register tools and prompts from various Python packages, making them available to AI assistants and other MCP clients. ## Key Features - **Simplified Architecture**: Direct function registration without complex abstractions -- **Configurable Tool Loading**: Register tools via string specifications (`module:function`) -- **Automatic Tool Discovery**: Python packages can expose tools via entrypoints +- **Configurable Tool & Prompt Loading**: Register tools and prompts via string specifications (`module:function`) +- **Automatic Discovery**: Python packages can expose tools and prompts 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 @@ -42,15 +42,21 @@ c.MCPExtensionApp.mcp_tools = [ "os:getcwd", "json:dumps", "time:time", - - # Jupyter AI Tools - Notebook operations + + # Jupyter AI Tools - Notebook operations "jupyter_ai_tools.toolkits.notebook:read_notebook", "jupyter_ai_tools.toolkits.notebook:edit_cell", - + # JupyterLab Commands Toolkit "jupyterlab_commands_toolkit.tools:list_all_commands", "jupyterlab_commands_toolkit.tools:execute_command", ] + +# Register prompts for code assistance +c.MCPExtensionApp.mcp_prompts = [ + "my_package.prompts:code_review_prompt", + "my_package.prompts:documentation_prompt", +] ``` ### 2. Start Jupyter Server @@ -121,9 +127,12 @@ await server.start_server() ``` **Key Methods:** -- `register_tool(func, name=None, description=None)` - Register a Python function +- `register_tool(func, name=None, description=None)` - Register a Python function as a tool - `register_tools(tools)` - Register multiple functions (list or dict) - `list_tools()` - Get list of registered tools +- `register_prompt(func, name=None, description=None)` - Register a Python function as a prompt +- `register_prompts(prompts)` - Register multiple prompt functions (list or dict) +- `list_prompts()` - Get list of registered prompts - `start_server(host=None)` - Start the HTTP MCP server #### MCPExtensionApp (`jupyter_server_mcp.extension.MCPExtensionApp`) @@ -134,7 +143,8 @@ Jupyter Server extension that manages the MCP server lifecycle: - `mcp_name` - Server name (default: "Jupyter MCP Server") - `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) +- `mcp_prompts` - List of prompts to register (format: "module:function") +- `use_tool_discovery` - Enable automatic discovery via entrypoints (default: True) ### Tool Registration @@ -185,6 +195,95 @@ Tools from entrypoints are discovered automatically when the extension starts. T c.MCPExtensionApp.use_tool_discovery = False ``` +### Prompt Registration + +Prompts can be registered in two ways: + +#### 1. Manual Configuration + +Specify prompts directly in your Jupyter configuration using `module:function` format: + +```python +c.MCPExtensionApp.mcp_prompts = [ + "my_package.prompts:code_review_prompt", + "my_package.prompts:documentation_prompt", +] +``` + +#### 2. Automatic Discovery via Entrypoints + +Python packages can expose prompts automatically using the `jupyter_server_mcp.prompts` entrypoint group. + +**In your package's `pyproject.toml`:** + +```toml +[project.entry-points."jupyter_server_mcp.prompts"] +my_package_prompts = "my_package.prompts:PROMPTS" +``` + +**In `my_package/prompts.py`:** + +```python +# Option 1: Define as a list +PROMPTS = [ + "my_package.prompts:code_review_prompt", + "my_package.prompts:documentation_prompt", + "my_package.prompts:test_generation_prompt", +] + +# Option 2: Define as a function +def get_prompts(): + return [ + "my_package.prompts:code_review_prompt", + "my_package.prompts:documentation_prompt", + ] + +# Example prompt functions +def code_review_prompt(file_path: str, focus_area: str = "general") -> str: + """Generate a code review prompt for a specific file. + + Args: + file_path: Path to the file to review + focus_area: Aspect to focus on (e.g., 'security', 'performance', 'general') + + Returns: + A formatted prompt for code review + """ + return f"""Please review the code in {file_path} with a focus on {focus_area}. + +Provide feedback on: +- Code quality and best practices +- Potential bugs or issues +- Performance considerations +- Security concerns +- Documentation and readability +""" + +def documentation_prompt(module_name: str, target_audience: str = "developers") -> str: + """Generate a documentation prompt for a module. + + Args: + module_name: Name of the module to document + target_audience: Who will read this (e.g., 'developers', 'users', 'contributors') + + Returns: + A formatted prompt for documentation generation + """ + return f"""Create comprehensive documentation for the {module_name} module. + +Target audience: {target_audience} + +Include: +- Overview and purpose +- Key features and capabilities +- Usage examples with code +- API reference +- Common patterns and best practices +""" +``` + +Prompts from entrypoints are discovered automatically when the extension starts, using the same `use_tool_discovery` setting as tools. + ## Configuration Examples ### Minimal Setup @@ -203,36 +302,50 @@ c.MCPExtensionApp.mcp_port = 8080 c.MCPExtensionApp.mcp_tools = [ # File system operations (jupyter-ai-tools) "jupyter_ai_tools.toolkits.file_system:read", - "jupyter_ai_tools.toolkits.file_system:write", + "jupyter_ai_tools.toolkits.file_system:write", "jupyter_ai_tools.toolkits.file_system:edit", "jupyter_ai_tools.toolkits.file_system:ls", "jupyter_ai_tools.toolkits.file_system:glob", - + # Notebook operations (jupyter-ai-tools) "jupyter_ai_tools.toolkits.notebook:read_notebook", "jupyter_ai_tools.toolkits.notebook:edit_cell", - "jupyter_ai_tools.toolkits.notebook:add_cell", + "jupyter_ai_tools.toolkits.notebook:add_cell", "jupyter_ai_tools.toolkits.notebook:delete_cell", "jupyter_ai_tools.toolkits.notebook:create_notebook", - + # Git operations (jupyter-ai-tools) "jupyter_ai_tools.toolkits.git:git_status", "jupyter_ai_tools.toolkits.git:git_add", "jupyter_ai_tools.toolkits.git:git_commit", "jupyter_ai_tools.toolkits.git:git_push", - + # JupyterLab operations (jupyterlab-commands-toolkit) "jupyterlab_commands_toolkit.tools:clear_all_outputs_in_notebook", "jupyterlab_commands_toolkit.tools:open_document", "jupyterlab_commands_toolkit.tools:open_markdown_file_in_preview_mode", "jupyterlab_commands_toolkit.tools:show_diff_of_current_notebook", - - # Utility functions + + # Utility functions "os:getcwd", "json:dumps", "time:time", "platform:system", ] + +c.MCPExtensionApp.mcp_prompts = [ + # Code analysis prompts + "my_package.prompts:code_review_prompt", + "my_package.prompts:refactoring_suggestions_prompt", + + # Documentation prompts + "my_package.prompts:documentation_prompt", + "my_package.prompts:api_documentation_prompt", + + # Testing prompts + "my_package.prompts:test_generation_prompt", + "my_package.prompts:test_coverage_analysis_prompt", +] ``` ### Running Tests @@ -258,17 +371,94 @@ jupyter_server_mcp/ │ └── extension.py # Jupyter Server extension ├── tests/ │ ├── test_mcp_server.py # MCPServer tests -│ └── test_extension.py # Extension tests -├── demo/ -│ ├── jupyter_config.py # Example configuration -│ └── *.py # Debug/diagnostic scripts +│ └── test_extension.py # Extension tests └── pyproject.toml # Package configuration ``` +## Example: Creating a Package with Prompts + +Here's how to create a package that exposes prompts via entrypoints: + +**my_package/prompts.py:** +```python +"""Example prompt functions for code assistance.""" + +def code_review_prompt(file_path: str, focus_area: str = "general") -> str: + """Generate a code review prompt for a specific file. + + Args: + file_path: Path to the file to review + focus_area: Aspect to focus on (e.g., 'security', 'performance', 'general') + + Returns: + A formatted prompt for code review + """ + return f"""Please review the code in {file_path} with a focus on {focus_area}. + +Provide feedback on: +- Code quality and best practices +- Potential bugs or issues +- Performance considerations +- Security concerns +- Documentation and readability +""" + +def documentation_prompt(module_name: str, target_audience: str = "developers") -> str: + """Generate a documentation prompt for a module.""" + return f"""Create comprehensive documentation for the {module_name} module. + +Target audience: {target_audience} + +Include: +- Overview and purpose +- Key features and capabilities +- Usage examples with code +- API reference +- Common patterns and best practices +""" + +def test_generation_prompt(function_name: str, function_code: str) -> str: + """Generate a prompt for creating unit tests.""" + return f"""Generate comprehensive unit tests for the following function: + +Function name: {function_name} + +Code: +```python +{function_code} +``` + +Please create: +- Test cases for normal operation +- Edge cases and boundary conditions +- Error handling tests +- Mocking examples if needed +""" + +# Define prompts list for entrypoint discovery +PROMPTS = [ + "my_package.prompts:code_review_prompt", + "my_package.prompts:documentation_prompt", + "my_package.prompts:test_generation_prompt", +] +``` + +**pyproject.toml:** +```toml +[project] +name = "my-package" +version = "0.1.0" + +[project.entry-points."jupyter_server_mcp.prompts"] +my_package_prompts = "my_package.prompts:PROMPTS" +``` + +After installing your package, the prompts will be automatically discovered and registered when the Jupyter Server MCP extension starts. + ## Contributing 1. Fork the repository 2. Create a feature branch -3. Add tests for new functionality +3. Add tests for new functionality 4. Ensure all tests pass: `pytest tests/` 5. Submit a pull request diff --git a/jupyter_server_mcp/extension.py b/jupyter_server_mcp/extension.py index de4e3f0..73362db 100644 --- a/jupyter_server_mcp/extension.py +++ b/jupyter_server_mcp/extension.py @@ -39,11 +39,22 @@ class MCPExtensionApp(ExtensionApp): ), ).tag(config=True) + mcp_prompts = List( + trait=Unicode(), + default_value=[], + help=( + "List of prompts to register with the MCP server. " + "Format: 'module_path:function_name' " + "(e.g., 'mypackage.prompts:code_review_prompt')" + ), + ).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" + "Whether to automatically discover and register tools and prompts from " + "Python entrypoints in the 'jupyter_server_mcp.tools' and " + "'jupyter_server_mcp.prompts' groups" ), ).tag(config=True) @@ -107,6 +118,29 @@ def _register_tools(self, tool_specs: list[str], source: str = "configuration"): ) continue + def _register_prompts(self, prompt_specs: list[str], source: str = "configuration"): + """Register prompts from a list of prompt specifications. + + Args: + prompt_specs: List of prompt specifications in 'module:function' format + source: Description of where prompts came from (for logging) + """ + if not prompt_specs: + return + + logger.info(f"Registering {len(prompt_specs)} prompts from {source}") + + for prompt_spec in prompt_specs: + try: + function = self._load_function_from_string(prompt_spec) + self.mcp_server_instance.register_prompt(function) + logger.info(f"✅ Registered prompt from {source}: {prompt_spec}") + except Exception as e: + logger.error( + f"❌ Failed to register prompt '{prompt_spec}' from {source}: {e}" + ) + continue + def _discover_entrypoint_tools(self) -> list[str]: """Discover tools from Python entrypoints in the 'jupyter_server_mcp.tools' group. @@ -176,6 +210,75 @@ def _discover_entrypoint_tools(self) -> list[str]: return discovered_tools + def _discover_entrypoint_prompts(self) -> list[str]: + """Discover prompts from Python entrypoints in the 'jupyter_server_mcp.prompts' group. + + Returns: + List of prompt specifications in 'module:function' format + """ + if not self.use_tool_discovery: + return [] + + discovered_prompts = [] + + 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"): + prompts_group = entrypoints.select(group="jupyter_server_mcp.prompts") + else: + prompts_group = entrypoints.get("jupyter_server_mcp.prompts", []) + + for entry_point in prompts_group: + try: + # Load the entrypoint value (can be a list or a function that returns a list) + loaded_value = entry_point.load() + + # Get prompt specs from either a list or callable + if isinstance(loaded_value, list): + prompt_specs = loaded_value + elif callable(loaded_value): + prompt_specs = loaded_value() + if not isinstance(prompt_specs, list): + logger.warning( + f"Entrypoint '{entry_point.name}' function returned " + f"{type(prompt_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 prompt specs + valid_specs = [spec for spec in prompt_specs if isinstance(spec, str)] + invalid_count = len(prompt_specs) - len(valid_specs) + + if invalid_count > 0: + logger.warning( + f"Skipped {invalid_count} non-string prompt specs from '{entry_point.name}'" + ) + + discovered_prompts.extend(valid_specs) + logger.info( + f"Discovered {len(valid_specs)} prompts 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_prompts: + logger.info("No prompts discovered from entrypoints") + + return discovered_prompts + def initialize(self): """Initialize the extension.""" super().initialize() @@ -200,11 +303,15 @@ async def start_extension(self): parent=self, name=self.mcp_name, port=self.mcp_port ) - # Register tools from entrypoints, then from configuration + # Register tools and prompts 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") + entrypoint_prompts = self._discover_entrypoint_prompts() + self._register_prompts(entrypoint_prompts, source="entrypoints") + self._register_prompts(self.mcp_prompts, source="configuration") + # Start the MCP server in a background task self.mcp_server_task = asyncio.create_task( self.mcp_server_instance.start_server() @@ -213,9 +320,11 @@ 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) + registered_tools_count = len(self.mcp_server_instance._registered_tools) + registered_prompts_count = len(self.mcp_server_instance._registered_prompts) self.log.info(f"✅ MCP server started on port {self.mcp_port}") - self.log.info(f"Total registered tools: {registered_count}") + self.log.info(f"Total registered tools: {registered_tools_count}") + self.log.info(f"Total registered prompts: {registered_prompts_count}") except Exception as e: self.log.error(f"Failed to start MCP server: {e}") diff --git a/jupyter_server_mcp/mcp_server.py b/jupyter_server_mcp/mcp_server.py index 7b8e070..592a47f 100644 --- a/jupyter_server_mcp/mcp_server.py +++ b/jupyter_server_mcp/mcp_server.py @@ -265,9 +265,10 @@ def __init__(self, **kwargs): """ super().__init__(**kwargs) - # Initialize FastMCP and tools registry + # Initialize FastMCP and tools/prompts registry self.mcp = FastMCP(self.name) self._registered_tools = {} + self._registered_prompts = {} self.log.info( f"Initialized MCP server '{self.name}' on {self.host}:{self.port}" ) @@ -344,6 +345,67 @@ def get_tool_info(self, tool_name: str) -> dict[str, Any] | None: """Get information about a specific tool.""" return self._registered_tools.get(tool_name) + def register_prompt( + self, + func: Callable, + name: str | None = None, + description: str | None = None, + ): + """Register a Python function as an MCP prompt. + + Args: + func: Python function to register as a prompt + name: Optional prompt name (defaults to function name) + description: Optional prompt description (defaults to function + docstring) + """ + prompt_name = name or func.__name__ + prompt_description = description or func.__doc__ or f"Prompt: {prompt_name}" + + self.log.info(f"Registering prompt: {prompt_name}") + self.log.debug( + f"Prompt details - Name: {prompt_name}, " + f"Description: {prompt_description}, Async: {iscoroutinefunction(func)}" + ) + + # Register with FastMCP using the @mcp.prompt decorator + self.mcp.prompt(func) + + # Keep track for listing + self._registered_prompts[prompt_name] = { + "name": prompt_name, + "description": prompt_description, + "function": func, + "is_async": iscoroutinefunction(func), + } + + def register_prompts(self, prompts: list[Callable] | dict[str, Callable]): + """Register multiple Python functions as MCP prompts. + + Args: + prompts: List of functions or dict mapping names to functions + """ + if isinstance(prompts, list): + for func in prompts: + self.register_prompt(func) + elif isinstance(prompts, dict): + for name, func in prompts.items(): + self.register_prompt(func, name=name) + else: + msg = "prompts must be a list of functions or dict mapping names to functions" + raise ValueError(msg) + + def list_prompts(self) -> list[dict[str, Any]]: + """List all registered prompts.""" + return [ + {"name": prompt["name"], "description": prompt["description"]} + for prompt in self._registered_prompts.values() + ] + + def get_prompt_info(self, prompt_name: str) -> dict[str, Any] | None: + """Get information about a specific prompt.""" + return self._registered_prompts.get(prompt_name) + async def start_server(self, host: str | None = None): """Start the MCP server on the specified host and port.""" server_host = host or self.host diff --git a/tests/test_extension.py b/tests/test_extension.py index bec6535..4515203 100644 --- a/tests/test_extension.py +++ b/tests/test_extension.py @@ -35,6 +35,16 @@ def test_extension_trait_configuration(self): extension.mcp_tools = ["os:getcwd", "math:sqrt"] assert extension.mcp_tools == ["os:getcwd", "math:sqrt"] + # Test prompts configuration + extension.mcp_prompts = [ + "package.prompts:code_review", + "package.prompts:documentation", + ] + assert extension.mcp_prompts == [ + "package.prompts:code_review", + "package.prompts:documentation", + ] + def test_initialize_handlers(self): """Test handler initialization (should be no-op).""" extension = MCPExtensionApp() @@ -298,6 +308,57 @@ def test_register_configured_tools_with_errors(self): "Could not import module 'invalid': No module named 'invalid'" ) + def test_register_configured_prompts_empty(self): + """Test registering prompts when mcp_prompts is empty.""" + extension = MCPExtensionApp() + extension.mcp_server_instance = Mock() + extension.mcp_prompts = [] + + # Should not call register_prompt + extension._register_prompts(extension.mcp_prompts, source="configuration") + extension.mcp_server_instance.register_prompt.assert_not_called() + + def test_register_configured_prompts_valid(self): + """Test registering valid configured prompts.""" + extension = MCPExtensionApp() + extension.mcp_server_instance = Mock() + extension.mcp_prompts = [ + "package.prompts:code_review", + "package.prompts:documentation", + ] + + # Capture log output + with patch("jupyter_server_mcp.extension.logger") as mock_logger: + extension._register_prompts(extension.mcp_prompts, source="configuration") + + # Should register both prompts + assert extension.mcp_server_instance.register_prompt.call_count == 2 + + # Check log messages + mock_logger.info.assert_any_call("Registering 2 prompts from configuration") + + def test_register_configured_prompts_with_errors(self): + """Test registering prompts when some fail to load.""" + extension = MCPExtensionApp() + extension.mcp_server_instance = Mock() + extension.mcp_prompts = [ + "os:getcwd", # Valid import but might not be a good prompt + "invalid:function", + "json:dumps", + ] + + with patch("jupyter_server_mcp.extension.logger") as mock_logger: + extension._register_prompts(extension.mcp_prompts, source="configuration") + + # Should register 2 valid prompts (os:getcwd and json:dumps) + assert extension.mcp_server_instance.register_prompt.call_count == 2 + + # Check error logging for the invalid one + mock_logger.error.assert_any_call( + "❌ Failed to register prompt 'invalid:function' from configuration: " + "Could not import module 'invalid': No module named 'invalid'" + ) + class TestExtensionWithTools: """Test extension lifecycle with configured tools.""" @@ -443,3 +504,147 @@ async def test_start_extension_with_entrypoints_and_config(self): # Should register both entrypoint (1) and configured (1) tools = 2 total assert mock_server.register_tool.call_count == 2 + + def test_discover_entrypoint_prompts_multiple_types(self): + """Test discovering prompts from both list and function entrypoints.""" + extension = MCPExtensionApp() + + # Create mock entrypoints - one list, one function + mock_ep1 = Mock() + mock_ep1.name = "package1_prompts" + mock_ep1.value = "package1.prompts:PROMPTS" + mock_ep1.load.return_value = [ + "package1.prompts:code_review", + "package1.prompts:documentation", + ] + + mock_ep2 = Mock() + mock_ep2.name = "package2_prompts" + mock_ep2.value = "package2.prompts:get_prompts" + mock_function = Mock( + return_value=[ + "package2.prompts:test_gen", + "package2.prompts:refactor", + ] + ) + 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]) + + prompts = extension._discover_entrypoint_prompts() + assert len(prompts) == 4 + assert set(prompts) == { + "package1.prompts:code_review", + "package1.prompts:documentation", + "package2.prompts:test_gen", + "package2.prompts:refactor", + } + mock_function.assert_called_once() # Function was called + + def test_discover_entrypoint_prompts_error_handling(self): + """Test that prompt 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 = ["package.prompts:valid_prompt"] + + 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"): + prompts = extension._discover_entrypoint_prompts() + # Should only get the valid one + assert prompts == ["package.prompts:valid_prompt"] + + def test_discover_entrypoint_prompts_disabled(self): + """Test that prompt discovery returns empty list when disabled.""" + extension = MCPExtensionApp() + extension.use_tool_discovery = False + + # Should return empty without trying to discover + prompts = extension._discover_entrypoint_prompts() + assert prompts == [] + + @pytest.mark.asyncio + async def test_start_extension_with_prompts_entrypoints_and_config(self): + """Test extension startup with both entrypoint and configured prompts.""" + extension = MCPExtensionApp() + extension.mcp_port = 3087 + extension.use_tool_discovery = True + extension.mcp_prompts = ["package.prompts:config_prompt"] + + discovered_prompts = ["package.prompts:discovered_prompt"] + + with patch("jupyter_server_mcp.extension.MCPServer") as mock_mcp_class: + mock_server = Mock() + mock_server.start_server = AsyncMock() + mock_server._registered_tools = {} + mock_server._registered_prompts = { + "discovered_prompt": {}, + "config_prompt": {}, + } + mock_mcp_class.return_value = mock_server + + with ( + patch.object( + extension, + "_discover_entrypoint_prompts", + return_value=discovered_prompts, + ), + patch.object( + extension, "_discover_entrypoint_tools", return_value=[] + ), + ): + await extension.start_extension() + + # Should register both entrypoint (1) and configured (1) prompts = 2 total + assert mock_server.register_prompt.call_count == 2 + + @pytest.mark.asyncio + async def test_start_extension_with_tools_and_prompts(self): + """Test extension startup with both tools and prompts.""" + extension = MCPExtensionApp() + extension.mcp_port = 3088 + extension.mcp_tools = ["os:getcwd"] + extension.mcp_prompts = ["package.prompts:review"] + + with patch("jupyter_server_mcp.extension.MCPServer") as mock_mcp_class: + mock_server = Mock() + mock_server.start_server = AsyncMock() + mock_server._registered_tools = {"getcwd": {}} + mock_server._registered_prompts = {"review": {}} + mock_mcp_class.return_value = mock_server + + with ( + patch.object(extension, "_discover_entrypoint_tools", return_value=[]), + patch.object( + extension, "_discover_entrypoint_prompts", return_value=[] + ), + ): + await extension.start_extension() + + # Should register 1 tool and 1 prompt + assert mock_server.register_tool.call_count == 1 + assert mock_server.register_prompt.call_count == 1 diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 6dfa781..31da7bf 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -27,6 +27,26 @@ def function_with_docstring(message: str) -> str: return f"Printed: {message}" +def simple_prompt(topic: str) -> str: + """Generate a prompt about a topic.""" + return f"Please explain the concept of '{topic}' in detail." + + +async def async_prompt(task: str, language: str = "Python") -> str: + """Generate a code generation prompt.""" + await asyncio.sleep(0.001) # Small delay + return f"Write a {language} function that performs the following task: {task}" + + +def prompt_with_detailed_docstring(file_path: str, focus: str = "general") -> str: + """Generate a code review prompt. + + This prompt helps reviewers focus on specific aspects of code quality. + Supports multiple focus areas including security, performance, and readability. + """ + return f"Review the code in {file_path} with focus on {focus}." + + class TestMCPServer: """Test MCPServer functionality.""" @@ -39,6 +59,7 @@ def test_server_creation(self): assert server.host == "localhost" assert server.mcp is not None assert len(server._registered_tools) == 0 + assert len(server._registered_prompts) == 0 def test_server_creation_with_params(self): """Test server creation with custom parameters.""" @@ -166,6 +187,135 @@ def test_get_tool_info(self): assert info["name"] == "simple_function" assert info["function"] == simple_function + def test_register_single_prompt(self): + """Test registering a single prompt.""" + server = MCPServer() + + server.register_prompt(simple_prompt) + + # Check prompt was registered + assert len(server._registered_prompts) == 1 + assert "simple_prompt" in server._registered_prompts + + prompt_info = server._registered_prompts["simple_prompt"] + assert prompt_info["name"] == "simple_prompt" + assert prompt_info["description"] == "Generate a prompt about a topic." + assert prompt_info["function"] == simple_prompt + assert prompt_info["is_async"] is False + + def test_register_prompt_with_custom_name(self): + """Test registering a prompt with custom name.""" + server = MCPServer() + + server.register_prompt(simple_prompt, name="topic_explainer") + + assert "topic_explainer" in server._registered_prompts + assert "simple_prompt" not in server._registered_prompts + + prompt_info = server._registered_prompts["topic_explainer"] + assert prompt_info["name"] == "topic_explainer" + + def test_register_prompt_with_custom_description(self): + """Test registering a prompt with custom description.""" + server = MCPServer() + + server.register_prompt(simple_prompt, description="Custom prompt description") + + prompt_info = server._registered_prompts["simple_prompt"] + assert prompt_info["description"] == "Custom prompt description" + + def test_register_async_prompt(self): + """Test registering an async prompt.""" + server = MCPServer() + + server.register_prompt(async_prompt) + + prompt_info = server._registered_prompts["async_prompt"] + assert prompt_info["is_async"] is True + + def test_register_prompts_as_list(self): + """Test registering multiple prompts as a list.""" + server = MCPServer() + + server.register_prompts([simple_prompt, async_prompt]) + + assert len(server._registered_prompts) == 2 + assert "simple_prompt" in server._registered_prompts + assert "async_prompt" in server._registered_prompts + + def test_register_prompts_as_dict(self): + """Test registering multiple prompts as a dict.""" + server = MCPServer() + + prompts = {"explainer": simple_prompt, "code_gen": async_prompt} + server.register_prompts(prompts) + + assert len(server._registered_prompts) == 2 + assert "explainer" in server._registered_prompts + assert "code_gen" in server._registered_prompts + + assert server._registered_prompts["explainer"]["function"] == simple_prompt + assert server._registered_prompts["code_gen"]["function"] == async_prompt + + def test_register_prompts_invalid_type(self): + """Test registering prompts with invalid type.""" + server = MCPServer() + + with pytest.raises( + ValueError, match="prompts must be a list of functions or dict" + ): + server.register_prompts("invalid") + + def test_list_prompts(self): + """Test listing registered prompts.""" + server = MCPServer() + + # Initially empty + assert server.list_prompts() == [] + + # After registering prompts + server.register_prompt(simple_prompt) + server.register_prompt(prompt_with_detailed_docstring) + + prompts = server.list_prompts() + assert len(prompts) == 2 + + prompt_names = [p["name"] for p in prompts] + assert "simple_prompt" in prompt_names + assert "prompt_with_detailed_docstring" in prompt_names + + # Check structure + for prompt in prompts: + assert "name" in prompt + assert "description" in prompt + + def test_get_prompt_info(self): + """Test getting prompt information.""" + server = MCPServer() + + # Non-existent prompt + assert server.get_prompt_info("nonexistent") is None + + # Existing prompt + server.register_prompt(simple_prompt) + info = server.get_prompt_info("simple_prompt") + + assert info is not None + assert info["name"] == "simple_prompt" + assert info["function"] == simple_prompt + + def test_tools_and_prompts_together(self): + """Test registering both tools and prompts in the same server.""" + server = MCPServer() + + server.register_tool(simple_function) + server.register_prompt(simple_prompt) + + assert len(server._registered_tools) == 1 + assert len(server._registered_prompts) == 1 + assert "simple_function" in server._registered_tools + assert "simple_prompt" in server._registered_prompts + class TestMCPServerDirect: """Test direct MCPServer instantiation.""" From d373b345ce4040fffd4dfd0aa7954b333b0b503f Mon Sep 17 00:00:00 2001 From: Jeremy Tuloup Date: Mon, 27 Oct 2025 14:14:39 +0100 Subject: [PATCH 2/4] lint --- jupyter_server_mcp/extension.py | 4 +++- jupyter_server_mcp/mcp_server.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/jupyter_server_mcp/extension.py b/jupyter_server_mcp/extension.py index 73362db..d47715f 100644 --- a/jupyter_server_mcp/extension.py +++ b/jupyter_server_mcp/extension.py @@ -254,7 +254,9 @@ def _discover_entrypoint_prompts(self) -> list[str]: continue # Validate and collect prompt specs - valid_specs = [spec for spec in prompt_specs if isinstance(spec, str)] + valid_specs = [ + spec for spec in prompt_specs if isinstance(spec, str) + ] invalid_count = len(prompt_specs) - len(valid_specs) if invalid_count > 0: diff --git a/jupyter_server_mcp/mcp_server.py b/jupyter_server_mcp/mcp_server.py index 592a47f..0b8ea42 100644 --- a/jupyter_server_mcp/mcp_server.py +++ b/jupyter_server_mcp/mcp_server.py @@ -392,7 +392,9 @@ def register_prompts(self, prompts: list[Callable] | dict[str, Callable]): for name, func in prompts.items(): self.register_prompt(func, name=name) else: - msg = "prompts must be a list of functions or dict mapping names to functions" + msg = ( + "prompts must be a list of functions or dict mapping names to functions" + ) raise ValueError(msg) def list_prompts(self) -> list[dict[str, Any]]: From df514e876516a8af6e70b9d919d47aeaf92e9a9b Mon Sep 17 00:00:00 2001 From: Jeremy Tuloup Date: Mon, 27 Oct 2025 14:17:30 +0100 Subject: [PATCH 3/4] more linting --- tests/test_extension.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_extension.py b/tests/test_extension.py index 4515203..d675ca0 100644 --- a/tests/test_extension.py +++ b/tests/test_extension.py @@ -613,9 +613,7 @@ async def test_start_extension_with_prompts_entrypoints_and_config(self): "_discover_entrypoint_prompts", return_value=discovered_prompts, ), - patch.object( - extension, "_discover_entrypoint_tools", return_value=[] - ), + patch.object(extension, "_discover_entrypoint_tools", return_value=[]), ): await extension.start_extension() From b19d8f0673e30ea0df99cc88e48022d029873161 Mon Sep 17 00:00:00 2001 From: Jeremy Tuloup Date: Mon, 27 Oct 2025 17:57:27 +0100 Subject: [PATCH 4/4] fix tests --- tests/test_extension.py | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/tests/test_extension.py b/tests/test_extension.py index d675ca0..6ad8b8f 100644 --- a/tests/test_extension.py +++ b/tests/test_extension.py @@ -72,7 +72,8 @@ async def test_start_extension_success(self): with patch("jupyter_server_mcp.extension.MCPServer") as mock_mcp_class: mock_server = Mock() mock_server.start_server = AsyncMock() - mock_server._registered_tools = [] # Use list instead of Mock + mock_server._registered_tools = {} + mock_server._registered_prompts = {} mock_mcp_class.return_value = mock_server await extension.start_extension() @@ -163,7 +164,8 @@ async def test_full_lifecycle(self): with patch("jupyter_server_mcp.extension.MCPServer") as mock_mcp_class: mock_server = Mock() mock_server.start_server = AsyncMock() - mock_server._registered_tools = [] # Use list instead of Mock + mock_server._registered_tools = {} + mock_server._registered_prompts = {} mock_mcp_class.return_value = mock_server # Start extension @@ -323,8 +325,8 @@ def test_register_configured_prompts_valid(self): extension = MCPExtensionApp() extension.mcp_server_instance = Mock() extension.mcp_prompts = [ - "package.prompts:code_review", - "package.prompts:documentation", + "os:getcwd", + "json:dumps", ] # Capture log output @@ -336,6 +338,12 @@ def test_register_configured_prompts_valid(self): # Check log messages mock_logger.info.assert_any_call("Registering 2 prompts from configuration") + mock_logger.info.assert_any_call( + "✅ Registered prompt from configuration: os:getcwd" + ) + mock_logger.info.assert_any_call( + "✅ Registered prompt from configuration: json:dumps" + ) def test_register_configured_prompts_with_errors(self): """Test registering prompts when some fail to load.""" @@ -378,6 +386,7 @@ async def test_start_extension_with_tools(self): "getcwd": {}, "sqrt": {}, } # Mock registered tools + mock_server._registered_prompts = {} mock_mcp_class.return_value = mock_server await extension.start_extension() @@ -401,6 +410,7 @@ async def test_start_extension_no_tools(self): mock_server = Mock() mock_server.start_server = AsyncMock() mock_server._registered_tools = {} + mock_server._registered_prompts = {} mock_mcp_class.return_value = mock_server await extension.start_extension() @@ -495,6 +505,7 @@ async def test_start_extension_with_entrypoints_and_config(self): mock_server = Mock() mock_server.start_server = AsyncMock() mock_server._registered_tools = {"getcwd": {}, "dumps": {}} + mock_server._registered_prompts = {} mock_mcp_class.return_value = mock_server with patch.object( @@ -593,17 +604,17 @@ async def test_start_extension_with_prompts_entrypoints_and_config(self): extension = MCPExtensionApp() extension.mcp_port = 3087 extension.use_tool_discovery = True - extension.mcp_prompts = ["package.prompts:config_prompt"] + extension.mcp_prompts = ["json:dumps"] - discovered_prompts = ["package.prompts:discovered_prompt"] + discovered_prompts = ["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 = {} mock_server._registered_prompts = { - "discovered_prompt": {}, - "config_prompt": {}, + "getcwd": {}, + "dumps": {}, } mock_mcp_class.return_value = mock_server @@ -626,13 +637,13 @@ async def test_start_extension_with_tools_and_prompts(self): extension = MCPExtensionApp() extension.mcp_port = 3088 extension.mcp_tools = ["os:getcwd"] - extension.mcp_prompts = ["package.prompts:review"] + extension.mcp_prompts = ["json:dumps"] with patch("jupyter_server_mcp.extension.MCPServer") as mock_mcp_class: mock_server = Mock() mock_server.start_server = AsyncMock() mock_server._registered_tools = {"getcwd": {}} - mock_server._registered_prompts = {"review": {}} + mock_server._registered_prompts = {"dumps": {}} mock_mcp_class.return_value = mock_server with (