Initial Implementation of Multi-Notebook Management Support

[ChengJiale150]

Summary

The current implementation of jupyter_mcp_server/server.py only supports the management of a single Notebook, with all operations hardcoded to the "default" notebook name. To enhance the system's flexibility and scalability, this proposal outlines the addition of comprehensive multi-notebook management features. This includes full lifecycle management for Notebooks and adding multi-notebook support to existing Cell tools.

Current Problems

1. Single Notebook Limitation

• The system can only manage one Notebook, using the hardcoded identifier "default".
• All tool functions access the Notebook via notebook_manager.get_notebook_connection("default").
• It's impossible to operate on multiple different Notebook files simultaneously.
• There is a lack of isolation and independent management capabilities between Notebooks.

2. Lack of Lifecycle Management

The current system is missing the following key features:

• Connection Management: No ability to dynamically connect to a specified Notebook file.
• Creation: No functionality to create new Notebook files.
• Listing: No way to view all currently connected Notebooks.
• Restarting: No feature to restart the Kernel of a specific Notebook.
• Disconnection: No mechanism to gracefully disconnect from a Notebook.

Improvement Plan

1. connect_notebook Tool

@mcp.tool()
async def connect_notebook(
    notebook_name: str,       # Unique identifier
    notebook_path: str,       # Path to the Notebook file
    mode: Literal["connect", "create"] = "connect",
    kernel_id: Optional[str] = None
) -> str:

Functionality:

• connect mode: Connect to an existing Notebook.
• create mode: Create a new Notebook and connect to it.
• kernel_id: Specify a Kernel ID. If not provided, a new Kernel will be created.

2. list_notebook Tool (Replaces get_notebook_info)

@mcp.tool()
async def list_notebook() -> str:

Returns information for all currently connected Notebooks in a TSV format, including:

• Notebook Name (unique identifier)
• Notebook File Path
• Number of Cells
• Status (whether it is currently active)

3. restart_notebook Tool

@mcp.tool()
async def restart_notebook(notebook_name: str) -> str:

Restarts the Kernel for the specified Notebook, clearing its memory state and imported packages.

4. disconnect_notebook Tool

@mcp.tool()
async def disconnect_notebook(notebook_name: str) -> str:

Disconnects the Kernel for the specified Notebook and releases associated resources.

5. switch_notebook Tool

@mcp.tool()
async def switch_notebook(notebook_name: str) -> str:

Switches the currently active Notebook. All subsequent Cell operations will be performed on this Notebook.

Functionality:

• Set the specified Notebook as the current active one.
• Verify that the target Notebook is already connected.
• Return the result of the switch and information about the new active Notebook.

Phase Two: Implement the Active Notebook Mechanism

The switch_notebook tool will maintain the state of the currently active Notebook. Existing Cell tools will not require parameter modifications.

Implementation Points:

• Add a current_notebook attribute to the NotebookManager.
• Replace the hardcoded "default" with the dynamic active Notebook.
• Add validation to ensure the active Notebook exists.
• Provide a default Notebook mechanism for backward compatibility.

Affected Tools (No parameter changes needed; only internal logic requires modification):

1. insert_cell() - Inserts a cell in the active Notebook.
2. insert_execute_code_cell() - Inserts and executes a cell in the active Notebook.
3. overwrite_cell_source() - Modifies a cell in the active Notebook.
4. execute_cell_with_progress() - Executes a cell in the active Notebook.
5. execute_cell_simple_timeout() - Executes a cell in the active Notebook.
6. execute_cell_streaming() - Executes a cell in the active Notebook with streaming output.
7. read_all_cells() - Reads all cells from the active Notebook.
8. list_cell() - Lists cells in the active Notebook.
9. read_cell() - Reads a specific cell from the active Notebook.
10. delete_cell() - Deletes a cell from the active Notebook.

Technical Implementation Details

1. Leverage Existing Infrastructure

jupyter_mcp_server already has a basic NotebookManager class that includes:

• Support for multi-notebook management.
• A connection context manager.
• Kernel lifecycle management.

2. NotebookManager Enhancements

The following functionality needs to be added to the NotebookManager class:

class NotebookManager:
    def __init__(self):
        self._notebooks: Dict[str, Dict[str, Any]] = {}
        self._current_notebook: Optional[str] = None  # New: The currently active notebook

    def set_current_notebook(self, name: str) -> bool:
        """Sets the currently active notebook."""
        if name in self._notebooks:
            self._current_notebook = name
            return True
        return False

    def get_current_notebook(self) -> Optional[str]:
        """Gets the name of the currently active notebook."""
        return self._current_notebook

    def get_current_connection(self) -> NotebookConnection:
        """Gets the connection for the currently active notebook."""
        current = self._current_notebook or "default"  # Backward compatibility
        return self.get_notebook_connection(current)

3. Utilize jupyter-server-client

• Check if a Notebook exists.
• Create a new Notebook.

4. Backward Compatibility

• The DOCUMENT_URL, DOCUMENT_TOKEN, DOCUMENT_ID, RUNTIME_URL, and RUNTIME_TOKEN configuration items will be retained for backward compatibility.
• The system will default to connecting to "default" as the active notebook.
• Existing configurations and HTTP interfaces will remain unchanged.

Final Remarks

This is an initial implementation to introduce multi-notebook management. Further improvements are still needed.