feat(jupyter): forward extra HTTP headers and honor an explicit token

code_sandboxes/jupyter_sandbox.py

@@ -50,7 +50,15 @@
 
 
 class JupyterSandbox(Sandbox):
-    """Jupyter Server sandbox using a persistent kernel."""
+    """Jupyter Server sandbox using a persistent kernel.
+
+    Pass ``headers`` to send extra HTTP headers on every request to an external
+    Jupyter Server, for deployments whose credentials are not a token — for
+    example a session ``Cookie`` plus ``X-XSRFToken`` from a password login.
+    Such a deployment has no token to send, so pass ``token=None`` and let the
+    headers authenticate: a token is only generated for a server this sandbox
+    starts itself.
+    """
 
     def __init__(
         self,
@@ -65,6 +73,7 @@ def __init__(
         kernel_path: str | None = None,
         client_kwargs: dict | None = None,
         reuse_kernel: bool = True,
+        headers: dict[str, str] | None = None,
         **kwargs,
     ):
         super().__init__(config)
@@ -81,7 +90,19 @@ def __init__(
         if parsed_url and parsed_token:
             cleaned = parsed_url._replace(query="", fragment="")
             self._server_url = urlunparse(cleaned)
-        self._token = token or uuid.uuid4().hex
+        self._headers = dict(headers) if headers else {}
+        # A generated token only means something for a server this sandbox
+        # starts itself, where it is passed as --ServerApp.token. When talking to
+        # an external server, honor the caller's token as given — including None,
+        # so that other credentials (for example the session cookie and XSRF
+        # header supplied via ``headers``) are what authenticate the requests
+        # instead of a fabricated token the server has never seen.
+        if token:
+            self._token = token
+        elif server_url is None:
+            self._token = uuid.uuid4().hex
+        else:
+            self._token = None
         self._host = host
         self._port = port
         self._python_executable = python_executable or os.environ.get("PYTHON", "python")
@@ -283,6 +304,7 @@ def _wait_for_server(self, timeout: float = DEFAULT_STARTUP_TIMEOUT) -> None:
                 response = requests.get(
                     f"{self._server_url}/api/status",
                     params={"token": self._token},
+                    headers=self._headers or None,
                     timeout=2,
                 )
                 if response.ok:
@@ -308,7 +330,11 @@ def _find_existing_kernel(self) -> str | None:
             return None
 
         try:
-            jsc = JupyterServerClient(base_url=self._server_url, token=self._token)
+            jsc = JupyterServerClient(
+                base_url=self._server_url,
+                token=self._token,
+                headers=self._headers or None,
+            )
             kernels = jsc.kernels.list_kernels()
             if not kernels:
                 logger.info("No existing kernels found, will create a new one")
@@ -364,12 +390,17 @@ def start(self) -> None:
         else:
             kernel_id = None
 
-        self._client = JupyterKernelClient(
-            server_url=self._server_url,
-            token=self._token,
-            kernel_id=kernel_id,
-            client_kwargs=self._client_kwargs or None,
-        )
+        client_kwargs: dict = {
+            "server_url": self._server_url,
+            "token": self._token,
+            "kernel_id": kernel_id,
+            "client_kwargs": self._client_kwargs or None,
+        }
+        # Only forward headers when the caller supplied some, so the call stays
+        # byte-identical for the common token/anonymous case.
+        if self._headers:
+            client_kwargs["headers"] = self._headers
+        self._client = JupyterKernelClient(**client_kwargs)
 
         self._client.start(path=self._kernel_path)
 
@@ -469,6 +500,7 @@ def _do_interrupt(self) -> bool:
                 resp = requests.post(
                     f"{self._server_url}/api/kernels/{kernel_id}/interrupt",
                     params={"token": self._token},
+                    headers=self._headers or None,
                     timeout=5,
                 )
                 return resp.ok

tests/test_jupyter.py

@@ -7,6 +7,7 @@
 import os
 import sys
 import types
+import uuid
 from pathlib import Path
 
 import pytest
@@ -169,3 +170,99 @@ def test_local_jupyter_persistence(self, tmp_path: Path):
             assert "8" in execution.results[0].data.get("text/plain", "")
         finally:
             sandbox.stop()
+
+
+def _kernel_client_stub(captured: dict):
+    """Build a JupyterKernelClient stub that records the kwargs it was built with."""
+
+    class _KernelClientStub:
+        def __init__(self, server_url, token, kernel_id, client_kwargs=None, **kwargs):
+            captured["server_url"] = server_url
+            captured["token"] = token
+            captured["kernel_id"] = kernel_id
+            captured["client_kwargs"] = client_kwargs
+            captured["headers"] = kwargs.get("headers")
+
+        def start(self, path=None):
+            captured["start_path"] = path
+
+        def stop(self):
+            return None
+
+    return _KernelClientStub
+
+
+def test_headers_are_forwarded_to_the_kernel_client(monkeypatch):
+    """Extra headers reach the kernel client, for cookie/XSRF authenticated servers."""
+
+    captured: dict[str, object] = {}
+    monkeypatch.setitem(
+        sys.modules,
+        "jupyter_kernel_client",
+        types.SimpleNamespace(JupyterKernelClient=_kernel_client_stub(captured)),
+    )
+
+    auth_headers = {"Cookie": "username-localhost=abc; _xsrf=tok", "X-XSRFToken": "tok"}
+    sandbox = JupyterSandbox(
+        server_url="http://localhost:8888",
+        token=None,
+        kernel_id="kernel-1",
+        reuse_kernel=False,
+        headers=auth_headers,
+    )
+
+    monkeypatch.setattr(sandbox, "_wait_for_server", lambda timeout=None: None)
+
+    sandbox.start()
+    try:
+        assert captured["headers"] == auth_headers
+    finally:
+        sandbox.stop()
+
+
+def test_no_headers_kwarg_when_none_supplied(monkeypatch):
+    """Without headers the kernel client is built exactly as before."""
+
+    captured: dict[str, object] = {}
+    monkeypatch.setitem(
+        sys.modules,
+        "jupyter_kernel_client",
+        types.SimpleNamespace(JupyterKernelClient=_kernel_client_stub(captured)),
+    )
+
+    credential = uuid.uuid4().hex
+    sandbox = JupyterSandbox(
+        server_url="http://localhost:8888",
+        token=credential,
+        kernel_id="kernel-1",
+        reuse_kernel=False,
+    )
+
+    monkeypatch.setattr(sandbox, "_wait_for_server", lambda timeout=None: None)
+
+    sandbox.start()
+    try:
+        assert captured["headers"] is None
+        assert captured["token"] == credential
+    finally:
+        sandbox.stop()
+
+
+def test_external_server_keeps_token_none():
+    """An external server with no token must not get a fabricated one.
+
+    Password-authenticated deployments have no token to send; generating one
+    would put a credential the server never issued on every request.
+    """
+
+    sandbox = JupyterSandbox(server_url="http://localhost:8888", token=None)
+
+    assert sandbox._token is None
+
+
+def test_owned_server_still_generates_a_token():
+    """A sandbox that starts its own server still needs a token to secure it."""
+
+    sandbox = JupyterSandbox()
+
+    assert sandbox._token