feat: Unify cell insertion functionality into insert_cell and insert_execute_code_cell

docs/docs/tools/index.mdx

@@ -1,78 +1,65 @@
 # Tools
 
-The server currently offers 13 tools:
+The server currently offers 11 tools:
 
-#### 1. `append_execute_code_cell`
+#### 1. `insert_cell`
 
-- Append at the end of the notebook a code cell with the provided source and execute it.
+- Insert a cell to specified position with unified API.
 - Input:
-  - `cell_source`(string): Code source.
-- Returns: List of outputs from the executed cell.
-
-#### 2. `append_markdown_cell`
-
-- Append at the end of the notebook a markdown cell with the provided source.
-- Input:
-  - `cell_source`(string): Markdown source.
-- Returns: Success message.
+  - `cell_index`(int): Target index for insertion (0-based). Use -1 to append at end.
+  - `cell_type`(string): Type of cell to insert ("code" or "markdown").
+  - `cell_source`(string): Source content for the cell.
+- Returns: Success message and the structure of its surrounding cells (up to 5 cells above and 5 cells below).
 
-#### 3. `insert_execute_code_cell`
+#### 2. `insert_execute_code_cell`
 
-- Insert a code cell at a specific index in the notebook and execute it.
+- Insert and execute a code cell in a Jupyter notebook.
 - Input:
-  - `cell_index`(int): Index where the cell should be inserted (0-based).
-  - `cell_source`(string): Code to be executed.
-- Returns: Cell output.
-
-#### 4. `insert_markdown_cell`
-
-- Insert a markdown cell at a specific index in the notebook.
-- Input:
-  - `cell_index`(int): Index where the cell should be inserted (0-based).
-  - `cell_source`(string): Markdown source.
-- Returns: Success message.
+  - `cell_index`(int): Index of the cell to insert (0-based). Use -1 to append at end and execute.
+  - `cell_source`(string): Code source.
+- Returns: List of outputs from the executed cell.
 
-#### 5. `delete_cell`
+#### 3. `delete_cell`
 
 - Delete a specific cell from the notebook.
 - Input:
   - `cell_index`(int): Index of the cell to delete (0-based).
 - Returns: Success message.
 
-#### 6. `get_notebook_info`
+#### 4. `get_notebook_info`
 
 - Get basic information about the notebook.
 - Returns: Dictionary with notebook path, total cells, and cell type counts.
 
-#### 7. `read_cell`
+#### 5. `read_cell`
 
 - Read a specific cell from the notebook.
 - Input:
   - `cell_index`(int): Index of the cell to read (0-based).
 - Returns: Dictionary with cell index, type, source, and outputs (for code cells).
 
-#### 8. `read_all_cells`
+#### 6. `read_all_cells`
 
 - Read all cells from the notebook.
 - Returns:  List of cell information including index, type, source, and outputs (for code cells).
 
-#### 9. `list_cell`
+#### 7. `list_cell`
 
 - List the basic information of all cells in the notebook.
 - Returns a formatted table showing the index, type, execution count (for code cells), and first line of each cell.
 - Provides a quick overview of the notebook structure and is useful for locating specific cells for operations.
 - Input: None
 - Returns: Formatted table string with cell information (Index, Type, Count, First Line).
 
-#### 10. `overwrite_cell_source`
+#### 8. `overwrite_cell_source`
 
 - Overwrite the source of an existing cell.
 - Input:
   - `cell_index`(int): Index of the cell to overwrite (0-based).
   - `cell_source`(string): New cell source - must match existing cell type.
 - Returns: Success message and diff style.
 
-#### 11. `execute_cell_streaming`
+#### 9. `execute_cell_streaming`
 
 - Execute cell with streaming progress updates. To be used for long-running cells.
 - Input:
@@ -82,7 +69,7 @@ The server currently offers 13 tools:
 - Returns:
   - `list[str]`: List of outputs including progress updates
 
-#### 12. `execute_cell_simple_timeout`
+#### 10. `execute_cell_simple_timeout`
 
 - Execute a cell with simple timeout (no forced real-time sync). To be used for short-running cells. This won't force real-time updates but will work reliably.
 - Input:
@@ -91,7 +78,7 @@ The server currently offers 13 tools:
 - Returns:
   - `list[str]`: List of outputs from the executed cell
 
-#### 13. `execute_cell_with_progress`
+#### 11. `execute_cell_with_progress`
 
 - Execute a specific cell with timeout and progress monitoring.
 - Input:

jupyter_mcp_server/server.py

@@ -24,8 +24,9 @@
 from starlette.middleware.cors import CORSMiddleware
 
 from jupyter_mcp_server.models import DocumentRuntime, CellInfo
-from jupyter_mcp_server.utils import extract_output, safe_extract_outputs, format_cell_list
+from jupyter_mcp_server.utils import extract_output, safe_extract_outputs, format_cell_list, get_surrounding_cells_info
 from jupyter_mcp_server.config import get_config, set_config
+from typing import Literal
 
 
 ###############################################################################
@@ -357,58 +358,109 @@ async def health_check(request: Request):
 
 
 @mcp.tool()
-async def append_markdown_cell(cell_source: str) -> str:
-    """Append at the end of the notebook a markdown cell with the provided source.
+async def insert_cell(
+    cell_index: int,
+    cell_type: Literal["code", "markdown"],
+    cell_source: str,
+) -> str:
+    """Insert a cell to specified position.
 
     Args:
-        cell_source: Markdown source
+        cell_index: target index for insertion (0-based). Use -1 to append at end.
+        cell_type: Type of cell to insert ("code" or "markdown")
+        cell_source: Source content for the cell
 
     Returns:
-        str: Success message
+        str: Success message and the structure of its surrounding cells (up to 5 cells above and 5 cells below)
     """
-    async def _append_markdown():
+    async def _insert_cell():
         notebook = None
         try:
             notebook = _get_notebook_client()
             await notebook.start()
-            notebook.add_markdown_cell(cell_source)
-            return "Jupyter Markdown cell added."
+
+            ydoc = notebook._doc
+            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."
+                )
+
+            if cell_type == "code":
+                if actual_index == total_cells:
+                    notebook.add_code_cell(cell_source)
+                else:
+                    notebook.insert_code_cell(actual_index, cell_source)
+            elif cell_type == "markdown":
+                if actual_index == total_cells:
+                    notebook.add_markdown_cell(cell_source)
+                else:
+                    notebook.insert_markdown_cell(actual_index, cell_source)
+
+            # Get surrounding cells info
+            new_total_cells = len(ydoc._ycells)
+            surrounding_info = get_surrounding_cells_info(notebook, actual_index, new_total_cells)
+
+            return f"Cell inserted successfully at index {actual_index} ({cell_type})!\n\nCurrent Surrounding Cells:\n{surrounding_info}"
         finally:
             if notebook:
                 try:
                     await notebook.stop()
                 except Exception as e:
-                    logger.warning(f"Error stopping notebook in append_markdown_cell: {e}")
+                    logger.warning(f"Error stopping notebook in insert_cell: {e}")
 
-    return await __safe_notebook_operation(_append_markdown)
+    return await __safe_notebook_operation(_insert_cell)
 
 
 @mcp.tool()
-async def insert_markdown_cell(cell_index: int, cell_source: str) -> str:
-    """Insert a markdown cell in a Jupyter notebook.
+async def insert_execute_code_cell(cell_index: int, cell_source: str) -> list[str]:
+    """Insert and execute a code cell in a Jupyter notebook.
 
     Args:
-        cell_index: Index of the cell to insert (0-based)
-        cell_source: Markdown source
+        cell_index: Index of the cell to insert (0-based). Use -1 to append at end and execute.
+        cell_source: Code source
 
     Returns:
-        str: Success message
+        list[str]: List of outputs from the executed cell
     """
-    async def _insert_markdown():
+    async def _insert_execute():
+        __ensure_kernel_alive()
         notebook = None
         try:
             notebook = _get_notebook_client()
             await notebook.start()
-            notebook.insert_markdown_cell(cell_index, cell_source)
-            return f"Jupyter Markdown cell {cell_index} inserted."
+
+            ydoc = notebook._doc
+            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."
+                )
+
+            if actual_index == total_cells:
+                notebook.add_code_cell(cell_source)
+            else:
+                notebook.insert_code_cell(actual_index, cell_source)
+
+            notebook.execute_cell(actual_index, kernel)
+
+            ydoc = notebook._doc
+            outputs = ydoc._ycells[actual_index]["outputs"]
+            return safe_extract_outputs(outputs)
         finally:
             if notebook:
                 try:
                     await notebook.stop()
                 except Exception as e:
-                    logger.warning(f"Error stopping notebook in insert_markdown_cell: {e}")
+                    logger.warning(f"Error stopping notebook in insert_execute_code_cell: {e}")
 
-    return await __safe_notebook_operation(_insert_markdown)
+    return await __safe_notebook_operation(_insert_execute)
 
 
 @mcp.tool()
@@ -479,72 +531,6 @@ async def _overwrite_cell():
 
     return await __safe_notebook_operation(_overwrite_cell)
 
-
-@mcp.tool()
-async def append_execute_code_cell(cell_source: str) -> list[str]:
-    """Append at the end of the notebook a code cell with the provided source and execute it.
-
-    Args:
-        cell_source: Code source
-
-    Returns:
-        list[str]: List of outputs from the executed cell
-    """
-    async def _append_execute():
-        __ensure_kernel_alive()
-        notebook = None
-        try:
-            notebook = _get_notebook_client()
-            await notebook.start()
-            cell_index = notebook.add_code_cell(cell_source)
-            notebook.execute_cell(cell_index, kernel)
-
-            ydoc = notebook._doc
-            outputs = ydoc._ycells[cell_index]["outputs"]
-            return safe_extract_outputs(outputs)
-        finally:
-            if notebook:
-                try:
-                    await notebook.stop()
-                except Exception as e:
-                    logger.warning(f"Error stopping notebook in append_execute_code_cell: {e}")
-
-    return await __safe_notebook_operation(_append_execute)
-
-
-@mcp.tool()
-async def insert_execute_code_cell(cell_index: int, cell_source: str) -> list[str]:
-    """Insert and execute a code cell in a Jupyter notebook.
-
-    Args:
-        cell_index: Index of the cell to insert (0-based)
-        cell_source: Code source
-
-    Returns:
-        list[str]: List of outputs from the executed cell
-    """
-    async def _insert_execute():
-        __ensure_kernel_alive()
-        notebook = None
-        try:
-            notebook = _get_notebook_client()
-            await notebook.start()
-            notebook.insert_code_cell(cell_index, cell_source)
-            notebook.execute_cell(cell_index, kernel)
-
-            ydoc = notebook._doc
-            outputs = ydoc._ycells[cell_index]["outputs"]
-            return safe_extract_outputs(outputs)
-        finally:
-            if notebook:
-                try:
-                    await notebook.stop()
-                except Exception as e:
-                    logger.warning(f"Error stopping notebook in insert_execute_code_cell: {e}")
-
-    return await __safe_notebook_operation(_insert_execute)
-
-
 @mcp.tool()
 async def execute_cell_with_progress(cell_index: int, timeout_seconds: int = 300) -> list[str]:
     """Execute a specific cell with timeout and progress monitoring.

jupyter_mcp_server/utils.py

@@ -159,4 +159,48 @@ def format_cell_list(ydoc_cells: Any) -> str:
         # Add to table
         lines.append(f"{i}\t{cell_type}\t{execution_count}\t{first_line}")
 
+    return "\n".join(lines)
+
+def get_surrounding_cells_info(notebook, cell_index: int, total_cells: int) -> str:
+    """Get information about surrounding cells for context."""
+    start_index = max(0, cell_index - 5)
+    end_index = min(total_cells, cell_index + 6)
+
+    if total_cells == 0:
+        return "Notebook is now empty, no cells remaining"
+
+    lines = ["Index\tType\tCount\tFirst Line"]
+    lines.append("-" * 60)
+
+    for i in range(start_index, end_index):
+        if i >= total_cells:
+            break
+
+        ydoc = notebook._doc
+        cell_data = ydoc._ycells[i]
+        cell_type = cell_data.get("cell_type", "unknown")
+
+        # Get execution count for code cells
+        if cell_type == "code":
+            execution_count = cell_data.get("execution_count") or "None"
+        else:
+            execution_count = "N/A"
+
+        # Get first line of source
+        source = cell_data.get("source", "")
+        if isinstance(source, list):
+            first_line = source[0] if source else ""
+        else:
+            first_line = str(source)
+
+        # Get just the first line and truncate if too long
+        first_line = first_line.split('\n')[0]
+        if len(first_line) > 50:
+            first_line = first_line[:47] + "..."
+
+        # Mark the target cell
+        marker = " ← inserted" if i == cell_index else ""
+
+        lines.append(f"{i}\t{cell_type}\t{execution_count}\t{first_line}{marker}")
+
     return "\n".join(lines)