Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.

Expand Down Expand Up @@ -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 <remote>` immediately before pushing or merging. If `git log main..<remote>/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.
Expand All @@ -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 <file.py>`).
- **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 `</dev/null` if my shell tool runs commands in a terminal); where ssh can't work, use `echo '!cmd' | colab exec -s s` (each `!` line is its own shell and IPython expands `{expr}` and `$name` in it).
- Specially crafted mock scripts that simulate timeouts or API calls.
- **I CANNOT Run (Requires User Assistance):**
- **`colab auth`**: This command relies on the traditional Gcloud fallback `input_request` (via `USE_AUTH_EPHEM='0'`), which prompts the user via Python's `input()` to click a URL, sign in, and paste back an authorization code. My shell tool will hang indefinitely on this `input()`.
- **`colab drivemount`**: This command prompts the user via `sys.stdin.readline()` (specifically querying `/dev/tty` to ensure input is captured) to press `Enter` after granting OAuth consent in the browser. My shell tool will timeout/hang waiting for `Enter`.
- **Interactive (TTY) `colab repl` / `colab console`**: When stdin is a real terminal these commands drop into interactive raw-TTY modes that require real-time keystroke streaming. My shell tools cannot support this. (Piped stdin is fine — see above.)
- **Interactive (TTY) `colab repl` / `colab ssh`**: When stdin is a real terminal these commands open interactive sessions that require real-time keystroke streaming. My shell tools cannot support this. (Piped stdin to `colab repl` or `colab exec` is fine — see above.)

Whenever working on interactive commands, I must build the core logic, write mock tests, and explicitly ask the user to run the live test in their terminal to verify success.
10 changes: 10 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,16 @@ below corresponds to a tag of the same name.

- **ssh:** `colab ssh [-s NAME] -- CMD` runs one command on an existing runtime and exits with its status. Output is byte-exact when piped (a PTY only when stdin and stdout are terminals), and it never creates a runtime. Its handshake is tagged `ssh-mode-cmd`.

### Removed

- **console:** **BREAKING:** Remove the `colab console` subcommand. Use
`colab ssh -s NAME` for an interactive shell on the runtime and
`colab ssh -s NAME -- CMD` for one-shot commands; where ssh can't work, pipe
`!`-prefixed shell commands to `colab exec -s NAME`. 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. Breaking, so
it ships in a new major version.

### Fixed

- **ssh:** `colab --config FILE ssh` and `colab --auth adc ssh` now keep those options in the ProxyCommand that `colab ssh` starts, so it uses the same sessions file and auth provider instead of the defaults.
Expand Down
7 changes: 3 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ Designed to support seamless developer productivity, headless automation, and AI
## Key Features

* **Instant VM Provisioning:** Spin up CPU, GPU (T4, L4, G4, H100, A100), or TPU (v5e1, v6e1) runtimes in seconds.
* **Robust Code Execution:** Run local Python scripts, Jupyter Notebooks (`.ipynb`), or piped `stdin` code; launch interactive REPLs or raw TTY console shells.
* **Robust Code Execution:** Run local Python scripts, Jupyter Notebooks (`.ipynb`), or piped `stdin` code; launch interactive REPLs or SSH shells (`colab ssh`).
* **Ephemeral Job Runner (`colab run`):** Provision a fresh VM, execute a local script with forwarded arguments, retrieve output files, and automatically tear down the runtime in a single command.
* **Activity-Based Liveness:** Sessions stay alive as long as the kernel is active, with the Colab backend automatically maintaining liveness.
* **Seamless Workspace Automation:** Mount Google Drive, authenticate Google Cloud Platform (GCP) credentials, and install dependencies with high-performance `uv` package management.
Expand Down Expand Up @@ -81,8 +81,7 @@ Run `colab <command> --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 |
Expand Down Expand Up @@ -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.

Expand Down
3 changes: 2 additions & 1 deletion docs/01_session_management.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
log:
2026-11-04: `--gpu` and `--tpu` now reject unknown values instead of silently selecting `A100` (GPU) or `v6e1` (TPU). `resolve_runtime_options()` uses the new `GPU_ACCELERATORS` / `TPU_ACCELERATORS` maps and raises a `typer.BadParameter` that names the flag and the supported list; `--help` is generated from the same maps so it cannot drift. Validation runs before any assignment request, in both `colab new` and `colab run`.
2026-11-02: Added `_check_linux_keyring_locked` in `auth.py` to check if the active Linux `SecretService` keyring collection is locked via `secretstorage` before reading or writing OAuth tokens, logging and printing an immediate warning with the collection label, timeout, and `--auth=adc` hint while still attempting the keyring operation.
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-29: The debug log (`colab.log` and `--logtostderr`) masks access tokens, runtime proxy tokens, OAuth codes, client secrets, auth headers and cookies, including in percent-encoded URLs; `jupyter_kernel_client` logs at INFO so kernel stdin replies are not recorded.
2026-10-29: `sessions.json`, `settings.json`, `colab.log` and session history files are created with mode 0600, and existing copies are tightened to 0600 the next time the CLI uses them (history files whenever the history logger starts).
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).
Expand Down Expand Up @@ -123,7 +124,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`.
Expand Down
Loading
Loading