| 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 | from typing import Optional |
| 16 | |
| 17 | import typer |
| 18 | from typing_extensions import Annotated |
| 19 | |
| 20 | from colab_cli import auto_update |
| 21 | from colab_cli.auto_update import get_app_version |
| 22 | from colab_cli.common import state |
| 23 | |
| 24 | |
| 25 | def pay(): |
| 26 | """Open the Colab signup page to manage compute units""" |
| 27 | import webbrowser |
| 28 | |
| 29 | url = "https://colab.research.google.com/signup" |
| 30 | typer.echo(f"[colab] Opening {url}...") |
| 31 | webbrowser.open(url) |
| 32 | |
| 33 | |
| 34 | def url( |
| 35 | session: Annotated[ |
| 36 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 37 | ] = None, |
| 38 | host: Annotated[ |
| 39 | str, |
| 40 | typer.Option( |
| 41 | "--host", |
| 42 | help=( |
| 43 | "Colab frontend host (origin) to use for the URL. The Colab " |
| 44 | "frontend resolves `dbu` against `window.location.origin`, " |
| 45 | "so this only changes the page origin, not the embedded " |
| 46 | "backend path." |
| 47 | ), |
| 48 | ), |
| 49 | ] = "https://colab.research.google.com", |
| 50 | open_browser: Annotated[ |
| 51 | bool, |
| 52 | typer.Option( |
| 53 | "--open", |
| 54 | help=( |
| 55 | "After printing the URL, also open it in the system browser. " |
| 56 | "Off by default so the command remains pipeable " |
| 57 | "(e.g. `colab url -s s1 | xclip`)." |
| 58 | ), |
| 59 | ), |
| 60 | ] = False, |
| 61 | ): |
| 62 | """Print a browser URL that connects to an existing session. |
| 63 | |
| 64 | Format: ``https://<host>/notebooks/empty.ipynb?dbu=<urlencoded path>#datalabBackendUrl=<host>/tun/m/<endpoint>``, |
| 65 | where the path is ``/tun/m/<endpoint>``. When opened, the Colab frontend |
| 66 | skips ``/tun/m/assign`` and attaches the kernel to our existing VM. |
| 67 | |
| 68 | Two backend-URL signals are embedded: |
| 69 | |
| 70 | - ``?dbu=<urlencoded path>`` — the ``datalab_backend_url`` development |
| 71 | query flag. The frontend resolves the value against |
| 72 | ``window.location.origin``. |
| 73 | |
| 74 | - ``#datalabBackendUrl=<full URL>`` — the hash-fragment form. Some |
| 75 | frontend code paths consult this first and ignore ``dbu``, so we |
| 76 | emit both for robustness. The fragment value is a FULL URL (with |
| 77 | scheme + host) and is intentionally NOT URL-encoded — browsers do |
| 78 | not decode the fragment before passing ``location.hash`` to page |
| 79 | JS, and Colab's hash parser expects the raw string. |
| 80 | |
| 81 | The fragment's host always matches ``--host`` (the page origin), so |
| 82 | Colab's same-origin enforcement on the embedded backend URL doesn't |
| 83 | block the connection, and sandbox/dev users get a sandbox fragment |
| 84 | automatically. |
| 85 | """ |
| 86 | # Imported here (not at module top) to mirror the lazy-state pattern used |
| 87 | # elsewhere in this module and avoid a circular import via colab_cli.common. |
| 88 | from urllib.parse import quote |
| 89 | |
| 90 | from colab_cli.common import state |
| 91 | |
| 92 | name = state.resolve_session(session) |
| 93 | s = state.store.get(name) |
| 94 | if not s: |
| 95 | typer.echo(f"[colab] Session '{name}' not found.", err=True) |
| 96 | raise typer.Exit(code=1) |
| 97 | |
| 98 | # Strip a trailing slash so we don't produce `https://host//notebooks/...` |
| 99 | # or `https://host//tun/m/...` in the fragment URL. |
| 100 | host_clean = host.rstrip("/") |
| 101 | backend_path = f"/tun/m/{s.endpoint}" |
| 102 | # `dbu` value is the backend path. URL-encode it (incl. the slashes via |
| 103 | # `safe=""`) so the value survives any downstream non-strict query-string |
| 104 | # re-parsing — this is also the form shown in real Colab connect URLs. |
| 105 | dbu_value = quote(backend_path, safe="") |
| 106 | # `#datalabBackendUrl=` value is the FULL backend URL, raw (un-encoded): |
| 107 | # the browser does not decode the fragment before passing it to page JS, |
| 108 | # and Colab's hash parser calls `new URL(rawString)` directly. Pinning |
| 109 | # the host to `host_clean` (not hardcoding research.google.com) keeps |
| 110 | # this aligned with the page origin so same-origin enforcement passes |
| 111 | # for sandbox / dev hosts too. |
| 112 | fragment_value = f"{host_clean}{backend_path}" |
| 113 | connect_url = ( |
| 114 | f"{host_clean}/notebooks/empty.ipynb" |
| 115 | f"?dbu={dbu_value}" |
| 116 | f"#datalabBackendUrl={fragment_value}" |
| 117 | ) |
| 118 | |
| 119 | # Print the URL on its own line with no `[colab]` prefix so the output |
| 120 | # is pipeable (`colab url -s s1 | xclip`, etc.). |
| 121 | typer.echo(connect_url) |
| 122 | |
| 123 | if open_browser: |
| 124 | import webbrowser |
| 125 | |
| 126 | webbrowser.open(connect_url) |
| 127 | |
| 128 | |
| 129 | def log( |
| 130 | session: Annotated[ |
| 131 | Optional[str], |
| 132 | typer.Option( |
| 133 | "-s", |
| 134 | "--session", |
| 135 | help="Session name (if omitted, lists all sessions with logs)", |
| 136 | ), |
| 137 | ] = None, |
| 138 | lines: Annotated[ |
| 139 | Optional[int], |
| 140 | typer.Option( |
| 141 | "-n", "--lines", help="Number of lines to show/export (default: all)" |
| 142 | ), |
| 143 | ] = None, |
| 144 | type: Annotated[ |
| 145 | Optional[str], |
| 146 | typer.Option( |
| 147 | "-t", |
| 148 | "--type", |
| 149 | help="Filter by event type (e.g., execution, file_operation)", |
| 150 | ), |
| 151 | ] = None, |
| 152 | output: Annotated[ |
| 153 | Optional[str], |
| 154 | typer.Option( |
| 155 | "-o", |
| 156 | "--output", |
| 157 | help="Output file path (suffix determines format: .ipynb, .md, .txt, .jsonl)", |
| 158 | ), |
| 159 | ] = None, |
| 160 | ): |
| 161 | """Manage and view session history logs""" |
| 162 | if not session: |
| 163 | sessions_with_logs = state.history.list_sessions() |
| 164 | if not sessions_with_logs: |
| 165 | typer.echo("[colab] No session history found.") |
| 166 | else: |
| 167 | typer.echo("[colab] Sessions with history logs:") |
| 168 | for n in sorted(sessions_with_logs): |
| 169 | typer.echo(f" {n}") |
| 170 | return |
| 171 | |
| 172 | events = state.history.get_history(session) |
| 173 | if not events: |
| 174 | typer.echo(f"[colab] No history found for session '{session}'.") |
| 175 | return |
| 176 | |
| 177 | if type: |
| 178 | events = [e for e in events if e.get("event_type") == type] |
| 179 | |
| 180 | if lines: |
| 181 | events = events[-lines:] |
| 182 | |
| 183 | if output: |
| 184 | from colab_cli.converter import export_history |
| 185 | |
| 186 | export_history(events, session, output) |
| 187 | else: |
| 188 | for event in events: |
| 189 | ts = event.get("timestamp", "").split(".")[0].replace("T", " ") |
| 190 | etype = event.get("event_type", "unknown") |
| 191 | |
| 192 | if etype == "execution": |
| 193 | preview = event.get("code", "").strip().split("\n")[0][:60] |
| 194 | typer.echo(f"[{ts}] EXEC: {preview}...") |
| 195 | elif etype == "file_operation": |
| 196 | typer.echo( |
| 197 | f"[{ts}] FILE: {event.get('op')} {event.get('path', event.get('remote', ''))}" |
| 198 | ) |
| 199 | elif etype == "automation": |
| 200 | typer.echo(f"[{ts}] AUTO: {event.get('op')}") |
| 201 | elif etype == "stdin_request": |
| 202 | typer.echo(f"[{ts}] INPT: {event.get('prompt', '').strip()}") |
| 203 | elif etype == "input_reply": |
| 204 | typer.echo(f"[{ts}] RPLY: {event.get('value', '').strip()}") |
| 205 | elif etype == "keep_alive_started": |
| 206 | typer.echo( |
| 207 | f"[{ts}] KEEP: started endpoint={event.get('endpoint')} pid={event.get('pid')}" |
| 208 | ) |
| 209 | elif etype == "keep_alive_error": |
| 210 | msg = ( |
| 211 | f"[{ts}] KEEP: error iter={event.get('iteration')} " |
| 212 | f"status={event.get('status_code')} " |
| 213 | f"type={event.get('error_type')} " |
| 214 | f"msg={event.get('error', '')[:120]}" |
| 215 | ) |
| 216 | body = event.get("response_body") |
| 217 | if body: |
| 218 | msg += f" body={body[:300]}" |
| 219 | typer.echo(msg) |
| 220 | elif etype == "keep_alive_stopped": |
| 221 | msg = ( |
| 222 | f"[{ts}] KEEP: stopped reason={event.get('reason')} " |
| 223 | f"iters={event.get('iterations')} " |
| 224 | f"duration={event.get('duration_seconds')}s" |
| 225 | ) |
| 226 | last_err = event.get("last_error") |
| 227 | if last_err: |
| 228 | msg += ( |
| 229 | f" last_error=[status={last_err.get('status_code')} " |
| 230 | f"type={last_err.get('error_type')} " |
| 231 | f"msg={str(last_err.get('error', ''))[:120]}]" |
| 232 | ) |
| 233 | if event.get("expected_endpoint") or event.get("actual_endpoint"): |
| 234 | msg += ( |
| 235 | f" expected={event.get('expected_endpoint')} " |
| 236 | f"actual={event.get('actual_endpoint')}" |
| 237 | ) |
| 238 | typer.echo(msg) |
| 239 | else: |
| 240 | typer.echo(f"[{ts}] EVENT: {etype}") |
| 241 | |
| 242 | |
| 243 | def whoami(): |
| 244 | """[debug] Print the active credentials' identity, scopes, and expiry. |
| 245 | |
| 246 | Mints an access token using the same path the rest of the CLI uses |
| 247 | (`auth.get_credentials(...)` honoring the global `--auth=...` flag), |
| 248 | then queries Google's tokeninfo endpoint and prints a human-readable |
| 249 | summary. Useful when debugging "why is my call to |
| 250 | colab.pa.googleapis.com 403-ing" — the answer is almost always a |
| 251 | missing scope or a token whose `email` doesn't match what you |
| 252 | expected. |
| 253 | |
| 254 | Hidden from `colab --help` because end users shouldn't need it; reach |
| 255 | it via `colab whoami --help` or by knowing the name. |
| 256 | """ |
| 257 | import json |
| 258 | import urllib.error |
| 259 | import urllib.parse |
| 260 | import urllib.request |
| 261 | |
| 262 | from colab_cli.auth import get_credentials |
| 263 | |
| 264 | provider = state.auth_provider |
| 265 | |
| 266 | # Mint a fresh token. Some credential types (service-account, GCE, some |
| 267 | # impersonated creds) don't populate `.token` until refresh() is called, |
| 268 | # so we always refresh — cheap, ~1 RPC, and avoids a confusing |
| 269 | # `creds.token is None` failure mode for valid credentials. |
| 270 | sess = get_credentials(state.client_oauth_config, provider=provider) |
| 271 | creds = sess.credentials |
| 272 | try: |
| 273 | from google.auth.transport.requests import Request as _GoogleAuthRequest |
| 274 | |
| 275 | creds.refresh(_GoogleAuthRequest()) |
| 276 | except Exception as e: |
| 277 | typer.echo(f"[colab] whoami: failed to refresh credentials: {e}", err=True) |
| 278 | raise typer.Exit(code=1) |
| 279 | |
| 280 | token = creds.token |
| 281 | if not token: |
| 282 | typer.echo( |
| 283 | "[colab] whoami: credentials have no access token after refresh; " |
| 284 | "the auth provider may have failed silently.", |
| 285 | err=True, |
| 286 | ) |
| 287 | raise typer.Exit(code=1) |
| 288 | |
| 289 | # Hit Google's tokeninfo endpoint. We use stdlib urllib (rather than the |
| 290 | # already-authorized `sess`) deliberately: tokeninfo accepts the token as |
| 291 | # a query parameter and does NOT want a Bearer header alongside it. |
| 292 | qs = urllib.parse.urlencode({"access_token": token}) |
| 293 | url = f"https://oauth2.googleapis.com/tokeninfo?{qs}" |
| 294 | try: |
| 295 | with urllib.request.urlopen(url, timeout=10) as resp: |
| 296 | body = resp.read().decode("utf-8") |
| 297 | info = json.loads(body) |
| 298 | except urllib.error.HTTPError as e: |
| 299 | # tokeninfo returns 400 for invalid/expired/revoked tokens with a |
| 300 | # JSON body like {"error":"invalid_token","error_description":"..."}. |
| 301 | # Surface that body so the developer can see *why* it was rejected. |
| 302 | try: |
| 303 | err_body = e.read().decode("utf-8") |
| 304 | except Exception: |
| 305 | err_body = "" |
| 306 | typer.echo( |
| 307 | f"[colab] whoami: tokeninfo returned HTTP {e.code}: {err_body or e.reason}", |
| 308 | err=True, |
| 309 | ) |
| 310 | raise typer.Exit(code=1) |
| 311 | except Exception as e: |
| 312 | typer.echo(f"[colab] whoami: tokeninfo request failed: {e}", err=True) |
| 313 | raise typer.Exit(code=1) |
| 314 | |
| 315 | # Format. Provider name from the AuthProvider enum (e.g. "adc"); email |
| 316 | # may be missing for tokens scoped without `userinfo.email`, in which |
| 317 | # case we say so explicitly rather than printing "Email: None". |
| 318 | email = info.get("email") or "<unavailable: token has no userinfo.email scope>" |
| 319 | expires_in = info.get("expires_in") |
| 320 | try: |
| 321 | expires_min = int(expires_in) // 60 |
| 322 | expires_str = f"{expires_min}m" |
| 323 | except (TypeError, ValueError): |
| 324 | expires_str = str(expires_in) if expires_in else "<unknown>" |
| 325 | |
| 326 | audience = info.get("audience") or info.get("aud") or "<none>" |
| 327 | scopes = (info.get("scope") or "").split() |
| 328 | |
| 329 | typer.echo(f"Auth provider: {provider.value}") |
| 330 | typer.echo(f"Email: {email}") |
| 331 | typer.echo(f"Audience: {audience}") |
| 332 | typer.echo(f"Expires in: {expires_str}") |
| 333 | if scopes: |
| 334 | typer.echo("Scopes:") |
| 335 | for s in sorted(scopes): |
| 336 | typer.echo(f" - {s}") |
| 337 | else: |
| 338 | typer.echo("Scopes: <none>") |
| 339 | |
| 340 | |
| 341 | def version_command(): |
| 342 | """Show the version of the Colab CLI""" |
| 343 | typer.echo(f"Version: {get_app_version()}") |
| 344 | |
| 345 | |
| 346 | def update_command( |
| 347 | install: Annotated[ |
| 348 | bool, |
| 349 | typer.Option( |
| 350 | "--install", |
| 351 | help=( |
| 352 | "After checking, run 'pip install -U google-colab-cli' to " |
| 353 | "upgrade the CLI in place. No-op if already up to date. " |
| 354 | "Linux only." |
| 355 | ), |
| 356 | ), |
| 357 | ] = False, |
| 358 | ): |
| 359 | """Check for latest version and print if an update is available""" |
| 360 | auto_update.check_for_updates(quiet=False) |
| 361 | if not install: |
| 362 | return |
| 363 | |
| 364 | if not auto_update.is_self_install_supported(): |
| 365 | typer.echo( |
| 366 | "[colab] '--install' self-install is only supported on Linux and macOS.", |
| 367 | err=True, |
| 368 | ) |
| 369 | raise typer.Exit(code=1) |
| 370 | |
| 371 | # Skip the install when the current version already matches (or exceeds) |
| 372 | # the latest known version, to avoid an unnecessary subprocess call. |
| 373 | settings = state.settings_store.load() |
| 374 | if settings.latest_version and not auto_update._is_newer( |
| 375 | settings.latest_version, auto_update.get_app_version() |
| 376 | ): |
| 377 | return |
| 378 | |
| 379 | auto_update.self_install() |
| 380 | |
| 381 | |
| 382 | def _print_resource(filename: str) -> None: |
| 383 | import importlib.resources |
| 384 | import os |
| 385 | |
| 386 | content = None |
| 387 | try: |
| 388 | # Try reading from package resources |
| 389 | ref = importlib.resources.files("colab_cli").joinpath(filename) |
| 390 | if ref.is_file(): |
| 391 | content = ref.read_text(encoding="utf-8") |
| 392 | except Exception: |
| 393 | pass |
| 394 | |
| 395 | if not content: |
| 396 | # Fallback to local file for development |
| 397 | local_path = os.path.abspath( |
| 398 | os.path.join(os.path.dirname(__file__), f"../../../{filename}") |
| 399 | ) |
| 400 | if os.path.exists(local_path): |
| 401 | try: |
| 402 | with open(local_path, "r", encoding="utf-8") as f: |
| 403 | content = f.read() |
| 404 | except Exception: |
| 405 | pass |
| 406 | |
| 407 | if content: |
| 408 | typer.echo(content) |
| 409 | else: |
| 410 | typer.echo(f"[colab] {filename} content not available.", err=True) |
| 411 | raise typer.Exit(code=1) |
| 412 | |
| 413 | |
| 414 | def readme(): |
| 415 | """Print the bundled README.md file""" |
| 416 | _print_resource("README.md") |
| 417 | |
| 418 | |
| 419 | def skill(): |
| 420 | """Print the bundled SKILL.md file""" |
| 421 | _print_resource("SKILL.md") |
| 422 | |
| 423 | |
| 424 | def register(app: typer.Typer): |
| 425 | app.command()(pay) |
| 426 | app.command()(log) |
| 427 | app.command(name="url")(url) |
| 428 | app.command(name="version")(version_command) |
| 429 | app.command(name="update")(update_command) |
| 430 | # Developer-only debugging aid; hidden from `colab --help` but still |
| 431 | # reachable via `colab whoami` / `colab whoami --help`. |
| 432 | app.command(name="whoami", hidden=True)(whoami) |
| 433 | app.command(name="readme")(readme) |
| 434 | app.command(name="README", hidden=True)(readme) |
| 435 | app.command(name="skill")(skill) |
| 436 | app.command(name="SKILL", hidden=True)(skill) |