| 1 | --- |
| 2 | log: |
| 3 | 2026-08-09: Added `--high-mem` passthrough when `colab ssh` auto-creates a runtime (forwards to `colab new --high-mem`). |
| 4 | 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). |
| 5 | 2026-07-22: Bare `colab ssh` now auto-creates a runtime (via `colab new`) when you have no active session, with `--gpu/--tpu` passthrough and `--rm` to stop an auto-created runtime on exit. Fixed two client bugs: the dead 403 branch (feature-off returns 404, not 403) and the RSA guidance (all `ssh-rsa` keys are server-rejected, so `id_rsa` is no longer auto-scanned and the 400 message no longer advertises `rsa-sha2`). Added `tests/test_ssh_wire_contract.py` (real loopback-server wire assertions) and `tests/test_ssh_autocreate.py`. |
| 6 | 2026-07-22: Interactive `colab ssh` now starts in `/content` (Colab's working dir) instead of `/root`, via a forced PTY (`-t`) plus a remote `cd /content 2>/dev/null; exec $SHELL -l`. A missing `/content` falls back to the login home. Added `tests/test_ssh_workdir.py`. |
| 7 | 2026-07-22: `--proxy-mode` now honors every `colab ssh` flag: with `-s NAME` it creates the session if missing (creation output routed to stderr so stdout stays the clean ssh byte stream), `--gpu/--tpu` set the accelerator, and `--rm` stops the bridged session on disconnect — so `~/.ssh/config` hosts work on first connect and can be made ephemeral. Removed the `--drive` subfeature entirely (code + tests). |
| 8 | 2026-07-22: Fixed `--proxy-mode --rm` not tearing down on disconnect. OpenSSH sends the ProxyCommand SIGHUP (verified empirically) when the session ends — not just stdin EOF — and Python's default SIGHUP action terminated the process before the teardown `finally` ran, leaking the runtime + keep-alive daemon. Now `--rm` installs SIGHUP/SIGTERM/SIGINT handlers that run the stop (idempotent with the `finally`). |
| 9 | 2026-07-23: Applied go/pystyle readability to the ssh code (80-col reflow, Args/Returns/Raises docstrings) and upgraded the integration test from a `--help` smoke into a real end-to-end: it drives a live remote command over `colab ssh --proxy-mode` used as an OpenSSH ProxyCommand (handshake + pubkey-header auth + bridge + remote exec), asserts the RSA-key rejection, and verifies no orphan VM — plus an always-on offline check (help flags + unknown-session exit 2). The live part now auto-runs when auth is present instead of being `RUN_LIVE`-gated. |
| 10 | 2026-07-24: Refactored `ssh()` into intent-named helpers (`_select_proxy_session`, `_select_interactive_session`, `_run_proxy_bridge`, `_run_interactive_shell`, `_install_rm_signal_handlers`, `_warn_accelerator_ignored`) — behavior-preserving — and unified the two `--gpu/--tpu ignored` messages into one. Added `tests/test_ssh_lifecycle.py` pinning lifecycle guarantees: `--rm` teardown survives an exception (try/finally), ssh/bridge exit-code propagation, `--proxy-mode` stdout cleanliness (create/`--rm` chatter stays on stderr), auto-create failure aborts before connect, `--gpu`+`--tpu` both forwarded to `colab new`, `--rm` idempotency across the signal + finally paths, and reused-session `--rm` teardown. Trimmed prose in this doc + the integration README. |
| 11 | --- |
| 12 | |
| 13 | # Design: `colab ssh` — SSH-over-WebSocket Runtime Access |
| 14 | |
| 15 | ## Motivation |
| 16 | Users want a real shell on their Colab runtime and, more importantly, IDE |
| 17 | remote-development (VS Code Remote-SSH, JetBrains Gateway, plain `ssh`). `colab ssh` is |
| 18 | the client that allows sshing into Colab, reusing the CLI's existing session resolution and |
| 19 | runtime-proxy token so no separate credential handling is needed. |
| 20 | |
| 21 | ## User Surface |
| 22 | |
| 23 | ``` |
| 24 | colab ssh [OPTIONS] |
| 25 | ``` |
| 26 | |
| 27 | | Flag | Type | Default | Purpose | |
| 28 | |---|---|---|---| |
| 29 | | `-s`, `--session` | str | auto | Session to connect to. If omitted, uses your only active session, auto-creates one when you have none, or errors when you have several. | |
| 30 | | `--proxy-mode` | bool | False | Act as an OpenSSH `ProxyCommand`-compatible WebSocket↔stdio bridge (reads stdin, writes stdout) for `~/.ssh/config`. Every other flag still applies. | |
| 31 | | `-i`, `--identity` | str | auto | Private key for the public key sent to Colab. Default: first of `~/.ssh/id_ed25519`, `id_ecdsa`. | |
| 32 | | `--gpu` | str | None | GPU accelerator for a runtime this command creates (T4, L4, G4, H100, A100). | |
| 33 | | `--tpu` | str | None | TPU accelerator for a runtime this command creates (v5e1, v6e1). | |
| 34 | | `--high-mem` | bool | False | Request high-RAM when this command auto-creates a runtime (ignored when connecting to an existing session). | |
| 35 | | `--rm` | bool | False | Stop the runtime when the session ends. Interactive: only a runtime `colab ssh` auto-created (a reused session is never removed). `--proxy-mode`: the bridged session, on disconnect. | |
| 36 | |
| 37 | ### `~/.ssh/config` usage |
| 38 | `--proxy-mode` turns `colab ssh` into a transport any SSH-based tool can drive: |
| 39 | |
| 40 | ``` |
| 41 | Host <alias> |
| 42 | ProxyCommand <abs-path-to>/colab ssh --proxy-mode -s <name> [--gpu T4] [--rm] |
| 43 | User root |
| 44 | StrictHostKeyChecking no |
| 45 | UserKnownHostsFile /dev/null |
| 46 | ``` |
| 47 | |
| 48 | Because every flag applies in `--proxy-mode`, `-s <name>` creates the session on |
| 49 | first connect, `--gpu/--tpu` size it, and `--rm` makes the host ephemeral. Use an |
| 50 | **absolute** `colab` path: `ssh` runs the `ProxyCommand` in a non-login shell |
| 51 | where a bare `colab` may not be on `PATH`. External SSH tools run their own |
| 52 | remote command, so to also land in `/content` add `RequestTTY yes` and |
| 53 | `RemoteCommand cd /content 2>/dev/null; exec bash -l`. |
| 54 | |
| 55 | ## Behavior |
| 56 | |
| 57 | 1. **Session resolution / auto-create**: With `-s NAME`, resolves that session |
| 58 | (via `state.resolve_session`, the same helper the other commands use). Bare |
| 59 | `colab ssh` uses your only active session; with **no** session it auto-creates |
| 60 | one (mirrors `colab new` end-to-end: assign → keep-alive pre-flight → spawn |
| 61 | keep-alive daemon → persist `SessionState`); with **multiple** it errors and |
| 62 | asks you to pick one with `-s`. |
| 63 | 2. **Connect**: Opens the WebSocket to `wss://<netloc>/colab/ssh?colab-runtime-proxy-token=<token>` |
| 64 | and sends the resolved public key verbatim in the `X-Colab-Ssh-Pubkey` header |
| 65 | (no transformation -- the bytes the user controls are exactly what the server |
| 66 | receives). Only `ssh-ed25519` / `ecdsa-sha2-nistp{256,384,521}` keys are |
| 67 | accepted. |
| 68 | 3. **Interactive shell**: Spawns the system `ssh` binary with the CLI re-invoked |
| 69 | as its own `ProxyCommand` (`python -m colab_cli.cli ssh --proxy-mode`), so the |
| 70 | WebSocket bridge and the interactive shell share one code path. It forces a |
| 71 | PTY (`-t`) and runs `cd /content 2>/dev/null; exec $SHELL -l` so you land in |
| 72 | `/content` (Colab's working dir) rather than root's home; a missing `/content` |
| 73 | falls back to the login home. |
| 74 | 4. **`--proxy-mode` bridge**: Bridges the WebSocket ↔ stdin/stdout for use as an |
| 75 | OpenSSH `ProxyCommand`. Honors every flag: `-s NAME` creates the session if |
| 76 | missing (creation output routed to stderr so stdout stays the clean ssh byte |
| 77 | stream); bare `--proxy-mode` with no `-s` just resolves an existing session. |
| 78 | 5. **`--rm` teardown**: Stops the runtime when the session ends. In `--proxy-mode` |
| 79 | this must survive how OpenSSH ends a `ProxyCommand`: on disconnect it sends |
| 80 | **SIGHUP** (verified), not just stdin EOF, and Python's default SIGHUP action |
| 81 | would terminate the process before the teardown `finally` ran — leaking the |
| 82 | runtime and its keep-alive daemon. `--rm` therefore installs |
| 83 | SIGHUP/SIGTERM/SIGINT handlers that run the stop, idempotent with the |
| 84 | `finally`. `SIGKILL` cannot be intercepted, so a `kill -9`/hard crash can |
| 85 | still leak; a normal disconnect is SIGHUP and is handled. |
| 86 | 6. **Error handling**: The WebSocket upgrade maps each common HTTP status to an |
| 87 | actionable message: |
| 88 | |
| 89 | | Status | Meaning surfaced to the user | |
| 90 | | --- | --- | |
| 91 | | 400 | Bad/unsupported/missing pubkey, with remediation (`ssh-keygen -t ed25519`) | |
| 92 | | 401 | Token invalid/expired — try `colab new` | |
| 93 | | 403 | Forbidden — token lacks permission for this action (feature-off returns 404, not 403) | |
| 94 | | 404 | SSH not exposed on this runtime — SSH is baked in at creation, so run `colab new` | |
| 95 | | 429 | Another `colab ssh` is already connected — disconnect first | |
| 96 | | 502 | Runtime `sshd` unreachable — runtime may be unhealthy | |
| 97 | | other / none | Raw status or a network-check hint | |
| 98 | |
| 99 | |
| 100 | ## Testing Strategy (TDD) |
| 101 | |
| 102 | ### Unit tests (`tests/test_ssh.py`) |
| 103 | 1. WebSocket URL construction (`wss` for https, `ws` for http; token query param). |
| 104 | 2. Pubkey resolution — `--identity` (via `ssh-keygen -y -f`) and the `~/.ssh` |
| 105 | default scan; missing-key and missing-identity exit paths. |
| 106 | 3. The full status→message map (400/401/403/404/429/502/other/none). |
| 107 | 4. Shell quoting for the `ProxyCommand` string. |
| 108 | 5. Session resolution (existing vs missing). |
| 109 | 6. End-to-end dispatch: interactive vs `--proxy-mode`, including a |
| 110 | verbatim-pubkey pass-through assertion and the actionable-400 message. |
| 111 | |
| 112 | ### Wire-contract tests (`tests/test_ssh_wire_contract.py`) |
| 113 | Stands up a loopback WebSocket server and drives the real connect path (no mock) |
| 114 | to assert the request path, the `colab-runtime-proxy-token` query param, and the |
| 115 | `X-Colab-Ssh-Pubkey` header reach the wire verbatim. Includes mutation tests that |
| 116 | fail if `_SSH_PATH`/`_PUBKEY_HEADER` drift, plus real HTTP 400/429 mapping via a |
| 117 | genuine `WebSocketBadStatusException`. |
| 118 | |
| 119 | ### Auto-create & proxy-mode tests (`tests/test_ssh_autocreate.py`) |
| 120 | Bare `colab ssh` create vs reuse vs ambiguous; `--gpu/--tpu` passthrough; `--rm` |
| 121 | stop-on-exit; and the `--proxy-mode` matrix — create-if-missing with `-s NAME`, |
| 122 | reuse of an existing session, `--gpu` passthrough, `--rm` teardown, and the |
| 123 | SIGHUP cleanup handler being installed only under `--rm`. |
| 124 | |
| 125 | ### Working-directory tests (`tests/test_ssh_workdir.py`) |
| 126 | Interactive `ssh` forces a PTY (`-t`) and runs a `cd /content` remote command |
| 127 | (host before the command, `2>/dev/null` tolerance for a missing directory). |
| 128 | |
| 129 | ### Integration test (`integration/repro_ssh/`) |
| 130 | Two parts. An offline smoke that always runs (no VM): ``--help`` advertises the |
| 131 | documented flags, and an unknown session exits 2 with an actionable message. A |
| 132 | live end-to-end that runs when auth is present (allocates a CPU VM): it uses |
| 133 | ``colab ssh --proxy-mode`` as an OpenSSH ProxyCommand to run a real remote |
| 134 | command over the WebSocket bridge -- exercising the same connect -> |
| 135 | pubkey-header auth -> handshake -> bridge -> remote-exec path as the interactive |
| 136 | shell, minus the TTY -- asserts the RSA-key rejection, and verifies ``colab |
| 137 | stop`` leaves no orphan VM. |