| 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 os |
| 16 | import subprocess |
| 17 | import sys |
| 18 | import time |
| 19 | import uuid |
| 20 | from typing import Any, Dict, Optional |
| 21 | import typer |
| 22 | from typing_extensions import Annotated |
| 23 | |
| 24 | from colab_cli.client import ( |
| 25 | Accelerator, |
| 26 | ColabRequestError, |
| 27 | HIGH_MEM_ONLY_ACCELERATORS, |
| 28 | PostAssignmentResponse, |
| 29 | Shape, |
| 30 | Variant, |
| 31 | resolve_assign_shape, |
| 32 | shape_display_label, |
| 33 | ) |
| 34 | from colab_cli.utils import get_status_code |
| 35 | from colab_cli.state import SessionState |
| 36 | from colab_cli.runtime import ColabRuntime |
| 37 | |
| 38 | |
| 39 | def _is_scope_error(e: Exception) -> bool: |
| 40 | """True if a ColabRequestError's response body indicates a missing OAuth scope. |
| 41 | |
| 42 | The frontend returns a `google.rpc.Status` with `code=7` (PERMISSION_DENIED) |
| 43 | and a `DebugInfo` payload mentioning `SCOPE_NOT_PERMITTED` / |
| 44 | "insufficient authentication scopes". Match on either substring so we |
| 45 | don't depend on the exact wording of one of them. |
| 46 | """ |
| 47 | body = getattr(e, "response_body", None) or "" |
| 48 | body_str = str(body) |
| 49 | return ( |
| 50 | "SCOPE_NOT_PERMITTED" in body_str |
| 51 | or "insufficient authentication scopes" in body_str |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | def _scope_remediation_message(provider) -> str: |
| 56 | """User-facing remediation hint, tailored per auth provider. |
| 57 | |
| 58 | Keep-alive is a Tunnel Frontend ping against the Colab session backend |
| 59 | (colab.research.google.com), authenticated with the user's own Gaia bearer |
| 60 | token — the same credential and host used to assign the VM. A missing-scope |
| 61 | error here is rare (assignment would normally have failed first), but if it |
| 62 | happens the fix is to re-authenticate with the standard Colab scopes. |
| 63 | """ |
| 64 | # Importing locally to avoid a circular import at module load time. |
| 65 | from colab_cli.auth import AuthProvider |
| 66 | |
| 67 | common = ( |
| 68 | "Keeping the session alive requires valid Colab credentials for " |
| 69 | "colab.research.google.com." |
| 70 | ) |
| 71 | if provider == AuthProvider.ADC: |
| 72 | return ( |
| 73 | f"{common}\n" |
| 74 | "Re-authenticate ADC with the standard Colab scopes (the " |
| 75 | "cloud-platform and openid scopes are required by gcloud itself):\n" |
| 76 | " gcloud auth application-default login \\\n" |
| 77 | " --scopes=openid," |
| 78 | "https://www.googleapis.com/auth/cloud-platform," |
| 79 | "https://www.googleapis.com/auth/userinfo.email," |
| 80 | "https://www.googleapis.com/auth/colaboratory\n" |
| 81 | "Then re-run `colab new`." |
| 82 | ) |
| 83 | # OAuth2 (and any future provider) fallback. |
| 84 | return ( |
| 85 | f"{common}\n" |
| 86 | "Delete the cached token at ~/.config/colab-cli/token.json and " |
| 87 | "re-run `colab new` to trigger a fresh consent flow." |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | def _hardware_label(accelerator: str) -> str: |
| 92 | """`NONE` -> `CPU`; everything else passes through.""" |
| 93 | return "CPU" if accelerator == "NONE" else accelerator |
| 94 | |
| 95 | |
| 96 | def _format_session_line( |
| 97 | name: str, |
| 98 | endpoint: str, |
| 99 | accelerator: str, |
| 100 | variant: str, |
| 101 | status: Optional[str] = None, |
| 102 | machine_shape: Optional[str] = None, |
| 103 | ) -> str: |
| 104 | """Single source of truth for session display lines. |
| 105 | |
| 106 | Format: ``[name] endpoint | Hardware: X | Shape: Y | Variant: Z[ | Status: W]``. |
| 107 | Use ``"?"`` as the name for orphaned server-side assignments with no local |
| 108 | state. |
| 109 | """ |
| 110 | parts = [ |
| 111 | f"[{name}] {endpoint}", |
| 112 | f"Hardware: {_hardware_label(accelerator)}", |
| 113 | f"Shape: {shape_display_label(machine_shape)}", |
| 114 | f"Variant: {variant}", |
| 115 | ] |
| 116 | if status is not None: |
| 117 | parts.append(f"Status: {status}") |
| 118 | return " | ".join(parts) |
| 119 | |
| 120 | |
| 121 | def resolve_runtime_options( |
| 122 | gpu: Optional[str] = None, |
| 123 | tpu: Optional[str] = None, |
| 124 | *, |
| 125 | high_mem: bool = False, |
| 126 | ) -> tuple[Variant, Accelerator, Optional[Shape]]: |
| 127 | """Map CLI flags to backend variant, accelerator, and optional shape.""" |
| 128 | if tpu: |
| 129 | variant = Variant.TPU |
| 130 | accelerator = Accelerator.V5E1 if tpu.lower() == "v5e1" else Accelerator.V6E1 |
| 131 | elif gpu: |
| 132 | variant = Variant.GPU |
| 133 | mapping = { |
| 134 | "a100": Accelerator.A100, |
| 135 | "h100": Accelerator.H100, |
| 136 | "l4": Accelerator.L4, |
| 137 | "t4": Accelerator.T4, |
| 138 | "g4": Accelerator.G4, |
| 139 | } |
| 140 | accelerator = mapping.get(gpu.lower(), Accelerator.A100) |
| 141 | else: |
| 142 | variant = Variant.DEFAULT |
| 143 | accelerator = Accelerator.NONE |
| 144 | |
| 145 | shape = resolve_assign_shape(accelerator, high_mem=high_mem) |
| 146 | return variant, accelerator, shape |
| 147 | |
| 148 | |
| 149 | def new( |
| 150 | session: Annotated[ |
| 151 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 152 | ] = None, |
| 153 | tpu: Annotated[ |
| 154 | Optional[str], |
| 155 | typer.Option( |
| 156 | help="TPU accelerator variant. Supported: v5e1, v6e1.", |
| 157 | ), |
| 158 | ] = None, |
| 159 | gpu: Annotated[ |
| 160 | Optional[str], |
| 161 | typer.Option( |
| 162 | help=( |
| 163 | "GPU accelerator variant. Supported: T4, L4, G4, H100, A100." |
| 164 | "\n\nIf omitted (along with --tpu), a CPU runtime is created." |
| 165 | "\n\nAvailability varies by Colab subscription tier." |
| 166 | ), |
| 167 | ), |
| 168 | ] = None, |
| 169 | high_mem: Annotated[ |
| 170 | bool, |
| 171 | typer.Option( |
| 172 | "--high-mem", |
| 173 | help=( |
| 174 | "Request a high-RAM machine shape (CPU, T4, A100, etc.). " |
| 175 | "Requires Colab Pro or Pro+ entitlement. Ignored for " |
| 176 | "accelerators that only offer a single shape (L4, v5e1, v6e1)." |
| 177 | ), |
| 178 | ), |
| 179 | ] = False, |
| 180 | ): |
| 181 | """Create a new session""" |
| 182 | from colab_cli.common import state |
| 183 | |
| 184 | name = session or uuid.uuid4().hex[:6] |
| 185 | variant, accelerator, shape = resolve_runtime_options(gpu, tpu, high_mem=high_mem) |
| 186 | |
| 187 | if high_mem and accelerator in HIGH_MEM_ONLY_ACCELERATORS: |
| 188 | typer.echo( |
| 189 | "[colab] --high-mem ignored: this accelerator only offers one " |
| 190 | "machine shape.", |
| 191 | err=True, |
| 192 | ) |
| 193 | |
| 194 | typer.echo(f"[colab] Creating session '{name}'...") |
| 195 | try: |
| 196 | res = state.client.assign( |
| 197 | uuid.uuid4(), variant=variant, accelerator=accelerator, shape=shape |
| 198 | ) |
| 199 | except ColabRequestError as e: |
| 200 | # The Colab backend returns 400 when the caller is not entitled to the |
| 201 | # requested accelerator (e.g. no A100 quota). Translate that to a |
| 202 | # friendly, actionable message instead of a raw traceback. We only |
| 203 | # interpret it this way when an accelerator was actually requested; |
| 204 | # otherwise we re-raise so the user sees the real cause. |
| 205 | if get_status_code(e) == 400 and accelerator != Accelerator.NONE: |
| 206 | typer.echo( |
| 207 | f"[colab] Backend rejected accelerator '{accelerator.value}'. " |
| 208 | "You may not have quota or entitlement for this accelerator on " |
| 209 | "your account. Try a different one (e.g. --gpu T4) or omit " |
| 210 | "--gpu/--tpu for a CPU runtime.", |
| 211 | err=True, |
| 212 | ) |
| 213 | raise typer.Exit(code=1) |
| 214 | raise |
| 215 | |
| 216 | if isinstance(res, PostAssignmentResponse): |
| 217 | token = res.runtime_proxy_info.token |
| 218 | url = res.runtime_proxy_info.url |
| 219 | endpoint = res.endpoint |
| 220 | else: |
| 221 | token = ( |
| 222 | res.runtime_proxy_info.token |
| 223 | if hasattr(res, "runtime_proxy_info") |
| 224 | else getattr(res, "runtime_proxy_token", "") |
| 225 | ) |
| 226 | url = res.runtime_proxy_info.url if hasattr(res, "runtime_proxy_info") else "" |
| 227 | endpoint = res.endpoint |
| 228 | |
| 229 | # Importing locally to avoid a top-level circular import via auth. |
| 230 | |
| 231 | s = SessionState( |
| 232 | name=name, |
| 233 | token=token, |
| 234 | url=url, |
| 235 | endpoint=endpoint, |
| 236 | variant=variant.value, |
| 237 | accelerator=accelerator.value, |
| 238 | machine_shape=( |
| 239 | Shape.HIGH_RAM.name if shape == Shape.HIGH_RAM else Shape.STANDARD.name |
| 240 | ), |
| 241 | ) |
| 242 | |
| 243 | # Pre-flight the keep-alive ping once. If it returns a 403 caused by |
| 244 | # missing OAuth scopes we know the daemon will fail and the VM would be |
| 245 | # idle-pruned. Catch it now so we (a) never leak a billable assignment, |
| 246 | # (b) surface an actionable remediation instead of a session that quietly |
| 247 | # disappears a few minutes later. |
| 248 | try: |
| 249 | state.client.keep_alive_assignment(endpoint) |
| 250 | except ColabRequestError as e: |
| 251 | if get_status_code(e) == 403 and _is_scope_error(e): |
| 252 | typer.echo( |
| 253 | "[colab] Keep-alive pre-flight failed: your credentials " |
| 254 | "are missing an OAuth scope required by Colab.\n", |
| 255 | err=True, |
| 256 | ) |
| 257 | typer.echo(_scope_remediation_message(state.auth_provider), err=True) |
| 258 | # Don't leak the assignment we just created. |
| 259 | try: |
| 260 | state.client.unassign(endpoint) |
| 261 | except Exception: |
| 262 | pass |
| 263 | raise typer.Exit(code=1) |
| 264 | # Other failures: don't block session creation — the daemon will |
| 265 | # retry and log via the existing keep_alive_error event path. |
| 266 | |
| 267 | # Persist the session BEFORE spawning the daemon so the daemon's |
| 268 | # initial `state.store.get(session_name)` check doesn't race and |
| 269 | # exit with `reason=session_not_found`. We re-persist below to also |
| 270 | # capture the daemon PID. |
| 271 | state.store.add(s) |
| 272 | s.keep_alive_pid = spawn_keep_alive( |
| 273 | endpoint, |
| 274 | name, |
| 275 | auth_provider=state.auth_provider, |
| 276 | config_path=state.config_path, |
| 277 | ) |
| 278 | |
| 279 | state.store.add(s) |
| 280 | state.history.log_event( |
| 281 | name, |
| 282 | "session_created", |
| 283 | { |
| 284 | "endpoint": endpoint, |
| 285 | "variant": variant.value, |
| 286 | "accelerator": accelerator.value, |
| 287 | "machine_shape": s.machine_shape, |
| 288 | }, |
| 289 | ) |
| 290 | typer.echo("[colab] Session READY.") |
| 291 | |
| 292 | |
| 293 | def restart_kernel( |
| 294 | session: Annotated[ |
| 295 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 296 | ] = None, |
| 297 | ): |
| 298 | """Restart a session's kernel""" |
| 299 | from colab_cli.common import state |
| 300 | |
| 301 | name = state.resolve_session(session) |
| 302 | s = state.store.get(name) |
| 303 | |
| 304 | def on_started(kid): |
| 305 | s.kernel_id = kid |
| 306 | state.store.add(s) |
| 307 | |
| 308 | def on_sess_started(sid): |
| 309 | s.session_id = sid |
| 310 | state.store.add(s) |
| 311 | |
| 312 | runtime = ColabRuntime( |
| 313 | s.url, |
| 314 | s.token, |
| 315 | kernel_id=s.kernel_id, |
| 316 | session_id=s.session_id, |
| 317 | on_kernel_started=on_started, |
| 318 | on_session_started=on_sess_started, |
| 319 | ) |
| 320 | |
| 321 | try: |
| 322 | runtime.restart() |
| 323 | finally: |
| 324 | runtime.stop() |
| 325 | |
| 326 | |
| 327 | def sessions_command(): |
| 328 | """List all active sessions""" |
| 329 | from colab_cli.common import state |
| 330 | |
| 331 | sessions, assignments = state.sync_sessions() |
| 332 | if not assignments: |
| 333 | typer.echo("[colab] No active sessions found on server.") |
| 334 | return |
| 335 | |
| 336 | # Build endpoint -> local-name lookup so we can lead with the friendly name. |
| 337 | name_by_endpoint = {s.endpoint: s.name for s in sessions.values()} |
| 338 | for a in assignments: |
| 339 | name = name_by_endpoint.get(a.endpoint, "?") |
| 340 | # `a.variant` is an int-valued AssignmentVariant (DEFAULT=0/GPU=1/TPU=2); |
| 341 | # its `.name` matches the user-facing string Variant enum, which is what |
| 342 | # `status` shows for locally-tracked sessions. |
| 343 | typer.echo( |
| 344 | _format_session_line( |
| 345 | name=name, |
| 346 | endpoint=a.endpoint, |
| 347 | accelerator=a.accelerator.value, |
| 348 | variant=a.variant.name, |
| 349 | machine_shape=a.machine_shape.name, |
| 350 | ) |
| 351 | ) |
| 352 | |
| 353 | |
| 354 | def _print_status_for(s: SessionState) -> None: |
| 355 | """Print one session's status line plus optional last-execution detail.""" |
| 356 | status = f"BUSY ({s.running})" if s.running else "IDLE" |
| 357 | typer.echo( |
| 358 | _format_session_line( |
| 359 | name=s.name, |
| 360 | endpoint=s.endpoint, |
| 361 | accelerator=s.accelerator, |
| 362 | variant=s.variant, |
| 363 | status=status, |
| 364 | machine_shape=s.machine_shape, |
| 365 | ) |
| 366 | ) |
| 367 | if s.last_execution: |
| 368 | exec_file, exec_cell, exec_time = s.last_execution |
| 369 | cell_str = f" | Cell: {exec_cell}" if exec_cell else "" |
| 370 | typer.echo(f" Last Execution: {exec_file}{cell_str} at {exec_time}") |
| 371 | |
| 372 | |
| 373 | def status( |
| 374 | session: Annotated[ |
| 375 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 376 | ] = None, |
| 377 | ): |
| 378 | """Show session status""" |
| 379 | from colab_cli.common import state |
| 380 | |
| 381 | local_sessions, _ = state.sync_sessions() |
| 382 | if session: |
| 383 | s = state.store.get(session) |
| 384 | if s: |
| 385 | _print_status_for(s) |
| 386 | else: |
| 387 | typer.echo(f"[colab] Session '{session}' not found.") |
| 388 | return |
| 389 | |
| 390 | if not local_sessions: |
| 391 | typer.echo("[colab] No active sessions.") |
| 392 | return |
| 393 | for s in local_sessions.values(): |
| 394 | _print_status_for(s) |
| 395 | |
| 396 | |
| 397 | def stop( |
| 398 | session: Annotated[ |
| 399 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 400 | ] = None, |
| 401 | ): |
| 402 | """Stop a session""" |
| 403 | from colab_cli.common import state |
| 404 | |
| 405 | name = state.resolve_session(session) |
| 406 | s = state.store.get(name) |
| 407 | if not s: |
| 408 | typer.echo(f"[colab] Session '{name}' not found.") |
| 409 | return |
| 410 | |
| 411 | typer.echo(f"[colab] Stopping session '{name}'...") |
| 412 | if s.keep_alive_pid: |
| 413 | from colab_cli.common import kill_process |
| 414 | |
| 415 | kill_process(s.keep_alive_pid) |
| 416 | |
| 417 | try: |
| 418 | runtime = ColabRuntime(s.url, s.token, kernel_id=s.kernel_id) |
| 419 | runtime.stop(shutdown_kernel=True) |
| 420 | except Exception: |
| 421 | pass |
| 422 | |
| 423 | state.client.unassign(s.endpoint) |
| 424 | state.store.remove(name) |
| 425 | state.history.log_event(name, "session_terminated", {"reason": "user_requested"}) |
| 426 | typer.echo("[colab] Session terminated.") |
| 427 | |
| 428 | |
| 429 | def spawn_keep_alive( |
| 430 | endpoint: str, session_name: str, auth_provider=None, config_path=None |
| 431 | ): |
| 432 | """Spawns a detached keep-alive process. |
| 433 | |
| 434 | Both `auth_provider` and `config_path` are propagated as global flags |
| 435 | so the detached child uses the same authentication strategy AND the |
| 436 | same session state file as the parent that invoked `colab new`. |
| 437 | Without this, the child inherits Typer's defaults (`--auth=oauth2`, |
| 438 | `--config=~/.config/colab-cli/sessions.json`), which causes: |
| 439 | (a) wrong auth backend, and |
| 440 | (b) the daemon's `state.store.get(session_name)` check finds nothing |
| 441 | and exits with `reason=session_not_found` when the parent used |
| 442 | `--config` to write to a non-default path. |
| 443 | """ |
| 444 | cmd = [sys.executable, "-m", "colab_cli.cli"] |
| 445 | if auth_provider is not None: |
| 446 | cmd.append(f"--auth={auth_provider.value}") |
| 447 | if config_path is not None: |
| 448 | cmd.extend(["--config", config_path]) |
| 449 | cmd.extend(["keep-alive", endpoint, session_name]) |
| 450 | # Detach process |
| 451 | kwargs = {} |
| 452 | if sys.platform != "win32": |
| 453 | kwargs["start_new_session"] = True |
| 454 | else: |
| 455 | # https://stackoverflow.com/questions/1356540/how-can-i-make-a-python-script-run-in-the-background-as-a-service-on-windows |
| 456 | CREATE_NEW_PROCESS_GROUP = 0x00000200 |
| 457 | DETACHED_PROCESS = 0x00000008 |
| 458 | kwargs["creationflags"] = DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP |
| 459 | |
| 460 | p = subprocess.Popen( |
| 461 | cmd, |
| 462 | stdout=subprocess.DEVNULL, |
| 463 | stderr=subprocess.DEVNULL, |
| 464 | stdin=subprocess.DEVNULL, |
| 465 | **kwargs, |
| 466 | ) |
| 467 | return p.pid |
| 468 | |
| 469 | |
| 470 | def keep_alive( |
| 471 | endpoint: Annotated[str, typer.Argument(help="Endpoint ID")], |
| 472 | session_name: Annotated[str, typer.Argument(help="Session name")], |
| 473 | ): |
| 474 | """Hidden command to run keep-alive loop. Terminate after 24h.""" |
| 475 | from colab_cli.common import state |
| 476 | |
| 477 | state.history.log_event( |
| 478 | session_name, |
| 479 | "keep_alive_started", |
| 480 | {"endpoint": endpoint, "pid": os.getpid()}, |
| 481 | ) |
| 482 | |
| 483 | start_time = time.time() |
| 484 | # 24 hours limit |
| 485 | max_duration = 24 * 3600 |
| 486 | consecutive_4xx = 0 |
| 487 | iterations = 0 |
| 488 | last_error: Optional[Dict[str, Any]] = None |
| 489 | |
| 490 | reason = "time_limit_reached" |
| 491 | extra: Dict[str, Any] = {} |
| 492 | while time.time() - start_time < max_duration: |
| 493 | iterations += 1 |
| 494 | # Check if session still exists in local state |
| 495 | s = state.store.get(session_name) |
| 496 | if not s: |
| 497 | reason = "session_not_found" |
| 498 | break |
| 499 | if s.endpoint != endpoint: |
| 500 | reason = "endpoint_mismatch" |
| 501 | extra["expected_endpoint"] = endpoint |
| 502 | extra["actual_endpoint"] = s.endpoint |
| 503 | break |
| 504 | |
| 505 | try: |
| 506 | state.client.keep_alive_assignment(endpoint) |
| 507 | consecutive_4xx = 0 |
| 508 | last_error = None |
| 509 | except Exception as e: |
| 510 | code = get_status_code(e) |
| 511 | response_body = getattr(e, "response_body", None) |
| 512 | err_info = { |
| 513 | "status_code": code, |
| 514 | "error_type": type(e).__name__, |
| 515 | "error": str(e)[:500], |
| 516 | "response_body": (str(response_body)[:1000] if response_body else None), |
| 517 | } |
| 518 | last_error = err_info |
| 519 | state.history.log_event( |
| 520 | session_name, |
| 521 | "keep_alive_error", |
| 522 | { |
| 523 | **err_info, |
| 524 | "iteration": iterations, |
| 525 | "consecutive_4xx": consecutive_4xx |
| 526 | + (1 if code is not None and 400 <= code < 500 else 0), |
| 527 | }, |
| 528 | ) |
| 529 | if code is not None and 400 <= code < 500: |
| 530 | consecutive_4xx += 1 |
| 531 | if consecutive_4xx >= 2: |
| 532 | reason = "consecutive_4xx_errors" |
| 533 | break |
| 534 | else: |
| 535 | # For other errors (network), we retry and don't count as 4xx |
| 536 | pass |
| 537 | |
| 538 | time.sleep(60) |
| 539 | |
| 540 | payload: Dict[str, Any] = { |
| 541 | "reason": reason, |
| 542 | "iterations": iterations, |
| 543 | "duration_seconds": round(time.time() - start_time, 2), |
| 544 | } |
| 545 | if last_error is not None: |
| 546 | payload["last_error"] = last_error |
| 547 | payload.update(extra) |
| 548 | state.history.log_event(session_name, "keep_alive_stopped", payload) |
| 549 | |
| 550 | |
| 551 | def register(app: typer.Typer): |
| 552 | app.command()(new) |
| 553 | app.command(name="sessions")(sessions_command) |
| 554 | app.command(name="restart-kernel")(restart_kernel) |
| 555 | app.command()(status) |
| 556 | app.command()(stop) |
| 557 | app.command(hidden=True)(keep_alive) |