From 9abb3bfb58a3e2b58757cc12fdd4fc26cea65a3b Mon Sep 17 00:00:00 2001 From: EnyMan Date: Wed, 29 Jul 2026 16:11:52 +0200 Subject: [PATCH] feat(jupyter): forward extra HTTP headers and honor an explicit token MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JupyterSandbox could only authenticate to an external Jupyter Server with a token. Deployments that authenticate with a password have no token: the credentials are a session Cookie plus a matching X-XSRFToken header, and there was no way to pass those through the sandbox. Two things blocked it: - Headers could not reach the requests. `client_kwargs` only reaches the kernel client's websocket, while the kernel manager sends its own HTTP requests (kernel GET/POST/DELETE, interrupt) using its `headers` argument, which nothing populated. `**kwargs` were collected into `_extra_kwargs` and never forwarded. - `token=None` was replaced by `uuid.uuid4().hex`, putting a credential the server never issued on every request. A generated token is only meaningful for a server this sandbox starts itself, where it becomes --ServerApp.token. So add a `headers` argument, forwarded to the kernel client, to the readiness and interrupt requests, and to the JupyterServerClient used for kernel reuse. Only pass it on when the caller supplied headers, so the common token/anonymous path is unchanged. And keep an explicit token as given — including None — when talking to a server the sandbox does not own. Co-Authored-By: Claude Opus 4.8 (1M context) --- code_sandboxes/jupyter_sandbox.py | 50 +++++++++++++--- tests/test_jupyter.py | 97 +++++++++++++++++++++++++++++++ 2 files changed, 138 insertions(+), 9 deletions(-) diff --git a/code_sandboxes/jupyter_sandbox.py b/code_sandboxes/jupyter_sandbox.py index 09e998a..74658ce 100644 --- a/code_sandboxes/jupyter_sandbox.py +++ b/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 diff --git a/tests/test_jupyter.py b/tests/test_jupyter.py index 3fba7c1..25ed0df 100644 --- a/tests/test_jupyter.py +++ b/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