Fix: sync error (#148)

* add thread lock and transaction to fix sync error

* simple api

Author: ChengJiale150

jupyter_mcp_server/__version__.py

@@ -4,4 +4,4 @@
 
 """Jupyter MCP Server."""
 
-__version__ = "0.18.0"
+__version__ = "0.18.1"

jupyter_mcp_server/tools/delete_cell_tool.py

@@ -10,12 +10,12 @@
 from jupyter_server_client import JupyterServerClient
 from jupyter_mcp_server.tools._base import BaseTool, ServerMode
 from jupyter_mcp_server.notebook_manager import NotebookManager
-from jupyter_mcp_server.utils import get_current_notebook_context, get_jupyter_ydoc, clean_notebook_outputs
+from jupyter_mcp_server.utils import get_current_notebook_context, get_notebook_model, clean_notebook_outputs
 
 
 class DeleteCellTool(BaseTool):
     """Tool to delete a specific cell from a notebook."""
-    
+
     def _get_cell_source(self, cell: Any) -> str:
         """Get the cell source from the cell"""
         cell_source = cell.get("source", "")
@@ -38,35 +38,21 @@ async def _delete_cell_ydoc(
             cell_index: Index of cell to delete
             
         Returns:
-            Success message
+            NotebookNode
         """
-        # Get file_id from file_id_manager
-        file_id_manager = serverapp.web_app.settings.get("file_id_manager")
-        if file_id_manager is None:
-            raise RuntimeError("file_id_manager not available in serverapp")
-        
-        file_id = file_id_manager.get_id(notebook_path)
-        
-        # Try to get YDoc
-        ydoc = await get_jupyter_ydoc(serverapp, file_id)
-        
-        if ydoc:
-            if cell_index >= len(ydoc.ycells):
+        nb = await get_notebook_model(serverapp, notebook_path)
+        if nb:
+            if cell_index >= len(nb):
                 raise ValueError(
-                    f"Cell index {cell_index} is out of range. Notebook has {len(ydoc.ycells)} cells."
+                    f"Cell index {cell_index} is out of range. Notebook has {len(nb)} cells."
                 )
 
-            cell = ydoc.ycells[cell_index]
-            result = {
+            cell = nb.delete_cell(cell_index)
+            return {
                 "index": cell_index,
-                "cell_type": cell.get("cell_type", "unknown"),
+                "cell_type": cell.cell_type,
                 "source": self._get_cell_source(cell),
             }
-            
-            # Delete the cell from YDoc
-            del ydoc.ycells[cell_index]
-            
-            return result
         else:
             # YDoc not available, use file operations
             return await self._delete_cell_file(notebook_path, cell_index)

jupyter_mcp_server/tools/insert_cell_tool.py

@@ -10,72 +10,81 @@
 from jupyter_server_client import JupyterServerClient
 from jupyter_mcp_server.tools._base import BaseTool, ServerMode
 from jupyter_mcp_server.notebook_manager import NotebookManager
-from jupyter_mcp_server.utils import get_current_notebook_context, get_jupyter_ydoc, clean_notebook_outputs
+from jupyter_mcp_server.utils import get_current_notebook_context, get_notebook_model, clean_notebook_outputs
 from jupyter_mcp_server.models import Notebook
 
 
+
 class InsertCellTool(BaseTool):
     """Tool to insert a cell at a specified position."""
 
+    def _validate_cell_insertion_params(
+        self,
+        cell_index: int,
+        total_cells: int,
+        cell_type: str
+    ) -> int:
+        """Validate and normalize cell insertion parameters.
+        
+        Args:
+            cell_index: Target index for insertion (-1 for append)
+            total_cells: Total number of cells in the notebook
+            cell_type: Type of cell to insert
+            
+        Returns:
+            Normalized actual_index for insertion
+            
+        Raises:
+            IndexError: When cell_index is out of valid range
+            ValueError: When cell_type is invalid
+        """
+        if cell_index < -1 or cell_index > total_cells:
+            raise IndexError(
+                f"Index {cell_index} is outside valid range [-1, {total_cells}]. "
+                f"Use -1 to append at end."
+            )
+        
+        # Normalize -1 to append position
+        actual_index = cell_index if cell_index != -1 else total_cells
+        return actual_index
+    
     async def _insert_cell_ydoc(
         self,
         serverapp: Any,
         notebook_path: str,
         cell_index: int,
         cell_type: Literal["code", "markdown"],
         cell_source: str
-    ) -> tuple[int, int]:
+    ) -> tuple[Notebook, int, int]:
         """Insert cell using YDoc (collaborative editing mode).
         
         Args:
             serverapp: Jupyter ServerApp instance
             notebook_path: Path to the notebook
             cell_index: Index to insert at (-1 for append)
-            cell_type: Type of cell to insert
+            cell_type: Type of cell to insert ("code", "markdown")
             cell_source: Source content for the cell
             
         Returns:
-            Success message with surrounding cells info
+            Tuple of (notebook, actual_index, total_cells_after_insertion)
+            
+        Raises:
+            IndexError: When cell_index is out of range
         """
-        # Get file_id from file_id_manager
-        file_id_manager = serverapp.web_app.settings.get("file_id_manager")
-        if file_id_manager is None:
-            raise RuntimeError("file_id_manager not available in serverapp")
+        nb = await get_notebook_model(serverapp, notebook_path)
 
-        file_id = file_id_manager.get_id(notebook_path)
-        
-        # Try to get YDoc
-        ydoc = await get_jupyter_ydoc(serverapp, file_id)
-        
-        if ydoc:
+        if nb:
             # Notebook is open in collaborative mode, use YDoc
-            total_cells = len(ydoc.ycells)
-            actual_index = cell_index if cell_index != -1 else total_cells
-            
-            if actual_index < 0 or actual_index > total_cells:
-                raise ValueError(
-                    f"Cell index {cell_index} is out of range. Notebook has {total_cells} cells. Use -1 to append at end."
-                )
-            
-            # Create the cell
-            cell = {
-                "cell_type": cell_type,
-                "source": "",
-            }
-            ycell = ydoc.create_ycell(cell)
+            total_cells = len(nb)
 
-            # Insert at the specified position
-            if actual_index >= total_cells:
-                ydoc.ycells.append(ycell)
-            else:
-                ydoc.ycells.insert(actual_index, ycell)
+            # Validate insertion parameters
+            actual_index = self._validate_cell_insertion_params(
+                cell_index, total_cells, cell_type
+            )
 
-            # Write content to the cell collaboratively
-            if cell_source:
-                # Set the source directly on the ycell
-                ycell["source"] = cell_source
+            nb.insert_cell(actual_index, cell_source, cell_type)
 
-            return actual_index, len(ydoc.ycells)
+            return Notebook(**nb.as_dict()), actual_index, len(nb)
         else:
             # YDoc not available, use file operations
             return await self._insert_cell_file(notebook_path, cell_index, cell_type, cell_source)
@@ -86,17 +95,21 @@ async def _insert_cell_file(
         cell_index: int,
         cell_type: Literal["code", "markdown"],
         cell_source: str
-    ) -> tuple[int, int]:
+    ) -> tuple[Notebook, int, int]:
         """Insert cell using file operations (non-collaborative mode).
-        
+                
         Args:
             notebook_path: Absolute path to the notebook
             cell_index: Index to insert at (-1 for append)
-            cell_type: Type of cell to insert
+            cell_type: Type of cell to insert ("code", "markdown")
             cell_source: Source content for the cell
             
         Returns:
-            Success message with surrounding cells info
+            Tuple of (notebook, actual_index, total_cells_after_insertion)
+            
+        Raises:
+            IndexError: When cell_index is out of range
+            ValueError: When cell_type is invalid
         """
         # Read notebook file
         with open(notebook_path, "r", encoding="utf-8") as f:
@@ -107,55 +120,63 @@ async def _insert_cell_file(
         clean_notebook_outputs(notebook)
 
         total_cells = len(notebook.cells)
-        actual_index = cell_index if cell_index != -1 else total_cells
 
-        if actual_index < 0 or actual_index > total_cells:
-            raise ValueError(
-                f"Cell index {cell_index} is out of range. Notebook has {total_cells} cells. Use -1 to append at end."
-            )
+        # Validate insertion parameters
+        actual_index = self._validate_cell_insertion_params(
+            cell_index, total_cells, cell_type
+        )
 
-        # Create and insert the cell
+        # Create and insert the cell using unified method
+         # Create and insert the cell
         if cell_type == "code":
             new_cell = nbformat.v4.new_code_cell(source=cell_source or "")
         elif cell_type == "markdown":
             new_cell = nbformat.v4.new_markdown_cell(source=cell_source or "")
-        else:
-            raise ValueError(f"Invalid cell_type: {cell_type}. Must be 'code' or 'markdown'.")
-        
         notebook.cells.insert(actual_index, new_cell)
 
         # Write back to file
         with open(notebook_path, "w", encoding="utf-8") as f:
             nbformat.write(notebook, f)
 
-        return actual_index, len(notebook.cells)
+        notebook = Notebook(**notebook)
+        
+        return notebook, actual_index, len(notebook.cells)
 
     async def _insert_cell_websocket(
         self,
         notebook_manager: NotebookManager,
         cell_index: int,
         cell_type: Literal["code", "markdown"],
         cell_source: str
-    ) -> tuple[int, Notebook]:
+    ) -> tuple[Notebook, int, int]:
         """Insert cell using WebSocket connection (MCP_SERVER mode).
         
         Args:
             notebook_manager: Notebook manager instance
             cell_index: Index to insert at (-1 for append)
-            cell_type: Type of cell to insert
+            cell_type: Type of cell to insert ("code", "markdown")
             cell_source: Source content for the cell
             
         Returns:
-            Success message with surrounding cells info
+            Tuple of (notebook, actual_index, total_cells_after_insertion)
+            
+        Raises:
+            IndexError: When cell_index is out of range
+            ValueError: When cell_type is invalid
         """
         async with notebook_manager.get_current_connection() as notebook:
-            actual_index = cell_index if cell_index != -1 else len(notebook)
-            if actual_index < 0 or actual_index > len(notebook):
-                raise ValueError(f"Cell index {cell_index} out of range")
+            total_cells = len(notebook)
+            
+            # Validate insertion parameters
+            actual_index = self._validate_cell_insertion_params(
+                cell_index, total_cells, cell_type
+            )
 
+            # Use the unified insert_cell method pattern
+            # The remote notebook should have: insert_cell(index, source, cell_type)
             notebook.insert_cell(actual_index, cell_source, cell_type)
 
-            return actual_index, Notebook(**notebook.as_dict())
+            return Notebook(**notebook.as_dict()), actual_index, len(notebook)
 
     async def execute(
         self,
@@ -174,33 +195,44 @@ async def execute(
     ) -> str:
         """Execute the insert_cell tool.
         
-        This tool supports three modes of operation:
+        This tool supports three modes of operation following a unified insertion pattern:
         
         1. JUPYTER_SERVER mode with YDoc (collaborative):
            - Checks if notebook is open in a collaborative session
            - Uses YDoc for real-time collaborative editing
            - Changes are immediately visible to all connected users
+           - Operations protected by thread locks and YDoc transactions
            
         2. JUPYTER_SERVER mode without YDoc (file-based):
            - Falls back to direct file operations using nbformat
            - Suitable when notebook is not actively being edited
            
         3. MCP_SERVER mode (WebSocket):
            - Uses WebSocket connection to remote Jupyter server
-           - Accesses YDoc through NbModelClient
+           - Delegates to remote notebook's unified insert_cell method
+        
+        Thread Safety:
+        - YDoc mode: Protected by thread lock + YDoc transaction (atomic)
+        - File mode: No synchronization needed (single-threaded file I/O)
+        - WebSocket mode: Remote server handles synchronization
         
         Args:
             mode: Server mode (MCP_SERVER or JUPYTER_SERVER)
             server_client: HTTP client for MCP_SERVER mode
             contents_manager: Direct API access for JUPYTER_SERVER mode
             notebook_manager: Notebook manager instance
             cell_index: Target index for insertion (0-based, -1 to append)
-            cell_type: Type of cell ("code" or "markdown")
+            cell_type: Type of cell ("code", "markdown")
             cell_source: Source content for the cell
             **kwargs: Additional parameters
             
         Returns:
             Success message with surrounding cells info
+            
+        Raises:
+            ValueError: When mode is invalid or required clients are missing
+            IndexError: When cell_index is out of range
+            ValueError: When cell_type is invalid
         """
         if mode == ServerMode.JUPYTER_SERVER and contents_manager is not None:
             # JUPYTER_SERVER mode: Try YDoc first, fall back to file operations
@@ -216,29 +248,27 @@ async def execute(
                 notebook_path = str(Path(root_dir) / notebook_path)
 
             if serverapp:
-                # Try YDoc approach first
-                actual_index, new_total_cells = await self._insert_cell_ydoc(serverapp, notebook_path, cell_index, cell_type, cell_source)
+                # Try YDoc approach first (with thread safety and transactions)
+                notebook, actual_index, new_total_cells = await self._insert_cell_ydoc(
+                    serverapp, notebook_path, cell_index, cell_type, cell_source
+                )
             else:
                 # Fall back to file operations
-                actual_index, new_total_cells = await self._insert_cell_file(notebook_path, cell_index, cell_type, cell_source)
-            
-            # Load notebook using same API
-            notebook_path = notebook_manager.get_current_notebook_path()
-            model = await contents_manager.get(notebook_path, content=True, type='notebook')
-            if 'content' not in model:
-                raise ValueError(f"Could not read notebook content from {notebook_path}")
-            notebook = Notebook(**model['content'])
+                notebook, actual_index, new_total_cells = await self._insert_cell_file(
+                    notebook_path, cell_index, cell_type, cell_source
+                )
 
         elif mode == ServerMode.MCP_SERVER and notebook_manager is not None:
-            # MCP_SERVER mode: Use WebSocket connection
-            actual_index, notebook = await self._insert_cell_websocket(notebook_manager, cell_index, cell_type, cell_source)
-            new_total_cells = len(notebook)
+            # MCP_SERVER mode: Use WebSocket connection with unified insert_cell pattern
+            notebook, actual_index, new_total_cells = await self._insert_cell_websocket(
+                notebook_manager, cell_index, cell_type, cell_source
+            )
         else:
             raise ValueError(f"Invalid mode or missing required clients: mode={mode}")
 
         info_list = [f"Cell inserted successfully at index {actual_index} ({cell_type})!"]
         info_list.append(f"Notebook now has {new_total_cells} cells, showing surrounding cells:")
-        # near to end
+        # Show context near the insertion
         if new_total_cells - actual_index < 5:
             start_index = max(0, new_total_cells - 10)
         else: