| 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 datetime |
| 16 | import os |
| 17 | import sys |
| 18 | import json |
| 19 | from typing import Optional, List |
| 20 | import typer |
| 21 | from rich.console import Console |
| 22 | from typing_extensions import Annotated |
| 23 | |
| 24 | from colab_cli.runtime import ColabRuntime |
| 25 | from colab_cli.contents import ContentsClient |
| 26 | from colab_cli.auth import get_credentials |
| 27 | from colab_cli.utils import get_status_code, render_display_data |
| 28 | |
| 29 | _console = Console() |
| 30 | |
| 31 | |
| 32 | # Default execute() timeout for human-in-the-loop automations (auth / |
| 33 | # drivemount). The kernel goes silent while the user completes a browser |
| 34 | # OAuth flow, which can routinely take 30s+; the upstream 10s default |
| 35 | # raises ``TimeoutError`` mid-flow even though the mount actually succeeds. |
| 36 | # 10 minutes is long enough for any realistic interactive auth ceremony |
| 37 | # without leaving CI hangs unbounded. |
| 38 | INTERACTIVE_AUTOMATION_TIMEOUT_SEC = 600 |
| 39 | |
| 40 | |
| 41 | |
| 42 | def run_automation( |
| 43 | name: str, |
| 44 | op: str, |
| 45 | code: str, |
| 46 | allow_stdin: bool = False, |
| 47 | path: str = None, |
| 48 | timeout: Optional[float] = None, |
| 49 | ): |
| 50 | from colab_cli.common import state |
| 51 | |
| 52 | s = state.store.get(name) |
| 53 | runtime = ColabRuntime(s.url, s.token, session_name=s.name, history=state.history) |
| 54 | |
| 55 | def drivefs_hook(deserialize_msg, wsclient): |
| 56 | content = deserialize_msg.get("content", {}) |
| 57 | if content.get("request", {}).get("authType") == "dfs_ephemeral": |
| 58 | msg_id = deserialize_msg.get("metadata", {}).get("colab_msg_id") |
| 59 | state.history.log_event( |
| 60 | s.name, |
| 61 | "colab_request", |
| 62 | {"type": "dfs_ephemeral", "colab_msg_id": msg_id}, |
| 63 | ) |
| 64 | url = f"{state.client.colab_domain}/tun/m/credentials-propagation/{s.endpoint}" |
| 65 | params = { |
| 66 | "authuser": "0", |
| 67 | "authtype": "dfs_ephemeral", |
| 68 | "version": "2", |
| 69 | "dryrun": "true", |
| 70 | "propagate": "true", |
| 71 | "record": "false", |
| 72 | } |
| 73 | typer.echo( |
| 74 | f"\n[colab] Intercepted Drive Auth Request. Connecting to {state.client.colab_domain}..." |
| 75 | ) |
| 76 | |
| 77 | creds = get_credentials( |
| 78 | state.client_oauth_config, provider=state.auth_provider |
| 79 | ) |
| 80 | resp = creds.request("GET", url, params=params) |
| 81 | token = ( |
| 82 | json.loads(resp.text.split("\n", 1)[-1]).get("token") |
| 83 | if get_status_code(resp) == 200 |
| 84 | else None |
| 85 | ) |
| 86 | |
| 87 | headers = {"x-goog-colab-token": token} |
| 88 | resp = creds.request( |
| 89 | "POST", |
| 90 | url, |
| 91 | params=params, |
| 92 | headers=headers, |
| 93 | files={"file_id": (None, "empty.ipynb")}, |
| 94 | ) |
| 95 | data = json.loads(resp.text.split("\n", 1)[-1]) |
| 96 | |
| 97 | if not data.get("success"): |
| 98 | uri = data.get("unauthorized_redirect_uri") |
| 99 | typer.echo( |
| 100 | f"\n[colab] REQUIRED: Google Drive Authorization needed.\nPlease visit:\n\n{uri}\n" |
| 101 | ) |
| 102 | state.history.log_event(s.name, "drive_auth_needed", {"uri": uri}) |
| 103 | sys.stdout.write("Press Enter after you have granted access... ") |
| 104 | sys.stdout.flush() |
| 105 | with open("/dev/tty") as tty: |
| 106 | tty.readline() |
| 107 | |
| 108 | typer.echo("[colab] Authorizing VM...") |
| 109 | params["dryrun"] = "false" |
| 110 | resp = creds.request( |
| 111 | "POST", |
| 112 | url, |
| 113 | params=params, |
| 114 | headers=headers, |
| 115 | files={"file_id": (None, "empty.ipynb")}, |
| 116 | ) |
| 117 | if get_status_code(resp) == 200: |
| 118 | typer.echo("[colab] Credentials propagated. Resuming mount...") |
| 119 | state.history.log_event(s.name, "drive_auth_success", {}) |
| 120 | reply = wsclient.session.msg( |
| 121 | "input_reply", |
| 122 | {"value": {"type": "colab_reply", "colab_msg_id": msg_id}}, |
| 123 | ) |
| 124 | if "header" in deserialize_msg: |
| 125 | reply["parent_header"] = deserialize_msg["header"] |
| 126 | wsclient.stdin_channel.send(reply) |
| 127 | else: |
| 128 | typer.echo( |
| 129 | f"[colab] Error propagating: {get_status_code(resp)} {resp.text}" |
| 130 | ) |
| 131 | return True |
| 132 | return False |
| 133 | |
| 134 | runtime.colab_request_hook = drivefs_hook |
| 135 | try: |
| 136 | s.running = f"automation({op})" |
| 137 | s.last_execution = ( |
| 138 | f"automation:{op}", |
| 139 | None, |
| 140 | datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), |
| 141 | ) |
| 142 | state.store.add(s) |
| 143 | |
| 144 | if op == "drivemount": |
| 145 | state.history.log_event( |
| 146 | name, "automation", {"op": "drivemount", "path": path, "code": code} |
| 147 | ) |
| 148 | else: |
| 149 | state.history.log_event(name, "automation", {"op": op, "code": code}) |
| 150 | |
| 151 | outputs = runtime.execute_code(code, allow_stdin=allow_stdin, timeout=timeout) |
| 152 | state.history.log_event( |
| 153 | name, "automation_result", {"op": op, "outputs": outputs} |
| 154 | ) |
| 155 | |
| 156 | for out in outputs: |
| 157 | if "text" in out: |
| 158 | sys.stdout.write(out["text"]) |
| 159 | elif "data" in out: |
| 160 | text = render_display_data(out["data"]) |
| 161 | if text is not None: |
| 162 | _console.print(text) |
| 163 | elif out.get("output_type") == "error": |
| 164 | ename = out.get("ename", "Error") |
| 165 | evalue = out.get("evalue", "") |
| 166 | tb = out.get("traceback", []) |
| 167 | if tb: |
| 168 | sys.stderr.write("".join(tb) + "\n") |
| 169 | else: |
| 170 | sys.stderr.write(f"{ename}: {evalue}\n") |
| 171 | finally: |
| 172 | s.running = None |
| 173 | state.store.add(s) |
| 174 | runtime.stop() |
| 175 | |
| 176 | |
| 177 | def auth( |
| 178 | session: Annotated[ |
| 179 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 180 | ] = None, |
| 181 | ): |
| 182 | """Authenticate with Google on the VM""" |
| 183 | from colab_cli.common import state |
| 184 | |
| 185 | name = state.resolve_session(session) |
| 186 | code = "import os\nos.environ['USE_AUTH_EPHEM'] = '0'\nfrom google.colab import auth\nauth.authenticate_user()" |
| 187 | typer.echo(f"[colab] Starting Google Auth flow on {name}...") |
| 188 | run_automation( |
| 189 | name, |
| 190 | "auth", |
| 191 | code, |
| 192 | allow_stdin=True, |
| 193 | timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC, |
| 194 | ) |
| 195 | |
| 196 | |
| 197 | def drivemount( |
| 198 | session: Annotated[ |
| 199 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 200 | ] = None, |
| 201 | path: Annotated[str, typer.Argument(help="Mount path")] = "/content/drive", |
| 202 | ): |
| 203 | """Mount Google Drive at path""" |
| 204 | from colab_cli.common import state |
| 205 | |
| 206 | name = state.resolve_session(session) |
| 207 | code = f"from google.colab import drive\ndrive.mount('{path}')" |
| 208 | typer.echo(f"[colab] Mounting Google Drive to '{path}' on {name}...") |
| 209 | run_automation( |
| 210 | name, |
| 211 | "drivemount", |
| 212 | code, |
| 213 | allow_stdin=True, |
| 214 | path=path, |
| 215 | timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC, |
| 216 | ) |
| 217 | |
| 218 | |
| 219 | def install( |
| 220 | session: Annotated[ |
| 221 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 222 | ] = None, |
| 223 | packages: Annotated[ |
| 224 | Optional[List[str]], typer.Argument(help="Packages to install") |
| 225 | ] = None, |
| 226 | requirement: Annotated[ |
| 227 | Optional[str], typer.Option("-r", "--requirement", help="Requirements file") |
| 228 | ] = None, |
| 229 | ): |
| 230 | """Install python packages on the VM""" |
| 231 | from colab_cli.common import state |
| 232 | |
| 233 | name = state.resolve_session(session) |
| 234 | if not packages and not requirement: |
| 235 | typer.echo("[colab] No packages or requirements specified.") |
| 236 | raise typer.Exit(1) |
| 237 | |
| 238 | commands = [] |
| 239 | if requirement: |
| 240 | if not os.path.isfile(requirement): |
| 241 | typer.echo(f"[colab] Requirements file '{requirement}' not found locally.") |
| 242 | raise typer.Exit(1) |
| 243 | contents = ContentsClient(state.store.get(name)) |
| 244 | remote_path = f"content/{os.path.basename(requirement)}" |
| 245 | contents.upload(requirement, remote_path) |
| 246 | commands.extend(["-r", f"/{remote_path}"]) |
| 247 | if packages: |
| 248 | commands.extend(packages) |
| 249 | |
| 250 | cmd_str = ", ".join(f"'{c}'" for c in commands) |
| 251 | code = f""" |
| 252 | import subprocess, sys |
| 253 | def install(): |
| 254 | packages = [{cmd_str}] |
| 255 | try: |
| 256 | subprocess.check_call(['uv', 'pip', 'install', '--system'] + packages) |
| 257 | print('Installation Complete (via uv)!') |
| 258 | except: |
| 259 | subprocess.check_call([sys.executable, '-m', 'pip', 'install'] + packages) |
| 260 | print('Installation Complete (via pip)!') |
| 261 | install() |
| 262 | """ |
| 263 | typer.echo(f"[colab] Installing packages on {name} (preferring uv)...") |
| 264 | run_automation(name, "install", code) |
| 265 | |
| 266 | |
| 267 | def register(app: typer.Typer): |
| 268 | app.command(hidden=True)(auth) |
| 269 | app.command()(drivemount) |
| 270 | app.command()(install) |