From acf800cf2c682aa11a117b40da606cbb51eac4a4 Mon Sep 17 00:00:00 2001 From: Drake Aiman Date: Thu, 29 Oct 2026 19:02:16 +0000 Subject: [PATCH 1/2] refactor!: remove `colab console` in favour of `colab ssh` `colab ssh` supersedes the `/colab/tty` console, whose piped output carried terminal-control bytes. This deletes the console module, its tests, the piped-console integration test and converter.py's stale console comment. Use `colab ssh -s NAME` for an interactive shell and `colab ssh -s NAME -- CMD` for one-shot commands, or `echo '!cmd' | colab exec -s NAME` where ssh can't work. BREAKING CHANGE: `colab console` is removed; use `colab ssh -s NAME`, or `colab ssh -s NAME -- CMD` for one-shot commands (`!`-prefixed `colab exec` where ssh can't work). --- AGENTS.md | 9 +- CHANGELOG.md | 10 + README.md | 7 +- docs/01_session_management.md | 3 +- docs/02_execution_and_interactive.md | 12 +- docs/06_ssh_access.md | 5 +- docs/demos.md | 16 +- integration/README.md | 1 - integration/repro_piped_console/test.sh | 96 --------- skills/colab-operator/SKILL.md | 4 +- src/colab_cli/commands/execution.py | 30 --- src/colab_cli/console.py | 197 ------------------ src/colab_cli/converter.py | 2 - tests/test_cli.py | 43 +++- tests/test_console.py | 261 ------------------------ 15 files changed, 64 insertions(+), 632 deletions(-) delete mode 100755 integration/repro_piped_console/test.sh delete mode 100644 src/colab_cli/console.py delete mode 100644 tests/test_console.py diff --git a/AGENTS.md b/AGENTS.md index 668ff1e2..7275050f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -37,7 +37,7 @@ ## Subcommand Workflows - **Session Management**: `new`, `sessions`, `status`, `stop`. -- **Execution**: `repl`, `exec`, `console`. +- **Execution**: `repl`, `exec`, `ssh`. - **Files**: `ls`, `rm`, `upload`, `download`, `edit`. - **Automation**: `auth`, `drivemount`, `install`, `log`, `pay`, `version`, `update`. @@ -69,7 +69,7 @@ The package version is derived from the git tag via `hatch-vcs` (see `pyproject. 10. **Live Probes Allocate Real Resources**: Probing the Colab API to debug an issue creates real, billable assignments — every successful POST `/tun/m/assign` reserves a VM. Prefer GET-only (read) probes whenever possible. For any state-mutating call, (a) record every endpoint you create as you go, and (b) clean up via `client.unassign(endpoint)` (or `colab stop`) before declaring the investigation done. Then verify with `colab sessions` that nothing was orphaned. 11. **Push Freshness**: The remote may have advanced during a session (other contributors land commits while you work). ALWAYS `git fetch ` immediately before pushing or merging. If `git log main../main` is non-empty, reset local `main` to the remote, rebase feature branches onto it, retest, then push. NEVER force-push `main` to recover from divergence. 12. **Amend Safety**: Before `git commit --amend`, explicitly verify all three preconditions: (a) the user requested amend OR a pre-commit hook auto-modified files for an otherwise-successful commit, (b) HEAD was created by you in this conversation (`git log -1 --format='%an %ae'`), (c) the commit has not been pushed to a remote. If any precondition fails, create a new commit instead. NEVER amend a failed/rejected commit — fix the issue and create a new one. -13. **Run Integration Tests Yourself**: Re-read AGENTS.md "Agent Execution Limitations" before claiming you can't run a test. The CANNOT-run list is exclusively interactive commands (`colab auth`, `colab drivemount`, `colab repl`, `colab console`). Tests built on `colab new` / `colab stop` / `colab log` are non-interactive and you MUST run them yourself before declaring a fix complete. Encoded after running through a full implement-and-document cycle for a fix that the integration test would have falsified in 30 seconds — the cure was `uv run bash integration/repro_run_command/test.sh`. +13. **Run Integration Tests Yourself**: Re-read AGENTS.md "Agent Execution Limitations" before claiming you can't run a test. The CANNOT-run list is exclusively interactive commands (`colab auth`, `colab drivemount`, unpiped `colab repl`, interactive `colab ssh`). Tests built on `colab new` / `colab stop` / `colab log` are non-interactive and you MUST run them yourself before declaring a fix complete. Encoded after running through a full implement-and-document cycle for a fix that the integration test would have falsified in 30 seconds — the cure was `uv run bash integration/repro_run_command/test.sh`. 14. **Heed Research Caveats**: When a research subagent surfaces a caveat ("the proto allows it but the policy may reject"), treat it as a TODO to verify, not as a footnote. The validation pattern: shell out to the actual service with the proposed inputs and confirm the response matches expectations BEFORE writing the code that depends on it. For policy-gated paths, run a one-shot probe before committing to a design. 15. **Two distinct auths — never confuse them**: This codebase has two unrelated authentication concerns: (a) **CLI-to-Colab-control-plane**: how `Client` authenticates HTTP requests to `colab.research.google.com` and `colab.pa.googleapis.com`. Driven by the global `--auth={oauth2,adc}` flag and `auth.py:get_credentials`. The OAuth2 path is bootstrapped automatically by `google_auth_oauthlib`'s `InstalledAppFlow` on first invocation when no token is stored in `keyring` (or `keyring` is inaccessible); **no separate command is needed** — any `--auth=oauth2` invocation (e.g. `colab --auth=oauth2 sessions`) triggers the browser consent flow. (b) **VM-side credentials**: how the `colab auth` subcommand injects user GCP credentials *into* the running Colab kernel so notebook code (e.g. `gcloud`, BigQuery client) can make authenticated calls from inside the VM. This uses the `USE_AUTH_EPHEM='0'` gcloud-fallback path executed on the kernel via `ColabRuntime`. **NEVER tell a user to "run `colab auth`" as a prerequisite for fixing CLI-side auth issues** — they're orthogonal layers and the suggestion misleads. 16. **Pydantic validation requires a schema**: `_issue_request` accepts an optional `schema=` and historically called `TypeAdapter(schema).validate_python(...)` unconditionally. When `schema=None` (a caller that doesn't care about the response body — e.g. unassign or fire-and-forget endpoints that return `[]`), this raises `pydantic.ValidationError: Input should be None`. Always guard with `if schema is None: return` after the empty-body short-circuit. @@ -83,11 +83,12 @@ As an AI agent operating via non-interactive shell tools (`run_shell_command`), - **I CAN Run:** - Automated tests (`pytest`), linting (`ruff`), and headless execution scripts. - Subcommands that don't pause for user input (e.g., `colab new`, `colab status`, `colab stop`, `colab ls`, `colab install`, `colab exec `). - - **Piped `colab repl` and `colab console`**: as of 2026-05-07 both commands support piped stdin and exit cleanly on EOF (`echo 'cmd' | colab console -s s` returns in ~1.2s). The unpiped, interactive variants still cannot be run autonomously. + - **Piped `colab repl`**: as of 2026-05-07 it supports piped stdin and exits cleanly on EOF (`echo 'print(1)' | colab repl -s s`). The unpiped, interactive variant still cannot be run autonomously. + - **One-shot shell commands**: `colab ssh -s s -- cmd` runs `cmd` non-interactively and exits with its status (add ` --help` to view specific options, defaults, and detailed he | `colab run [--gpu GPU] [--tpu TPU] [--high-mem] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it | | `colab exec [-s NAME] [-f FILE] [--output-image PATH]` | Execute Python code from stdin, a local `.py` file, or a `.ipynb` notebook | | `colab repl [-s NAME] [--output-image PATH]` | Start an interactive Python REPL on the VM (exits cleanly on piped EOF) | -| `colab console [-s NAME]` | Connect to a raw interactive TTY shell (tmux) on the remote VM | -| `colab ssh [-s NAME] [--proxy-mode] [-i KEY] [--gpu GPU] [--tpu TPU] [--high-mem]` | Open an SSH shell to the runtime over WebSocket, or act as an OpenSSH `ProxyCommand` bridge for IDE remote-dev | +| `colab ssh [-s NAME] [--proxy-mode] [-i KEY] [--gpu GPU] [--tpu TPU] [--high-mem] [-- CMD ...]` | Open an SSH shell to the runtime over WebSocket, run a one-shot remote command (`-- CMD`), or act as an OpenSSH `ProxyCommand` bridge for IDE remote-dev | ### File Operations | Command | Description | @@ -144,7 +143,7 @@ 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:** `colab repl` and interactive `colab ssh` expect a local TTY. When running inside automated scripts or pipelines, pipe stdin to `colab repl` (e.g., `echo "print(1)" | colab repl`) to trigger its non-interactive mode, or run one-shot shell commands with `colab ssh -s NAME -- CMD` (e.g., `colab ssh -s NAME -- ls /content`); pass untrusted values as separate arguments (after `--` where the program accepts it), since a single argument is a remote shell line. `colab ssh` needs a local `ssh` client and key, so where it can't work (no client or key, a tool not on the minimal ssh environment's `PATH`, or the runtime's concurrent SSH connections in use), pipe the commands to `colab exec` with IPython's `!` prefix (e.g., `echo '!ls /content' | colab exec`); each `!` line is its own shell and IPython expands `{expr}` and `$name` in it, so start a multi-line script with `%%bash` and keep untrusted text out of `!` lines. * **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/01_session_management.md b/docs/01_session_management.md index 78e8244f..86871d0b 100644 --- a/docs/01_session_management.md +++ b/docs/01_session_management.md @@ -1,5 +1,6 @@ --- log: +2026-10-29: Removed the `colab console` subcommand; the long-lived connections that skip mid-session token refresh are now `repl` and `ssh`. 2026-10-21: Stored OAuth2 credentials in the OS native keyring via the `keyring` library instead of `~/.config/colab-cli/token.json`, falling back to `_run_remote_flow` if keyring is inaccessible. Keyring failures are reported on stderr, and an existing `token.json` is migrated into the keyring and then deleted (or deleted outright if the keyring already has a token). 2026-10-14: Added configurable HTTP request timeouts and bounded allocation retries. Uncertain allocation POSTs are reconciled by GET using the same notebook hash before resubmission. 2026-09-25: Runtime proxy tokens are now refreshed before expiry. `SessionState` persists `token_expires_at`, and `State.get_session` re-fetches the token via `GET /tun/m/assignments` when it's within 5 minutes of expiring. @@ -112,7 +113,7 @@ end to end; mocked timeouts alone do not establish backend compatibility. VM liveness is automatically maintained by the Colab backend based on active kernel executions and interactions. The CLI runs no keep-alive daemon. A session persists until explicitly terminated via `colab stop` or until the backend reclaims it due to inactivity. ### 6. Runtime Proxy Token Refresh -The runtime proxy token (`runtimeProxyInfo.token`) authenticates every request to the VM and expires after `tokenExpiresInSeconds` (1 hour). `SessionState.token_expires_at` stores the absolute expiry, computed when the token is received. Commands that talk to the VM load their session through `State.get_session(name)`, which, when the token is missing an expiry or is within 5 minutes of it, calls `GET /tun/m/assignments`, takes the fresh token/URL for the matching endpoint, and persists it. If the endpoint is no longer assigned, the session is pruned. `sync_sessions` also writes back the fresh tokens it already fetches. Long-lived connections (`repl`, `console`, `ssh`) only present the token on initial connect and aren't refreshed mid-session. +The runtime proxy token (`runtimeProxyInfo.token`) authenticates every request to the VM and expires after `tokenExpiresInSeconds` (1 hour). `SessionState.token_expires_at` stores the absolute expiry, computed when the token is received. Commands that talk to the VM load their session through `State.get_session(name)`, which, when the token is missing an expiry or is within 5 minutes of it, calls `GET /tun/m/assignments`, takes the fresh token/URL for the matching endpoint, and persists it. If the endpoint is no longer assigned, the session is pruned. `sync_sessions` also writes back the fresh tokens it already fetches. Long-lived connections (`repl`, `ssh`) only present the token on initial connect and aren't refreshed mid-session. ## TODO / Future Work - **Backend Sync**: Implement a way to reconcile the local `sessions.json` with the output of `colab sessions`. diff --git a/docs/02_execution_and_interactive.md b/docs/02_execution_and_interactive.md index ed1d782a..8f80e7d1 100644 --- a/docs/02_execution_and_interactive.md +++ b/docs/02_execution_and_interactive.md @@ -5,9 +5,11 @@ log: 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. 2026-06-04: Bumped the default `--timeout` for `colab exec` from 10s to 30s (and the matching `colab run` default) so brief silent tasks are less likely to hit a premature `TimeoutError`. Explicit `--timeout` overrides are unaffected. + +2026-10-29: Removed the `colab console` subcommand and its `/colab/tty` transport (`console.py`, its tests, and `integration/repro_piped_console/`). Use `colab ssh -s NAME` for an interactive shell and `colab ssh -s NAME -- CMD` for one-shot commands, or pipe `!`-prefixed shell commands to `colab exec` where ssh can't work. --- -# Design: Execution and Interactive Interaction (`repl`, `exec`, `console`) +# Design: Execution and Interactive Interaction (`repl`, `exec`) ## Overview Execution involves sending Python code (or shell commands) to the Jupyter kernel running on the Colab VM and processing the stream of output messages. @@ -30,12 +32,6 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel - **Multi-Modal Output**: Handle `display_data` messages (e.g., `image/png`, `text/html`). For the CLI, we'll save images to temporary files and print their paths, or if the terminal supports it (e.g., iTerm2), inline them. - **Timeout Configuration**: Exposes a `--timeout` flag (default 30s) to allow long-running silent tasks (like model compilation or data downloading) to execute without being prematurely killed. -### 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. -- **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 - **Kernel Management**: `ColabRuntime` (from `colab-agent`) already handles message signing and message types. - **Output Streaming**: Continuous polling or asynchronous message handling to provide real-time output. @@ -52,6 +48,4 @@ TDD is mandatory for all execution features. ### 2. TTY and Piping - **Test Case**: Mock `sys.stdin.isatty()` to verify `colab repl` correctly switches between interactive mode and one-shot piped execution. - **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**: `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/06_ssh_access.md b/docs/06_ssh_access.md index bc12e400..15eb731f 100644 --- a/docs/06_ssh_access.md +++ b/docs/06_ssh_access.md @@ -1,5 +1,6 @@ --- log: +2026-10-29: Removed the `colab console` subcommand, so `colab ssh` is the interactive shell on the runtime. Dropped console behavior from the Windows validation caveat. 2026-10-29: Added `colab ssh [-s NAME] -- CMD`, which runs one command on an existing session and exits with its status (a PTY only when stdin and stdout are terminals; neither it nor its ProxyCommand child creates a runtime). Its handshake carries `ssh-mode-cmd`, and the update banner is skipped for `ssh` when stdout is not a terminal. Added `tests/test_ssh_cmd.py`. 2026-10-29: HTTP 429 now shows the server's SSH session limit ("N of N SSH sessions in use") when the body carries one. 2026-10-29: The ProxyCommand that `colab ssh` starts now keeps a non-default global `--auth` and `--config` (as an absolute path), so `colab --config ALT ssh -s S` bridges S from ALT instead of the default sessions file. @@ -191,8 +192,8 @@ Portable tests also check registration without SIGHUP and the local null device. Existing POSIX shell quoting and teardown tests remain in the suite. These and the existing SSH tests passed on Windows, Linux, and macOS with -Python 3.12. This covers SSH compatibility only; interactive console behavior, -local file editing, and self-update require separate platform validation. +Python 3.12. This covers SSH compatibility only; local file editing and +self-update require separate platform validation. `uv run python integration/repro_windows_ssh/test.py` provides a native Windows live test using cached OAuth2 credentials. It creates a CPU runtime, executes a diff --git a/docs/demos.md b/docs/demos.md index 69a2f415..84faf8ed 100644 --- a/docs/demos.md +++ b/docs/demos.md @@ -8,7 +8,7 @@ Eleven scenarios that exercise common workflows, plus a final "bridging back to **Methodology** - Auth: `--auth=adc`. To set up: `gcloud auth application-default login --scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory`. - Accelerator: every session uses **CPU**. Provisioning real accelerators is gated by per-account quota and would not work for most readers; the workflows themselves are accelerator-agnostic, so where a demo's narrative mentions a GPU or TPU the prose flags the substitution. -- Interactive subcommands — `colab auth`, `colab drivemount`, and unpiped `colab repl` / `colab console` — are **not run** here because they require human interaction at a TTY. Demos that would normally use them include an inline note explaining what they do and the workflow continues with the non-interactive parts. +- Interactive subcommands — `colab auth`, `colab drivemount`, interactive `colab ssh`, and unpiped `colab repl` — are **not run** here because they require human interaction at a TTY. Demos that would normally use them include an inline note explaining what they do and the workflow continues with the non-interactive parts. - `enable_update_check` is set to `false` in `~/.config/colab-cli/settings.json` for the duration of recording so the daily upgrade banner doesn't pollute output. - `PYTHONWARNINGS=ignore` is set in the environment to suppress the ADC quota-project warning that `google.auth` emits on every call from end-user credentials. @@ -138,7 +138,7 @@ uv run colab --auth=adc stop -s trainer ## Demo 3: Interactive troubleshooting (piped) -Both `colab console` and `colab repl` accept piped stdin and exit on EOF, so they compose well with shell pipelines and other CLI tools. This demo investigates remote disk usage with a one-shot shell command, lists `/content`, creates and removes a scratch file, and then queries free space from a one-shot REPL. +`colab exec` and `colab repl` accept piped stdin and exit on EOF, so they compose well with shell pipelines and other CLI tools. This demo lists `/content`, creates and removes a scratch file, and then queries free space from a one-shot REPL. ```bash uv run colab --auth=adc new -s debug @@ -149,15 +149,7 @@ uv run colab --auth=adc new -s debug [colab] Session READY. ``` -*Note:* `colab console` connects to a tmux-wrapped pty on the VM, so even when stdin is piped the raw stdout contains terminal-control bytes (cursor moves, status-line repaints, ANSI color). For programmatic consumption, pipe the output through `grep -a` (force binary-safe) and a regex matching the line(s) you care about, as shown below. - -```bash -echo 'df -h /content' | uv run colab --auth=adc console -s debug 2>&1 | grep -aE 'overlay|/dev/' -``` - -```output -overlay 108G 21G 87G 20% / -``` +*Note:* for an interactive shell on the VM use `colab ssh -s debug` (interactive, so not run here); one-shot shell commands run with `colab ssh -s debug -- df -h /content`, or, where ssh can't work, piped to `colab exec` with IPython's `!` prefix, e.g. `echo '!df -h /content' | colab exec -s debug` (neither is part of this recording). ```bash uv run colab --auth=adc ls -s debug /content @@ -272,7 +264,7 @@ uv run colab --auth=adc stop -s reporter ## Demo 5: Bulk data via GCS -Pull a batch of objects down from a Google Cloud Storage bucket, transform them on the VM, and pull the results back. A full-fidelity workflow is `colab new --gpu L4` -> `colab auth` (so VM-side `gcloud` works) -> `gcloud storage cp gs://bucket/raw/*.jpg /content/images/` (via piped `colab console`) -> install pillow/torchvision -> process -> download. We skip the auth step here (interactive; the user has to click through OAuth) and substitute synthetic image generation in place of the GCS pull, which keeps the input -> install -> batch-process -> download shape intact. +Pull a batch of objects down from a Google Cloud Storage bucket, transform them on the VM, and pull the results back. A full-fidelity workflow is `colab new --gpu L4` -> `colab auth` (so VM-side `gcloud` works) -> `gcloud storage cp gs://bucket/raw/*.jpg /content/images/` (as a `!`-prefixed shell command piped to `colab exec`) -> install pillow/torchvision -> process -> download. We skip the auth step here (interactive; the user has to click through OAuth) and substitute synthetic image generation in place of the GCS pull, which keeps the input -> install -> batch-process -> download shape intact. ```bash uv run colab --auth=adc new -s data-proc diff --git a/integration/README.md b/integration/README.md index 3588f5ec..5ad646d2 100644 --- a/integration/README.md +++ b/integration/README.md @@ -14,7 +14,6 @@ End-to-end tests that run against a **live Colab backend** (unlike the mocked un | `repro_plot_redirection/` | `colab exec` of a matplotlib script with `--output-image` redirection. | | `repro_run_command/` | `colab run ` runs a local script on a fresh VM and cleans up; `--keep` leaves the session alive. | | `repro_variable_persistence/` | Variables persist across `colab exec` calls in the same session. | -| `repro_piped_console/` | Fast smoke test (~5s including session creation): `echo cmd \| colab console -s s` runs the command and exits within 30s. Regression test for the 2026-05-07 EOF-handler fix. | | `repro_bundled_oauth/` | Fast smoke test (~5s): verifies that the fallback OAuth configuration is loaded and starts the OAuth flow with the default client ID when local config is missing. | | `repro_legacy_token_migration/` | Fast smoke test (~5s, no VM): a pre-keyring `~/.config/colab-cli/token.json` is migrated into the keyring and deleted, deleted when the keyring already has a token, and kept (with a stderr warning) when the keyring is inaccessible. Uses a throwaway HOME and file-backed test keyring. | | `repro_ssh/` | Fast smoke test (~5s): `--help` advertises the flags and an unknown session exits. Slow soak test (~95s): Live e2e allocates a CPU VM, runs a real remote command over `colab ssh --proxy-mode` | diff --git a/integration/repro_piped_console/test.sh b/integration/repro_piped_console/test.sh deleted file mode 100755 index fe21ca80..00000000 --- a/integration/repro_piped_console/test.sh +++ /dev/null @@ -1,96 +0,0 @@ -#!/bin/bash -# 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. -# -# Integration Test: Piped console exits cleanly on EOF -# Verifies that `echo cmd | colab console -s s` runs the command on the remote -# /colab/tty endpoint and then exits within a few seconds. -# -# Regression: prior to 2026-05-07 this hung indefinitely because the EOF -# handler sent \x04 (Ctrl-D), which the remote tmux-wrapped bash treated as a -# literal character rather than a session terminator. The fix sends "exit\n" -# and closes the websocket from the client side after a short grace period. - -set -e - -TMP_DIR=$(mktemp -d) -SESSION_FILE="$TMP_DIR/sessions.json" -SESSION_NAME="test-piped-console" - -# OAuth2 is usable if a token is cached in the OS keyring. -has_oauth2_token() { - timeout 10 uv run python -c "import sys, keyring; sys.exit(0 if keyring.get_password('google-colab-cli', 'oauth_token') else 1)" > /dev/null 2>&1 -} -# Auth selection (OAuth2 if token cached, else ADC). -if has_oauth2_token; then - AUTH_FLAGS="--auth=oauth2" -elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then - ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) - ADC_SCOPES=$(curl -s "https://www.googleapis.com/oauth2/v3/tokeninfo?access_token=$ADC_TOKEN" | python3 -c "import json,sys; print(json.load(sys.stdin).get('scope',''))" 2>/dev/null) - if echo "$ADC_SCOPES" | grep -q "userinfo.email"; then - AUTH_FLAGS="--auth=adc" - else - echo "Error: ADC token lacks the userinfo.email scope." - exit 1 - fi -else - echo "Error: No usable auth provider found." >&2 - exit 1 -fi - -cleanup_session() { - uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" 2>/dev/null || true -} -trap "cleanup_session; rm -rf $TMP_DIR" EXIT - -echo "[*] Creating session for piped-console test using $AUTH_FLAGS..." -uv run colab $AUTH_FLAGS --config "$SESSION_FILE" new -s "$SESSION_NAME" - -# Marker string we'll grep for in the captured output. Made unique enough that -# accidental matches in tmux status lines or shell prompts are impossible. -MARKER="PIPED-CONSOLE-OK-$(date +%s)-$$" - -OUTPUT_FILE="$TMP_DIR/console-output.txt" - -echo "[*] Running: echo 'echo $MARKER' | colab console (must exit within 30s)..." -START=$(date +%s) -# 30s upper bound guards against the prior hang regression. The fix typically -# completes in ~1-2s; anything beyond that means the EOF/close path broke. -if ! timeout 30 bash -c "echo 'echo $MARKER' | uv run colab $AUTH_FLAGS --config '$SESSION_FILE' console -s '$SESSION_NAME'" > "$OUTPUT_FILE" 2>&1; then - EXIT=$? - if [ $EXIT -eq 124 ]; then - echo "[FAILURE] Piped console hung past the 30s timeout (regression — was the EOF handler reverted?)." - else - echo "[FAILURE] Piped console exited non-zero: $EXIT" - fi - cat "$OUTPUT_FILE" - exit 1 -fi -ELAPSED=$(($(date +%s) - START)) -echo "[*] Piped console returned in ${ELAPSED}s." - -if ! grep -q "$MARKER" "$OUTPUT_FILE"; then - echo "[FAILURE] Marker '$MARKER' not found in console output (the command did not actually execute on the VM)." - cat "$OUTPUT_FILE" - exit 1 -fi -echo "[*] Confirmed: the piped command actually ran on the remote VM." - -# Sanity: the doc claims ~1.2s round-trip. Anything > 10s is suspicious even -# if it eventually exited (probably means the grace window was bumped wildly). -if [ "$ELAPSED" -gt 10 ]; then - echo "[WARNING] Piped console took ${ELAPSED}s — slow, but not a hang." -fi - -echo "[SUCCESS] Piped console integration test passed (elapsed=${ELAPSED}s)." diff --git a/skills/colab-operator/SKILL.md b/skills/colab-operator/SKILL.md index c8892d70..fe41419b 100644 --- a/skills/colab-operator/SKILL.md +++ b/skills/colab-operator/SKILL.md @@ -53,8 +53,8 @@ by running `uv tool install google-colab-cli` or `pip install google-colab-cli`. - **Piped code**: `echo "print(1)" | colab exec -s ` or `cat script.py | colab exec -s `. - **Notebooks**: `colab exec -s -f nb.ipynb` runs each code cell and writes results to `_output.ipynb` next to the input. A `# @title Foo` first line labels the cell in progress output. - **Plots/images**: PNG/JPEG outputs are intercepted. Use `--output-image ` on `exec`/`repl` to save to a known location (otherwise a temp path is printed). Inline terminal-image escapes are auto-suppressed when stdout isn't a TTY, so piped/captured output stays clean. -- **Shell**: `echo "cmd" | colab console -s ` for batch shell. Console wraps bash in tmux, so even piped output contains terminal-control bytes — filter with `grep -a` for a specific line. `exec` is faster when you don't need a real shell. -- **Never run `colab repl`, `colab console`, `colab auth`, or `colab drivemount` interactively from an agent** — they expect a TTY and will hang. `repl`/`console` accept piped stdin and exit on EOF; `auth`/`drivemount` genuinely require a human at the terminal. +- **Shell**: run a one-shot shell command with `colab ssh -s -- `: it runs in `/content` and exits with the command's status. Output is byte-exact when stdout is piped or redirected; ssh allocates a PTY only when stdin and stdout are both terminals, so a shell tool that runs commands in a terminal should add ` -- 'nvidia-smi | head -5'`); several arguments are an exact argv, so pass untrusted values as separate arguments after `--` where the program accepts it (e.g. `colab ssh -s -- ls -- "$f"`), never inside a shell line or a piped script. Pipe a multi-line script to `colab ssh -s -- bash`. Use `colab exec` for Python code. When ssh can't work (no local `ssh` client, no `~/.ssh/id_ed25519` or `id_ecdsa` key and no `-i KEY`, a tool not on the minimal ssh environment's `PATH` such as `gcloud` or `node`, or the runtime's limited concurrent SSH connections in use, which the server answers with 429, e.g. while a person or an IDE is connected), fall back to `exec` with IPython's `!` prefix: `echo '!df -h /content' | colab exec -s ` (single quotes stop your local shell expanding `!`). Each `!` line runs in its own shell with no input (a command that prompts hangs), and IPython evaluates `{expr}` as Python even inside quotes and expands `$name` outside single quotes. Start a multi-line script with a `%%bash` line (not with `--env`, which prepends Python lines), use non-interactive flags such as `apt-get -y`, and never put untrusted text in a `%%bash` script or a `!` line (even via `{var}` or `$var`): pass it with `--env NAME=VALUE` and run `subprocess.run(["cmd", "--", os.environ["NAME"]])` from Python. For an interactive shell on the VM, a human runs `colab ssh -s `. +- **Never run `colab repl`, `colab ssh` without `-- CMD`, `colab auth`, or `colab drivemount` interactively from an agent** — `repl`, `auth` and `drivemount` expect a TTY and will hang or fail, and bare `colab ssh` with no saved session creates a billable runtime; an interactive shell is for a human. `repl` accepts piped stdin and exits on EOF; `auth`/`drivemount` genuinely require a human at the terminal. ### Ephemeral one-shot jobs (`colab run`) - `colab run [--gpu T4] [--tpu v6e1] [--keep] [-s NAME] script.py [args...]` = `new` + `exec` + `stop` in one command. It provisions a fresh VM, runs the script with `sys.argv` and `__name__ == "__main__"` set like native `python script.py args`, then tears the VM down (unless `--keep`). diff --git a/src/colab_cli/commands/execution.py b/src/colab_cli/commands/execution.py index 71784195..71c0cd91 100644 --- a/src/colab_cli/commands/execution.py +++ b/src/colab_cli/commands/execution.py @@ -26,7 +26,6 @@ from colab_cli.runtime import ColabRuntime from colab_cli.utils import handle_image, is_terminal_error, render_display_data -from colab_cli.console import connect_console _console = Console() @@ -375,35 +374,6 @@ def on_sess_started(sid): state.store.add(s) -def console( - session: Annotated[ - Optional[str], typer.Option("-s", "--session", help="Session name") - ] = None, -): - """Connect to raw TTY console""" - from colab_cli.common import state - - name = state.resolve_session(session) - s = state.get_session(name) - state.history.log_event(s.name, "console_started", {}) - s.running = "console" - state.store.add(s) - try: - connect_console(s) - except Exception as e: - if is_terminal_error(e): - typer.echo( - f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." - ) - state.prune_session(name) - raise typer.Exit(1) - raise e - finally: - s.running = None - state.store.add(s) - - def register(app: typer.Typer): app.command(name="exec")(exec_command) app.command()(repl) - app.command()(console) diff --git a/src/colab_cli/console.py b/src/colab_cli/console.py deleted file mode 100644 index ae2786bb..00000000 --- a/src/colab_cli/console.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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. - -import json -import logging -import os -import signal -import sys -import threading -import time - -try: - import termios - import tty -except ImportError: - # Windows has neither termios nor tty. connect_console() below degrades - # gracefully to line-buffered input when these are unavailable. - termios = None # type: ignore - tty = None # type: ignore -from urllib.parse import urlparse - -import websocket - -from colab_cli.state import SessionState - -logger = logging.getLogger(__name__) - -# Global flag to stop the read thread when the websocket closes -_is_running = False -_last_error = None - -# When stdin is piped and reaches EOF, we send "exit\n" to the remote shell and -# then wait this many seconds for any remaining output (the shell's goodbye, -# tmux teardown messages, etc.) to flush before closing the websocket from the -# client side. Empirically 0.5s is enough for the typical /colab/tty backend -# wrapped in tmux + bash; bumping it just delays exit, lowering it risks -# truncating tail output. -PIPED_EOF_GRACE_SECONDS = 0.5 - - -def on_message(ws, message): - """Callback for when a message is received from the server.""" - try: - data = json.loads(message) - if "data" in data: - # The backend sends raw ANSI escape sequences and string content. - # We write it directly to stdout buffer to avoid python print() formatting. - sys.stdout.buffer.write(data["data"].encode("utf-8")) - sys.stdout.buffer.flush() - except Exception as e: - logger.debug(f"Error parsing message: {e}") - - -def on_error(ws, error): - """Callback for when a websocket error occurs.""" - global _last_error - _last_error = error - logger.error(f"WebSocket Error: {error}") - - -def on_close(ws, close_status_code, close_msg): - """Callback for when the websocket is closed.""" - global _is_running - _is_running = False - - -def send_terminal_size(ws): - """Sends the current terminal size to the remote backend.""" - try: - size = os.get_terminal_size() - payload = json.dumps({"cols": size.columns, "rows": size.lines}) - ws.send(payload) - except Exception as e: - logger.debug(f"Failed to send terminal size: {e}") - - -def on_open(ws): - """Callback for when the websocket connection is opened.""" - global _is_running - _is_running = True - - # Send initial terminal size - send_terminal_size(ws) - - # Setup the background thread to read from stdin - def read_stdin(): - is_tty = sys.stdin.isatty() - 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})) - except Exception: - break - - thread = threading.Thread(target=read_stdin, daemon=True) - thread.start() - - -def connect_console(session: SessionState): - """ - Connects to the Colab TTY endpoint and sets up a raw terminal session. - """ - global _is_running, _last_error - _last_error = None - - # Construct the WebSocket URL from the base URL - parsed = urlparse(session.url) - ws_scheme = "wss" if parsed.scheme == "https" else "ws" - ws_url = f"{ws_scheme}://{parsed.netloc}/colab/tty?colab-runtime-proxy-token={session.token}" - - is_tty = sys.stdin.isatty() - fd = None - old_settings = None - can_raw = termios is not None and tty is not None and is_tty - if can_raw: - try: - fd = sys.stdin.fileno() - old_settings = termios.tcgetattr(fd) - except Exception: - fd = None - old_settings = None - can_raw = False - - ws = websocket.WebSocketApp( - url=ws_url, - on_open=on_open, - on_message=on_message, - on_error=on_error, - on_close=on_close, - ) - - def handle_sigwinch(signum, frame): - """Handle window resize events.""" - if _is_running: - 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) - - # 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: - if can_raw: - # Always ensure the terminal is restored to its original state - try: - termios.tcsetattr(fd, termios.TCSANOW, old_settings) - except Exception: - pass - # Restore the default signal handler for resize - if sigwinch is not None: - try: - signal.signal(sigwinch, signal.SIG_DFL) - except Exception: - pass - print("\r\nConnection closed.") diff --git a/src/colab_cli/converter.py b/src/colab_cli/converter.py index e19aa2de..c015a970 100644 --- a/src/colab_cli/converter.py +++ b/src/colab_cli/converter.py @@ -105,8 +105,6 @@ def convert_history_to_ipynb( elif etype == "execution": code = event.get("code", "") - # Check for shell commands (starting with ! or from piped console) - # If it's a raw shell command from a console pipe, wrap it in %%bash if it doesn't have ! if event.get("source") == "piped" and not code.startswith("!"): code = "%%bash\n" + code diff --git a/tests/test_cli.py b/tests/test_cli.py index b7d08ef2..d10ef2db 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -16,7 +16,9 @@ from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch +import click import pytest +import typer from typer.testing import CliRunner from colab_cli.cli import app @@ -384,7 +386,7 @@ def _extract_command_names(help_output: str) -> list[str]: inside the `╭─ Commands ─...` rich box. Returns names in the order they appear. """ - lines = help_output.splitlines() + lines = click.unstyle(help_output).splitlines() in_commands = False names = [] for line in lines: @@ -437,18 +439,37 @@ def test_cli_no_args(): assert result.exit_code in [0, 2] -def test_cli_console(mock_store, mock_common_state): - mock_session_state = MagicMock() - mock_session_state.name = "s1" - mock_session_state.token = "t1" - mock_session_state.url = "http://test.com" - mock_store.get.return_value = mock_session_state +def test_cli_console_removed(): + """`colab console` stays removed: unregistered, absent from --help, rejected.""" + assert "console" not in typer.main.get_command(app).commands + + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0, result.output + names = _extract_command_names(result.output) + assert names, result.output + assert "ssh" in names, result.output + assert "console" not in names, result.output + + result = runner.invoke(app, ["console"]) + assert result.exit_code == 2, result.output + assert "No such command" in result.output, result.output + +def test_cli_exec_stdin_forwards_shell_escape(mock_store, mock_common_state): + mock_session = MagicMock() + mock_session.name = "s1" + mock_session.kernel_id = None + mock_session.session_id = None + mock_store.get.return_value = mock_session mock_common_state.resolve_session.return_value = "s1" - with patch("colab_cli.commands.execution.connect_console") as mock_connect: - result = runner.invoke(app, ["console", "-s", "s1"]) - assert result.exit_code == 0 - mock_connect.assert_called_once_with(mock_session_state) + + with patch("colab_cli.commands.execution.ColabRuntime") as mock_runtime_class: + execute_code = mock_runtime_class.return_value.execute_code + execute_code.return_value = [] + result = runner.invoke(app, ["exec", "-s", "s1"], input="!echo hi") + + assert result.exit_code == 0, result.output + assert execute_code.call_args.args[0] == "!echo hi" @patch("colab_cli.commands.files.ContentsClient") diff --git a/tests/test_console.py b/tests/test_console.py deleted file mode 100644 index 13f20961..00000000 --- a/tests/test_console.py +++ /dev/null @@ -1,261 +0,0 @@ -# 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. - -import json -import os -import sys -from unittest.mock import MagicMock, patch - -try: - import termios -except ImportError: - # Windows has no termios; tests requiring raw-terminal mode are skipped. - termios = None # type: ignore - -from colab_cli.console import connect_console, on_message, on_open -from colab_cli.state import SessionState -import pytest - -needs_termios = pytest.mark.skipif( - termios is None, reason="requires termios (POSIX only)" -) - - -@pytest.fixture -def mock_session(): - return SessionState( - name="test-session", - token="test-token", - url="https://8080-m-s-kkb-usc1f1.us-central1-1.colab.dev", - endpoint="some-endpoint", - ) - - -@needs_termios -@patch("colab_cli.console.websocket.WebSocketApp") -@patch("colab_cli.console.tty.setraw") -@patch("colab_cli.console.termios.tcgetattr") -@patch("colab_cli.console.termios.tcsetattr") -@patch("colab_cli.console.os.get_terminal_size") -@patch("colab_cli.console.sys.stdin.fileno") -@patch("colab_cli.console.sys.stdin.isatty") -def test_console_initialization( - mock_isatty, - mock_fileno, - mock_get_term_size, - mock_tcsetattr, - mock_tcgetattr, - mock_setraw, - mock_ws_app, - mock_session, -): - # Setup mocks - mock_isatty.return_value = True - mock_fileno.return_value = 0 - mock_get_term_size.return_value = os.terminal_size((80, 24)) - mock_tcgetattr.return_value = ["fake_attrs"] - mock_ws_instance = MagicMock() - mock_ws_app.return_value = mock_ws_instance - - # We don't want run_forever to actually block or start threads in the test - mock_ws_instance.run_forever.return_value = None - - with patch("colab_cli.console.threading.Thread"): - connect_console(mock_session) - - # 1. Verify URL transformation - expected_url = "wss://8080-m-s-kkb-usc1f1.us-central1-1.colab.dev/colab/tty?colab-runtime-proxy-token=test-token" - mock_ws_app.assert_called_once() - assert mock_ws_app.call_args[1]["url"] == expected_url - - # 2. Verify raw mode setup and teardown - mock_tcgetattr.assert_called_once_with(sys.stdin.fileno()) - mock_setraw.assert_called_once_with(sys.stdin.fileno(), termios.TCSANOW) - - # Teardown should happen in a finally block - mock_tcsetattr.assert_called_once_with( - sys.stdin.fileno(), termios.TCSANOW, ["fake_attrs"] - ) - - -@needs_termios -@patch("colab_cli.console.websocket.WebSocketApp") -@patch("colab_cli.console.tty.setraw") -@patch("colab_cli.console.termios.tcgetattr") -@patch("colab_cli.console.termios.tcsetattr") -@patch("colab_cli.console.sys.stdin.isatty") -def test_console_piped_input( - mock_isatty, - mock_tcsetattr, - mock_tcgetattr, - mock_setraw, - mock_ws_app, - mock_session, -): - mock_isatty.return_value = False - mock_ws_instance = MagicMock() - mock_ws_app.return_value = mock_ws_instance - mock_ws_instance.run_forever.return_value = None - - with patch("colab_cli.console.threading.Thread"): - connect_console(mock_session) - - # In a piped environment, we should not attempt to use termios or tty - mock_tcgetattr.assert_not_called() - mock_setraw.assert_not_called() - mock_tcsetattr.assert_not_called() - - -@patch("colab_cli.console.os.get_terminal_size") -def test_on_open_sends_terminal_size(mock_get_term_size): - mock_ws = MagicMock() - mock_get_term_size.return_value = os.terminal_size((100, 40)) - - on_open(mock_ws) - - # Verify that the initial terminal size is sent - mock_ws.send.assert_called_once() - payload = json.loads(mock_ws.send.call_args[0][0]) - assert payload == {"cols": 100, "rows": 40} - - -@patch("colab_cli.console.sys.stdout.buffer.write") -@patch("colab_cli.console.sys.stdout.buffer.flush") -def test_on_message_writes_to_stdout(mock_flush, mock_write): - mock_ws = MagicMock() - test_data = "Hello \x1b[34mWorld\x1b[0m" - message_json = json.dumps({"data": test_data}) - - on_message(mock_ws, message_json) - - # Verify that the data is written exactly as received - mock_write.assert_called_once_with(test_data.encode("utf-8")) - mock_flush.assert_called_once() - - -@patch("colab_cli.console.os.get_terminal_size") -@patch("colab_cli.console.sys.stdin.isatty") -@patch("colab_cli.console.sys.stdin") -def test_read_stdin_eof_piped_sends_exit_and_closes_ws( - mock_stdin, mock_isatty, mock_get_term_size -): - """When stdin is piped and reaches EOF, the read thread should send 'exit\\n' - to the remote shell and then close the websocket from the client side. - - The remote shell at /colab/tty is wrapped in tmux which swallows the bare - \\x04 (Ctrl-D) we used to send, so EOF used to leave the websocket open - indefinitely. Sending 'exit\\n' + ws.close() guarantees clean termination. - """ - import colab_cli.console as console_mod - - mock_isatty.return_value = False - # Simulate piped stdin: returns one line then EOF - mock_stdin.read.side_effect = ["e", "c", "h", "o", " ", "h", "i", "\n", ""] - mock_get_term_size.return_value = os.terminal_size((80, 24)) - - mock_ws = MagicMock() - - # on_open spawns the read thread; we want it to run synchronously here - # so we patch threading.Thread to call target immediately and join(). - real_thread = [] - - class SyncThread: - def __init__(self, target, daemon=None): - self.target = target - real_thread.append(self) - - def start(self): - self.target() - - console_mod._is_running = True - with patch("colab_cli.console.threading.Thread", SyncThread): - # Use a tiny grace period for the test - with patch("colab_cli.console.PIPED_EOF_GRACE_SECONDS", 0.01): - on_open(mock_ws) - - # Collect what was sent to the websocket - sent_payloads = [json.loads(c.args[0]) for c in mock_ws.send.call_args_list] - - # Initial send is the terminal size; everything after is stdin chars or our exit string. - # Verify "exit\n" was sent on EOF (one send per character) - assert {"data": "exit\n"} in sent_payloads, ( - f"Expected 'exit\\n' to be sent on piped EOF, got: {sent_payloads}" - ) - - # Verify we closed the websocket from the client side - mock_ws.close.assert_called_once() - - -@patch("colab_cli.console.os.get_terminal_size") -@patch("colab_cli.console.sys.stdin.isatty") -@patch("colab_cli.console.sys.stdin") -def test_read_stdin_eof_tty_does_not_close_ws( - mock_stdin, mock_isatty, mock_get_term_size -): - """When stdin is a real TTY and read() returns empty (which happens on - Ctrl-D in raw mode), we should NOT inject 'exit\\n' or close the websocket - \u2014 the user is in interactive mode and may have intended Ctrl-D as a literal - char. The websocket lifecycle is owned by the remote shell in this case. - """ - import colab_cli.console as console_mod - - mock_isatty.return_value = True - # TTY EOF is rare but possible; should be passed through transparently - mock_stdin.read.side_effect = [""] - mock_get_term_size.return_value = os.terminal_size((80, 24)) - - 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("colab_cli.console.threading.Thread", SyncThread): - on_open(mock_ws) - - sent_payloads = [json.loads(c.args[0]) for c in mock_ws.send.call_args_list] - assert {"data": "exit\n"} not in sent_payloads - mock_ws.close.assert_not_called() - - -@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 -): - """Windows has no termios/tty/SIGWINCH. connect_console must not crash. - - Regression test for `ModuleNotFoundError: No module named 'termios'` - which broke every `colab` command on Windows via - cli -> execution -> console imports. - """ - 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 - - with ( - patch.object(console_mod, "termios", None), - patch.object(console_mod, "tty", None), - patch("colab_cli.console.threading.Thread"), - ): - connect_console(mock_session) - - mock_ws_instance.run_forever.assert_called_once() From 7006c186b7d8e39995c7e228cbb70570303cb0e1 Mon Sep 17 00:00:00 2001 From: Drake Aiman Date: Tue, 3 Nov 2026 17:24:40 +0000 Subject: [PATCH 2/2] test: move the exec `!` pin to test_exec.py and use Click's usage-error code The `!echo hi` passthrough test now sits next to the other `colab exec` stdin tests and uses their runtime fixture. The console regression test drops a docstring that restated it and checks Click's named usage-error exit code instead of a bare 2. --- tests/test_cli.py | 20 +------------------- tests/test_exec.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index d10ef2db..28df3026 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -440,7 +440,6 @@ def test_cli_no_args(): def test_cli_console_removed(): - """`colab console` stays removed: unregistered, absent from --help, rejected.""" assert "console" not in typer.main.get_command(app).commands result = runner.invoke(app, ["--help"]) @@ -451,27 +450,10 @@ def test_cli_console_removed(): assert "console" not in names, result.output result = runner.invoke(app, ["console"]) - assert result.exit_code == 2, result.output + assert result.exit_code == click.UsageError.exit_code, result.output assert "No such command" in result.output, result.output -def test_cli_exec_stdin_forwards_shell_escape(mock_store, mock_common_state): - mock_session = MagicMock() - mock_session.name = "s1" - mock_session.kernel_id = None - mock_session.session_id = None - mock_store.get.return_value = mock_session - mock_common_state.resolve_session.return_value = "s1" - - with patch("colab_cli.commands.execution.ColabRuntime") as mock_runtime_class: - execute_code = mock_runtime_class.return_value.execute_code - execute_code.return_value = [] - result = runner.invoke(app, ["exec", "-s", "s1"], input="!echo hi") - - assert result.exit_code == 0, result.output - assert execute_code.call_args.args[0] == "!echo hi" - - @patch("colab_cli.commands.files.ContentsClient") def test_cli_ls(mock_contents_class, mock_store, mock_common_state): mock_session_state = MagicMock() diff --git a/tests/test_exec.py b/tests/test_exec.py index 203405c2..bd0d6477 100644 --- a/tests/test_exec.py +++ b/tests/test_exec.py @@ -87,6 +87,24 @@ def test_cli_exec_stdin(mock_store, mock_runtime_class, mock_common_state): ) +def test_cli_exec_stdin_forwards_shell_escape( + mock_store, mock_runtime_class, mock_common_state +): + mock_session = MagicMock() + mock_session.name = "s1" + mock_session.kernel_id = None + mock_session.session_id = None + mock_store.get.return_value = mock_session + + mock_common_state.resolve_session.return_value = "s1" + mock_runtime = mock_runtime_class.return_value + mock_runtime.execute_code.return_value = [] + + result = runner.invoke(app, ["exec", "-s", "s1"], input="!echo hi") + assert result.exit_code == 0 + mock_runtime.execute_code.assert_any_call("!echo hi", output_hook=ANY, timeout=30.0) + + def test_cli_exec_env_injects_prelude( mock_store, mock_runtime_class, mock_common_state ):