google-colab-cli 0.6.0__tar.gz → 0.7.2__tar.gz

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/AGENTS.md

@@ -42,6 +42,21 @@
 - **Files**: `ls`, `rm`, `upload`, `download`, `edit`.
 - **Automation**: `auth`, `drivemount`, `install`, `log`, `pay`, `version`, `update`.
 
+## Release Tagging Workflow
+The package version is derived from the git tag via `hatch-vcs` (see `pyproject.toml`), so a release is just (1) a `CHANGELOG.md` entry and (2) a `vX.Y.Z` tag on the merged commit. Follow this workflow ONLY when the user explicitly requests a release (e.g. "cut a release", "tag v0.7.0", "release the keep-alive fix"). NEVER propose or perform a release proactively.
+
+1. **Verify a clean, current `main`**: `git fetch origin && git checkout main && git pull --ff-only origin main`. Confirm `git status` is clean. If `main` has diverged or there are unmerged feature branches the user expects in the release, stop and ask.
+2. **Identify unreleased commits**: Find the last release tag with `git describe --tags --abbrev=0`, then list commits since then with `git log <last-tag>..HEAD --oneline`. These are what the new changelog section must cover. Cross-reference each commit's PR number (the `(#NN)` suffix in the merge commit subject) — every entry in the changelog should cite at least one PR.
+3. **Infer the next version (semver)**: Categorize each commit using the Keep-a-Changelog buckets already present in `CHANGELOG.md` (`Added`, `Changed`, `Fixed`, `Removed`, `Deprecated`, `Security`). Pick the bump:
+   - **Major** (`X.0.0`): any breaking change to the CLI surface, flag semantics, on-disk state schema (`sessions.json` / `settings.json`), or persisted token format.
+   - **Minor** (`0.X.0`): at least one `Added` entry (new subcommand, new flag, new capability) with no breaking changes.
+   - **Patch** (`0.0.X`): only `Fixed`, `Changed`, `Removed` (non-breaking cleanup), `Deprecated`, or `Security` entries.
+   Confirm the inferred version with the user before proceeding — never tag without that confirmation, even when the bump seems obvious.
+4. **Draft and commit the CHANGELOG entry on a branch**: Create a release branch (e.g. `release-v0.7.0`), then prepend a new `## [X.Y.Z] - YYYY-MM-DD` section above the previous release. Group entries by category in the order already established in `CHANGELOG.md` (`Changed`, `Added`, `Fixed`, `Removed`). Each bullet should be one short paragraph naming the user-visible behavior change (not the implementation detail) and ending with the PR reference `(#NN)`. Append a new link reference at the bottom: `[X.Y.Z]: https://github.com/googlecolab/google-colab-cli/compare/v<PREV>...vX.Y.Z`. Commit with message `docs: add CHANGELOG.md for vX.Y.Z release`, push the branch, and open a PR with `gh pr create`.
+5. **Wait for the changelog PR to merge**: Do NOT tag the pre-merge commit on your branch. The tag must land on the squash-merge commit that appears on `main` (this is what `hatch-vcs` will see and what users will `git checkout`). After the user reports the PR is merged (or you observe it via `gh pr view <num> --json state,mergeCommit`), `git fetch origin && git checkout main && git pull --ff-only origin main`.
+6. **Tag the merged changelog commit and push the tag**: Verify `git log -1 --oneline` is the merged changelog commit (subject matches `docs: add CHANGELOG.md for vX.Y.Z release (#NN)`). Create an annotated tag whose message is the changelog section body: `git tag -a vX.Y.Z -m "vX.Y.Z" -m "<changelog section body>"`. Push with `git push origin vX.Y.Z`. Do NOT create a GitHub Release entry via `gh release create` — the tag alone is sufficient because `hatch-vcs` reads it directly, and the canonical release notes already live in `CHANGELOG.md`.
+7. **Verify**: Run `git describe --tags --exact-match HEAD` (should print `vX.Y.Z`) and `uv run colab version` (the version string should match). Confirm with `gh api repos/googlecolab/google-colab-cli/tags --jq '.[0].name'` that the tag is visible on the remote.
+
 ## Implementation Principles
 1.  **Direct Execution**: Code for `auth`, `drivemount`, etc., should be injected and executed on the VM kernel.
 2.  **Contents API**: Use the Jupyter Contents API for file management as seen in the browser traces.

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/CHANGELOG.md

@@ -8,6 +8,35 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 The package version is derived from the git tag via `hatch-vcs`; each release
 below corresponds to a tag of the same name.
 
+## [0.7.1] - 2026-09-14
+
+### Fixed
+
+- **build:** Pin core metadata to version 2.4 for both the wheel and the sdist.
+  hatchling 1.32.0 bumped its default to 2.5, which our publishing pipeline
+  rejects, so 0.7.0 was built but never reached PyPI. `colab ssh` (added in
+  0.7.0) was therefore unavailable to anyone installing from PyPI.
+
+## [0.7.0] - 2026-09-03
+
+### Changed
+
+- **deps:** Bump cryptography from 49.0.0 to 50.0.0 (#98), and pyasn1 from 0.6.3 to 0.6.4 (#87).
+- **lockfile:** Upgrade lockfile dependencies via `uv lock --upgrade`. (#67)
+- **skill:** Move the CLI help skill details into a dedicated `SKILL.md`. (#55)
+- **docs:** Update `AGENTS.md` with instructions for the release tagging workflow. (#66)
+
+### Added
+
+- **ssh:** Add `colab ssh` command to provide secure SSH-over-WebSocket direct access to the Colab runtime VM. (#88)
+- **execution:** Add `--env KEY=VALUE` flag to `colab run` and `colab exec` commands to allow injecting custom environment variables into the running kernel. (#65)
+- **session:** Add `--high-mem` flag to machine shape selection when creating a new session to support high-RAM resources. (#105)
+
+### Fixed
+
+- **session:** Surface a friendly error message on 412 GPU/TPU allocation failure instead of printing a raw traceback. (#112)
+- **runtime:** Support both `ColabKernelClient` and standard `KernelClient` in the runtime layer to improve compatibility across environments. (#95)
+
 ## [0.6.0] - 2026-06-16
 
 ### Changed
@@ -47,4 +76,5 @@ below corresponds to a tag of the same name.
   `colaboratory`-scope / `pa.googleapis.com` pre-flight remediation messaging,
   superseded by the TFE keep-alive ping. (#61)
 
+[0.7.0]: https://github.com/googlecolab/google-colab-cli/compare/v0.6.0...v0.7.0
 [0.6.0]: https://github.com/googlecolab/google-colab-cli/compare/v0.5.11...v0.6.0

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-colab-cli
-Version: 0.6.0
+Version: 0.7.2
 Summary: CLI for interacting with Colab.
 Project-URL: Homepage, https://github.com/googlecolab/google-colab-cli
 Project-URL: Repository, https://github.com/googlecolab/google-colab-cli
@@ -21,7 +21,7 @@ Requires-Dist: filelock>=3.29.2
 Requires-Dist: google-auth-oauthlib>=1.3.0
 Requires-Dist: google-auth>=2.49.1
 Requires-Dist: html2text>=2024.2.26
-Requires-Dist: jupyter-kernel-client
+Requires-Dist: jupyter-kernel-client==0.8
 Requires-Dist: nbformat>=5.10.4
 Requires-Dist: packaging>=24.0
 Requires-Dist: prompt-toolkit>=3.0.52
@@ -104,9 +104,9 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he
 ### Session Management
 | Command | Description |
 | --- | --- |
-| `colab new [-s NAME] [--gpu GPU] [--tpu TPU]` | Allocate a new CPU, GPU, or TPU VM runtime |
+| `colab new [-s NAME] [--gpu GPU] [--tpu TPU] [--high-mem]` | Allocate a new CPU, GPU, or TPU VM runtime (optionally high-RAM) |
 | `colab sessions` | List all active sessions currently active on the backend |
-| `colab status [-s NAME]` | Display hardware, status, and local metadata for active sessions |
+| `colab status [-s NAME]` | Display hardware, machine shape, status, and local metadata for active sessions |
 | `colab restart-kernel [-s NAME]` | Restart the active session's Jupyter kernel |
 | `colab stop [-s NAME]` | Terminate a session VM and tear down its keep-alive daemon |
 | `colab url [-s NAME] [--open]` | Print or open a browser URL connecting to the active session |
@@ -114,10 +114,11 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he
 ### Execution
 | Command | Description |
 | --- | --- |
-| `colab run [--gpu GPU] [--tpu TPU] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it |
+| `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 |
 
 ### File Operations
 | Command | Description |
@@ -135,6 +136,7 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he
 | `colab drivemount [-s NAME] [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 |
 | `colab pay` | Open the Colab subscription page to manage compute units |
 | `colab version` | Print the installed version of the CLI |
 | `colab update [--install]` | Check for a newer release (and optionally upgrade the CLI in place) |
@@ -177,6 +179,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.
 * **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.

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/README.md

@@ -68,9 +68,9 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he
 ### Session Management
 | Command | Description |
 | --- | --- |
-| `colab new [-s NAME] [--gpu GPU] [--tpu TPU]` | Allocate a new CPU, GPU, or TPU VM runtime |
+| `colab new [-s NAME] [--gpu GPU] [--tpu TPU] [--high-mem]` | Allocate a new CPU, GPU, or TPU VM runtime (optionally high-RAM) |
 | `colab sessions` | List all active sessions currently active on the backend |
-| `colab status [-s NAME]` | Display hardware, status, and local metadata for active sessions |
+| `colab status [-s NAME]` | Display hardware, machine shape, status, and local metadata for active sessions |
 | `colab restart-kernel [-s NAME]` | Restart the active session's Jupyter kernel |
 | `colab stop [-s NAME]` | Terminate a session VM and tear down its keep-alive daemon |
 | `colab url [-s NAME] [--open]` | Print or open a browser URL connecting to the active session |
@@ -78,10 +78,11 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he
 ### Execution
 | Command | Description |
 | --- | --- |
-| `colab run [--gpu GPU] [--tpu TPU] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it |
+| `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 |
 
 ### File Operations
 | Command | Description |
@@ -99,6 +100,7 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he
 | `colab drivemount [-s NAME] [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 |
 | `colab pay` | Open the Colab subscription page to manage compute units |
 | `colab version` | Print the installed version of the CLI |
 | `colab update [--install]` | Check for a newer release (and optionally upgrade the CLI in place) |
@@ -141,6 +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.
 * **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.

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/cloudbuild.yaml

@@ -38,6 +38,27 @@ steps:
         echo "TAG_NAME=${TAG_NAME}"
         echo "git describe: $(git describe --tags --always)"
 
+  # 1b. Delete any artifacts a previous run left staged.
+  #     Staged artifacts are only cleaned up after a successful release; a
+  #     failed one leaves them in place so the release can be retried. Since
+  #     step 4 publishes with `publish_all: true`, a later release would pick
+  #     those up too and republish them alongside the new ones.
+  #     This has to run as part of the build, which holds the delete
+  #     permission that individual accounts do not.
+  #     `|| true` so a first-ever release (nothing to delete) isn't a failure.
+  - id: purge-stale-artifacts
+    name: gcr.io/google.com/cloudsdktool/cloud-sdk:slim
+    entrypoint: bash
+    args:
+      - -c
+      - |
+        set -e
+        gcloud artifacts packages delete "${_EG_PROJECT_NAME}" \
+          --project="${_AR_PROJECT}" \
+          --repository="${_AR_REPOSITORY}" \
+          --location="${_AR_LOCATION}" \
+          --quiet || true
+
   # 2. Build sdist + wheel into dist/ using uv.
   #    Use the non-slim bookworm variant: hatch-vcs derives the version
   #    by shelling out to `git describe`, which requires a git binary
@@ -57,6 +78,8 @@ steps:
   #    google-artifactregistry-auth keyring plugin (uses ADC from the
   #    Cloud Build service account, which is registered as a builder in
   #    the project's project.txtpb).
+  #    Delete any existing artifact for this tag which can happen if the
+  #    build failed to publish in step 4.
   - id: publish-to-ar
     name: python:3.13-slim
     entrypoint: bash
@@ -66,6 +89,12 @@ steps:
         set -e
         pip install --quiet --root-user-action=ignore \
           twine keyrings.google-artifactregistry-auth
+        gcloud artifacts versions delete ${TAG_NAME}  \
+          --package={_EG_PROJECT_NAME} \
+          --project={_AR_PROJECT} \
+          --repository=${_AR_REPOSITORY} \
+          --location=us \
+          --quiet
         twine upload \
           --repository-url "https://${_AR_LOCATION}-python.pkg.dev/${_AR_PROJECT}/${_AR_REPOSITORY}/" \
           --verbose \

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/docs/01_session_management.md

@@ -1,5 +1,7 @@
 ---
 log:
+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-08-09: Added `--high-mem` to `colab new`, `colab run`, and `colab ssh` (auto-create). Assign requests now send `shape=hm` when high-RAM is requested; `colab sessions` and `colab status` display machine shape.
 2026-06-15: Switched the keep-alive daemon from the `colab.pa.googleapis.com` `RuntimeService/KeepAliveAssignment` RPC to a Tunnel Frontend HTTP ping (`GET /tun/m/<endpoint>/keep-alive/` with `X-Colab-Tunnel: Google`) on `colab.research.google.com`. The RPC required `serviceusage` consumer access to Colab's internal project `1014160490159`, which ordinary user accounts lack, so every external user hit HTTP 403 `USER_PROJECT_DENIED` and their CLI sessions were idle-pruned within minutes (issue #14). Reproduced live with a third-party account; verified the tunnel ping is accepted by the same bearer-token credential that already works for `assign`. A `ReadTimeout` on the ping is treated as success (TFE records activity before forwarding to the often-non-responding VM). Generalized the pre-flight remediation messaging away from the now-irrelevant `colaboratory`/`pa.googleapis.com` framing, and removed the dead grpc-web client-registry/API-key code.
 2026-06-10: Replaced the POSIX-only `fcntl.flock` file locking in `_LockedFileStore` with the cross-platform `filelock` library (reported broken on Windows). Reads use `ReadWriteLock.read_lock()` (shared) and writes use `write_lock()` (exclusive), preserving the original `LOCK_SH`/`LOCK_EX` semantics. The lock is constructed with `is_singleton=False` so two `StateStore` instances for the same path in one process don't collapse into a single reentrant lock (which would raise `RuntimeError` on multi-threaded write contention). Added shared-read, cross-process exclusion, and multi-thread/multi-process regression tests.
 ---
@@ -32,11 +34,20 @@ Defines the specific hardware model.
     - `V5E1`: TPU v5e (1 core, optimized for inference/efficient training).
     - `V6E1`: TPU v6e (1 core, high performance).
 
-### 3. CLI Mapping
+### 3. Machine shape (`shape`)
+Defines the RAM profile for runtimes that support a choice (CPU, T4, A100, etc.).
+- `STANDARD` (default): omit the `shape` query param on assign.
+- `HIGH_RAM`: send `shape=hm` on assign (requires Colab Pro/Pro+ entitlement).
+
+Accelerators with only one shape (L4, v5e1, v6e1) ignore `--high-mem`.
+
+### 4. CLI Mapping
 The CLI maps user flags to these backend parameters:
 - `colab new my-session` -> `variant=DEFAULT`, `accelerator=NONE`
-- `colab new my-session -gpu=L4` -> `variant=GPU`, `accelerator=L4`
-- `colab new my-session -tpu=v5e1` -> `variant=TPU`, `accelerator=V5E1`
+- `colab new my-session --gpu=L4` -> `variant=GPU`, `accelerator=L4`
+- `colab new my-session --tpu=v5e1` -> `variant=TPU`, `accelerator=V5E1`
+- `colab new my-session --high-mem` -> adds `shape=hm` (when supported)
+- `colab new my-session --gpu A100 --high-mem` -> `variant=GPU`, `accelerator=A100`, `shape=hm`
 
 ## Approach
 
@@ -51,8 +62,8 @@ The CLI maps user flags to these backend parameters:
     - Format: `{ "session_name": { "token": "...", "backend_url": "...", "hardware": "..." } }`
 
 ### 2. Session Status (`colab status`)
-- **API**: `/api/sessions` or querying the kernel for resource usage via a special "status" message.
-- **Metric Collection**: Execute a small snippet on the VM to get memory/CPU usage if the backend API doesn't provide it directly.
+- **Session lines**: Local metadata per tracked session — hardware, shape, variant, IDLE/BUSY (from CLI bookkeeping), and optional last-execution detail. Syncs with `GET /tun/m/assignments` to prune stale entries.
+- **Future**: Real-time VM resource usage (CPU/RAM/GPU) via a kernel diagnostic snippet is still TODO. For account-level compute-unit balance and usage rate, use `colab usage`.
 
 ### 3. Stop Session (`colab stop`)
 - **API**: `POST https://colab.sandbox.google.com/tun/m/unassign/<endpoint>` (based on `tpu-v5e1-unassign.har`).
@@ -84,7 +95,7 @@ To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 mi
 
 ## TODO / Future Work
 - **Backend Sync**: Implement a way to reconcile the local `sessions.json` with the output of `colab sessions`.
-- **Resource Usage**: Add real-time resource usage (CPU/RAM/GPU) to the `status` output by executing a diagnostic snippet on the VM.
+- **VM Resource Usage**: Add real-time resource usage (CPU/RAM/GPU) to the `status` output by executing a diagnostic snippet on the VM (distinct from account-level CU info, which is available via `colab usage`).
 
 ## Implementation Details
 - **Authentication**: Uses `google-auth-oauthlib` to perform a local server OAuth flow.

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/docs/04_automation_and_utility.md

@@ -1,5 +1,6 @@
 ---
 log:
+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.
 2026-06-01: Enabled `colab update --install` self-update on macOS in addition to Linux. Refactored platform check logic to keep the implementation DRY and updated both tests and documentation. Also, on these platforms, an additional message is shown recommending `colab update --install` to upgrade in place, positioned above the standard `pip`/`uv` installation command.
 2026-05-29: Added default OAuth2 client config (`oauth_config.json`) as a bundled package resource and restored fallback loading logic in `get_credentials()`. The CLI now falls back to using these default credentials when no explicit local config is found. Added `integration/repro_bundled_oauth` integration test.
@@ -12,7 +13,7 @@ log:
 2026-05-12: Added an optional `timeout=` parameter to `ColabRuntime.execute_code` that flows through to both the `execute()` and `execute_interactive()` branches. `colab auth` and `colab drivemount` now pass `timeout=600` (10 min) via a shared `INTERACTIVE_AUTOMATION_TIMEOUT_SEC` constant in `commands/automation.py`. Background: `jupyter_kernel_client` defaults to a 10s wall-clock timeout that is consumed even when the kernel is idle waiting on `input_request`. With the drivefs hook intercepting that request and prompting the user to OAuth in their browser, any user that takes >10s to click through (essentially everyone) hit `TimeoutError` and saw "drivemount failed" even though the mount had actually succeeded server-side. The fix is scoped narrowly to the two human-in-the-loop subcommands; non-interactive paths (`colab exec`, `colab run`, `colab install`, `colab repl --pipe`, `colab console --pipe`) keep the upstream default since they receive continuous iopub traffic that resets the practical inactivity ceiling.
 ---
 
-# Design: Automation and Utility (`auth`, `install`, `log`, `pay`, `version`, `update`, `whoami`)
+# Design: Automation and Utility (`auth`, `install`, `log`, `pay`, `usage`, `version`, `update`, `whoami`)
 
 ## Overview
 
@@ -151,13 +152,22 @@ remediation guidance) rather than silently after ~1 minute via the daemon.
 -   **Conversion (Planned)**: Future expansion to convert history logs to
     `.ipynb` or `.html`.
 
-### 5. Subscription Management (`colab pay`)
+### 5. Compute-Unit Usage (`colab usage`)
+
+-   **Action**: Print account-level compute-unit (CU) usage rate and balance.
+-   **Data source**: `GET https://colab.research.google.com/tun/m/ccu-info`
+    (balance, hourly rate, assignment count). Same bearer token and session
+    backend as `colab new`.
+-   **Output**: `Current balance: … compute units`, `Usage rate: {rate}/hr`
+    (aggregate CU/hour across assigned VMs), `Active assignments: N`.
+
+### 6. Subscription Management (`colab pay`)
 
 -   **Action**: Open the Colab signup page in the user's browser.
 -   **Implementation**: Uses
     `webbrowser.open("https://colab.research.google.com/signup")`.
 
-### 6. Version Information (`colab version`)
+### 7. Version Information (`colab version`)
 
 -   **Action**: Show the current version of the Colab CLI.
 -   **Implementation**:
@@ -167,7 +177,7 @@ remediation guidance) rather than silently after ~1 minute via the daemon.
         Git commit hash using `git rev-parse --short HEAD`.
     -   Dynamic versioning is supported in the build system via `hatch-vcs`.
 
-### 7. Auto-Update (`colab update`)
+### 8. Auto-Update (`colab update`)
 
 -   **Action**: Check if a new version of the Colab CLI is available.
 -   **Auto-check**: The CLI automatically checks for updates once every 24 hours
@@ -218,7 +228,7 @@ remediation guidance) rather than silently after ~1 minute via the daemon.
     automation. If the upgrade command exits non-zero, `colab update --install`
     propagates the same exit code.
 
-### 8. Identity Inspection (`colab whoami`) [developer-only]
+### 9. Identity Inspection (`colab whoami`) [developer-only]
 
 -   **Action**: Resolve the active credentials, mint an access token, and
     print the email, audience, scopes, and expiry of that token.
@@ -260,7 +270,7 @@ remediation guidance) rather than silently after ~1 minute via the daemon.
       - openid
     ```
 
-### 9. README and AGENT (`colab README`, `colab AGENT`)
+### 10. README and AGENT (`colab README`, `colab AGENT`)
 
 -   **Action**: Print the bundled `README.md` or `AGENTS.md` file.
 -   **Implementation**:

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/docs/05_run_command.md

@@ -1,5 +1,6 @@
 ---
 log:
+2026-08-09: Added `--high-mem` flag (passthrough to session creation; sends `shape=hm` on assign when supported).
 2026-05-12: Initial design and implementation of `colab run <script.py> [args...]`. Combines `colab new` + `colab exec` + `colab stop` into a single fire-and-forget invocation so a Python file can use `#!/usr/bin/env -S colab run` as a shebang line and execute on a freshly-allocated Colab VM. Adds `--keep` (skip auto-stop), `--gpu` / `--tpu` (passthrough to session creation), `-s/--session` (name the ephemeral session), and propagates the script's exit status (non-zero on any uncaught exception in the kernel). The script's `sys.argv` is re-set inside the kernel to mirror native `python script.py arg1 arg2` semantics, and `__name__` is set to `"__main__"`.
 2026-05-12: Native CPython exit-code semantics for `sys.exit()` / `raise SystemExit(...)` from the script body. The Colab kernel reports a `SystemExit` as `output_type=='error'`, which under the previous logic would have (a) printed the IPython traceback (`An exception has occurred, use %tb...`) and (b) flagged the run as a failure regardless of the integer exit code. Now: `sys.exit()` / `sys.exit(0)` exit 0 silently; `sys.exit(N)` exits N; `sys.exit('msg')` exits 1 (matching CPython). The IPython "To exit: use 'exit', 'quit', or Ctrl-D." UserWarning is filtered via the prelude. Encoded after running `examples/gpu_hello.py` end-to-end and seeing the noisy `SystemExit: 0` traceback at the end of an otherwise-successful GPU run.
 2026-06-04: Bumped the default value of the `--timeout` flag from 10.0s to 30.0s so short-but-silent tasks aren't prematurely killed out of the box. Mirrors the same change for `colab exec`.
@@ -29,6 +30,7 @@ colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...
 | `-s`, `--session` | str | auto | Name the ephemeral session (helpful with `--keep`). Auto-generated as `run-<6 hex>` if omitted. |
 | `--gpu` | str | None | Same set as `colab new --gpu` (T4, L4, G4, H100, A100). |
 | `--tpu` | str | None | Same set as `colab new --tpu` (v5e1, v6e1). |
+| `--high-mem` | bool | False | Same as `colab new --high-mem` — request high-RAM when supported. |
 | `--keep` | bool | False | Do **not** stop the session after the script finishes. |
 | `--timeout` | float | 30.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. |
 

google_colab_cli-0.7.2/docs/06_ssh_access.md

@@ -0,0 +1,137 @@
+---
+log:
+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).
+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`.
+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`.
+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).
+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`).
+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.
+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.
+---
+
+# Design: `colab ssh` — SSH-over-WebSocket Runtime Access
+
+## Motivation
+Users want a real shell on their Colab runtime and, more importantly, IDE
+remote-development (VS Code Remote-SSH, JetBrains Gateway, plain `ssh`). `colab ssh` is
+the client that allows sshing into Colab, reusing the CLI's existing session resolution and
+runtime-proxy token so no separate credential handling is needed.
+
+## User Surface
+
+```
+colab ssh [OPTIONS]
+```
+
+| Flag | Type | Default | Purpose |
+|---|---|---|---|
+| `-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. |
+| `--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. |
+| `-i`, `--identity` | str | auto | Private key for the public key sent to Colab. Default: first of `~/.ssh/id_ed25519`, `id_ecdsa`. |
+| `--gpu` | str | None | GPU accelerator for a runtime this command creates (T4, L4, G4, H100, A100). |
+| `--tpu` | str | None | TPU accelerator for a runtime this command creates (v5e1, v6e1). |
+| `--high-mem` | bool | False | Request high-RAM when this command auto-creates a runtime (ignored when connecting to an existing session). |
+| `--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. |
+
+### `~/.ssh/config` usage
+`--proxy-mode` turns `colab ssh` into a transport any SSH-based tool can drive:
+
+```
+Host <alias>
+  ProxyCommand <abs-path-to>/colab ssh --proxy-mode -s <name> [--gpu T4] [--rm]
+  User root
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+```
+
+Because every flag applies in `--proxy-mode`, `-s <name>` creates the session on
+first connect, `--gpu/--tpu` size it, and `--rm` makes the host ephemeral. Use an
+**absolute** `colab` path: `ssh` runs the `ProxyCommand` in a non-login shell
+where a bare `colab` may not be on `PATH`. External SSH tools run their own
+remote command, so to also land in `/content` add `RequestTTY yes` and
+`RemoteCommand cd /content 2>/dev/null; exec bash -l`.
+
+## Behavior
+
+1. **Session resolution / auto-create**: With `-s NAME`, resolves that session
+   (via `state.resolve_session`, the same helper the other commands use). Bare
+   `colab ssh` uses your only active session; with **no** session it auto-creates
+   one (mirrors `colab new` end-to-end: assign → keep-alive pre-flight → spawn
+   keep-alive daemon → persist `SessionState`); with **multiple** it errors and
+   asks you to pick one with `-s`.
+2. **Connect**: Opens the WebSocket to `wss://<netloc>/colab/ssh?colab-runtime-proxy-token=<token>`
+   and sends the resolved public key verbatim in the `X-Colab-Ssh-Pubkey` header
+   (no transformation -- the bytes the user controls are exactly what the server
+   receives). Only `ssh-ed25519` / `ecdsa-sha2-nistp{256,384,521}` keys are
+   accepted.
+3. **Interactive shell**: Spawns the system `ssh` binary with the CLI re-invoked
+   as its own `ProxyCommand` (`python -m colab_cli.cli ssh --proxy-mode`), so the
+   WebSocket bridge and the interactive shell share one code path. It forces a
+   PTY (`-t`) and runs `cd /content 2>/dev/null; exec $SHELL -l` so you land in
+   `/content` (Colab's working dir) rather than root's home; a missing `/content`
+   falls back to the login home.
+4. **`--proxy-mode` bridge**: Bridges the WebSocket ↔ stdin/stdout for use as an
+   OpenSSH `ProxyCommand`. Honors every flag: `-s NAME` creates the session if
+   missing (creation output routed to stderr so stdout stays the clean ssh byte
+   stream); bare `--proxy-mode` with no `-s` just resolves an existing session.
+5. **`--rm` teardown**: Stops the runtime when the session ends. In `--proxy-mode`
+   this must survive how OpenSSH ends a `ProxyCommand`: on disconnect it sends
+   **SIGHUP** (verified), not just stdin EOF, and Python's default SIGHUP action
+   would terminate the process before the teardown `finally` ran — leaking the
+   runtime and its keep-alive daemon. `--rm` therefore installs
+   SIGHUP/SIGTERM/SIGINT handlers that run the stop, idempotent with the
+   `finally`. `SIGKILL` cannot be intercepted, so a `kill -9`/hard crash can
+   still leak; a normal disconnect is SIGHUP and is handled.
+6. **Error handling**: The WebSocket upgrade maps each common HTTP status to an
+   actionable message:
+
+   | Status | Meaning surfaced to the user |
+   | --- | --- |
+   | 400 | Bad/unsupported/missing pubkey, with remediation (`ssh-keygen -t ed25519`) |
+   | 401 | Token invalid/expired — try `colab new` |
+   | 403 | Forbidden — token lacks permission for this action (feature-off returns 404, not 403) |
+   | 404 | SSH not exposed on this runtime — SSH is baked in at creation, so run `colab new` |
+   | 429 | Another `colab ssh` is already connected — disconnect first |
+   | 502 | Runtime `sshd` unreachable — runtime may be unhealthy |
+   | other / none | Raw status or a network-check hint |
+
+
+## Testing Strategy (TDD)
+
+### Unit tests (`tests/test_ssh.py`)
+1. WebSocket URL construction (`wss` for https, `ws` for http; token query param).
+2. Pubkey resolution — `--identity` (via `ssh-keygen -y -f`) and the `~/.ssh`
+   default scan; missing-key and missing-identity exit paths.
+3. The full status→message map (400/401/403/404/429/502/other/none).
+4. Shell quoting for the `ProxyCommand` string.
+5. Session resolution (existing vs missing).
+6. End-to-end dispatch: interactive vs `--proxy-mode`, including a
+   verbatim-pubkey pass-through assertion and the actionable-400 message.
+
+### Wire-contract tests (`tests/test_ssh_wire_contract.py`)
+Stands up a loopback WebSocket server and drives the real connect path (no mock)
+to assert the request path, the `colab-runtime-proxy-token` query param, and the
+`X-Colab-Ssh-Pubkey` header reach the wire verbatim. Includes mutation tests that
+fail if `_SSH_PATH`/`_PUBKEY_HEADER` drift, plus real HTTP 400/429 mapping via a
+genuine `WebSocketBadStatusException`.
+
+### Auto-create & proxy-mode tests (`tests/test_ssh_autocreate.py`)
+Bare `colab ssh` create vs reuse vs ambiguous; `--gpu/--tpu` passthrough; `--rm`
+stop-on-exit; and the `--proxy-mode` matrix — create-if-missing with `-s NAME`,
+reuse of an existing session, `--gpu` passthrough, `--rm` teardown, and the
+SIGHUP cleanup handler being installed only under `--rm`.
+
+### Working-directory tests (`tests/test_ssh_workdir.py`)
+Interactive `ssh` forces a PTY (`-t`) and runs a `cd /content` remote command
+(host before the command, `2>/dev/null` tolerance for a missing directory).
+
+### Integration test (`integration/repro_ssh/`)
+Two parts. An offline smoke that always runs (no VM): ``--help`` advertises the
+documented flags, and an unknown session exits 2 with an actionable message. A
+live end-to-end that runs when auth is present (allocates a CPU VM): it uses
+``colab ssh --proxy-mode`` 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, minus the TTY -- asserts the RSA-key rejection, and verifies ``colab
+stop`` leaves no orphan VM.

{google_colab_cli-0.6.0 → google_colab_cli-0.7.2}/integration/README.md

@@ -17,6 +17,7 @@ End-to-end tests that run against a **live Colab backend** (unlike the mocked un
 | `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_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` |
 
 
 ## Running