| 1 | # Copyright 2026 Google LLC |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # |
| 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | # |
| 9 | # Unless required by applicable law or agreed to in writing, software |
| 10 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. |
| 14 | |
| 15 | import logging |
| 16 | import time |
| 17 | from typing import Any, Callable, Dict, List, Optional |
| 18 | |
| 19 | import jupyter_kernel_client |
| 20 | import requests |
| 21 | |
| 22 | |
| 23 | class ColabRuntime: |
| 24 | def __init__( |
| 25 | self, |
| 26 | url: str, |
| 27 | token: str, |
| 28 | session_name: Optional[str] = None, |
| 29 | history: Optional[Any] = None, |
| 30 | kernel_id: Optional[str] = None, |
| 31 | session_id: Optional[str] = None, |
| 32 | on_kernel_started: Optional[Callable[[str], None]] = None, |
| 33 | on_session_started: Optional[Callable[[str], None]] = None, |
| 34 | ): |
| 35 | self.url = url |
| 36 | self.token = token |
| 37 | self.session_name = session_name |
| 38 | self.history = history |
| 39 | self.kernel_id = kernel_id |
| 40 | self.session_id = session_id |
| 41 | self.on_kernel_started = on_kernel_started |
| 42 | self.on_session_started = on_session_started |
| 43 | self._kernel_client = None |
| 44 | self.colab_request_hook: Optional[Callable[[Dict[str, Any], Any], None]] = None |
| 45 | |
| 46 | def _apply_ws_hook(self): |
| 47 | wsclient = self._kernel_client._manager.client |
| 48 | original_on_message = wsclient.kernel_socket.on_message |
| 49 | |
| 50 | def hooked_on_message(s_ws, message): |
| 51 | if not self.colab_request_hook: |
| 52 | return original_on_message(s_ws, message) |
| 53 | |
| 54 | try: |
| 55 | from jupyter_kernel_client.wsclient import JupyterSubprotocol |
| 56 | |
| 57 | if wsclient._subprotocol == JupyterSubprotocol.DEFAULT: |
| 58 | from jupyter_kernel_client.wsclient import ( |
| 59 | deserialize_msg_from_ws_default, |
| 60 | ) |
| 61 | |
| 62 | deserialize_msg = deserialize_msg_from_ws_default(message) |
| 63 | elif wsclient._subprotocol == JupyterSubprotocol.V1: |
| 64 | from jupyter_kernel_client.wsclient import ( |
| 65 | deserialize_msg_from_ws_v1, |
| 66 | ) |
| 67 | |
| 68 | channel, msg_list = deserialize_msg_from_ws_v1(message) |
| 69 | deserialize_msg = wsclient.session.deserialize(msg_list) |
| 70 | else: |
| 71 | deserialize_msg = None |
| 72 | |
| 73 | if deserialize_msg: |
| 74 | msg_type = deserialize_msg.get("msg_type") |
| 75 | if msg_type == "colab_request": |
| 76 | # We pass the deserialized msg and the wsclient to the hook |
| 77 | if self.colab_request_hook(deserialize_msg, wsclient): |
| 78 | # If the hook returns True, we intercept and do NOT pass to original |
| 79 | return |
| 80 | |
| 81 | except Exception as e: |
| 82 | logging.debug(f"Error in colab_request hook: {e}") |
| 83 | |
| 84 | # Call original for all other messages |
| 85 | original_on_message(s_ws, message) |
| 86 | |
| 87 | wsclient.kernel_socket.on_message = hooked_on_message |
| 88 | |
| 89 | @property |
| 90 | def kernel_client(self): |
| 91 | if not self._kernel_client: |
| 92 | retries = 3 |
| 93 | backoff = 2 |
| 94 | last_err = None |
| 95 | |
| 96 | for i in range(retries): |
| 97 | try: |
| 98 | client_kwargs = { |
| 99 | "subprotocol": jupyter_kernel_client.JupyterSubprotocol.DEFAULT, |
| 100 | "extra_params": {"colab-runtime-proxy-token": self.token}, |
| 101 | } |
| 102 | if self.session_id: |
| 103 | # WSSession (Session) expects 'session' for the ID |
| 104 | client_kwargs["session"] = self.session_id |
| 105 | |
| 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 |
| 131 | |
| 132 | self._kernel_client.start() |
| 133 | self._apply_ws_hook() |
| 134 | |
| 135 | # Capture IDs if we started fresh |
| 136 | if not self.kernel_id and self._kernel_client.id: |
| 137 | self.kernel_id = self._kernel_client.id |
| 138 | if self.on_kernel_started: |
| 139 | self.on_kernel_started(self.kernel_id) |
| 140 | |
| 141 | if ( |
| 142 | not self.session_id |
| 143 | and self._kernel_client._manager.client.session.session |
| 144 | ): |
| 145 | self.session_id = ( |
| 146 | self._kernel_client._manager.client.session.session |
| 147 | ) |
| 148 | if self.on_session_started: |
| 149 | self.on_session_started(self.session_id) |
| 150 | break |
| 151 | except ( |
| 152 | requests.exceptions.ReadTimeout, |
| 153 | requests.exceptions.ConnectTimeout, |
| 154 | ) as e: |
| 155 | last_err = e |
| 156 | if i < retries - 1: |
| 157 | sleep_time = backoff ** (i + 1) |
| 158 | logging.debug( |
| 159 | f"Kernel startup timeout, retrying in {sleep_time}s..." |
| 160 | f" ({i + 1}/{retries})" |
| 161 | ) |
| 162 | time.sleep(sleep_time) |
| 163 | else: |
| 164 | raise last_err |
| 165 | except Exception as e: |
| 166 | raise e |
| 167 | |
| 168 | return self._kernel_client |
| 169 | |
| 170 | def restart( |
| 171 | self, |
| 172 | timeout: Optional[float] = None, |
| 173 | ): |
| 174 | self.kernel_client.restart(timeout=timeout) |
| 175 | |
| 176 | def execute_code( |
| 177 | self, |
| 178 | code: str, |
| 179 | allow_stdin: bool = False, |
| 180 | stdin_hook: Any = None, |
| 181 | output_hook: Optional[Callable[[Dict[str, Any]], None]] = None, |
| 182 | timeout: Optional[float] = None, |
| 183 | ) -> List[Dict[str, Any]]: |
| 184 | # ``jupyter_kernel_client`` defaults ``timeout`` to ``REQUEST_TIMEOUT`` |
| 185 | # (10 seconds) on both ``execute`` and ``execute_interactive``. That |
| 186 | # value is a wall-clock budget that shrinks every time the poll loop |
| 187 | # iterates -- as long as iopub/stdin events arrive back-to-back the |
| 188 | # call survives, but a single >10s quiet stretch (e.g. a kernel |
| 189 | # blocked on ``input_request`` while the user OAuths in the browser) |
| 190 | # will raise ``TimeoutError`` even though the underlying execution is |
| 191 | # still healthy. Callers that know they need a longer ceiling can |
| 192 | # pass ``timeout=`` here; otherwise we forward whatever the upstream |
| 193 | # default is (currently 10s). |
| 194 | kwargs = {"allow_stdin": allow_stdin} |
| 195 | if timeout is not None: |
| 196 | kwargs["timeout"] = timeout |
| 197 | |
| 198 | # Wrap stdin_hook to log inputs |
| 199 | original_stdin_hook = stdin_hook |
| 200 | |
| 201 | def wrapped_stdin_hook(prompt): |
| 202 | if self.history and self.session_name: |
| 203 | self.history.log_event( |
| 204 | self.session_name, "stdin_request", {"prompt": prompt} |
| 205 | ) |
| 206 | |
| 207 | res = original_stdin_hook(prompt) if original_stdin_hook else input(prompt) |
| 208 | |
| 209 | if self.history and self.session_name: |
| 210 | self.history.log_event(self.session_name, "input_reply", {"value": res}) |
| 211 | return res |
| 212 | |
| 213 | if allow_stdin: |
| 214 | kwargs["stdin_hook"] = wrapped_stdin_hook |
| 215 | |
| 216 | if output_hook: |
| 217 | # If we have an output hook, we use execute_interactive and manage buffering ourselves |
| 218 | outputs = [] |
| 219 | |
| 220 | def wrapped_output_hook(msg): |
| 221 | from jupyter_kernel_client.client import ( |
| 222 | output_hook as default_output_hook, |
| 223 | ) |
| 224 | |
| 225 | # Update local outputs list using the default logic |
| 226 | new_indexes = default_output_hook(outputs, msg) |
| 227 | # If new outputs were added, call our streaming hook with the new data |
| 228 | if new_indexes: |
| 229 | for idx in sorted(new_indexes): |
| 230 | if idx < len(outputs): |
| 231 | output_hook(outputs[idx]) |
| 232 | |
| 233 | reply = self.kernel_client.execute_interactive( |
| 234 | code, output_hook=wrapped_output_hook, **kwargs |
| 235 | ) |
| 236 | # execute_interactive returns the raw reply message |
| 237 | reply_content = reply["content"] if reply else {"status": "error"} |
| 238 | else: |
| 239 | reply = self.kernel_client.execute(code, **kwargs) |
| 240 | if not reply: |
| 241 | return [] |
| 242 | outputs = reply.get("outputs", []) |
| 243 | reply_content = reply |
| 244 | |
| 245 | # If there's an error status but no error in outputs, synthesize one |
| 246 | if reply_content.get("status") == "error": |
| 247 | has_error_output = any(o.get("output_type") == "error" for o in outputs) |
| 248 | if not has_error_output: |
| 249 | outputs.append( |
| 250 | { |
| 251 | "output_type": "error", |
| 252 | "ename": reply_content.get("ename", "Error"), |
| 253 | "evalue": reply_content.get("evalue", "Unknown error"), |
| 254 | "traceback": reply_content.get("traceback", []), |
| 255 | } |
| 256 | ) |
| 257 | |
| 258 | return outputs |
| 259 | |
| 260 | def stop(self, shutdown_kernel: bool = False): |
| 261 | if self._kernel_client: |
| 262 | try: |
| 263 | # We manage kernel lifecycle explicitly. |
| 264 | # To prevent automatic shutdown, we bypass the manager's stop() and |
| 265 | # directly close the channels and socket. |
| 266 | client = self._kernel_client._manager.client |
| 267 | client.stop_channels() |
| 268 | if client.kernel_socket: |
| 269 | client.kernel_socket.close() |
| 270 | |
| 271 | if shutdown_kernel: |
| 272 | self._kernel_client._manager.shutdown_kernel(now=True) |
| 273 | except Exception: |
| 274 | logging.exception("Error stopping kernel client") |