fix(workbench): keep JupyterLab sessions fresh to stop launch/exec skip (#532)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: samcofer

selftests/test_workbench_jupyter_cleanup.py

@@ -0,0 +1,145 @@
+"""Selftests for the JupyterLab notebook-cleanup helpers on WorkbenchClient.
+
+These cover the pure URL/header builders and the ``delete_jupyter_notebook``
+method (via an ``httpx.MockTransport``), so the contents-API teardown that keeps
+JupyterLab sessions fresh is exercised without a live Workbench.
+"""
+
+from __future__ import annotations
+
+import httpx
+import pytest
+
+from vip.clients.workbench import (
+    WorkbenchClient,
+    jupyterlab_app_base,
+    jupyterlab_contents_delete_url,
+    jupyterlab_xsrf_headers,
+)
+
+# -- jupyterlab_app_base -----------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "page_url, expected",
+    [
+        # Standard per-session proxy prefix, notebook open under /lab/tree.
+        (
+            "https://wb.example.com/s/abc123/lab/tree/Untitled.ipynb",
+            "https://wb.example.com/s/abc123",
+        ),
+        # Bare /lab root.
+        ("https://wb.example.com/s/abc123/lab", "https://wb.example.com/s/abc123"),
+        # Query string and fragment are stripped.
+        (
+            "https://wb.example.com/s/xyz/lab/tree/Untitled1.ipynb?reset#cell",
+            "https://wb.example.com/s/xyz",
+        ),
+        # No /lab segment — return the URL trimmed rather than raising.
+        ("https://wb.example.com/s/abc123/", "https://wb.example.com/s/abc123"),
+        # Real Workbench shape (validated live via CDP on pwb.demo): JupyterLab
+        # inserts a /lab/workspaces/<id>/tree/<nb> segment. The base must still
+        # cut at /lab and match PageConfig.baseUrl (".../s/<sid>").
+        (
+            "https://pwb.demo.soleng.posit.it/s/0da2921f2ed72297f2efa/lab/workspaces/auto-m/tree/Untitled7.ipynb",
+            "https://pwb.demo.soleng.posit.it/s/0da2921f2ed72297f2efa",
+        ),
+    ],
+)
+def test_jupyterlab_app_base(page_url, expected):
+    assert jupyterlab_app_base(page_url) == expected
+
+
+# -- jupyterlab_contents_delete_url ------------------------------------------
+
+
+def test_contents_delete_url_basic():
+    url = jupyterlab_contents_delete_url(
+        "https://wb.example.com/s/abc123/lab/tree/Untitled.ipynb", "Untitled.ipynb"
+    )
+    assert url == "https://wb.example.com/s/abc123/api/contents/Untitled.ipynb"
+
+
+def test_contents_delete_url_quotes_spaces_and_strips_slashes():
+    url = jupyterlab_contents_delete_url("https://wb.example.com/s/abc123/lab", "/_vip run 1.ipynb")
+    # Leading slash stripped, space percent-encoded.
+    assert url == "https://wb.example.com/s/abc123/api/contents/_vip%20run%201.ipynb"
+
+
+# -- jupyterlab_xsrf_headers -------------------------------------------------
+
+
+def test_xsrf_headers_present():
+    assert jupyterlab_xsrf_headers({"_xsrf": "tok123"}) == {"X-XSRFToken": "tok123"}
+
+
+def test_xsrf_headers_absent():
+    assert jupyterlab_xsrf_headers({"other": "v"}) == {}
+
+
+# -- delete_jupyter_notebook (MockTransport) ---------------------------------
+
+
+def _client_with_handler(handler) -> WorkbenchClient:
+    wc = WorkbenchClient("https://wb.example.com")
+    wc._client.close()
+    wc._client = httpx.Client(
+        base_url="https://wb.example.com",
+        transport=httpx.MockTransport(handler),
+    )
+    return wc
+
+
+def test_delete_notebook_success_sends_xsrf_and_targets_contents_api():
+    seen: dict[str, object] = {}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen["method"] = request.method
+        seen["path"] = request.url.path
+        seen["xsrf"] = request.headers.get("X-XSRFToken")
+        return httpx.Response(204)
+
+    wc = _client_with_handler(handler)
+    ok = wc.delete_jupyter_notebook(
+        "https://wb.example.com/s/abc123/lab/tree/Untitled.ipynb",
+        "Untitled.ipynb",
+        {"_xsrf": "tok123"},
+    )
+    assert ok is True
+    assert seen["method"] == "DELETE"
+    assert seen["path"] == "/s/abc123/api/contents/Untitled.ipynb"
+    assert seen["xsrf"] == "tok123"
+
+
+def test_delete_notebook_404_treated_as_success():
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(404)
+
+    wc = _client_with_handler(handler)
+    assert (
+        wc.delete_jupyter_notebook("https://wb.example.com/s/abc/lab", "Untitled.ipynb", {}) is True
+    )
+
+
+def test_delete_notebook_server_error_is_false():
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(500)
+
+    wc = _client_with_handler(handler)
+    assert (
+        wc.delete_jupyter_notebook(
+            "https://wb.example.com/s/abc/lab", "Untitled.ipynb", {"_xsrf": "t"}
+        )
+        is False
+    )
+
+
+def test_delete_notebook_transport_error_is_false():
+    def handler(request: httpx.Request) -> httpx.Response:
+        raise httpx.ConnectError("boom")
+
+    wc = _client_with_handler(handler)
+    assert (
+        wc.delete_jupyter_notebook("https://wb.example.com/s/abc/lab", "Untitled.ipynb", {})
+        is False
+    )

src/vip/clients/workbench.py

@@ -29,6 +29,57 @@ def is_vip_session(label: str) -> bool:
     return any(label.startswith(prefix) for prefix in _VIP_SESSION_PREFIXES)
 
 
+def jupyterlab_app_base(page_url: str) -> str:
+    """Return the JupyterLab app-root URL from a session page URL.
+
+    A Workbench JupyterLab session is served under a per-session proxy prefix,
+    e.g. ``https://wb.example.com/s/abc123/lab/tree/Untitled.ipynb``.  The
+    JupyterLab server (and its ``/api/contents`` REST surface) is mounted at the
+    segment ending in ``/lab`` — everything *before* ``/lab`` is the app base
+    (``https://wb.example.com/s/abc123``).  The returned value has no trailing
+    slash.  If the URL has no ``/lab`` segment (unexpected), the input is
+    returned with any query/fragment and trailing slash stripped, so callers
+    still get a usable base rather than raising mid-teardown.
+    """
+    # Drop query string and fragment first — the contents API path is built
+    # fresh, so anything after ``?``/``#`` is irrelevant and would corrupt the base.
+    core = page_url.split("?", 1)[0].split("#", 1)[0]
+    marker = "/lab"
+    idx = core.find(marker)
+    if idx != -1:
+        return core[:idx].rstrip("/")
+    return core.rstrip("/")
+
+
+def jupyterlab_contents_delete_url(page_url: str, notebook_name: str) -> str:
+    """Build the ``DELETE /api/contents/<notebook>`` URL for a JupyterLab session.
+
+    *page_url* is the current session page URL (see :func:`jupyterlab_app_base`);
+    *notebook_name* is the notebook's filename as shown on its dock tab (e.g.
+    ``"Untitled.ipynb"``).  The name is URL-quoted so spaces/special characters
+    in a renamed notebook do not break the path.  Any leading slashes on the
+    name are stripped so the result is always ``<app_base>/api/contents/<name>``.
+    """
+    from urllib.parse import quote
+
+    base = jupyterlab_app_base(page_url)
+    clean = notebook_name.strip().lstrip("/")
+    return f"{base}/api/contents/{quote(clean)}"
+
+
+def jupyterlab_xsrf_headers(cookies: dict[str, str]) -> dict[str, str]:
+    """Return the XSRF header JupyterLab requires for mutating contents-API calls.
+
+    JupyterLab protects non-GET ``/api`` requests with a double-submit token:
+    the ``_xsrf`` cookie value must be echoed back in the ``X-XSRFToken`` header.
+    Returns ``{"X-XSRFToken": <value>}`` when the cookie is present, or an empty
+    dict when it is absent (older/token-auth deployments) so the caller can still
+    attempt the request without sending a bogus header.
+    """
+    token = cookies.get("_xsrf")
+    return {"X-XSRFToken": token} if token else {}
+
+
 class WorkbenchClient(BaseClient):
     """Minimal Workbench HTTP wrapper."""
 
@@ -76,6 +127,39 @@ def set_cookies(self, cookies: dict[str, str]) -> None:
         """Set cookies on the client instance for authenticated requests."""
         self._client.cookies.update(cookies)
 
+    def delete_jupyter_notebook(
+        self, page_url: str, notebook_name: str, cookies: dict[str, str]
+    ) -> bool:
+        """Delete a notebook file from a live JupyterLab session's contents API.
+
+        Sends ``DELETE <app_base>/api/contents/<notebook_name>`` authenticated
+        with the browser session *cookies* and the ``X-XSRFToken`` header
+        JupyterLab requires for mutating calls (see
+        :func:`jupyterlab_contents_delete_url` / :func:`jupyterlab_xsrf_headers`).
+
+        This exists so a test that creates a notebook can remove it *before the
+        session is quit* — the contents API only exists while the session is
+        alive — so JupyterLab's layout-restorer stops reopening stale
+        ``Untitled*.ipynb`` tabs (and eagerly starting their kernels) on every
+        subsequent session launch, which widens the launch/exec interaction race.
+
+        Returns True when the server accepts the delete (HTTP < 400, including
+        the ``204`` success and a ``404`` for an already-absent file), False on
+        any other status or transport error.  Never raises — teardown cleanup is
+        a best-effort safety net, not an assertion.
+        """
+        url = jupyterlab_contents_delete_url(page_url, notebook_name)
+        headers = jupyterlab_xsrf_headers(cookies)
+        try:
+            # Set cookies on the client instance (per-request cookies= is
+            # deprecated in httpx and its persistence semantics are ambiguous).
+            self._client.cookies.update(cookies)
+            resp = self._client.request("DELETE", url, headers=headers)
+        except Exception:
+            return False
+        # 404 → already gone, treat as success (idempotent teardown).
+        return resp.status_code < 400 or resp.status_code == 404
+
     def list_sessions(self) -> list[dict[str, Any]]:
         """List active sessions for the authenticated user."""
         resp = self._client.get("/api/sessions")

src/vip_tests/workbench/pages/jupyterlab_session.py

@@ -26,6 +26,19 @@ class JupyterLabSession:
     DIALOG = ".jp-Dialog"
     DIALOG_ACCEPT = ".jp-Dialog .jp-Dialog-button.jp-mod-accept"
 
+    # Kernel execution-status indicator on the notebook toolbar. JupyterLab sets
+    # its ``data-status`` to "idle" once the kernel is connected and ready, and
+    # "busy" while a cell is running. Gating the first cell run on "idle" avoids
+    # typing/running before the kernel has finished connecting, which drops the
+    # input and produces the spurious "kernel did not produce output" timeout.
+    KERNEL_STATUS = ".jp-Notebook-ExecutionIndicator"
+    KERNEL_STATUS_IDLE = ".jp-Notebook-ExecutionIndicator[data-status='idle']"
+
+    # Current (active) notebook tab in the dock area. Its label text is the
+    # notebook's filename (e.g. "Untitled.ipynb"), used to build the contents-API
+    # path for teardown deletion.
+    CURRENT_TAB_LABEL = ".lm-DockPanel-tabBar .lm-TabBar-tab.jp-mod-current .lm-TabBar-tabLabel"
+
     # Sidebar
     FILE_BROWSER = ".jp-FileBrowser"
     SIDEBAR_LEFT = ".jp-SideBar.jp-mod-left"