From c19ce5075683c680e77bf12052e886af18e3d8d9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 7 Oct 2026 17:44:58 +0000 Subject: [PATCH] fix: Windows consoles and Drive mount consent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Interactive colab console on cmd.exe and PowerShell no longer depends on termios. It enables ANSI processing, switches the console to raw input, and translates Windows scan codes to the sequences a POSIX tty would send. colab ssh --proxy-mode reads stdin with a blocking os.read and quotes its ProxyCommand for cmd.exe. colab drivemount passes timeout_ms=600000 so the kernel does not abandon the browser consent step at 120s, runs the handshake off the websocket thread, opens the consent URL, and explains the multi-account HTTP 400 instead of dumping Google's HTML error page. Co-authored-by: János Rózsavölgyi --- AGENTS.md | 2 +- README.md | 7 +- docs/02_execution_and_interactive.md | 6 +- docs/04_automation_and_utility.md | 35 +++- docs/06_ssh_access.md | 2 + skills/colab-operator/SKILL.md | 2 +- src/colab_cli/cli.py | 17 +- src/colab_cli/commands/automation.py | 201 +++++++++--------- src/colab_cli/commands/ssh.py | 75 ++++--- src/colab_cli/console.py | 119 +++++++---- src/colab_cli/drive_auth.py | 303 +++++++++++++++++++++++++++ src/colab_cli/terminal.py | 242 +++++++++++++++++++++ tests/test_automation.py | 55 ++++- tests/test_console.py | 79 ++++++- tests/test_drive_auth.py | 271 ++++++++++++++++++++++++ tests/test_ssh.py | 56 +++-- tests/test_terminal.py | 146 +++++++++++++ 17 files changed, 1422 insertions(+), 196 deletions(-) create mode 100644 src/colab_cli/drive_auth.py create mode 100644 src/colab_cli/terminal.py create mode 100644 tests/test_drive_auth.py create mode 100644 tests/test_terminal.py diff --git a/AGENTS.md b/AGENTS.md index 6e84b75b..04777814 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,7 +87,7 @@ As an AI agent operating via non-interactive shell tools (`run_shell_command`), - Specially crafted mock scripts that simulate timeouts or API calls. - **I CANNOT Run (Requires User Assistance):** - **`colab auth`**: This command relies on the traditional Gcloud fallback `input_request` (via `USE_AUTH_EPHEM='0'`), which prompts the user via Python's `input()` to click a URL, sign in, and paste back an authorization code. My shell tool will hang indefinitely on this `input()`. - - **`colab drivemount`**: This command prompts the user via `sys.stdin.readline()` (specifically querying `/dev/tty` to ensure input is captured) to press `Enter` after granting OAuth consent in the browser. My shell tool will timeout/hang waiting for `Enter`. + - **`colab drivemount`**: This command prompts the user to press `Enter` after granting OAuth consent in the browser. The prompt reads the controlling terminal (`/dev/tty` on POSIX, `CONIN$` on Windows), falling back to stdin. My shell tool will timeout/hang waiting for `Enter`. The consent URL is also opened in a GUI browser when one is available. - **Interactive (TTY) `colab repl` / `colab console`**: When stdin is a real terminal these commands drop into interactive raw-TTY modes that require real-time keystroke streaming. My shell tools cannot support this. (Piped stdin is fine — see above.) Whenever working on interactive commands, I must build the core logic, write mock tests, and explicitly ask the user to run the live test in their terminal to verify success. diff --git a/README.md b/README.md index 1cbfedd6..c9ffe973 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ Designed to support seamless developer productivity, headless automation, and AI [Demo](https://github.com/user-attachments/assets/656226a9-af13-4fdb-8eda-d7de747336a2) > [!NOTE] -> **Platform support:** the Colab CLI currently supports **Linux and macOS** only. Windows is not supported at this time. +> **Platform support:** Linux, macOS, and Windows (PowerShell and cmd.exe). Interactive `colab console` uses the Windows console API on Windows (raw keystrokes and ANSI color). `colab ssh --proxy-mode` reads stdin with a blocking read so it works as an OpenSSH ProxyCommand under Windows OpenSSH. > [!TIP] > Looking for in-notebook, interactive agent-assisted coding instead of a terminal workflow? See the [Colab MCP Server](https://github.com/googlecolab/colab-mcp). @@ -97,7 +97,7 @@ Run `colab --help` to view specific options, defaults, and detailed he | Command | Description | | --- | --- | | `colab auth [-s NAME]` | Authenticate the VM for GCP services (BigQuery, GCS, etc.) | -| `colab drivemount [-s NAME] [PATH]` | Mount Google Drive on the VM (default: `/content/drive`) | +| `colab drivemount [-s NAME] [--authuser N] [PATH]` | Mount Google Drive on the VM (default: `/content/drive`) | | `colab install [-s NAME] [-r FILE \| PKG...]` | Install packages on the VM using `uv` (falls back to `pip`) | | `colab log [-s NAME] [-n N] [-o FILE]` | View or export session history (`.ipynb`, `.md`, `.txt`, `.jsonl`) | | `colab usage` | Show account compute-unit usage rate and balance | @@ -144,7 +144,8 @@ colab stop -s analysis ## Usage Notes * **Machine shape:** Use `--high-mem` with `colab new`, `colab run`, or `colab ssh` (when auto-creating a runtime) to request a high-RAM machine shape. Requires Colab Pro or Pro+ entitlement for supported accelerators (CPU, T4, A100, etc.). L4 and TPU runtimes ignore this flag because they only offer one shape. Machine shape is shown in `colab sessions` and `colab status`. -* **TTY Requirements:** The interactive commands `repl` and `console` require a local TTY. When running inside automated scripts or pipelines, make sure to pipe stdin (e.g., `echo "print(1)" | colab repl`) to trigger non-interactive execution modes. +* **TTY Requirements:** The interactive commands `repl` and `console` require a local TTY. When running inside automated scripts or pipelines, make sure to pipe stdin (e.g., `echo "print(1)" | colab repl`) to trigger non-interactive execution modes. On Windows, `colab console` switches the console out of line-editing mode so arrow keys and Ctrl-C reach the remote shell; colors need virtual-terminal processing, which the CLI turns on for cmd.exe and Windows PowerShell. +* **Drive mount and multiple Google accounts:** `colab drivemount` prints a consent URL (and opens it in the browser when it can) and waits up to 10 minutes. On the account chooser, select the same account the CLI authenticated as (`colab whoami` / the `login_hint` in the URL). A browser that is signed into several accounts otherwise attributes the Drive session to the active account, and Colab rejects the mount with HTTP 400. `--authuser N` appends that index to the consent URL when the chooser keeps opening the wrong account. * **Transparent Code Execution:** When calling `colab exec -f file.py`, the CLI reads the file locally and transmits its content to the remote kernel. You do not need to manually upload files before execution. * **Storage & State Paths:** Session tokens and metadata are stored at `~/.config/colab-cli/sessions.json`. Global CLI settings are located at `~/.config/colab-cli/settings.json`. These can be customized or isolated via the global `--config` flag. diff --git a/docs/02_execution_and_interactive.md b/docs/02_execution_and_interactive.md index ed1d782a..29c1f963 100644 --- a/docs/02_execution_and_interactive.md +++ b/docs/02_execution_and_interactive.md @@ -1,5 +1,7 @@ --- log: +2026-10-07: Windows console support for `colab console`. cmd.exe and PowerShell have no `termios`, `/dev/tty`, or `SIGWINCH`, so interactive console used to stay line-buffered (arrow keys never left the local line editor, Ctrl-C killed the CLI, ANSI from the remote PTY rendered as garbage in conhost). `console.py` now enables virtual-terminal processing and UTF-8, switches the input handle to raw mode via `SetConsoleMode`, reads keys with `msvcrt.getwch`, translates scan codes to ANSI, and polls the window size. POSIX raw mode and piped stdin are unchanged. `colab ssh --proxy-mode` reads stdin with a blocking `os.read` (Windows `select()` cannot wait on a pipe) and quotes the ProxyCommand for `cmd.exe`. + 2026-05-07: Fixed `colab console` piped-stdin handling. Previously a piped invocation (e.g. `echo 'cmd' | colab console -s s`) sent the command and then hung indefinitely because the previous EOF handler emitted a bare `\x04` (Ctrl-D), which the remote `tmux`-wrapped bash treats as a literal character rather than a session terminator. The new handler sends `exit\n` (which bash actually exits on) and then closes the websocket from the client side after a short grace period (`PIPED_EOF_GRACE_SECONDS = 0.5s`) so any tail output (bash `logout`, tmux `[exited]`) makes it back to the user. TTY mode is unchanged: real-terminal EOF is left to the remote shell. Verified live: `echo 'echo HELLO' | colab console -s s` now exits in ~1.2s instead of hanging. 2026-05-07: Fixed `print_kitty` (used by `colab exec --output-image` and any image-producing exec) to no-op when `sys.stdout.isatty()` is false. The Kitty Graphics Protocol escape sequence is meaningless when stdout is a file or pipe and was visually corrupting captured output (a multi-KB base64 PNG blob would land in log files, grep targets, or showboat captures). Image bytes are still saved to disk via `handle_image`'s file-write path; only the inline-render attempt is suppressed. @@ -33,7 +35,7 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel ### 3. Console (`colab console`) - **Implementation**: Connects directly to the backend terminal endpoint (`/colab/tty`) via WebSockets using `websocket-client`. - **Interactive**: Bypasses the Jupyter kernel entirely to provide a raw, PTY-backed bash session on the Colab VM. -- **Terminal Management**: Configures `sys.stdin` to raw mode using `termios` and `tty`, passing single characters to the socket and writing raw ANSI escape sequences directly to `sys.stdout.buffer`. Hooks into `SIGWINCH` to communicate local terminal dimensions (`cols`/`rows`) to the remote bash environment so output rendering works perfectly during resizing. +- **Terminal Management**: On Linux and macOS, configures `sys.stdin` to raw mode using `termios` and `tty`, passing single characters to the socket and writing raw ANSI escape sequences directly to `sys.stdout.buffer`. Hooks into `SIGWINCH` to communicate local terminal dimensions (`cols`/`rows`) to the remote bash environment so output rendering works perfectly during resizing. On Windows (cmd.exe and PowerShell), `termios` does not exist: the CLI enables `ENABLE_VIRTUAL_TERMINAL_PROCESSING` and UTF-8 on the console, sets raw input mode with `SetConsoleMode` (so the local line editor and Ctrl-C handler get out of the way), reads keys with `msvcrt.getwch`, maps scan codes to ANSI sequences, and polls `os.get_terminal_size` because there is no `SIGWINCH`. - **Piped stdin**: Detected via `sys.stdin.isatty()`. When piped, the input characters are forwarded one at a time to the remote pty, and on EOF the client sends `exit\n` and then closes the websocket itself after `PIPED_EOF_GRACE_SECONDS` (0.5s) so the user's shell goodbye text drains back. The remote `/colab/tty` endpoint wraps bash in tmux, which intercepts a bare `\x04` as a literal character — that is why we send `exit\n` rather than Ctrl-D. ## Implementation Details @@ -54,4 +56,6 @@ TDD is mandatory for all execution features. - **Test Case**: Verify large piped inputs are handled without buffer overflow or truncation. - **Test Case**: `colab console` with piped stdin sends `exit\n` and calls `ws.close()` on EOF (regression: previously sent `\x04` only and hung). - **Test Case**: `colab console` in TTY mode does not synthesize an exit on EOF (the user owns the session lifecycle). +- **Test Case**: With `termios` unavailable and `os.name == "nt"`, `connect_console` enters the Windows raw-console context manager and does not call `tty.setraw`. +- **Test Case**: Windows key translation maps scan codes to ANSI and the stdin thread forwards those characters one at a time. - **Test Case**: `print_kitty` is a no-op when `sys.stdout.isatty()` is false (regression: previously emitted ANSI/base64 into pipes and files). diff --git a/docs/04_automation_and_utility.md b/docs/04_automation_and_utility.md index 1a0e2287..009204df 100644 --- a/docs/04_automation_and_utility.md +++ b/docs/04_automation_and_utility.md @@ -1,5 +1,7 @@ --- log: +2026-10-07: Drive mount fixes for upstream issues #113, #103, and #108. `drive.mount` is now called with `timeout_ms=600000` so the kernel's `blocking_request` does not give up at the 120s default while the user is still in the browser (that produced `ValueError: mount failed` with an empty reason). The propagation handshake runs on a daemon thread so it does not block the websocket recv loop. The consent URL is opened in the system browser when one is available (`os.startfile` on Windows, `open` on macOS, `xdg-open` elsewhere) and the prompt names the `login_hint` account. HTTP 400 is explained as a multi-account browser mismatch instead of dumping Google's HTML error page, and an error `colab_reply` unblocks the kernel immediately. Optional `--authuser N` rewrites the consent URL only when the user asks; the propagation request's `authuser` parameter was shown in #103 not to fix the 400 by itself, and the URL is not forced to `authuser=0`. The Enter prompt reads `/dev/tty` on POSIX and `CONIN$` on Windows. + 2026-09-25: Cleaned up references to keep-alive background daemon and pre-flight pings. VM liveness is automatically maintained by the Colab backend based on active kernel activity. 2026-08-10: Added `colab usage` for account-level compute-unit rate/balance via `GET /tun/m/ccu-info` on the session backend (same bearer token as `colab new`). 2026-06-11: Replaced the `oauth2` provider's `run_local_server()` (localhost redirect) with a remote copy-paste flow (`_run_remote_flow` in `auth.py`). The CLI now prints an authorization URL built with `redirect_uri=https://sdk.cloud.google.com/applicationdefaultauthcode.html` and `token_usage=remote`, then reads the pasted authorization code via `input()` and exchanges it with `flow.fetch_token(code=...)`. This is the same flow `gcloud auth application-default login` uses and works identically in local and remote/headless/container environments, removing the heuristic of whether to auto-open a browser. Confirmed server-side acceptance with a live GET-only check against the bundled cloud-SDK client (`764086051850-...`); the OOB redirect and a non-bundled client id were both verified to be rejected (`OOB flow has been blocked` / `redirect_uri_mismatch`). Unit tests in `tests/test_auth.py` assert no localhost server is started, the redirect URI + `token_usage=remote` are set, and the pasted code is exchanged. @@ -116,7 +118,7 @@ How each provider supplies the scope: - **Action**: Execute `drive.mount()` and transparently proxy Colab's proprietary credential propagation flow. - **Code**: `python from google.colab import drive - drive.mount('/content/drive')` + drive.mount('/content/drive', timeout_ms=600000)` - **Handling**: Because `drivefs` enforces the ephemeral side-channel propagation (`colab_request` over websocket), the CLI intercepts these messages using `ColabRuntime.colab_request_hook`. When intercepted, the CLI @@ -128,7 +130,27 @@ How each provider supplies the scope: user is OAuthing in their browser. To avoid the upstream 10s `jupyter_kernel_client` default raising `TimeoutError` mid-flow, this subcommand passes `timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC` (600s) to - `ColabRuntime.execute_code`. Same applies to `colab auth`. + `ColabRuntime.execute_code`. Same applies to `colab auth`. Separately, + `google.colab.drive.mount` defaults `timeout_ms` to 120000 and that budget + is the `blocking_request` wait, which includes the browser step. The + injected call passes `timeout_ms=600000` so the kernel does not abandon + the mount at two minutes and then fail with `ValueError: mount failed`. +- **Where the prompt runs**: The `colab_request` hook returns immediately + and finishes propagation on a daemon thread. The websocket recv loop stays + free to deliver the kernel's later output. +- **Consent UX**: The CLI prints the `unauthorized_redirect_uri`, highlights + the `login_hint` account, and tries to open that URL in the GUI browser. + The user presses Enter on the controlling terminal (`/dev/tty` or, on + Windows, `CONIN$`). +- **Multi-account 400**: Propagation HTTP 400 means the browser's active + Google account was not the CLI account. The CLI says so and replies to + the kernel with an error `colab_reply` instead of dumping the HTML body + and leaving `drive.mount` to time out. `--authuser N` is an opt-in that + appends `authuser` to the consent URL; it is not applied by default. +- **What still depends on Google**: The CLI cannot make the browser + attribute the gsession to `login_hint` when a different account is active. + The user has to pick the matching account on the chooser. There is no + fully non-interactive Drive mount; the ephemeral consent page is required. ### 4. Logging and Notebook Capture (`colab log`) @@ -291,8 +313,13 @@ TDD is mandatory for all automation features. - **Test Case**: Verify `colab install` correctly injects `pip install` or `uv install` commands to the remote VM kernel. - **Test Case**: Verify `colab drivemount` correctly injects `drive.mount()` - commands and registers the `colab_request_hook` to intercept credential - propagation events. + commands with `timeout_ms` of at least 10 minutes and registers the + `colab_request_hook` to intercept credential propagation events. +- **Test Case**: A propagation HTTP 400 produces an account-mismatch message + (no raw HTML dump) and a `colab_reply` that carries `error`, so the kernel + does not sit until `timeout_ms`. +- **Test Case**: `--authuser` rewrites the consent URL; the default path + leaves the server-built URL unchanged. ### 2. History Capture diff --git a/docs/06_ssh_access.md b/docs/06_ssh_access.md index 312ada8a..0d085eb8 100644 --- a/docs/06_ssh_access.md +++ b/docs/06_ssh_access.md @@ -1,5 +1,7 @@ --- log: +2026-10-07: Windows ProxyCommand. `select()` on the stdin fd raises on Windows (only sockets are selectable), which killed the `--proxy-mode` input pump. The pump now blocks in `os.read`. The ProxyCommand string is quoted with `subprocess.list2cmdline` on Windows because OpenSSH there re-parses it with cmd.exe, and `UserKnownHostsFile` uses `os.devnull` (`nul`) instead of `/dev/null`. `--rm` signal installation skips `SIGHUP` when the interpreter does not define it. + 2026-09-25: Cleaned up keep-alive daemon references in session creation and teardown documentation. 2026-08-09: Added `--high-mem` passthrough when `colab ssh` auto-creates a runtime (forwards to `colab new --high-mem`). 2026-07-17: Initial design and implementation of `colab ssh` — client side of SSH-over-WebSocket runtime access. Adds three modes (interactive shell, `-s SESSION`, and `--proxy-mode` OpenSSH ProxyCommand bridge), `--identity/-i` key selection, and per-HTTP-status handshake error messages. Server side is out of scope for this repo; the subcommand is a no-op against runtimes that do not expose the `/colab/ssh` endpoint (surfaces an actionable HTTP 404 message). diff --git a/skills/colab-operator/SKILL.md b/skills/colab-operator/SKILL.md index b5f2a687..9ac30ed6 100644 --- a/skills/colab-operator/SKILL.md +++ b/skills/colab-operator/SKILL.md @@ -65,7 +65,7 @@ by running `uv tool install google-colab-cli` or `pip install google-colab-cli`. ### Automate - `colab auth -s ` — VM-side GCP creds, needed before in-VM GCS/BigQuery calls (interactive; not agent-runnable). -- `colab drivemount -s [PATH]` — mounts Drive at `/content/drive` by default (interactive; not agent-runnable). +- `colab drivemount -s [--authuser N] [PATH]` — mounts Drive at `/content/drive` by default (interactive; not agent-runnable). The user must approve the printed consent URL in a browser as the same account `colab whoami` shows; a multi-account browser rejects the mount with HTTP 400 unless that account is the one selected on the chooser. The kernel waits 10 minutes, not 2. - `colab install -s pkg1 pkg2` — installs via `uv pip install --system`, falling back to `pip`. Also `colab install -s -r requirements.txt`. ### Inspect & report diff --git a/src/colab_cli/cli.py b/src/colab_cli/cli.py index 5f837eaf..77b4c8f7 100644 --- a/src/colab_cli/cli.py +++ b/src/colab_cli/cli.py @@ -23,7 +23,16 @@ from colab_cli import auto_update from colab_cli.auth import AuthProvider from colab_cli.common import state, setup_logging -from colab_cli.commands import session, execution, files, automation, run, ssh, utility, usage +from colab_cli.commands import ( + session, + execution, + files, + automation, + run, + ssh, + utility, + usage, +) class AlphabeticalGroup(TyperGroup): @@ -80,6 +89,12 @@ def callback( """ Colab CLI global configuration. """ + # ANSI colors and cursor control in cmd.exe / Windows PowerShell 5.1. + # No-op on Linux, macOS, and consoles that already have virtual-terminal + # processing enabled. + from colab_cli.terminal import enable_windows_virtual_terminal + + enable_windows_virtual_terminal() state.client_oauth_config = client_oauth_config state.config_path = config state.logtostderr = logtostderr diff --git a/src/colab_cli/commands/automation.py b/src/colab_cli/commands/automation.py index d2600091..ff0797b6 100644 --- a/src/colab_cli/commands/automation.py +++ b/src/colab_cli/commands/automation.py @@ -15,16 +15,21 @@ import datetime import os import sys -import json -from typing import Optional, List +import threading +from typing import Callable, Optional, List import typer from rich.console import Console from typing_extensions import Annotated from colab_cli.runtime import ColabRuntime from colab_cli.contents import ContentsClient -from colab_cli.auth import get_credentials -from colab_cli.utils import get_status_code, render_display_data +from colab_cli.drive_auth import ( + DRIVE_AUTH_TIMEOUT_SEC, + mount_code, + perform_drive_authorization, + send_colab_reply, +) +from colab_cli.utils import render_display_data _console = Console() @@ -33,11 +38,76 @@ # drivemount). The kernel goes silent while the user completes a browser # OAuth flow, which can routinely take 30s+; the upstream 10s default # raises ``TimeoutError`` mid-flow even though the mount actually succeeds. -# 10 minutes is long enough for any realistic interactive auth ceremony -# without leaving CI hangs unbounded. -INTERACTIVE_AUTOMATION_TIMEOUT_SEC = 600 +# 10 minutes matches ``drive.mount(timeout_ms=...)`` so the kernel-side +# ``blocking_request`` does not give up at 120s while the user is still in +# the browser (upstream issue #113). +INTERACTIVE_AUTOMATION_TIMEOUT_SEC = DRIVE_AUTH_TIMEOUT_SEC +def make_drivefs_hook( + state, + session_state, + *, + authuser: str = "0", + rewrite_auth_url: bool = False, + runner: Optional[Callable[[Callable[[], None]], None]] = None, +): + """Build the ``colab_request`` hook that answers ``dfs_ephemeral`` mounts. + + The handshake (HTTP propagation, browser prompt, stdin reply) runs off the + websocket thread. Doing it inside ``on_message`` blocks the recv loop for + the whole consent wait, so the kernel's execute reply cannot be delivered + and a late Enter looks like a hang. ``runner`` exists so tests can run the + worker inline; production starts a daemon thread. + """ + + def drivefs_hook(deserialize_msg, wsclient): + content = deserialize_msg.get("content") or {} + request = content.get("request") or {} + if request.get("authType") != "dfs_ephemeral": + return False + msg_id = (deserialize_msg.get("metadata") or {}).get("colab_msg_id") + state.history.log_event( + session_state.name, + "colab_request", + {"type": "dfs_ephemeral", "colab_msg_id": msg_id}, + ) + + def work(): + try: + from colab_cli.auth import get_credentials + + http = get_credentials( + state.client_oauth_config, provider=state.auth_provider + ) + perform_drive_authorization( + colab_domain=state.client.colab_domain, + endpoint=session_state.endpoint, + http=http, + wsclient=wsclient, + deserialize_msg=deserialize_msg, + msg_id=msg_id, + authuser=authuser, + rewrite_auth_url=rewrite_auth_url, + on_event=lambda ev, payload: state.history.log_event( + session_state.name, ev, payload + ), + ) + except Exception as exc: + typer.echo(f"[colab] Drive authorization failed: {exc}", err=True) + try: + send_colab_reply(wsclient, deserialize_msg, msg_id, error=str(exc)) + except Exception: + pass + + if runner is not None: + runner(work) + else: + threading.Thread(target=work, daemon=True, name="colab-drive-auth").start() + return True + + return drivefs_hook + def run_automation( name: str, @@ -46,102 +116,19 @@ def run_automation( allow_stdin: bool = False, path: str = None, timeout: Optional[float] = None, + authuser: str = "0", + rewrite_auth_url: bool = False, ): from colab_cli.common import state s = state.get_session(name) runtime = ColabRuntime(s.url, s.token, session_name=s.name, history=state.history) - - def drivefs_hook(deserialize_msg, wsclient): - content = deserialize_msg.get("content", {}) - if content.get("request", {}).get("authType") == "dfs_ephemeral": - msg_id = deserialize_msg.get("metadata", {}).get("colab_msg_id") - state.history.log_event( - s.name, - "colab_request", - {"type": "dfs_ephemeral", "colab_msg_id": msg_id}, - ) - url = f"{state.client.colab_domain}/tun/m/credentials-propagation/{s.endpoint}" - params = { - "authuser": "0", - "authtype": "dfs_ephemeral", - "version": "2", - "dryrun": "true", - "propagate": "true", - "record": "false", - } - typer.echo( - f"\n[colab] Intercepted Drive Auth Request. Connecting to {state.client.colab_domain}..." - ) - - creds = get_credentials( - state.client_oauth_config, provider=state.auth_provider - ) - resp = creds.request("GET", url, params=params) - token = ( - json.loads(resp.text.split("\n", 1)[-1]).get("token") - if get_status_code(resp) == 200 - else None - ) - - headers = {"x-goog-colab-token": token} - resp = creds.request( - "POST", - url, - params=params, - headers=headers, - files={"file_id": (None, "empty.ipynb")}, - ) - data = json.loads(resp.text.split("\n", 1)[-1]) - - if not data.get("success"): - uri = data.get("unauthorized_redirect_uri") - typer.echo( - f"\n[colab] REQUIRED: Google Drive Authorization needed.\nPlease visit:\n\n{uri}\n" - ) - state.history.log_event(s.name, "drive_auth_needed", {"uri": uri}) - sys.stdout.write("Press Enter after you have granted access... ") - sys.stdout.flush() - try: - with open("/dev/tty") as tty: - tty.readline() - except (FileNotFoundError, OSError, IOError): - # Windows has no /dev/tty; fall back to stdin. - try: - sys.stdin.readline() - except Exception: - try: - input() - except Exception: - pass - - typer.echo("[colab] Authorizing VM...") - params["dryrun"] = "false" - resp = creds.request( - "POST", - url, - params=params, - headers=headers, - files={"file_id": (None, "empty.ipynb")}, - ) - if get_status_code(resp) == 200: - typer.echo("[colab] Credentials propagated. Resuming mount...") - state.history.log_event(s.name, "drive_auth_success", {}) - reply = wsclient.session.msg( - "input_reply", - {"value": {"type": "colab_reply", "colab_msg_id": msg_id}}, - ) - if "header" in deserialize_msg: - reply["parent_header"] = deserialize_msg["header"] - wsclient.stdin_channel.send(reply) - else: - typer.echo( - f"[colab] Error propagating: {get_status_code(resp)} {resp.text}" - ) - return True - return False - - runtime.colab_request_hook = drivefs_hook + runtime.colab_request_hook = make_drivefs_hook( + state, + s, + authuser=authuser, + rewrite_auth_url=rewrite_auth_url, + ) try: s.running = f"automation({op})" s.last_execution = ( @@ -208,13 +195,27 @@ def drivemount( session: Annotated[ Optional[str], typer.Option("-s", "--session", help="Session name") ] = None, + authuser: Annotated[ + Optional[str], + typer.Option( + "--authuser", + help=( + "Account index to select on the Drive consent screen. " + "Appended to the authorization URL. Set this only when the " + "chooser opens the wrong Google account. The propagation " + "request's authuser parameter does not by itself fix a " + "multi-account HTTP 400." + ), + ), + ] = None, path: Annotated[str, typer.Argument(help="Mount path")] = "/content/drive", ): """Mount Google Drive at path""" from colab_cli.common import state name = state.resolve_session(session) - code = f"from google.colab import drive\ndrive.mount('{path}')" + explicit_authuser = authuser is not None + code = mount_code(path, timeout_ms=INTERACTIVE_AUTOMATION_TIMEOUT_SEC * 1000) typer.echo(f"[colab] Mounting Google Drive to '{path}' on {name}...") run_automation( name, @@ -223,6 +224,8 @@ def drivemount( allow_stdin=True, path=path, timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC, + authuser=authuser if explicit_authuser else "0", + rewrite_auth_url=explicit_authuser, ) diff --git a/src/colab_cli/commands/ssh.py b/src/colab_cli/commands/ssh.py index 5bceea7a..2d37caae 100644 --- a/src/colab_cli/commands/ssh.py +++ b/src/colab_cli/commands/ssh.py @@ -37,8 +37,6 @@ import contextlib import os from pathlib import Path -import select -import shlex import signal import subprocess import sys @@ -48,6 +46,7 @@ import uuid from colab_cli.state import SessionState +from colab_cli.terminal import join_argv import typer from typing_extensions import Annotated import websocket @@ -93,9 +92,7 @@ def _pubkey_from_identity(identity: str) -> str: raise typer.Exit(code=2) pubkey = res.stdout.strip() if not pubkey: - typer.echo( - f"[colab] ssh-keygen produced no key for {identity}.", err=True - ) + typer.echo(f"[colab] ssh-keygen produced no key for {identity}.", err=True) raise typer.Exit(code=2) return pubkey @@ -346,6 +343,33 @@ def _close_quietly(ws: websocket.WebSocket) -> None: _DATA_OPCODES = (websocket.ABNF.OPCODE_BINARY, websocket.ABNF.OPCODE_TEXT) +def _pump_stdin_to_ws(ws: websocket.WebSocket, read=os.read, fileno=None) -> None: + """Copy stdin bytes onto the WebSocket until EOF. + + A blocking ``os.read`` is equivalent to ``select`` on a single fd with no + timeout, and it works on Windows. ``select.select`` only accepts sockets + there, so calling it on the stdin handle raises ``OSError`` and the pump + thread dies immediately — ``colab ssh --proxy-mode`` then accepts output + but never forwards keystrokes (upstream issues #85 / #100, PR #158). + + Args: + ws: The connected WebSocket to write binary frames to. + read: Byte reader, injectable in tests. Defaults to ``os.read``. + fileno: Stdin fd. Defaults to ``sys.stdin.buffer.fileno()``. + """ + stdin_fd = sys.stdin.buffer.fileno() if fileno is None else fileno + try: + while True: + data = read(stdin_fd, 8192) + if not data: + break + ws.send_binary(data) + except (OSError, websocket.WebSocketException): + pass + finally: + _close_quietly(ws) + + def _bridge_proxy_mode(ws: websocket.WebSocket) -> int: """Bridges the WebSocket <-> stdin/stdout as an OpenSSH ProxyCommand. @@ -355,24 +379,8 @@ def _bridge_proxy_mode(ws: websocket.WebSocket) -> int: Returns: 0 when either side closes. """ - stdin_fd = sys.stdin.buffer.fileno() - def stdin_to_ws(): - try: - while True: - ready, _, _ = select.select([stdin_fd], [], [], None) - if not ready: - continue - data = os.read(stdin_fd, 8192) - if not data: - break - ws.send_binary(data) - except (OSError, websocket.WebSocketException): - pass - finally: - _close_quietly(ws) - - threading.Thread(target=stdin_to_ws, daemon=True).start() + threading.Thread(target=_pump_stdin_to_ws, args=(ws,), daemon=True).start() try: while True: @@ -409,7 +417,9 @@ def _proxy_command(session: SessionState, identity: Optional[str]) -> str: ] if identity: self_cmd.extend(["--identity", identity]) - return shlex.join(self_cmd) + # POSIX OpenSSH re-parses this with /bin/sh. Windows OpenSSH uses cmd.exe, + # which does not honor single quotes. + return join_argv(self_cmd) def _ssh_base_args(proxy_command: str, identity: Optional[str]) -> list[str]: @@ -421,7 +431,9 @@ def _ssh_base_args(proxy_command: str, identity: Optional[str]) -> list[str]: "-o", "StrictHostKeyChecking=no", "-o", - "UserKnownHostsFile=/dev/null", + # ``/dev/null`` on POSIX, ``nul`` on Windows. Windows OpenSSH rejects + # a POSIX null path for the known-hosts file. + f"UserKnownHostsFile={os.devnull}", "-o", "LogLevel=ERROR", ] @@ -482,9 +494,7 @@ def _select_proxy_session( """ if session and not _session_exists(session): with contextlib.redirect_stdout(sys.stderr): - return _auto_create_session( - gpu, tpu, name=session, high_mem=high_mem - ), True + return _auto_create_session(gpu, tpu, name=session, high_mem=high_mem), True return _resolve_session(session), False @@ -546,16 +556,19 @@ def _on_signal(signum, frame): do_rm() os._exit(0) - for sig in (signal.SIGHUP, signal.SIGTERM, signal.SIGINT): + # Windows has SIGINT and SIGTERM but not SIGHUP. Looking the name up + # with getattr avoids AttributeError while evaluating the signal list. + for name in ("SIGHUP", "SIGTERM", "SIGINT"): + sig = getattr(signal, name, None) + if sig is None: + continue try: signal.signal(sig, _on_signal) except (ValueError, OSError): pass # e.g. not running in the main thread -def _run_proxy_bridge( - s: SessionState, identity: Optional[str], rm: bool -) -> int: +def _run_proxy_bridge(s: SessionState, identity: Optional[str], rm: bool) -> int: """Runs the ``--proxy-mode`` WebSocket-stdio bridge, honoring ``--rm``. Args: diff --git a/src/colab_cli/console.py b/src/colab_cli/console.py index ae2786bb..a73c165b 100644 --- a/src/colab_cli/console.py +++ b/src/colab_cli/console.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import json import logging import os @@ -33,6 +34,12 @@ import websocket from colab_cli.state import SessionState +from colab_cli.terminal import ( + enable_windows_virtual_terminal, + is_windows, + read_windows_key, + windows_raw_console, +) logger = logging.getLogger(__name__) @@ -96,30 +103,44 @@ def on_open(ws): # Setup the background thread to read from stdin def read_stdin(): is_tty = sys.stdin.isatty() + # Windows consoles are line-buffered unless we switched to raw mode + # in connect_console(). Read keystrokes through the console API and + # translate scan codes to ANSI so cmd.exe and PowerShell behave like + # a POSIX raw tty. Piped stdin stays on sys.stdin.read so EOF still + # works. + windows_tty = is_windows() and is_tty while _is_running: try: - # Read a single character (or escape sequence byte) - char = sys.stdin.read(1) - if not char: - if not is_tty: - # Piped input has reached EOF. The remote /colab/tty - # endpoint wraps bash in tmux which intercepts \x04 - # (Ctrl-D) as a literal character, so it never exits. - # Instead send "exit\n" so bash voluntarily terminates, - # wait a short grace period for the shell's goodbye - # output to drain back to us, then close the websocket - # ourselves to guarantee the client unblocks. - try: - ws.send(json.dumps({"data": "exit\n"})) - except Exception: - pass - time.sleep(PIPED_EOF_GRACE_SECONDS) - try: - ws.close() - except Exception: - pass - break - ws.send(json.dumps({"data": char})) + if windows_tty: + text = read_windows_key() + if not text: + continue + chunks = list(text) + else: + # Read a single character (or escape sequence byte) + char = sys.stdin.read(1) + if not char: + if not is_tty: + # Piped input has reached EOF. The remote /colab/tty + # endpoint wraps bash in tmux which intercepts \x04 + # (Ctrl-D) as a literal character, so it never exits. + # Instead send "exit\n" so bash voluntarily terminates, + # wait a short grace period for the shell's goodbye + # output to drain back to us, then close the websocket + # ourselves to guarantee the client unblocks. + try: + ws.send(json.dumps({"data": "exit\n"})) + except Exception: + pass + time.sleep(PIPED_EOF_GRACE_SECONDS) + try: + ws.close() + except Exception: + pass + break + chunks = [char] + for chunk in chunks: + ws.send(json.dumps({"data": chunk})) except Exception: break @@ -133,6 +154,10 @@ def connect_console(session: SessionState): """ global _is_running, _last_error _last_error = None + # cmd.exe and Windows PowerShell 5.1 do not enable ANSI processing by + # default. The remote PTY speaks ANSI, and we write those bytes straight + # to stdout.buffer, so the console mode has to be switched first. + enable_windows_virtual_terminal() # Construct the WebSocket URL from the base URL parsed = urlparse(session.url) @@ -143,6 +168,9 @@ def connect_console(session: SessionState): fd = None old_settings = None can_raw = termios is not None and tty is not None and is_tty + # Windows has no termios. Raw mode is the console API instead, so + # interactive `colab console` is not stuck in cooked line editing. + use_windows_raw = is_windows() and is_tty and not can_raw if can_raw: try: fd = sys.stdin.fileno() @@ -166,22 +194,43 @@ def handle_sigwinch(signum, frame): send_terminal_size(ws) sigwinch = getattr(signal, "SIGWINCH", None) - try: - if can_raw: - tty.setraw(fd, termios.TCSANOW) - if sigwinch is not None: - signal.signal(sigwinch, handle_sigwinch) + resize_stop = threading.Event() - # This is a blocking call until the connection is closed - ws.run_forever() + def _watch_resize(): + """Poll the console size. Windows has no SIGWINCH.""" + last = None + while not resize_stop.is_set(): + try: + size = os.get_terminal_size() + current = (size.columns, size.lines) + if current != last: + last = current + send_terminal_size(ws) + except Exception: + pass + resize_stop.wait(0.4) - if _last_error: - # Re-raise or wrap terminal errors - err_msg = str(_last_error) - if "404" in err_msg or "401" in err_msg: - # We raise a standard exception that the caller can recognize - raise RuntimeError(f"Connection failed: {err_msg}") + raw_cm = windows_raw_console() if use_windows_raw else contextlib.nullcontext() + try: + with raw_cm: + if use_windows_raw: + threading.Thread(target=_watch_resize, daemon=True).start() + if can_raw: + tty.setraw(fd, termios.TCSANOW) + if sigwinch is not None: + signal.signal(sigwinch, handle_sigwinch) + + # This is a blocking call until the connection is closed + ws.run_forever() + + if _last_error: + # Re-raise or wrap terminal errors + err_msg = str(_last_error) + if "404" in err_msg or "401" in err_msg: + # We raise a standard exception that the caller can recognize + raise RuntimeError(f"Connection failed: {err_msg}") finally: + resize_stop.set() if can_raw: # Always ensure the terminal is restored to its original state try: diff --git a/src/colab_cli/drive_auth.py b/src/colab_cli/drive_auth.py new file mode 100644 index 00000000..a9ae5dde --- /dev/null +++ b/src/colab_cli/drive_auth.py @@ -0,0 +1,303 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Google Drive credential propagation for ``colab drivemount``. + +The kernel's ``drive.mount()`` blocks in ``google.colab._message.blocking_request`` +until the frontend answers the ``dfs_ephemeral`` ``colab_request`` with an +``input_reply`` whose value is ``{"type": "colab_reply", "colab_msg_id": ...}``. +The default ``timeout_ms`` is 120000, and that budget includes the time the +user spends in the browser. This module keeps that wait aligned with the CLI's +interactive timeout and turns the multi-account HTTP 400 into an actionable +message instead of a raw HTML error page. +""" + +import json +import re +from typing import Any, Callable, Optional +from urllib.parse import parse_qs, parse_qsl, urlencode, urlparse, urlunparse + +from colab_cli.terminal import open_url_in_browser, read_controlling_line + +# google.colab.drive.mount defaults to 120s. That is shorter than a normal +# consent click-through, so the kernel gives up, starts DriveFS with no +# credentials, and raises ``ValueError: mount failed`` (upstream issue #113). +# 10 minutes matches the CLI execute() ceiling for this command. +DRIVE_AUTH_TIMEOUT_SEC = 600 +DRIVE_MOUNT_TIMEOUT_MS = DRIVE_AUTH_TIMEOUT_SEC * 1000 + +_XSSI_PREFIX = ")]}'" +_ISSUE_103 = "https://github.com/googlecolab/google-colab-cli/issues/103" +_ISSUE_113 = "https://github.com/googlecolab/google-colab-cli/issues/113" + +Echo = Callable[[str], None] +BrowserOpener = Callable[[str], bool] +LineReader = Callable[..., str] +EventLogger = Callable[[str, dict], None] + + +def mount_code(path: str, timeout_ms: int = DRIVE_MOUNT_TIMEOUT_MS) -> str: + """Python source that mounts Drive and waits long enough for browser consent. + + ``path`` is embedded with ``repr`` so quotes and backslashes round-trip. + """ + return ( + "from google.colab import drive\n" + f"drive.mount({path!r}, timeout_ms={int(timeout_ms)})" + ) + + +def login_hint(uri: str) -> Optional[str]: + """Return the ``login_hint`` query value, or None when it is absent.""" + if not uri: + return None + values = parse_qs(urlparse(uri).query).get("login_hint") or [] + if not values or not values[0]: + return None + return values[0] + + +def with_authuser(uri: str, authuser: Optional[str]) -> str: + """Return ``uri`` with its ``authuser`` query parameter set. + + ``authuser is None`` leaves the URL alone. The server-built consent URL + already carries ``login_hint``; appending ``authuser=0`` by default can + pin a multi-account browser to the wrong account. + """ + if not uri or authuser is None: + return uri + parts = urlparse(uri) + pairs = [ + (key, value) + for key, value in parse_qsl(parts.query, keep_blank_values=True) + if key != "authuser" + ] + pairs.append(("authuser", str(authuser))) + return urlunparse(parts._replace(query=urlencode(pairs))) + + +def propagation_params(authuser: str, *, dryrun: bool) -> dict: + """Query parameters for ``/tun/m/credentials-propagation``.""" + return { + "authuser": str(authuser), + "authtype": "dfs_ephemeral", + "version": "2", + "dryrun": "true" if dryrun else "false", + "propagate": "true", + "record": "false", + } + + +def parse_colab_json(text: str) -> dict: + """Parse a Colab JSON body, stripping the XSSI guard prefix when present.""" + if not text: + return {} + raw = text + if raw.startswith(_XSSI_PREFIX): + raw = raw[len(_XSSI_PREFIX) :] + if raw.startswith("\n"): + raw = raw[1:] + raw = raw.strip() + if not raw: + return {} + try: + data = json.loads(raw) + except json.JSONDecodeError: + return {} + return data if isinstance(data, dict) else {} + + +def _plain_snippet(body: str, limit: int = 300) -> str: + text = re.sub(r"<[^>]+>", " ", body or "") + text = re.sub(r"\s+", " ", text).strip() + if len(text) > limit: + text = text[:limit] + "..." + return text + + +def explain_propagation_failure(status: Optional[int], body: str) -> str: + """Human-readable propagation failure. + + A 400 from this endpoint is the multi-account mismatch described in + upstream issue #103: the consent screen follows ``login_hint``, but the + gsession created at the redirect is attributed to the browser's *active* + account. The response body is a full HTML error page; dumping it does not + tell the user what to do. + """ + snippet = _plain_snippet(body) + guidance = ( + "The browser account that approved access did not match the CLI " + "account. On the account chooser, select the account named in the " + "login_hint (printed above, and shown by `colab whoami`). " + f"{_ISSUE_103}" + ) + if status == 400 or (body and "Bad Request" in body): + message = ( + f"[colab] Drive credential propagation was rejected (HTTP {status}). " + f"{guidance}" + ) + if snippet: + message = f"{message}\n{snippet}" + return message + suffix = f" {snippet}" if snippet else "" + return f"[colab] Error propagating: {status}{suffix}" + + +def send_colab_reply( + wsclient: Any, + deserialize_msg: dict, + msg_id: Any, + *, + error: Optional[str] = None, +) -> None: + """Unblock ``drive.mount`` with a ``colab_reply`` on the stdin channel. + + The kernel matches ``content.value.colab_msg_id`` to the integer it stored + in the request metadata. An ``error`` field makes ``read_reply_from_input`` + raise ``MessageError`` immediately instead of waiting out the mount timeout + and then failing with an empty ``ValueError: mount failed``. + """ + value: dict = {"type": "colab_reply", "colab_msg_id": msg_id} + if error: + value["error"] = error + reply = wsclient.session.msg("input_reply", {"value": value}) + header = deserialize_msg.get("header") + if header: + reply["parent_header"] = header + wsclient.stdin_channel.send(reply) + + +def _json_ok(resp: Any) -> dict: + if getattr(resp, "status_code", None) != 200: + return {} + return parse_colab_json(getattr(resp, "text", "") or "") + + +def perform_drive_authorization( + *, + colab_domain: str, + endpoint: str, + http: Any, + wsclient: Any, + deserialize_msg: dict, + msg_id: Any, + authuser: str = "0", + rewrite_auth_url: bool = False, + open_browser: BrowserOpener = open_url_in_browser, + read_line: LineReader = read_controlling_line, + echo: Optional[Echo] = None, + on_event: Optional[EventLogger] = None, +) -> None: + """Run the Drive credential-propagation handshake and reply to the kernel. + + ``http`` is an authorized ``requests`` session (``get_credentials``). + ``rewrite_auth_url`` appends ``authuser`` to the consent URL. It is off + unless the user passed ``--authuser``: issue #103 showed that changing the + propagation *request* parameter does not fix a multi-account 400, and + forcing ``authuser=0`` onto the browser URL can select the wrong account. + """ + if echo is None: + import typer + + echo = typer.echo + + def emit(message: str) -> None: + echo(message) + + def event(name: str, payload: dict) -> None: + if on_event: + on_event(name, payload) + + url = f"{colab_domain}/tun/m/credentials-propagation/{endpoint}" + emit(f"\n[colab] Intercepted Drive Auth Request. Connecting to {colab_domain}...") + params = propagation_params(authuser, dryrun=True) + token_resp = http.request("GET", url, params=params) + token = _json_ok(token_resp).get("token") + headers = {"x-goog-colab-token": token} if token else {} + files = {"file_id": (None, "empty.ipynb")} + probe = http.request("POST", url, params=params, headers=headers, files=files) + if getattr(probe, "status_code", None) != 200: + message = explain_propagation_failure( + getattr(probe, "status_code", None), getattr(probe, "text", "") or "" + ) + emit(message) + send_colab_reply( + wsclient, + deserialize_msg, + msg_id, + error="Drive credential propagation failed before consent.", + ) + return + + data = _json_ok(probe) + if not data.get("success"): + uri = data.get("unauthorized_redirect_uri") or "" + if rewrite_auth_url: + uri = with_authuser(uri, authuser) + if not uri: + emit("[colab] Drive auth did not return a consent URL.") + emit( + explain_propagation_failure( + getattr(probe, "status_code", None), + getattr(probe, "text", "") or "", + ) + ) + send_colab_reply( + wsclient, + deserialize_msg, + msg_id, + error="Drive auth did not return a consent URL.", + ) + return + account = ( + login_hint(uri) or "the account this CLI authenticated as (colab whoami)" + ) + emit("[colab] Google Drive authorization needed.") + emit(f"Sign in as {account}.") + emit( + "If the browser is signed into more than one Google account, pick " + "this account on the chooser. login_hint alone is not enough: Colab " + "attributes the Drive session to the browser's active account and " + f"rejects a mismatch with HTTP 400 ({_ISSUE_103})." + ) + emit("") + emit(uri) + emit("") + event("drive_auth_needed", {"uri": uri}) + if open_browser(uri): + emit("[colab] Opened that link in your browser.") + else: + emit("[colab] Could not launch a browser here. Open the link above.") + emit(f"[colab] The kernel waits up to 10 minutes for this step ({_ISSUE_113}).") + read_line("Press Enter after you have granted access... ") + + emit("[colab] Authorizing VM...") + params = propagation_params(authuser, dryrun=False) + final = http.request("POST", url, params=params, headers=headers, files=files) + if getattr(final, "status_code", None) == 200: + emit("[colab] Credentials propagated. Resuming mount...") + event("drive_auth_success", {}) + send_colab_reply(wsclient, deserialize_msg, msg_id) + return + + message = explain_propagation_failure( + getattr(final, "status_code", None), getattr(final, "text", "") or "" + ) + emit(message) + send_colab_reply( + wsclient, + deserialize_msg, + msg_id, + error=message, + ) diff --git a/src/colab_cli/terminal.py b/src/colab_cli/terminal.py new file mode 100644 index 00000000..9a440932 --- /dev/null +++ b/src/colab_cli/terminal.py @@ -0,0 +1,242 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cross-platform terminal helpers. + +POSIX keeps using ``termios`` raw mode in ``console.py``. Windows consoles +(cmd.exe and PowerShell, both conhost and Windows Terminal) have no +``termios``, no ``/dev/tty``, and no ``SIGWINCH``. This module owns the +Windows console-mode and key-translation pieces, plus the controlling-terminal +line read used by interactive prompts. +""" + +import os +import shutil +import subprocess +import sys +from contextlib import contextmanager +from typing import Iterator, Optional + +# Win32 console mode flags. See +# https://learn.microsoft.com/en-us/windows/console/setconsolemode +_ENABLE_PROCESSED_OUTPUT = 0x0001 +_ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004 +_ENABLE_VIRTUAL_TERMINAL_INPUT = 0x0200 +_ENABLE_EXTENDED_FLAGS = 0x0080 + +_STD_INPUT_HANDLE = -10 +_STD_OUTPUT_HANDLE = -11 +_STD_ERROR_HANDLE = -12 + +# Second byte of a Windows console special key (after \x00 or \xe0) mapped to +# the ANSI sequence a POSIX raw tty would emit for the same key. +_SCAN_TO_ANSI = { + "H": "\x1b[A", # up + "P": "\x1b[B", # down + "M": "\x1b[C", # right + "K": "\x1b[D", # left + "G": "\x1b[H", # home + "O": "\x1b[F", # end + "I": "\x1b[5~", # page up + "Q": "\x1b[6~", # page down + "S": "\x1b[3~", # delete + "R": "\x1b[2~", # insert +} + + +def is_windows() -> bool: + """True on Windows, including cmd.exe and PowerShell.""" + return os.name == "nt" + + +def translate_windows_key(first: str, second: Optional[str] = None) -> str: + """Map one Windows console key event to bytes a POSIX raw tty would send. + + ``msvcrt.getwch`` returns a Unicode character. Special keys arrive as + ``\\x00`` or ``\\xe0`` followed by a scan-code character. Arrow keys must + become ANSI sequences so the remote shell's line editor understands them. + Enter arrives as ``\\r``, which is also what a POSIX raw terminal sends + (local ICRNL is off; the remote pty's ICRNL turns ``\\r`` into ``\\n``). + Unrecognized scan codes return an empty string so they are not forwarded + as garbage. + """ + if first in ("\x00", "\xe0"): + if not second: + return "" + return _SCAN_TO_ANSI.get(second, "") + return first + + +def enable_windows_virtual_terminal() -> bool: + """Enable ANSI processing and UTF-8 on the Windows console. + + cmd.exe and Windows PowerShell 5.1 leave + ``ENABLE_VIRTUAL_TERMINAL_PROCESSING`` off, so cursor movement and colors + written as raw bytes show up as garbage. Windows Terminal and PowerShell 7 + often already have the bit set; setting it again is harmless. No-op off + Windows. Returns True when an output mode was updated. + """ + if not is_windows(): + return False + try: + import ctypes + + kernel32 = ctypes.windll.kernel32 + kernel32.SetConsoleOutputCP(65001) + kernel32.SetConsoleCP(65001) + updated = False + for handle_id in (_STD_OUTPUT_HANDLE, _STD_ERROR_HANDLE): + handle = kernel32.GetStdHandle(handle_id) + mode = ctypes.c_uint() + if not kernel32.GetConsoleMode(handle, ctypes.byref(mode)): + continue + new_mode = ( + mode.value + | _ENABLE_PROCESSED_OUTPUT + | _ENABLE_VIRTUAL_TERMINAL_PROCESSING + ) + if kernel32.SetConsoleMode(handle, new_mode): + updated = True + return updated + except Exception: + return False + + +@contextmanager +def windows_raw_console() -> Iterator[None]: + """Put the Windows console input handle into raw (non-line) mode. + + Clears line-input, echo, and processed-input so keystrokes — including + Ctrl-C as ``\\x03`` — reach this process instead of being edited by + conhost or killing the local CLI. Enables virtual-terminal input so newer + consoles can emit ANSI sequences themselves. Restores the previous mode + on exit. No-op when this is not a Windows console (pipes, POSIX). + """ + if not is_windows(): + yield + return + import ctypes + + kernel32 = ctypes.windll.kernel32 + handle = kernel32.GetStdHandle(_STD_INPUT_HANDLE) + old = ctypes.c_uint() + if not kernel32.GetConsoleMode(handle, ctypes.byref(old)): + yield + return + # Leaving LINE/ECHO/PROCESSED unset is the raw mode. EXTENDED_FLAGS has + # to be set for the mode change to stick on modern conhost. + raw = _ENABLE_VIRTUAL_TERMINAL_INPUT | _ENABLE_EXTENDED_FLAGS + kernel32.SetConsoleMode(handle, raw) + try: + yield + finally: + kernel32.SetConsoleMode(handle, old.value) + + +def read_windows_key() -> str: + """Block for one Windows console key and return an ANSI-compatible string. + + Returns ``""`` for an unrecognized scan code. The caller should not + forward that. + """ + import msvcrt + + first = msvcrt.getwch() + if first in ("\x00", "\xe0"): + second = msvcrt.getwch() + return translate_windows_key(first, second) + return translate_windows_key(first) + + +def read_controlling_line(prompt: str = "") -> str: + """Read one line from the controlling terminal, not from a redirected pipe. + + POSIX opens ``/dev/tty``. Windows opens ``CONIN$``, which is the console + even when stdin is redirected (the same reason the POSIX path avoids + stdin). Falls back to ``sys.stdin`` and then ``input()`` when neither + device exists, which is the case under CI and some IDE consoles. + """ + if prompt: + sys.stdout.write(prompt) + sys.stdout.flush() + if is_windows(): + try: + with open("CONIN$", "r", encoding="utf-8", errors="replace") as con: + return con.readline() + except OSError: + pass + else: + try: + with open("/dev/tty", encoding="utf-8", errors="replace") as tty: + return tty.readline() + except OSError: + pass + try: + return sys.stdin.readline() + except Exception: + try: + return input() + "\n" + except Exception: + return "" + + +def open_url_in_browser(url: str) -> bool: + """Open ``url`` in the system GUI browser. + + Windows uses ``os.startfile``, the same association cmd and PowerShell use + for ``start``. macOS uses ``open``. Other platforms use ``xdg-open`` when + it is on ``PATH``. A missing opener returns False so the caller can keep + the printed URL. This never launches a TUI browser (lynx, w3m), which + would steal the terminal the CLI is prompting on. + """ + if is_windows(): + startfile = getattr(os, "startfile", None) + if startfile is None: + return False + try: + startfile(url) + return True + except OSError: + return False + opener_name = "open" if sys.platform == "darwin" else "xdg-open" + opener = shutil.which(opener_name) + if not opener: + return False + try: + subprocess.Popen( + [opener, url], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return True + except OSError: + return False + + +def join_argv(argv: list[str], *, windows: Optional[bool] = None) -> str: + """Quote an argv for the shell that will re-parse it. + + POSIX OpenSSH runs ``ProxyCommand`` via ``/bin/sh`` (``shlex.join``). + Windows OpenSSH runs it via ``cmd.exe`` (``subprocess.list2cmdline``). + cmd.exe does not treat single quotes as quoting, so a POSIX-quoted path + such as ``'C:\\Program Files\\Python\\python.exe'`` is passed through + literally and the bridge never starts. + """ + if windows is None: + windows = is_windows() + if windows: + return subprocess.list2cmdline(argv) + import shlex + + return shlex.join(argv) diff --git a/tests/test_automation.py b/tests/test_automation.py index 7407edc9..b3da1f45 100644 --- a/tests/test_automation.py +++ b/tests/test_automation.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from typer.testing import CliRunner from colab_cli.cli import app @@ -96,7 +96,11 @@ def test_cli_drivemount(mock_state, mock_runtime_class, mock_session): mock_runtime.execute_code.assert_called_once() called_code = mock_runtime.execute_code.call_args[0][0] - assert "drive.mount('/foo/bar')" in called_code + assert "drive.mount('/foo/bar'" in called_code + assert "timeout_ms=" in called_code + # The kernel-side default is 120s and that budget includes the browser + # consent wait. Anything under 5 minutes will still abort a normal mount. + assert "timeout_ms=120000" not in called_code assert mock_runtime.colab_request_hook is not None # Drivemount waits for the user to OAuth in their browser; the kernel # goes silent during that wait and the default 10s execute() timeout @@ -106,6 +110,53 @@ def test_cli_drivemount(mock_state, mock_runtime_class, mock_session): assert kwargs.get("timeout") is not None and kwargs["timeout"] >= 300 +@patch("colab_cli.commands.automation.perform_drive_authorization") +@patch("colab_cli.auth.get_credentials") +@patch("colab_cli.commands.automation.ColabRuntime") +@patch("colab_cli.common.state") +def test_drivemount_hook_runs_off_the_websocket_thread( + mock_state, mock_runtime_class, _mock_creds, mock_perform, mock_session +): + """The dfs_ephemeral handler must return without doing the browser prompt + on the caller (the websocket recv thread).""" + mock_state.get_session.return_value = mock_session + mock_state.resolve_session.return_value = "test-session" + mock_state.client.colab_domain = "https://colab.research.google.com" + mock_runtime = mock_runtime_class.return_value + mock_runtime.execute_code.return_value = [{"text": "Mounted"}] + + result = runner.invoke(app, ["drivemount", "-s", "test-session", "--authuser", "6"]) + assert result.exit_code == 0 + + hook = mock_runtime.colab_request_hook + msg = { + "content": {"request": {"authType": "dfs_ephemeral"}}, + "metadata": {"colab_msg_id": 3}, + "header": {"msg_id": "parent"}, + } + + class SyncThread: + def __init__(self, target=None, daemon=None, name=None): + self.target = target + + def start(self): + self.target() + + with patch("colab_cli.commands.automation.threading.Thread", SyncThread): + assert hook(msg, MagicMock()) is True + assert mock_perform.call_args.kwargs["authuser"] == "6" + assert mock_perform.call_args.kwargs["rewrite_auth_url"] is True + + other = {"content": {"request": {"authType": "something_else"}}} + assert hook(other, MagicMock()) is False + + +def test_drivemount_help_documents_authuser(): + result = runner.invoke(app, ["drivemount", "--help"]) + assert result.exit_code == 0 + assert "--authuser" in result.output + + @patch("colab_cli.commands.automation.ColabRuntime") @patch("colab_cli.common.state") def test_cli_auth_uses_long_timeout(mock_state, mock_runtime_class, mock_session): diff --git a/tests/test_console.py b/tests/test_console.py index 13f20961..55577eb9 100644 --- a/tests/test_console.py +++ b/tests/test_console.py @@ -235,9 +235,7 @@ def start(self): @patch("colab_cli.console.websocket.WebSocketApp") @patch("colab_cli.console.sys.stdin.isatty") -def test_console_no_termios_degrades_gracefully( - mock_isatty, mock_ws_app, mock_session -): +def test_console_no_termios_degrades_gracefully(mock_isatty, mock_ws_app, mock_session): """Windows has no termios/tty/SIGWINCH. connect_console must not crash. Regression test for `ModuleNotFoundError: No module named 'termios'` @@ -259,3 +257,78 @@ def test_console_no_termios_degrades_gracefully( connect_console(mock_session) mock_ws_instance.run_forever.assert_called_once() + + +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty") +def test_console_windows_uses_raw_console_not_termios( + mock_isatty, mock_ws_app, mock_session +): + """cmd.exe and PowerShell have no termios. Interactive console must switch + the Windows console into raw mode instead of staying line-buffered. + """ + import colab_cli.console as console_mod + + mock_isatty.return_value = True + mock_ws_instance = MagicMock() + mock_ws_app.return_value = mock_ws_instance + mock_ws_instance.run_forever.return_value = None + raw_cm = MagicMock() + + with ( + patch.object(console_mod, "termios", None), + patch.object(console_mod, "tty", None), + patch.object(console_mod, "is_windows", return_value=True), + patch.object( + console_mod, "windows_raw_console", return_value=raw_cm + ) as mock_raw, + patch("colab_cli.console.threading.Thread"), + ): + connect_console(mock_session) + + mock_raw.assert_called_once() + raw_cm.__enter__.assert_called_once() + raw_cm.__exit__.assert_called_once() + mock_ws_instance.run_forever.assert_called_once() + + +@patch("colab_cli.console.os.get_terminal_size") +@patch("colab_cli.console.sys.stdin.isatty") +def test_read_stdin_windows_tty_forwards_translated_keys( + mock_isatty, mock_get_term_size +): + """Arrow keys arrive as Windows scan codes. The remote shell needs the + ANSI sequence, one character per websocket frame, matching POSIX raw mode. + """ + import colab_cli.console as console_mod + + mock_isatty.return_value = True + mock_get_term_size.return_value = os.terminal_size((80, 24)) + keys = ["\x1b[A", "\r"] + + def fake_key(): + if not keys: + console_mod._is_running = False + return "" + return keys.pop(0) + + mock_ws = MagicMock() + + class SyncThread: + def __init__(self, target, daemon=None): + self.target = target + + def start(self): + self.target() + + console_mod._is_running = True + with ( + patch.object(console_mod, "is_windows", return_value=True), + patch.object(console_mod, "read_windows_key", side_effect=fake_key), + patch("colab_cli.console.threading.Thread", SyncThread), + ): + on_open(mock_ws) + + sent = [json.loads(call.args[0]) for call in mock_ws.send.call_args_list] + assert sent[0] == {"cols": 80, "rows": 24} + assert [item["data"] for item in sent[1:]] == ["\x1b", "[", "A", "\r"] diff --git a/tests/test_drive_auth.py b/tests/test_drive_auth.py new file mode 100644 index 00000000..419e96bc --- /dev/null +++ b/tests/test_drive_auth.py @@ -0,0 +1,271 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Drive credential propagation. + +Pins the fixes for upstream issues: +- #113 mount failed after exactly 120s (kernel-side drive.mount timeout) +- #103 HTTP 400 when the browser's active account is not the CLI account +- #108 open the consent URL instead of copy-paste only +""" + +import json +from unittest.mock import MagicMock + +from colab_cli.drive_auth import ( + DRIVE_MOUNT_TIMEOUT_MS, + explain_propagation_failure, + login_hint, + mount_code, + perform_drive_authorization, + with_authuser, +) + + +def test_mount_code_extends_kernel_timeout_past_120s(): + """google.colab.drive.mount defaults timeout_ms to 120000. That budget + covers the browser consent wait, so a normal OAuth click-through makes + the kernel give up and raise ``ValueError: mount failed`` (#113). + """ + code = mount_code("/content/drive") + assert "drive.mount('/content/drive'" in code + assert f"timeout_ms={DRIVE_MOUNT_TIMEOUT_MS}" in code + assert DRIVE_MOUNT_TIMEOUT_MS >= 600_000 + + +def test_mount_code_quotes_path_with_repr(): + code = mount_code("/tmp/drive's") + assert ( + "drive.mount('/tmp/drive\\'s'" in code or 'drive.mount("/tmp/drive\'s"' in code + ) + + +def test_login_hint_is_decoded(): + uri = ( + "https://accounts.google.com/o/oauth2/v2/auth" + "?login_hint=user%40example.com&response_type=none+gsession" + ) + assert login_hint(uri) == "user@example.com" + assert login_hint("") is None + assert login_hint("https://accounts.google.com/o/oauth2/v2/auth") is None + + +def test_with_authuser_appends_and_replaces(): + uri = "https://accounts.google.com/o/oauth2/v2/auth?login_hint=a@b.c&authuser=0" + rewritten = with_authuser(uri, "6") + assert "authuser=6" in rewritten + assert "authuser=0" not in rewritten + assert "login_hint=a%40b.c" in rewritten or "login_hint=a@b.c" in rewritten + assert with_authuser(uri, None) == uri + + +def test_explain_400_is_actionable_and_not_a_raw_html_dump(): + html = ( + "Error 400 (Bad Request)!!1" + "" + ("x" * 8000) + "" + ) + text = explain_propagation_failure(400, html) + assert "400" in text + assert "account" in text.lower() + assert "login_hint" in text or "chooser" in text.lower() + assert "https://github.com/googlecolab/google-colab-cli/issues/103" in text + assert "