Enhance: delete cell (#150)

* enhance delete cell

* fix the missing index

Author: ChengJiale150

docs/docs/tools/index.mdx

@@ -135,11 +135,12 @@ The server currently offers 14 tools organized into 3 categories:
 
 #### 13. `delete_cell`
 
-- Delete a specific cell from the currently activated notebook and return the cell source of deleted cell.
+- Delete a specific cell or multiple cells from the currently activated notebook and return the cell source of deleted cells (if include_source=True).
 - When deleting many cells, MUST delete them in descending order of their index to avoid index shifting.
 - Input:
-  - `cell_index`(int): Index of the cell to delete (0-based)
-- Returns: Success message with deletion confirmation and the source code of the deleted cell
+  - `cell_indices`(list[int]): List of indices of the cells to delete (0-based)
+  - `include_source`(bool, optional): Whether to include the source of deleted cells (default: true)
+- Returns: Success message with deletion confirmation and the source code of the deleted cells (if include_source=True)
 
 #### 14. `execute_code`
 

jupyter_mcp_server/server.py

@@ -462,18 +462,19 @@ async def read_cell(
 
 @mcp.tool()
 async def delete_cell(
-    cell_index: Annotated[int, Field(description="Index of the cell to delete (0-based)", ge=0)],
-) -> Annotated[str, Field(description="Success message and the cell source of deleted cell")]:
-    """Delete a specific cell from the currently activated notebook and return the cell source of deleted cell.
-    When deleting many cells, MUST delete them in descending order of their index to avoid index shifting."""
+    cell_indices: Annotated[list[int], Field(description="List of cell indices to delete (0-based)",min_items=1)],
+    include_source: Annotated[bool, Field(description="Whether to include the source of deleted cells")] = True,
+) -> Annotated[str, Field(description="Success message with list of deleted cells and their source (if include_source=True)")]:
+    """Delete specific cells from the currently activated notebook and return the cell source of deleted cells (if include_source=True)."""
     return await safe_notebook_operation(
         lambda: DeleteCellTool().execute(
             mode=server_context.mode,
             server_client=server_context.server_client,
             contents_manager=server_context.contents_manager,
             kernel_manager=server_context.kernel_manager,
             notebook_manager=notebook_manager,
-            cell_index=cell_index,
+            cell_indices=cell_indices,
+            include_source=include_source,
         )
     )
 

jupyter_mcp_server/tools/delete_cell_tool.py

@@ -14,7 +14,7 @@
 
 
 class DeleteCellTool(BaseTool):
-    """Tool to delete a specific cell from a notebook."""
+    """Tool to delete specific cells from a notebook."""
 
     def _get_cell_source(self, cell: Any) -> str:
         """Get the cell source from the cell"""
@@ -28,102 +28,98 @@ async def _delete_cell_ydoc(
         self,
         serverapp: Any,
         notebook_path: str,
-        cell_index: int
-    ) -> dict:
+        cell_indices: list[int]
+    ) -> list:
         """Delete cell using YDoc (collaborative editing mode).
         
         Args:
             serverapp: Jupyter ServerApp instance
             notebook_path: Path to the notebook
-            cell_index: Index of cell to delete
+            cell_indices: List of indices of cells to delete
             
         Returns:
             NotebookNode
         """
         nb = await get_notebook_model(serverapp, notebook_path)
         if nb:
-            if cell_index >= len(nb):
+            if max(cell_indices) >= len(nb):
                 raise ValueError(
-                    f"Cell index {cell_index} is out of range. Notebook has {len(nb)} cells."
+                    f"Cell index {max(cell_indices)} is out of range. Notebook has {len(nb)} cells."
                 )
 
-            cell = nb.delete_cell(cell_index)
-            return {
-                "index": cell_index,
-                "cell_type": cell.cell_type,
-                "source": self._get_cell_source(cell),
-            }
+            cells = nb.delete_many_cells(cell_indices)
+            return cells
         else:
             # YDoc not available, use file operations
-            return await self._delete_cell_file(notebook_path, cell_index)
+            return await self._delete_cell_file(notebook_path, cell_indices)
 
     async def _delete_cell_file(
         self,
         notebook_path: str,
-        cell_index: int
-    ) -> dict:
+        cell_indices: list[int]
+    ) -> list:
         """Delete cell using file operations (non-collaborative mode).
         
         Args:
             notebook_path: Absolute path to the notebook
-            cell_index: Index of cell to delete
+            cell_indices: List of indices of cells to delete
             
         Returns:
-            Success message
+            List of deleted cells
         """
         # Read notebook file as version 4 for consistency
         with open(notebook_path, "r", encoding="utf-8") as f:
             notebook = nbformat.read(f, as_version=4)
 
         clean_notebook_outputs(notebook)
 
-        if cell_index >= len(notebook.cells):
+        if max(cell_indices) >= len(notebook.cells):
             raise ValueError(
-                f"Cell index {cell_index} is out of range. Notebook has {len(notebook.cells)} cells."
+                f"Cell index {max(cell_indices)} is out of range. Notebook has {len(notebook.cells)} cells."
             )
 
-        cell = notebook.cells[cell_index]
-        result = {
-            "index": cell_index,
-            "cell_type": cell.cell_type,
-            "source": self._get_cell_source(cell),
-        }
+        deleted_cells = []
+        for cell_index in cell_indices:
+            cell = notebook.cells[cell_index]
+            result = {
+                "index": cell_index,
+                "cell_type": cell.cell_type,
+                "source": self._get_cell_source(cell),
+            }
+            deleted_cells.append(result)
 
         # Delete the cell
-        notebook.cells.pop(cell_index)
+        for cell_index in sorted(cell_indices, reverse=True):
+            notebook.cells.pop(cell_index)
 
         # Write back to file
         with open(notebook_path, "w", encoding="utf-8") as f:
             nbformat.write(notebook, f)
 
-        return result
+        return deleted_cells
 
     async def _delete_cell_websocket(
         self,
         notebook_manager: NotebookManager,
-        cell_index: int
-    ) -> dict:
+        cell_indices: list[int]
+    ) -> list:
         """Delete cell using WebSocket connection (MCP_SERVER mode).
         
         Args:
             notebook_manager: Notebook manager instance
-            cell_index: Index of cell to delete
+            cell_indices: List of indices of cells to delete
             
         Returns:
-            Success message
+            List of deleted cell information
         """
         async with notebook_manager.get_current_connection() as notebook:
-            if cell_index >= len(notebook):
+            if max(cell_indices) >= len(notebook):
                 raise ValueError(
-                    f"Cell index {cell_index} is out of range. Notebook has {len(notebook)} cells."
+                    f"Cell index {max(cell_indices)} is out of range. Notebook has {len(notebook)} cells."
                 )
 
-            cell = notebook.delete_cell(cell_index)
-            return {
-                "index": cell_index,
-                "cell_type": cell.cell_type,
-                "source": self._get_cell_source(cell),
-            }
+            cells = notebook.delete_many_cells(cell_indices)
+            return cells
 
     async def execute(
         self,
@@ -135,7 +131,8 @@ async def execute(
         kernel_spec_manager: Optional[Any] = None,
         notebook_manager: Optional[NotebookManager] = None,
         # Tool-specific parameters
-        cell_index: int = None,
+        cell_indices: list[int] = None,
+        include_source: bool = True,
         **kwargs
     ) -> str:
         """Execute the delete_cell tool.
@@ -181,17 +178,22 @@ async def execute(
 
             if serverapp:
                 # Try YDoc approach first
-                cell_info = await self._delete_cell_ydoc(serverapp, notebook_path, cell_index)
+                cells = await self._delete_cell_ydoc(serverapp, notebook_path, cell_indices)
             else:
                 # Fall back to file operations
-                cell_info = await self._delete_cell_file(notebook_path, cell_index)
+                cells = await self._delete_cell_file(notebook_path, cell_indices)
 
         elif mode == ServerMode.MCP_SERVER and notebook_manager is not None:
             # MCP_SERVER mode: Use WebSocket connection
-            cell_info = await self._delete_cell_websocket(notebook_manager, cell_index)
+            cells = await self._delete_cell_websocket(notebook_manager, cell_indices)
         else:
             raise ValueError(f"Invalid mode or missing required clients: mode={mode}")
 
-        info_list = [f"Cell {cell_info['index']} ({cell_info['cell_type']}) deleted successfully."]
-        info_list.append(f"deleted cell source:\n{cell_info['source']}")
+        info_list = []
+        for cell_index, cell_info in zip(cell_indices, cells):
+            info_list.append(f"Cell {cell_index} ({cell_info['cell_type']}) deleted successfully.")
+            if include_source:
+                info_list.append(f"deleted cell source:\n{cell_info['source']}")
+                info_list.append("\n---\n")
+        
         return "\n".join(info_list)

pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
 dependencies = [
     "jupyter-kernel-client>=0.7.3",
     "jupyter-mcp-tools>=0.1.4",
-    "jupyter-nbmodel-client>=0.14.2",
+    "jupyter-nbmodel-client>=0.14.4",
     "jupyter-server-nbmodel",
     "jupyter-server-client",
     "jupyter_server>=1.6,<3",

tests/test_common.py

@@ -336,8 +336,8 @@ async def list_kernels(self):
         return self._extract_text_content(result)
 
     @requires_session
-    async def delete_cell(self, cell_index):
-        result = await self._call_tool_safe("delete_cell", {"cell_index": cell_index})
+    async def delete_cell(self, cell_indices: list[int], include_source: bool = True):
+        result = await self._call_tool_safe("delete_cell", {"cell_indices": cell_indices, "include_source": include_source})
         return self._get_structured_content_safe(result) if result else None
 
     @requires_session

tests/test_tools.py

@@ -88,7 +88,7 @@ async def check_and_delete_cell(client: MCPClient, index, expected_type, content
         assert f"=====Cell {index} | type: {expected_type}" in cell_info['result'][0], "Cell metadata should be included"
         assert content in cell_info['result'][1], "Cell source should be included"
         # delete created cell
-        result = await client.delete_cell(index)
+        result = await client.delete_cell([index])
         assert result is not None, "delete_cell result should not be None"
         assert f"Cell {index} ({expected_type}) deleted successfully" in result["result"]
         assert f"deleted cell source:\n{content}" in result["result"]
@@ -158,7 +158,7 @@ async def test_multimodal_output(mcp_client_parametrized: MCPClient):
         assert isinstance(result['result'], list), "Result should be a list"
         assert isinstance(result['result'][0], dict)
         assert result['result'][0]['mimeType'] == "image/png", "Result should be a list of ImageContent"
-        await mcp_client_parametrized.delete_cell(1)
+        await mcp_client_parametrized.delete_cell([1])
 
 
 ###############################################################################