| 1 | --- |
| 2 | log: |
| 3 | 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. |
| 4 | 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. |
| 5 | 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. |
| 6 | --- |
| 7 | |
| 8 | # Design: Session Management (`new`, `status`, `stop`, `sessions`) |
| 9 | |
| 10 | ## Overview |
| 11 | Session management involves interacting with the Colab backend to allocate, monitor, and terminate runtimes. |
| 12 | |
| 13 | ## Runtime Parameters |
| 14 | |
| 15 | The `colab new` command supports selecting specific hardware and runtime environments. Based on the `tpu-v5e1.har` trace and `colab-agent` source code, the following parameters and values are identified: |
| 16 | |
| 17 | ### 1. Variants (`variant`) |
| 18 | Defines the general class of hardware requested. |
| 19 | - `DEFAULT`: Standard CPU-based runtime. |
| 20 | - `GPU`: Request a GPU-accelerated runtime. |
| 21 | - `TPU`: Request a TPU-accelerated runtime. |
| 22 | |
| 23 | ### 2. Accelerators (`accelerator`) |
| 24 | Defines the specific hardware model. |
| 25 | - **None**: For `DEFAULT` variant. |
| 26 | - **GPU Accelerators**: |
| 27 | - `T4`: NVIDIA T4 (standard free-tier GPU). |
| 28 | - `L4`: NVIDIA L4 (cost-effective modern GPU). |
| 29 | - `A100`: NVIDIA A100 (high-performance GPU). |
| 30 | - `H100`: NVIDIA H100 (latest-gen performance GPU). |
| 31 | - **TPU Accelerators**: |
| 32 | - `V2-8`: TPU v2 (8 cores). |
| 33 | - `V5E1`: TPU v5e (1 core, optimized for inference/efficient training). |
| 34 | - `V6E1`: TPU v6e (1 core, high performance). |
| 35 | |
| 36 | ### 3. Machine shape (`shape`) |
| 37 | Defines the RAM profile for runtimes that support a choice (CPU, T4, A100, etc.). |
| 38 | - `STANDARD` (default): omit the `shape` query param on assign. |
| 39 | - `HIGH_RAM`: send `shape=hm` on assign (requires Colab Pro/Pro+ entitlement). |
| 40 | |
| 41 | Accelerators with only one shape (L4, v5e1, v6e1) ignore `--high-mem`. |
| 42 | |
| 43 | ### 4. CLI Mapping |
| 44 | The CLI maps user flags to these backend parameters: |
| 45 | - `colab new my-session` -> `variant=DEFAULT`, `accelerator=NONE` |
| 46 | - `colab new my-session --gpu=L4` -> `variant=GPU`, `accelerator=L4` |
| 47 | - `colab new my-session --tpu=v5e1` -> `variant=TPU`, `accelerator=V5E1` |
| 48 | - `colab new my-session --high-mem` -> adds `shape=hm` (when supported) |
| 49 | - `colab new my-session --gpu A100 --high-mem` -> `variant=GPU`, `accelerator=A100`, `shape=hm` |
| 50 | |
| 51 | ## Approach |
| 52 | |
| 53 | ### 1. New Session (`colab new`) |
| 54 | - **API**: `GET https://colab.sandbox.google.com/tun/m/assign` (based on HAR). |
| 55 | - **Parameters**: |
| 56 | - `nbh`: Notebook hash. Generated from a unique UUID per CLI session/client instance, transformed to web-safe base64 with specific padding (44 characters total). |
| 57 | - `nsa`: 1 (Standard flag observed in browser traces, typically for "next-gen session architecture"). |
| 58 | - `variant`: Selected from the list above. |
| 59 | - `accelerator`: Selected from the list above. |
| 60 | - **State Persistence**: The response contains a `token` and potentially a backend URL or identifier. We will store this in a local JSON file (default `~/.config/colab-cli/sessions.json`). |
| 61 | - Format: `{ "session_name": { "token": "...", "backend_url": "...", "hardware": "..." } }` |
| 62 | |
| 63 | ### 2. Session Status (`colab status`) |
| 64 | - **API**: `/api/sessions` or querying the kernel for resource usage via a special "status" message. |
| 65 | - **Metric Collection**: Execute a small snippet on the VM to get memory/CPU usage if the backend API doesn't provide it directly. |
| 66 | |
| 67 | ### 3. Stop Session (`colab stop`) |
| 68 | - **API**: `POST https://colab.sandbox.google.com/tun/m/unassign/<endpoint>` (based on `tpu-v5e1-unassign.har`). |
| 69 | - **Flow**: |
| 70 | 1. Perform a `GET` request to the unassign URL to obtain a fresh XSRF token. |
| 71 | 2. Perform a `POST` request to the same URL with the `X-Goog-Colab-Token` header. |
| 72 | - **Parameters**: |
| 73 | - `authuser`: 0. |
| 74 | - `<endpoint>`: The unique session identifier returned during assignment (e.g., `tpu-v5e1-s-kkb-...`). |
| 75 | - **Cleanup**: Remove the session from the local state file upon successful 204 response. |
| 76 | |
| 77 | ### 4. Session Listing (`colab sessions`) |
| 78 | - **API**: `GET https://colab.research.google.com/tun/m/assignments` (based on `colab-agent` implementation). |
| 79 | - **Function**: Lists all active VM assignments for the user. This is useful for synchronizing local state with actual backend sessions. |
| 80 | |
| 81 | ### 5. Keep-Alive Protocol |
| 82 | To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 minutes), the CLI implements a background keep-alive mechanism. |
| 83 | - **Daemon Process**: Since the CLI is a fire-and-forget tool, `colab new` spawns a detached background process running a hidden `keep-alive` command. |
| 84 | - **Tunnel ping**: Every 60 seconds, the daemon issues `GET https://colab.research.google.com/tun/m/<endpoint>/keep-alive/` with the header `X-Colab-Tunnel: Google`, authenticated with the user's own Gaia bearer token (the same credential and host used for `/tun/m/assign`). The Tunnel Frontend (TFE) records `LastActiveTime` before forwarding the request, which refreshes the idle timer. This matches the official `colab-vscode` extension's `sendKeepAlive`. TFE notes the activity on arrival and then forwards to the VM, which often does not answer on this path — so the request commonly read-times-out even though the keep-alive succeeded; a `ReadTimeout` is therefore treated as success, while genuine HTTP errors (e.g. 404 for a deleted assignment) propagate. |
| 85 | - **Why not the RuntimeService RPC**: The previous implementation called `google.internal.colab.v1.RuntimeService/KeepAliveAssignment` at `colab.pa.googleapis.com` with `X-Goog-User-Project: 1014160490159`. That path requires the caller to be a `serviceusage` consumer of Colab's internal project `1014160490159`, which no ordinary user account is — so it returned HTTP 403 `USER_PROJECT_DENIED` for every external user, causing CLI sessions to be idle-pruned within minutes (issue #14). Dropping the header instead produced HTTP 400 `CONSUMER_INVALID` (public API-key project ≠ bearer-token quota project). The browser only succeeds because it rides the user's `google.com` cookie through an internal cookie-proxy (`colab.clients6.google.com`), which a headless bearer-token client cannot use. The TFE tunnel ping needs no project entitlement and works for any account that can assign a VM. |
| 86 | - **Pre-flight (`colab new`, OAuth2/ADC only)**: Immediately after a successful `assign`, the CLI invokes `keep_alive_assignment` once synchronously. If the response is 403 with a `SCOPE_NOT_PERMITTED` body, it unassigns the new VM (to avoid leaking a billable assignment) and prints a per-provider remediation message before exiting non-zero. Other errors are tolerated — the daemon will retry and surface them via the structured event log. (Because keep-alive now uses the same backend/credential as `assign`, a scope failure at this stage is rare — assignment would normally have failed first.) |
| 87 | - **Structured logging**: The daemon emits `keep_alive_started` (with `pid`, `endpoint`), one `keep_alive_error` per failed iteration (with `status_code`, `error_type`, truncated `error`, `response_body`, `iteration`, `consecutive_4xx`), and `keep_alive_stopped` (with `reason`, `iterations`, `duration_seconds`, optional `last_error`, optional `expected_endpoint`/`actual_endpoint`). All three are rendered specially by `colab log` so users get diagnostic context without parsing JSONL by hand. |
| 88 | - **Termination**: |
| 89 | - **Explicit**: `colab stop` terminates the daemon using its stored PID. |
| 90 | - **Implicit**: If a session is pruned (e.g., during `sync_sessions`), its daemon is also terminated. |
| 91 | - **Safety Fallback**: The daemon automatically terminates after 24 hours to prevent permanent zombie processes. |
| 92 | - **State Check**: The daemon periodically verifies that its session still exists in the local state store; if missing, it exits. |
| 93 | - **Repeated 4xx**: After two consecutive 4xx responses, the daemon exits with `reason=consecutive_4xx_errors`. With the TFE tunnel ping, a normal read-timeout is not counted as a 4xx (it is treated as success), so this branch is now reached only by genuine HTTP errors such as a 404 for a deleted/expired assignment. |
| 94 | |
| 95 | ## TODO / Future Work |
| 96 | - **Backend Sync**: Implement a way to reconcile the local `sessions.json` with the output of `colab sessions`. |
| 97 | - **Resource Usage**: Add real-time resource usage (CPU/RAM/GPU) to the `status` output by executing a diagnostic snippet on the VM. |
| 98 | |
| 99 | ## Implementation Details |
| 100 | - **Authentication**: Uses `google-auth-oauthlib` to perform a local server OAuth flow. |
| 101 | - **Global Flags**: |
| 102 | - `-c`, `--client-oauth-config`: Path to the client secrets JSON file (default: `~/.colab-cli-oauth-config.json`). |
| 103 | - `--config`: Path to the session state JSON file (default: `~/.config/colab-cli/sessions.json`). |
| 104 | - **Token Storage**: Credentials are persisted to `~/.config/colab-cli/token.json` after the initial flow. |
| 105 | - Use `requests` for robust HTTP interactions and `pydantic` for schema validation. |
| 106 | - Handle authentication headers (likely `Authorization: Bearer <token>` or cookies). |
| 107 | |
| 108 | ### State Persistence & File Locking |
| 109 | - **Stores**: `StateStore` (`sessions.json`) and `SettingsStore` (`settings.json`) both derive from `_LockedFileStore`, which guards concurrent access across independent `colab` invocations and the detached keep-alive daemon. |
| 110 | - **Cross-platform locking**: Locking uses the [`filelock`](https://pypi.org/project/filelock/) library rather than `fcntl.flock`. `fcntl` is POSIX-only and is unavailable on Windows, so the original implementation crashed on import there. `filelock` provides the same advisory cross-process locking on Linux, macOS, and Windows. |
| 111 | - **Shared vs. exclusive**: Each store owns a `filelock.ReadWriteLock` bound to a sidecar file (`<path>.lock`). Reads acquire `read_lock()` (shared — multiple concurrent readers allowed) and writes acquire `write_lock()` (exclusive). This preserves the `LOCK_SH`/`LOCK_EX` distinction of the previous `fcntl` implementation. |
| 112 | - **`is_singleton=False`**: The `ReadWriteLock` is created with `is_singleton=False`. With `filelock`'s default (`True`), two `ReadWriteLock` objects for the same path *within a single process* are deduplicated into one reentrant lock; its reentrancy guard then raises `RuntimeError` when two threads each construct their own `StateStore` and contend for the write lock. Disabling the singleton registry makes each store's lock independent so they serialize via the underlying file lock instead. |
| 113 | |
| 114 | ## Testing Strategy |
| 115 | TDD is mandatory for all session management features. |
| 116 | |
| 117 | ### 1. Mock Assignment API |
| 118 | - **Test Case**: Verify `colab new` correctly parses a `PostAssignmentResponse` and stores it in the local `StateStore`. |
| 119 | - **Test Case**: Verify `colab stop` sends a `POST` request with the correct XSRF token to the unassign endpoint. |
| 120 | - **Test Case**: Verify that the path provided via `-c` is correctly passed to the authentication flow. |
| 121 | - **Mocking**: Use `unittest.mock` to intercept `requests.Session.request` and return simulated XSSI-prefixed JSON payloads matching the HAR traces. |
| 122 | |
| 123 | ### 2. State Store Validation |
| 124 | - **Test Case**: Verify `StateStore` correctly handles file locking and multiple concurrent reads/writes. |
| 125 | - **Test Case**: Verify `--config` override correctly directs all operations to the specified file path. |
| 126 | - **Test Case (cross-platform locking)**: Verify the store locks via `filelock.ReadWriteLock` on the `<path>.lock` sidecar and does not import the POSIX-only `fcntl`. |
| 127 | - **Test Case (shared/exclusive semantics)**: Verify reads go through `read_lock()` and writes through `write_lock()`. |
| 128 | - **Test Case (cross-process exclusion)**: Hold the write lock from a separate process and confirm the store's in-process write blocks until release. |
| 129 | - **Test Case (concurrent readers)**: Hold a read lock from a separate process and confirm the store can still complete a read concurrently. |
| 130 | - **Test Case (multi-thread regression)**: Two `StateStore` instances writing from different threads must serialize without raising `RuntimeError` (guards the `is_singleton=False` choice). |