🔙 Backport Support talking to Colab from the JKC

Makefile

@@ -28,6 +28,9 @@ build:
 clean: ## clean
 	git clean -fdx
 
+test: ## run tests
+	uv run pytest
+
 jupyter-server: ## jupyter-server
 	jupyter server --port 8888 --ServerApp.port_retries 0 --IdentityProvider.token MY_TOKEN
 

jupyter_kernel_client/__init__.py

@@ -11,6 +11,7 @@
 from jupyter_kernel_client.models import VariableDescription
 from jupyter_kernel_client.snippets import SNIPPETS_REGISTRY, LanguageSnippets
 from jupyter_kernel_client.wsclient import KernelWebSocketClient
+from jupyter_kernel_client.wsclient import JupyterSubprotocol
 
 __all__ = [
     "SNIPPETS_REGISTRY",
@@ -20,5 +21,6 @@
     "KonsoleApp",
     "LanguageSnippets",
     "VariableDescription",
+    "JupyterSubprotocol",
     "__version__",
 ]

jupyter_kernel_client/client.py

@@ -172,9 +172,7 @@ def __del__(self) -> None:
         try:
             self.stop()
         except BaseException as e:
-            self.log.error(
-                "Failed to stop the kernel client for %s.", self._manager.kernel_url, exc_info=e
-            )
+            self.log.error("Failed to stop the kernel client", exc_info=e)
 
     def _set_variables(self, variables: dict[str, t.Any] | None) -> None:
         """Set variables in the kernel's globals dictionary.
@@ -230,6 +228,10 @@ def last_activity(self) -> datetime.datetime | None:
             else None
         )
 
+    def list_kernels(self) -> list[dict[str, t.Any]]:
+        """List the running kernels."""
+        return self._manager.list_kernels()
+
     @property
     def username(self) -> str:
         """Client owner username."""
@@ -429,7 +431,7 @@ def stop(
             timeout: Request timeout in seconds
         """
         self.log.info("Stopping the kernel client…")
-        if self._manager.has_kernel:
+        if self._manager and self._manager.has_kernel:
             self._manager.client.stop_channels()
             shutdown = self._own_kernel if shutdown_kernel is None else shutdown_kernel
             if shutdown:

jupyter_kernel_client/manager.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import datetime
+import logging
 import os
 import re
 import typing as t
@@ -30,13 +31,12 @@ def fetch(
     """Fetch a network resource as a context manager."""
     method = kwargs.pop("method", "GET")
     f = getattr(requests, method.lower())
-    headers = kwargs.pop("headers", {})
-    if len(headers) == 0:
-        headers = {
-            "Accept": "application/json",
-            "Content-Type": "application/json",
-            "User-Agent": "Jupyter kernels CLI",
-        }
+    headers = {
+        "Accept": "application/json",
+        "Content-Type": "application/json",
+        "User-Agent": "Jupyter Kernel Client"
+    }
+    headers.update(kwargs.pop("headers", {}))
     if token:
         headers["Authorization"] = f"Bearer {token}"
     if "timeout" not in kwargs:
@@ -66,10 +66,11 @@ class KernelHttpManager(LoggingConfigurable):
     def __init__(
         self,
         server_url: str,
-        token: str,
+        token: str | None,
         username: str = os.environ.get("USER", "username"),
         kernel_id: str | None = None,
         client_kwargs: dict[str, t.Any] | None = None,
+        headers: dict[str, t.Any] | None = None,
         **kwargs,
     ):
         """Initialize the kernel manager."""
@@ -79,7 +80,17 @@ def __init__(
         self.username = username
         self.__kernel: dict | None = None
         self.__client: t.Any | None = None
+        if not client_kwargs:
+            client_kwargs = {}
         self.__client_kwargs = client_kwargs
+        if headers:
+            self.__extra_headers = headers
+        else:
+            self.__extra_headers = {}
+
+        if 'headers' not in self.__client_kwargs:
+            self.__client_kwargs['headers'] = {}
+        self.__client_kwargs['headers'].update(self.__extra_headers)
 
         if kernel_id:
             self.__kernel = {
@@ -162,7 +173,7 @@ def refresh_model(self, timeout: float = REQUEST_TIMEOUT) -> dict[str, t.Any] |
 
         self.log.debug("Request kernel at: %s", self.kernel_url)
         try:
-            response = fetch(self.kernel_url, token=self.token, method="GET", timeout=timeout)
+            response = fetch(self.kernel_url, token=self.token, method="GET", timeout=timeout, headers=self.__extra_headers)
         except HTTPError as error:
             if error.response.status_code == 404:
                 self.log.warning("Kernel not found at: %s", self.kernel_url)
@@ -179,6 +190,26 @@ def refresh_model(self, timeout: float = REQUEST_TIMEOUT) -> dict[str, t.Any] |
             self.__client = None
         return model
 
+    def list_kernels(self, timeout: float = REQUEST_TIMEOUT) -> list[dict[str, t.Any]]:
+        """List the running kernels.
+
+        Returns
+        -------
+            The list of running kernels.
+        """
+        kernels_url = url_path_join(self.server_url, "api/kernels")
+        self.log.debug("Request kernels at: %s", kernels_url)
+        try:
+            response = fetch(kernels_url, token=self.token, method="GET", timeout=timeout, headers=self.__extra_headers)
+        except HTTPError as error:
+            self.log.error("Error fetching kernels: %s", error)
+            return []
+        else:
+            models = response.json()
+        self.log.debug("Kernels retrieved: %s", models)
+        return models
+
+
     # --------------------------------------------------------------------------
     # Kernel management
     # --------------------------------------------------------------------------
@@ -210,6 +241,7 @@ def start_kernel(
             method="POST",
             json={"name": name, "path": path},
             timeout=timeout,
+            headers=self.__extra_headers
         )
 
         self.__kernel = response.json()
@@ -239,7 +271,7 @@ def shutdown_kernel(
 
         # If not now and refreshing the model still returns it, try the http way
         try:
-            response = fetch(self.kernel_url, token=self.token, method="DELETE", timeout=timeout)
+            response = fetch(self.kernel_url, token=self.token, method="DELETE", timeout=timeout, headers=self.__extra_headers)
             self.log.debug(
                 "Shutdown kernel response: %d %s",
                 response.status_code,
@@ -263,7 +295,7 @@ def restart_kernel(self, timeout: float = REQUEST_TIMEOUT, **kw):
 
         kernel_url = self.kernel_url + "/restart"
         self.log.debug("Request restart kernel at: %s", kernel_url)
-        response = fetch(kernel_url, token=self.token, method="POST", timeout=timeout)
+        response = fetch(kernel_url, token=self.token, method="POST", timeout=timeout, headers=self.__extra_headers)
         self.log.debug("Restart kernel response: %d %s", response.status_code, response.reason)
 
     def interrupt_kernel(self, timeout: float = REQUEST_TIMEOUT):
@@ -278,6 +310,7 @@ def interrupt_kernel(self, timeout: float = REQUEST_TIMEOUT):
             token=self.token,
             method="POST",
             timeout=timeout,
+            headers=self.__extra_headers
         )
         self.log.debug(
             "Interrupt kernel response: %d %s",

jupyter_kernel_client/tests/test_client.py

@@ -61,6 +61,31 @@ def test_execution_no_context_manager(jupyter_server):
     assert reply["status"] == "ok"
 
 
+def test_list_kernels_client(jupyter_server):
+    port, token = jupyter_server
+
+    # Start a kernel to ensure the list is not empty
+    with KernelClient(server_url=f"http://localhost:{port}", token=token) as kernel:
+        kernel_id = kernel.id
+
+        # Use a new client to list the kernels
+        listing_client = KernelClient(server_url=f"http://localhost:{port}", token=token)
+        kernels = listing_client.list_kernels()
+
+        assert isinstance(kernels, list)
+        assert len(kernels) > 0
+
+        # Check that the kernel we started is in the list
+        found = False
+        for k in kernels:
+            assert "id" in k
+            assert "name" in k
+            if k["id"] == kernel_id:
+                found = True
+
+        assert found, f"Kernel with id {kernel_id} not found in the list of running kernels."
+
+
 def test_list_variables(jupyter_server):
     port, token = jupyter_server
 
@@ -212,15 +237,18 @@ def test_set_variables(jupyter_server, variable, set_variable, expected):
 async def test_multi_execution_in_event_loop(jupyter_server):
     port, token = jupyter_server
 
+    current_user = os.environ.get('USER', 'John Smith')
+    current_node = node()
+
     with KernelClient(server_url=f"http://localhost:{port}", token=token) as kernel:
         all = await asyncio.gather(
             asyncio.to_thread(
                 kernel.execute,
-                """import os
+                f"""import os
 from platform import node
 import time
 time.sleep(5)
-print(f"Hey {os.environ.get('USER', 'John Smith')} from {node()}.")
+print(f"Hey {{os.environ.get('USER', 'John Smith')}} from {{node()}}.")
 """
             ),
             asyncio.to_thread(
@@ -236,7 +264,7 @@ async def test_multi_execution_in_event_loop(jupyter_server):
             {
                 "output_type": "stream",
                 "name": "stdout",
-                "text": f"Hey {os.environ.get('USER', 'John Smith')} from {node()}.\n",
+                "text": f"Hey {current_user} from {current_node}.\n",
             }
         ]
         assert all[0]["status"] == "ok"

jupyter_kernel_client/tests/test_manager.py

@@ -0,0 +1,33 @@
+
+# Copyright (c) 2023-2024 Datalayer, Inc.
+# Copyright (c) 2025      Google
+#
+# BSD 3-Clause License
+
+from jupyter_kernel_client.manager import KernelHttpManager
+from jupyter_kernel_client import KernelClient
+
+
+def test_list_kernels(jupyter_server):
+    port, token = jupyter_server
+
+    # Start a kernel to ensure the list is not empty
+    with KernelClient(server_url=f"http://localhost:{port}", token=token) as kernel:
+        kernel_id = kernel.id
+        # The manager is created after the kernel is started to ensure we can list it.
+        manager = KernelHttpManager(server_url=f"http://localhost:{port}", token=token)
+        kernels = manager.list_kernels()
+
+        assert isinstance(kernels, list)
+        assert len(kernels) > 0
+
+        # Check that the kernel we started is in the list
+        found = False
+        for k in kernels:
+            assert "id" in k
+            assert "name" in k
+            if k["id"] == kernel_id:
+                found = True
+
+        assert found, f"Kernel with id {kernel_id} not found in the list of running kernels."
+