feat: add `colab ssh` — SSH-over-WebSocket runtime access (#88)

* feat: add colab ssh subcommand for SSH-over-WebSocket runtime access Implements the client side of SSH-over-WebSocket runtime access. Three modes: colab ssh Pick the only active session and open an interactive SSH shell. colab ssh -s SESSION Same, targeting SESSION explicitly. colab ssh --proxy-mode -s SESSION Act as an OpenSSH ProxyCommand-compatible WebSocket-stdio bridge, so a Host colab-runtime ProxyCommand colab ssh --proxy-mode -s SESSION block in ~/.ssh/config works with any SSH-based IDE / remote-dev tool. --identity/-i overrides the default key order (~/.ssh/id_ed25519, id_ecdsa, id_rsa); the public key is derived via `ssh-keygen -y -f` and sent verbatim in the X-Colab-Ssh-Pubkey header. Per-failure-mode error messages map each common HTTP status (400/401/403/404/ 429/502 + a network catch-all) to an actionable hint. Adds src/colab_cli/commands/ssh.py, registers it in cli.py, plus tests/test_ssh.py (24 tests), docs/06_ssh_access.md, and integration/repro_ssh/. The server side is out of scope for this repo; against a runtime that does not expose the SSH endpoint the command surfaces an actionable HTTP 404. * Add `colab ssh --drive` to mint and install a Drive backup token. Adds a --drive flag to `colab ssh`. When set, the CLI mints a drive.file-scoped OAuth token via gcloud (after the user's browser consent), reshapes it into the frozen drive_token.json schema, and installs it on the runtime at /root/.config/colab/drive_token.json with mode 0600 before opening the shell -- so in-VM code can back up notebooks to Drive. All of this is gated behind `if drive:`, so the default `colab ssh` path is byte-for-byte unchanged. - The token path, 0600 mode, drive.file scope, and JSON schema are the frozen cross-repo credential contract the in-VM reader depends on: there is no shared import across the OSS<->google3 boundary, so a later golden fixture pins the exact bytes on both sides. - gcloud is mocked in every test; no test invokes real gcloud or a browser, and the default (no --drive) path makes no gcloud or scp call. * feat(ssh): auto-create, proxy-mode flags, /content default; drop --drive Bare `colab ssh` now creates a runtime when you have none and drops you in /content; every flag also works in --proxy-mode so `~/.ssh/config` hosts (ssh colab / colab-gpu / colab-ephem) work on first connect. - Auto-create: bare `colab ssh` with no session runs `colab new` then connects; reuses a single existing session; errors on multiple. - Flags: --gpu/--tpu pick the accelerator for an auto-created runtime; --rm stops a runtime this command created when the session ends. - /content: interactive shells cd to /content (Colab's working dir) via -t and a remote `cd /content 2>/dev/null; exec $SHELL -l` (falls back to home if absent). - --proxy-mode honors all flags: `-s NAME` creates the session if missing (creation output routed to stderr so stdout stays the ssh byte stream); --rm stops the bridged session on disconnect. - --rm teardown survives OpenSSH SIGHUP: on disconnect ssh SIGHUPs the ProxyCommand (verified), which by default skipped the cleanup `finally` and leaked the runtime; now SIGHUP/SIGTERM/SIGINT handlers run the stop, idempotent with the finally. - Fix two client bugs: the 403 branch was dead (feature-off returns 404, not 403); RSA keys are server-rejected, so id_rsa is no longer auto-scanned and the 400 message no longer advertises rsa-sha2. - Remove the --drive subfeature entirely (code + tests + golden fixture). - Tests: add test_ssh_autocreate.py, test_ssh_wire_contract.py, test_ssh_workdir.py; drop the drive tests. Full suite: 280 passing. * docs(ssh): reformat 06_ssh_access to match sibling design-doc structure Restructure the ssh design doc to the same shape as the other command docs (esp. 05_run_command): Motivation -> User Surface (flags table) -> Behavior (numbered) -> AGENTS.md Constraints Honoured -> Testing Strategy (TDD). No content dropped; the frontmatter log is preserved verbatim. * style(ssh): apply go/pystyle readability to the ssh code Bring the ssh subcommand + its tests to the Google Python Style Guide (go/pystyle) readability bar: - 3.2 line length: reflow all ssh files to <=80 cols (ruff format is idempotent at --line-length 80). Repo enforces no line length, so this was the main gap. - 3.8 docstrings: add Args/Returns/Raises where it aids a caller (_resolve_pubkey, _resolve_session, _auto_create_session, _explain_handshake_failure, _connect_websocket, _bridge_proxy_mode, _run_interactive_ssh); keep summary lines <=80. - 3.8.5 comments: move overflowing trailing comments above their statements; annotate the best-effort except-cleanups (isolation points, 2.4.4). - 2.2 imports: lift stdlib signal/uuid to the top-level group (the colab_cli.* imports stay local to avoid circular imports). - Extract locals for two long f-string echoes (clearer + <=80). Unchanged (already compliant or out of scope): 4-space indentation (3.4), naming (3.16), and the from-x-import-Class convention (2.2 exemptions; repo-wide). Behavior-preserving; full suite: 280 passing. * test(ssh): turn repro_ssh into a real end-to-end integration test The old integration test only grepped `colab ssh --help`, so it could pass while the whole feature was broken. Replace it with a genuine e2e, following the repro_run_command / repro_keep_alive conventions (auth detection, isolated --config, trap cleanup, no-orphan-VM assertion). Key idea: `colab ssh --proxy-mode` is non-interactive, so it can be used as an OpenSSH ProxyCommand to run a real remote command over the WebSocket bridge -- exercising the same connect -> pubkey-header auth -> handshake -> bridge -> remote-exec path as the interactive shell, without a TTY. - Part A (offline, always runs, no VM): --help advertises the documented flags; an unknown session exits 2 with an actionable message. - Part B (live, auto-runs when auth is present; allocates a CPU VM): colab new -> substrate check (sshd up) -> `ssh root@colab-runtime "whoami"` over --proxy-mode (assert root + a unique marker) -> RSA key is rejected -> colab stop -> assert no orphan VM. The live part is no longer RUN_LIVE-gated. Interactive shell (/content, raw TTY) remains a documented manual step. Verified: offline part + full live e2e pass against prod (handshake, pubkey auth, bridge, remote exec, RSA rejection, clean teardown). Updates integration/README.md and docs/06 accordingly. * test(ssh): close coverage gaps and parameterize; fix empty-pubkey exit Fix: `_resolve_pubkey` raised a bare RuntimeError when `ssh-keygen -y` returned empty output, which escaped the surrounding `except (CalledProcessError, FileNotFoundError)` as an uncaught traceback. It now echoes an actionable message and `raise typer.Exit(code=2)`, matching every other failure path. Tests: add coverage for the previously-untested logic surfaced by a branch coverage run (ssh.py 75% -> 90%, BrPart 0): - _resolve_pubkey --identity failures (ssh-keygen error / missing / empty) - default ~/.ssh scan order incl. id_rsa NOT auto-selected + no-keys - _connect_websocket network-failure branch (address/timeout/refused/OSError) - _bridge_proxy_mode ws->stdout byte pump (binary/text/ignored-opcode/close) - _stop_session error handling (generic swallowed vs typer.Exit re-raised) - --gpu/--tpu "ignored" warnings (proxy-existing + interactive-reuse) - --identity threading through _proxy_command / _ssh_base_args Parameterize repetitive tests with readable ids: WS-URL scheme, shquote, session resolution, pubkey scan/failures, bytes-vs-str body, bare-ssh create/reuse/ambiguous, proxy-mode create/reuse, --rm teardown (proxy + interactive), signal-handler installation, and the two wire-contract mutations (-> one). Remaining uncovered lines are intentionally left: the stdin->ws thread body (covered by the live integration test), the os._exit signal-handler wrapper, and best-effort except-close cleanups. SSH tests 48 -> 68; full suite 300 passing; ruff clean; 80-col formatting held. * refactor(ssh): split ssh() into intent-named helpers; unify --gpu/--tpu warning Extract _select_proxy_session / _select_interactive_session, _run_proxy_bridge / _run_interactive_shell, _install_rm_signal_handlers, _warn_accelerator_ignored, _pubkey_from_identity, and _close_quietly so ssh() is a thin dispatcher, and convert the handshake-status mapping to a match statement. Behavior-preserving except the two divergent '--gpu/--tpu ignored' messages are unified into one; test_ssh_autocreate asserts the new wording. * test(ssh): add lifecycle guarantees; shorten test_ssh docstring New tests/test_ssh_lifecycle.py pins dispatch/lifecycle contracts the mocked suite missed: --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. * docs(ssh): trim SSH design doc + integration prose; reindent repro_ssh Simplify the 06 design-doc motivation/testing sections, drop the redundant 'AGENTS.md Constraints Honoured' block, and remove the stale 'manual step' framing now that the e2e is fully automated. Reindent integration/repro_ssh/test.sh to 2-space. Adds a 2026-07-24 frontmatter log entry. * refactor(ssh): use shlex.join for the ProxyCommand instead of a hand-rolled quoter Replaces _shquote with the stdlib. Verified equivalent: all cases round-trip identically through a real shell. shlex is also stricter on non-ASCII -- _shquote left 'ünïcode' unquoted because str.isalnum() is unicode-aware, while shlex.quote uses re.ASCII. Retargets the test at the real contract: _proxy_command output must re-parse (shlex.split) into the exact argv, including hostile session names (command substitution, semicolons, quotes). Addresses review comment on PR #88. * docs(ssh): document the ProxyCommand / IDE remote-dev flow Adds a README 'SSH / IDE Remote-Dev' section and a ready-to-paste ~/.ssh/config stanza in the design doc (PATH caveat for ssh's non-login ProxyCommand shell, ephemeral --rm host, auto-create on a -s NAME miss), plus VS Code Remote-SSH / plain-ssh / one-shot examples. Fixes the one-shot examples to connect as root@: the runtime authorizes root (_SSH_HOST = root@colab-runtime), so without it ssh would default to the local username and fail auth. %h expands to the host only, so -s %h is unaffected. * Revert "docs(ssh): document the ProxyCommand / IDE remote-dev flow" This reverts commit 7691d4ef5210595e08787a5a7ea1c871da0e22ba. * test(ssh): fold the ProxyCommand argv prefix into one assertion sethtroisi asked for this shape in review. Comparing the whole 7-element prefix at once also fixes the failure mode: the old argv[5]/argv[6] index accesses raised IndexError on a short argv instead of failing cleanly as an assertion. --------- Co-authored-by: Drake Aiman <drakeaiman@google.com>