Fix: Add RTC mode to reading tools (read_cell, read_cells, list_cells)

TESTING_RTC_FIX.md

@@ -0,0 +1,194 @@
+<!--
+  ~ Copyright (c) 2023-2024 Datalayer, Inc.
+  ~
+  ~ BSD 3-Clause License
+-->
+
+# Testing the RTC Mode Fix
+
+This document explains how to test that the WebSocket disconnect bug has been fixed.
+
+## The Bug
+
+**Original Issue:** MCP cell editing tools (execute_cell, overwrite_cell_source, insert_cell, delete_cell, insert_execute_code_cell) were using file mode instead of RTC mode, causing "Out-of-band changes" that broke the WebSocket connection in JupyterLab, requiring browser refresh.
+
+**Root Cause:** jupyter-mcp-server couldn't find `yroom_manager` because it wasn't added to `web_app.settings` by jupyter-collaboration. The code fell back to file mode.
+
+**Fix:** Access `ywebsocket_server` via `extension_manager.extension_points['jupyter_server_ydoc'].app.ywebsocket_server` and document via `room._document`.
+
+**Additional Fix:** Reading tools (read_cell, read_cells, list_cells) also needed RTC mode to see live unsaved changes instead of stale file data.
+
+## Test Coverage
+
+### Automated Tests
+
+Run the RTC mode tests:
+
+```bash
+# Install test dependencies
+uv pip install pytest pytest-asyncio
+
+# Run RTC-specific tests
+pytest tests/test_rtc_mode.py -v
+
+# Run with JUPYTER_SERVER mode only (the one with jupyter-collaboration)
+TEST_MCP_SERVER=false pytest tests/test_rtc_mode.py -v -m jupyter_server
+```
+
+### What the Tests Verify
+
+#### 1. `test_rtc_mode_for_cell_operations`
+- **Purpose**: Verifies all cell operations complete without causing WebSocket disconnects
+- **Tests**: insert_cell, overwrite_cell_source, execute_cell, read_cell, list_cells, delete_cell
+- **Pass Criteria**: All operations complete successfully, indicating RTC mode is used
+- **Failure Mode**: If file mode is used, operations would cause "Out-of-band changes"
+
+#### 2. `test_reading_tools_see_unsaved_changes`
+- **Purpose**: Verifies reading tools see live YDoc data, not stale file data
+- **Tests**: read_cell, read_cells, list_cells after making unsaved edits
+- **Pass Criteria**: Reading tools see modified content (modified_value=999), NOT initial content (initial_value=1)
+- **Failure Mode**: If reading from disk, would see stale initial_value
+
+**This is the KEY test for the new fix** - before our changes, reading tools would fail this test because they read from disk.
+
+#### 3. `test_jupyter_collaboration_extension_loaded`
+- **Purpose**: Confirms jupyter-collaboration infrastructure is available
+- **Pass Criteria**: Can perform basic operations that require the extension
+
+## Manual Testing
+
+To manually reproduce and verify the fix:
+
+### Prerequisites
+```bash
+# Start JupyterLab with jupyter-collaboration
+cd ~/workspace/tools_and_scripts/main/notebooks/jupyter
+./scripts/jupyter-server start
+
+# Or in this repo, start test server
+uv run jupyter lab --no-browser --port 8888 --NotebookApp.token=MY_TOKEN
+```
+
+### Test Scenario 1: No WebSocket Disconnect
+
+1. **Open notebook in JupyterLab browser UI**
+   - Navigate to http://localhost:8888
+   - Create or open a notebook
+   - Add a cell with: `x = 1`
+
+2. **Use MCP tools to edit the same notebook**
+   ```python
+   # Via MCP client or Claude Code
+   use_notebook("test.ipynb", kernel_id="<existing-kernel-id>")
+   overwrite_cell_source(0, "x = 2")
+   ```
+
+3. **Verify in JupyterLab UI**
+   - ✅ **Expected (FIXED)**: Cell updates to `x = 2`, no error, no need to refresh
+   - ❌ **Broken (ORIGINAL BUG)**: "Out-of-band changes" error, WebSocket disconnect, requires browser refresh
+
+### Test Scenario 2: Reading Tools See Live Changes
+
+1. **Create cell via MCP**
+   ```python
+   use_notebook("test.ipynb", kernel_id="<existing-kernel-id>")
+   insert_cell(0, "code", "initial_value = 1")
+   ```
+
+2. **Edit cell via MCP (unsaved change)**
+   ```python
+   overwrite_cell_source(0, "modified_value = 999")
+   ```
+
+3. **Read cell back immediately (before any save)**
+   ```python
+   cell = read_cell(0)
+   print(cell["source"])
+   ```
+
+4. **Verify**
+   - ✅ **Expected (FIXED)**: Shows `modified_value = 999` (live YDoc data)
+   - ❌ **Broken (BEFORE FIX)**: Shows `initial_value = 1` (stale file data)
+
+## Validation Script
+
+The `validate_fixes.py` script verifies code patterns without running tests:
+
+```bash
+python3 validate_fixes.py
+```
+
+This checks:
+- ✅ All 8 tools have `_get_jupyter_ydoc()` with RTC logic
+- ✅ All tools use `extension_manager.extension_points['jupyter_server_ydoc']`
+- ✅ All tools access document via `room._document`
+- ✅ No old broken `yroom_manager` patterns
+
+## What Makes a Test Pass vs Fail?
+
+### Indicators RTC Mode is Working (Test PASSES):
+1. No "Out-of-band changes" errors in logs
+2. Cell operations complete without WebSocket disconnect
+3. Reading tools see latest edits immediately (not stale file data)
+4. JupyterLab UI updates automatically without refresh
+
+### Indicators File Mode Fallback (Test FAILS):
+1. "Out-of-band changes" errors in logs
+2. WebSocket disconnects requiring browser refresh
+3. Reading tools see old content from disk
+4. JupyterLab UI requires refresh to see changes
+
+## Files Modified
+
+### Editing Tools (Fixed in PR #135)
+- `execute_cell_tool.py`
+- `overwrite_cell_source_tool.py`
+- `insert_cell_tool.py`
+- `insert_execute_code_cell_tool.py`
+- `delete_cell_tool.py`
+
+### Reading Tools (Fixed in follow-up)
+- `read_cell_tool.py`
+- `read_cells_tool.py`
+- `list_cells_tool.py`
+
+## Expected Test Results
+
+```bash
+$ pytest tests/test_rtc_mode.py -v
+
+tests/test_rtc_mode.py::test_rtc_mode_for_cell_operations PASSED
+tests/test_rtc_mode.py::test_reading_tools_see_unsaved_changes PASSED
+tests/test_rtc_mode.py::test_jupyter_collaboration_extension_loaded PASSED
+
+============================================================
+✅ All cell operations completed without WebSocket disconnect!
+
+This confirms:
+  - RTC mode is being used (via extension_points)
+  - No file-mode fallback occurred
+  - No 'Out-of-band changes' would occur
+  - WebSocket would remain stable in JupyterLab UI
+============================================================
+```
+
+## Troubleshooting
+
+### Test fails with "extension_points not found"
+- Ensure jupyter-collaboration is installed: `uv pip install jupyter-collaboration`
+- Check JupyterLab version compatibility (requires JupyterLab 4.x)
+
+### Test fails with "reading tools see stale data"
+- Verify the reading tools have been patched with RTC mode
+- Check that `_get_jupyter_ydoc()` method exists in read_cell_tool.py
+
+### WebSocket still disconnects
+- Check server logs for "Out-of-band changes" errors
+- Verify file_id_manager is accessible
+- Confirm ywebsocket_server is being found (should see RTC mode logs)
+
+## References
+
+- Original PR: https://github.com/datalayer/jupyter-mcp-server/pull/135
+- jupyter-collaboration docs: https://jupyterlab-realtime-collaboration.readthedocs.io/
+- Issue: WebSocket disconnects with "Out-of-band changes"

jupyter_mcp_server/tools/list_cells_tool.py

@@ -15,26 +15,102 @@
 
 class ListCellsTool(BaseTool):
     """Tool to list basic information of all cells."""
-
-    async def _list_cells_local(self, contents_manager: Any, path: str) -> str:
-        """List cells using local contents_manager (JUPYTER_SERVER mode)."""
-        # Read the notebook file directly
+
+    async def _get_jupyter_ydoc(self, serverapp, file_id: str):
+        """Get the YNotebook document if it's currently open in a collaborative session."""
+        try:
+            # Access ywebsocket_server from YDocExtension via extension_manager
+            # jupyter-collaboration doesn't add yroom_manager to web_app.settings
+            ywebsocket_server = None
+
+            if hasattr(serverapp, 'extension_manager'):
+                extension_points = serverapp.extension_manager.extension_points
+                if 'jupyter_server_ydoc' in extension_points:
+                    ydoc_ext_point = extension_points['jupyter_server_ydoc']
+                    if hasattr(ydoc_ext_point, 'app') and ydoc_ext_point.app:
+                        ydoc_app = ydoc_ext_point.app
+                        if hasattr(ydoc_app, 'ywebsocket_server'):
+                            ywebsocket_server = ydoc_app.ywebsocket_server
+
+            if ywebsocket_server is None:
+                return None
+
+            room_id = f"json:notebook:{file_id}"
+
+            # Get room and access document via room._document
+            # DocumentRoom stores the YNotebook as room._document, not via get_jupyter_ydoc()
+            try:
+                yroom = await ywebsocket_server.get_room(room_id)
+                if yroom and hasattr(yroom, '_document'):
+                    return yroom._document
+            except Exception:
+                pass
+
+        except Exception:
+            pass
+
+        return None
+
+    async def _list_cells_local(self, serverapp: Any, contents_manager: Any, path: str) -> str:
+        """List cells using local contents_manager (JUPYTER_SERVER mode).
+
+        First tries to read from YDoc (RTC mode) if notebook is open, otherwise falls back to file mode.
+        """
+        # Get file_id for YDoc lookup
+        file_id_manager = serverapp.web_app.settings.get("file_id_manager")
+        file_id = file_id_manager.get_id(path) if file_id_manager else None
+        if file_id is None and file_id_manager:
+            file_id = file_id_manager.index(path)
+
+        # Try to get YDoc if notebook is open (RTC mode)
+        if file_id:
+            ydoc = await self._get_jupyter_ydoc(serverapp, file_id)
+            if ydoc:
+                # Notebook is open - read from YDoc (live data)
+                cells = ydoc.ycells
+                if len(cells) == 0:
+                    return "Notebook is empty, no cells found."
+
+                headers = ["Index", "Type", "Count", "First Line"]
+                rows = []
+
+                for idx, cell in enumerate(cells):
+                    cell_type = cell.get('cell_type', 'unknown')
+                    execution_count = cell.get('execution_count', '-') if cell_type == 'code' else '-'
+
+                    # Get the first line of source
+                    source = cell.get('source', '')
+                    if isinstance(source, list):
+                        first_line = source[0] if source else ''
+                        lines = len(source)
+                    else:
+                        first_line = source.split('\n')[0] if source else ''
+                        lines = len(source.split('\n'))
+
+                    if lines > 1:
+                        first_line += f"...({lines - 1} lines hidden)"
+
+                    rows.append([idx, cell_type, execution_count, first_line])
+
+                return format_TSV(headers, rows)
+
+        # Fall back to file mode if notebook not open
         model = await contents_manager.get(path, content=True, type='notebook')
-        
+
         if 'content' not in model:
             raise ValueError(f"Could not read notebook content from {path}")
-        
+
         notebook_content = model['content']
         cells = notebook_content.get('cells', [])
-        
+
         # Format the cells into a table
         headers = ["Index", "Type", "Count", "First Line"]
         rows = []
-        
+
         for idx, cell in enumerate(cells):
             cell_type = cell.get('cell_type', 'unknown')
             execution_count = cell.get('execution_count', '-') if cell_type == 'code' else '-'
-            
+
             # Get the first line of source
             source = cell.get('source', '')
             if isinstance(source, list):
@@ -43,12 +119,12 @@ async def _list_cells_local(self, contents_manager: Any, path: str) -> str:
             else:
                 first_line = source.split('\n')[0] if source else ''
                 lines = len(source.split('\n'))
-            
+
             if lines > 1:
                 first_line += f"...({lines - 1} lines hidden)"
-            
+
             rows.append([idx, cell_type, execution_count, first_line])
-        
+
         return format_TSV(headers, rows)
 
     def _list_cells_websocket(self, notebook: NbModelClient) -> str:
@@ -129,7 +205,7 @@ async def execute(
                     # Path is not under root_dir, use as-is
                     pass
 
-            return await self._list_cells_local(contents_manager, notebook_path)
+            return await self._list_cells_local(serverapp, contents_manager, notebook_path)
         elif mode == ServerMode.MCP_SERVER and notebook_manager is not None:
             # Remote mode: use WebSocket connection to Y.js document
             async with notebook_manager.get_current_connection() as notebook: