fix: support both ColabKernelClient and KernelClient in runtime (#95)
Mahdi Seddigh committed
Jul 30, 2026 at 23:26 UTC
604aac1fda8b54a349a87c4ee61d540f28ff8cd3
2 files changed
+50
-35
src/colab_cli/runtime.py
+22
-10
@@ -103,16 +103,28 @@ class ColabRuntime:
103
# WSSession (Session) expects 'session' for the ID
104
client_kwargs["session"] = self.session_id
105
106
- self._kernel_client = jupyter_kernel_client.KernelClient(
107
- server_url=self.url,
108
- token=self.token,
109
- kernel_id=self.kernel_id,
110
- client_kwargs=client_kwargs,
111
- headers={
112
- "X-Colab-Client-Agent": "colab-cli",
113
- "X-Colab-Runtime-Proxy-Token": self.token,
114
- },
115
- )
106
+ if hasattr(jupyter_kernel_client, "ColabKernelClient"):
107
+ self._kernel_client = jupyter_kernel_client.ColabKernelClient(
108
+ server_url=self.url,
109
+ proxy_token=self.token,
110
+ kernel_id=self.kernel_id,
111
+ client_kwargs=client_kwargs,
112
+ headers={
113
+ "X-Colab-Client-Agent": "colab-cli",
114
+ "X-Colab-Runtime-Proxy-Token": self.token,
115
+ },
116
+ )
117
+ else:
118
+ self._kernel_client = jupyter_kernel_client.KernelClient(
119
+ server_url=self.url,
120
+ token=self.token,
121
+ kernel_id=self.kernel_id,
122
+ client_kwargs=client_kwargs,
123
+ headers={
124
+ "X-Colab-Client-Agent": "colab-cli",
125
+ "X-Colab-Runtime-Proxy-Token": self.token,
126
+ },
127
+ )
128
# Force _own_kernel to False. This prevents jupyter-kernel-client
129
# from automatically deleting the kernel when the client is closed or deleted.
130
self._kernel_client._own_kernel = False
tests/test_runtime.py
+28
-25
@@ -19,31 +19,34 @@ import jupyter_kernel_client
19
from colab_cli.runtime import ColabRuntime
20
21
22
-@patch("colab_cli.runtime.jupyter_kernel_client.KernelClient")
23
-def test_colab_runtime_kernel_client(mock_kc_cls):
24
- mock_kc = mock_kc_cls.return_value
25
-
26
- runtime = ColabRuntime("http://url", "token123")
27
-
28
- assert runtime._kernel_client is None
29
-
30
- kc = runtime.kernel_client
31
-
32
- mock_kc_cls.assert_called_once_with(
33
- server_url="http://url",
34
- token="token123",
35
- kernel_id=None,
36
- client_kwargs={
37
- "subprotocol": jupyter_kernel_client.JupyterSubprotocol.DEFAULT,
38
- "extra_params": {"colab-runtime-proxy-token": "token123"},
39
- },
40
- headers={
41
- "X-Colab-Client-Agent": "colab-cli",
42
- "X-Colab-Runtime-Proxy-Token": "token123",
43
- },
44
- )
45
- mock_kc.start.assert_called_once()
46
- assert kc == mock_kc
22
+def test_colab_runtime_kernel_client():
23
+ target_attr = "ColabKernelClient" if hasattr(jupyter_kernel_client, "ColabKernelClient") else "KernelClient"
24
+ token_param_name = "proxy_token" if hasattr(jupyter_kernel_client, "ColabKernelClient") else "token"
25
+
26
+ with patch.object(jupyter_kernel_client, target_attr) as mock_kc_cls:
27
+ mock_kc = mock_kc_cls.return_value
28
+ runtime = ColabRuntime("http://url", "token123")
29
+
30
+ assert runtime._kernel_client is None
31
+
32
+ kc = runtime.kernel_client
33
+
34
+ expected_kwargs = {
35
+ "server_url": "http://url",
36
+ token_param_name: "token123",
37
+ "kernel_id": None,
38
+ "client_kwargs": {
39
+ "subprotocol": jupyter_kernel_client.JupyterSubprotocol.DEFAULT,
40
+ "extra_params": {"colab-runtime-proxy-token": "token123"},
41
+ },
42
+ "headers": {
43
+ "X-Colab-Client-Agent": "colab-cli",
44
+ "X-Colab-Runtime-Proxy-Token": "token123",
45
+ },
46
+ }
47
+ mock_kc_cls.assert_called_once_with(**expected_kwargs)
48
+ mock_kc.start.assert_called_once()
49
+ assert kc == mock_kc
50
51
52
def test_colab_runtime_execute_code():