fix(keep-alive): use TFE tunnel ping instead of RuntimeService RPC (#14) (#61)
The keep-alive daemon called RuntimeService/KeepAliveAssignment on 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 and their CLI sessions were idle-pruned within minutes. No header permutation fixes it (dropping the project header yields 400 CONSUMER_INVALID); the browser only succeeds by riding the user's google.com cookie through an internal cookie-proxy (colab.clients6.google.com) that a headless bearer-token CLI cannot use. Switch keep-alive to the Tunnel Frontend HTTP ping used by the official colab-vscode extension: GET /tun/m/<endpoint>/keep-alive/ with header X-Colab-Tunnel: Google, on colab.research.google.com, authenticated with the user's own bearer token (the same host/credential as assign). TFE records LastActiveTime before forwarding, so it refreshes the idle timer with no project entitlement. The VM often doesn't answer on this path, so the request commonly read-times-out even on success; ReadTimeout is treated as success while genuine HTTP errors propagate. Also: generalize the pre-flight remediation messaging away from the now irrelevant colaboratory-scope/pa.googleapis.com framing, remove the dead grpc-web client-registry/API-key code, and update docs + AGENTS.md. Verified live with a third-party account: the old RPC 403'd with USER_PROJECT_DENIED; the tunnel ping succeeded and kept the VM alive with zero keep_alive_error events.
Tyler committed
Jun 15, 2026 at 12:06 UTC
05027b6a17b1fdb5442231adb5885cafb4597a1e
8 files changed
+157
-129
AGENTS.md
+4
-4
@@ -6,10 +6,10 @@
6
- **Client**: `ColabClient` handles API interactions (assignment, unassignment).
7
- **Auth**: `auth.py` exposes a single `get_credentials(config_path, provider)` facade that dispatches on the `AuthProvider` enum. Two providers are supported, selected via the global `--auth=oauth2|adc` flag (default `oauth2`):
8
- `oauth2`: public `google-auth-oauthlib` `InstalledAppFlow`, token cached at `~/.config/colab-cli/token.json`. Reads the client OAuth config from `-c/--client-oauth-config` (default `~/.colab-cli-oauth-config.json`), falling back to the **bundled** `src/colab_cli/oauth_config.json` resource (re-added in PR #41 / `9f44fe2`, 2026-05-29 — the earlier "removed in `20eb88e`" note was stale/incorrect; the file exists and `auth.py:_get_google_auth_credentials` loads it via `importlib.resources`). As of 2026-06-11 the flow is a **remote copy-paste flow**, not a localhost server: `_run_remote_flow` sets `redirect_uri=https://sdk.cloud.google.com/applicationdefaultauthcode.html` + `token_usage=remote`, prints the URL, and reads the pasted code via `input()`. NEVER revert to OOB (`urn:ietf:wg:oauth:2.0:oob`) — Google blocked it in 2022 ("OOB flow has been blocked"); the `sdk.cloud.google.com` redirect is registered only to the bundled cloud-SDK client (`764086051850-...`), so any other client id gets `redirect_uri_mismatch`. Server-side acceptance/rejection of these variants is verifiable GET-only by building the authorization URL and inspecting whether Google reaches sign-in vs. an OAuth error page (no resources allocated).
9
- - `adc`: Google Application Default Credentials via `google.auth.default()`. The CLI passes `scopes=PUBLIC_SCOPES` (which includes `colaboratory`) and re-applies via `creds.with_scopes()` for credential types that support it. **User credentials minted by `gcloud auth application-default login` ignore the `scopes=` kwarg AND raise `NotImplementedError` on `with_scopes`**: ADC users must explicitly re-authenticate with `gcloud auth application-default login --scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory`. `userinfo.email` is required by the session backend at `colab.research.google.com` (assign/unassign/sessions return 401 without it); `colaboratory` is required by the `RuntimeService` at `colab.pa.googleapis.com` (keep-alive returns 403 without it); `openid` and `cloud-platform` are mandated by `gcloud` itself, which rejects scope lists that omit `cloud-platform` with `Invalid value for [--scopes]`. Service-account / GCE / GKE / impersonated creds get the right scopes transparently via `with_scopes`.
9
+ - `adc`: Google Application Default Credentials via `google.auth.default()`. The CLI passes `scopes=PUBLIC_SCOPES` (which includes `colaboratory`) and re-applies via `creds.with_scopes()` for credential types that support it. **User credentials minted by `gcloud auth application-default login` ignore the `scopes=` kwarg AND raise `NotImplementedError` on `with_scopes`**: ADC users must explicitly re-authenticate with `gcloud auth application-default login --scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory`. `userinfo.email` is required by the session backend at `colab.research.google.com` (assign/unassign/sessions/keep-alive return 401 without it); `colaboratory` is retained for forward compatibility and other Colab features (keep-alive no longer uses `colab.pa.googleapis.com` — see the Keep-alive note below); `openid` and `cloud-platform` are mandated by `gcloud` itself, which rejects scope lists that omit `cloud-platform` with `Invalid value for [--scopes]`. Service-account / GCE / GKE / impersonated creds get the right scopes transparently via `with_scopes`.
10
- **Backend Hosts**: Two distinct backends with different requirements:
11
- - `colab.research.google.com` (session backend / `tun/m/...`): accepts the `userinfo.email` scope.
12
- - `colab.pa.googleapis.com` (`RuntimeService`, used by `KeepAliveAssignment`): requires (a) the `colaboratory` OAuth scope, AND (b) `X-Goog-Api-Client` header containing the substring `grpc-web`. Both are enforced server-side; missing either yields HTTP 400 / 403 with descriptive `google.rpc.DebugInfo` payloads. Always log `response_body` on failure for these RPCs to avoid silent debugging.
11
+ - `colab.research.google.com` (session backend / `tun/m/...`): accepts the `userinfo.email` scope. Handles assign, unassign, the contents API, **and keep-alive** (see below).
12
+ - **Keep-alive (2026-06-15, issue #14)**: keep-alive is a **Tunnel Frontend (TFE) HTTP ping** — `GET https://colab.research.google.com/tun/m/<endpoint>/keep-alive/` with header `X-Colab-Tunnel: Google`, authenticated by the user's own Gaia bearer token (same host/credential as `assign`). TFE records `LastActiveTime` before forwarding, refreshing the idle timer. The VM usually doesn't answer on this path, so the request commonly **read-times-out even on success** — `client.keep_alive_assignment` therefore catches `requests.exceptions.ReadTimeout` and treats it as success, while genuine HTTP errors (e.g. 404 for a deleted assignment) propagate. This mirrors the official `colab-vscode` extension's `sendKeepAlive` (`src/colab/client.ts`). **DO NOT revert to the `colab.pa.googleapis.com` `RuntimeService/KeepAliveAssignment` RPC**: that RPC 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 (issue #14) and silently idle-pruned their sessions within minutes. The browser only made that RPC work by riding the user's `google.com` cookie through an internal cookie-proxy (`colab.clients6.google.com`), which a headless bearer-token CLI cannot use. Verified live 2026-06-15 with a third-party account (the RPC 403'd; the tunnel ping succeeded and kept the VM alive).
13
- **Runtime**: `ColabRuntime` wraps `jupyter-kernel-client` for execution.
14
- **State**:
15
- `StateStore` persists session metadata in `~/.config/colab-cli/sessions.json`.
@@ -60,7 +60,7 @@
60
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 `~/.config/colab-cli/token.json` doesn't exist; **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.
61
16. **Detached daemons inherit nothing useful**: When spawning a detached child process via `subprocess.Popen` (e.g. `spawn_keep_alive`), the child does NOT inherit the parent's parsed Typer flags. It re-parses argv from scratch, so any global flag the parent saw via `--auth=adc` or `--config /tmp/foo.json` MUST be re-emitted as part of the child's command line. Forgetting this causes silent fallback to Typer defaults: in 2026-04-30 this manifested as the keep-alive daemon (a) using OAUTH2 instead of the parent's ADC, and (b) reading from `~/.config/colab-cli/sessions.json` instead of the parent's `--config` path. Always propagate every relevant global flag in `cmd = [sys.executable, "-m", ...]` BEFORE the subcommand name (Typer requires global flags before the subcommand).
62
17. **Persist-before-spawn for daemons that read shared state**: When a detached child reads from a state file the parent owns (e.g. `state.store.get(session_name)`), the parent MUST persist BEFORE spawning. Otherwise the child can race ahead of the parent's `add()` call and observe an empty store. Symptom in 2026-04-30: keep-alive daemon exits immediately with `keep_alive_stopped reason=session_not_found iters=1 duration=0.0s`. Fix: call `state.store.add(s)` once before `spawn_keep_alive`, and again after (to capture the PID).
63
-18. **`colab.pa.googleapis.com` rejects ADC user creds without `X-Goog-User-Project`**: When calling `colab.pa.googleapis.com` from ADC user credentials minted by `gcloud auth application-default login`, the bearer token carries the user's gcloud quota project. The colab API key sent alongside it belongs to a different project (Colab, `1014160490159`). The backend enforces project-match between the two and returns HTTP 400 `CONSUMER_INVALID` ("The API Key and the authentication credential are from different projects."). Fix: send `X-Goog-User-Project: 1014160490159` to pin the consumer project to Colab's. Any signed-in user has implicit access to the colab project via the public web client, so this works for ordinary user accounts. Service-account / GCE / GKE / impersonated creds don't hit this because their bearer token IS owned by a project that matches.
63
+18. **[SUPERSEDED 2026-06-15 — keep-alive no longer calls `colab.pa.googleapis.com`; see issue #14 / the Keep-alive note above] `colab.pa.googleapis.com` rejects ADC user creds without `X-Goog-User-Project`**: When calling `colab.pa.googleapis.com` from ADC user credentials minted by `gcloud auth application-default login`, the bearer token carries the user's gcloud quota project. The colab API key sent alongside it belongs to a different project (Colab, `1014160490159`). The backend enforces project-match between the two and returns HTTP 400 `CONSUMER_INVALID` ("The API Key and the authentication credential are from different projects."). The old fix was to send `X-Goog-User-Project: 1014160490159` to pin the consumer project — but that in turn required `serviceusage.serviceUsageConsumer` on project `1014160490159`, which ordinary users lack, yielding HTTP 403 `USER_PROJECT_DENIED` (the issue #14 root cause). Both failure modes are now moot because keep-alive uses the TFE tunnel ping on `colab.research.google.com` instead. Retained here as institutional knowledge: if any future code path must call `colab.pa.googleapis.com` with a bearer token, it will face this same project-entitlement wall for non-internal accounts.
64
19. **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. fire-and-forget RPCs like `KeepAliveAssignment` that return `[]`), this raises `pydantic.ValidationError: Input should be None`. Always guard with `if schema is None: return` after the empty-body short-circuit.
65
20. **Suggest the branch-diff review command after committing**: The user reviews changes with `git diff main..<branch-name>` (full cumulative diff against `main`, not just the latest commit). After landing one or more commits on a feature branch, ALWAYS suggest the exact command — e.g. "Review with `git diff main..sort-help-commands`" — instead of `git show <sha>` (which only shows a single commit and misses context when a branch has multiple commits). Encoded 2026-05-05 after suggesting `git show 9d9c7da` for a branch the user wanted to review holistically.
66
21. **Verify research-tool claims with primary sources**: Research tools and AI assistants can be confidently wrong, especially about edge cases or features outside their training corpus. When such a tool says "X is not used / not parsed / doesn't exist", treat it as a hypothesis to verify, not a fact. Always cross-check against the primary source (the actual code or config) — and when a tool names files, check whether the indirection chain it describes actually exists. The cost of believing the tool when it's wrong is shipping a non-functional feature; the cost of double-checking is small. Encoded 2026-05-05 after a `colab url` first-cut shipped the wrong URL format because of unverified output.
docs/01_session_management.md
+5
-3
@@ -1,5 +1,6 @@
1
---
2
log:
3
+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.
4
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.
5
---
6
@@ -70,15 +71,16 @@ The CLI maps user flags to these backend parameters:
71
### 5. Keep-Alive Protocol
72
To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 minutes), the CLI implements a background keep-alive mechanism.
73
- **Daemon Process**: Since the CLI is a fire-and-forget tool, `colab new` spawns a detached background process running a hidden `keep-alive` command.
73
-- **RPC**: Every 60 seconds, the daemon calls `google.internal.colab.v1.RuntimeService/KeepAliveAssignment` at `colab.pa.googleapis.com`. The wire format is grpc-web JSON: `Content-Type: application/json+protobuf`, body `["<endpoint>"]` (positional protojson), `X-Goog-Api-Client: grpc-web/0.1`, `x-user-agent: grpc-web-javascript/0.1`. **Both** the `colaboratory` OAuth scope (see `04_automation_and_utility.md`) and the `grpc-web` substring in `X-Goog-Api-Client` are server-enforced; missing either yields a descriptive 403/400 response.
74
-- **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.
74
+- **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.
75
+ - **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.
76
+- **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.)
77
- **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.
78
- **Termination**:
79
- **Explicit**: `colab stop` terminates the daemon using its stored PID.
80
- **Implicit**: If a session is pruned (e.g., during `sync_sessions`), its daemon is also terminated.
81
- **Safety Fallback**: The daemon automatically terminates after 24 hours to prevent permanent zombie processes.
82
- **State Check**: The daemon periodically verifies that its session still exists in the local state store; if missing, it exits.
81
- - **Repeated 4xx**: After two consecutive 4xx responses, the daemon exits with `reason=consecutive_4xx_errors`. The pre-flight in `colab new` now catches the most common cause (missing `colaboratory` scope) before it reaches this branch.
83
+ - **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.
84
85
## TODO / Future Work
86
- **Backend Sync**: Implement a way to reconcile the local `sessions.json` with the output of `colab sessions`.
docs/04_automation_and_utility.md
+17
-18
@@ -56,18 +56,18 @@ allowing the core `Client` to remain authentication-agnostic — it only sees a
56
57
### Required Scopes
58
59
-The CLI talks to two distinct backends, each with different scope demands:
60
-
61
-- `colab.research.google.com` (session assignment / unassignment /
62
- contents API): the `userinfo.email` scope is sufficient.
63
-- `colab.pa.googleapis.com` (`RuntimeService`, used by
64
- `KeepAliveAssignment`): **requires** the
65
- `https://www.googleapis.com/auth/colaboratory` scope. Without it, every
66
- request returns HTTP 403 with body `[7,"Request had insufficient
67
- authentication scopes.",...]` and a `DebugInfo` mentioning
68
- `SCOPE_NOT_PERMITTED`. (The frontend additionally requires
69
- `X-Goog-Api-Client` to contain `grpc-web` — see
70
- `01_session_management.md` §5.)
59
+The CLI talks to the Colab session backend at `colab.research.google.com`
60
+for assignment, unassignment, the contents API, **and keep-alive** (the TFE
61
+tunnel ping — see `01_session_management.md`). The `userinfo.email` scope is
62
+sufficient for this host.
63
+
64
+> Historical note: keep-alive previously used the `RuntimeService`
65
+> (`KeepAliveAssignment`) at `colab.pa.googleapis.com`, which required the
66
+> `https://www.googleapis.com/auth/colaboratory` scope **and** the caller to
67
+> be a `serviceusage` consumer of Colab's internal project `1014160490159`.
68
+> The latter is impossible for ordinary user accounts, which made keep-alive
69
+> fail with HTTP 403 `USER_PROJECT_DENIED` for all external users (issue #14).
70
+> Keep-alive no longer touches `colab.pa.googleapis.com`.
71
72
How each provider supplies the scope:
73
@@ -91,12 +91,11 @@ How each provider supplies the scope:
91
```
92
93
`userinfo.email` is required for the session backend at
94
- `colab.research.google.com` (otherwise assign/unassign/sessions return
95
- HTTP 401); `colaboratory` is required for the `RuntimeService` at
96
- `colab.pa.googleapis.com` (otherwise keep-alive returns HTTP 403);
97
- `openid` and `cloud-platform` are mandated by `gcloud` itself
98
- (`gcloud auth application-default login` rejects scope lists that
99
- omit `cloud-platform` with `Invalid value for [--scopes]`).
94
+ `colab.research.google.com` (otherwise assign/unassign/sessions/keep-alive
95
+ return HTTP 401); `colaboratory` is retained for forward compatibility and
96
+ other Colab features; `openid` and `cloud-platform` are mandated by
97
+ `gcloud` itself (`gcloud auth application-default login` rejects scope
98
+ lists that omit `cloud-platform` with `Invalid value for [--scopes]`).
99
100
`colab new` performs a one-shot keep-alive pre-flight after `assign`
101
succeeds so missing-scope failures surface immediately (with per-provider
integration/repro_keep_alive_scope/test.sh
+16
-10
@@ -16,16 +16,22 @@
16
# Integration Test: Keep-Alive Daemon Soak (OAuth Scope Regression Guard)
17
#
18
# Background:
19
-# On 2026-04-30, a regression was discovered where `colab new` would succeed
20
-# but the keep-alive daemon would silently die ~1 minute later, causing the
21
-# VM to be idle-pruned shortly after. Two underlying causes:
22
-# (a) The RuntimeService at colab.pa.googleapis.com requires
23
-# `X-Goog-Api-Client` to contain `grpc-web`. Missing this returns 400.
24
-# (b) The same service requires the
25
-# `https://www.googleapis.com/auth/colaboratory` OAuth scope.
26
-# Missing this returns 403 SCOPE_NOT_PERMITTED.
27
-# Both unit-test layers passed because they mock the network. The bug only
28
-# surfaces against the live backend.
19
+# The keep-alive daemon would `colab new` successfully but then silently die
20
+# ~1 minute later, causing the VM to be idle-pruned shortly after. The
21
+# dominant cause for EXTERNAL users (issue #14, fixed 2026-06-15) was that
22
+# keep-alive used the RuntimeService RPC at colab.pa.googleapis.com, which
23
+# requires the caller to be a serviceusage consumer of Colab's internal
24
+# project 1014160490159 — something no ordinary account is. That returned
25
+# HTTP 403 USER_PROJECT_DENIED. Keep-alive now uses the Tunnel Frontend ping
26
+# (GET /tun/m/<endpoint>/keep-alive/ with X-Colab-Tunnel: Google) on
27
+# colab.research.google.com, authenticated with the user's bearer token and
28
+# requiring no project entitlement.
29
+#
30
+# Earlier (2026-04-30) the RPC path also required `X-Goog-Api-Client` to
31
+# contain `grpc-web` (else 400) and the `colaboratory` OAuth scope (else 403
32
+# SCOPE_NOT_PERMITTED). Those are moot now but kept here for history.
33
+# Unit-test layers pass because they mock the network; these bugs only
34
+# surface against the live backend — which is why this soak test exists.
35
#
36
# What this test does:
37
# 1. Spawns a real Colab session via `colab new`.
src/colab_cli/client.py
+28
-47
@@ -32,28 +32,16 @@ COLAB_CLIENT_AGENT_HEADER = {
32
"value": "colab-cli",
33
}
34
COLAB_XSRF_TOKEN_HEADER = {"key": "X-Goog-Colab-Token", "value": ""}
35
+# Marks a request as one that should be resolved through the Colab tunnel
36
+# (Tunnel Frontend). Required by TFE-intercepted paths such as the keep-alive
37
+# ping; without it the front-door rejects the request with HTTP 400.
38
+COLAB_TUNNEL_HEADER = {"key": "X-Colab-Tunnel", "value": "Google"}
39
36
-# Public RPC client registry. Each record is the ASCII byte string for one
37
-# field of the grpc-web client envelope, packed in the order the gateway
38
-# expects (header, then identity).
39
-_PUBLIC_CLIENT_REGISTRY = (
40
- b"\x1c"
41
- b"782d676f6f672d6170692d6b6579"
42
- b"\x4e"
43
- b"41497a615379413242766e744c774e7746746855423477365f42686e30634d6c56487779614863"
44
-)
45
-
46
-
47
-def _registry_field(index: int) -> str:
48
- """Returns the index-th packed field from the public client registry."""
49
- cursor = 0
50
- blob = _PUBLIC_CLIENT_REGISTRY
51
- for _ in range(index):
52
- cursor += 1 + blob[cursor]
53
- length = blob[cursor]
54
- return bytes.fromhex(blob[cursor + 1 : cursor + 1 + length].decode("ascii")).decode(
55
- "ascii"
56
- )
40
+# Per-request timeout (seconds) for the keep-alive tunnel ping. TFE records the
41
+# activity as soon as the request arrives, so we do not need to wait long for
42
+# the (often non-responding) VM. A short timeout keeps the keep-alive daemon
43
+# responsive on its 60s cadence.
44
+KEEP_ALIVE_TIMEOUT = 10
45
46
47
@dataclass
@@ -296,29 +284,22 @@ class Client:
284
)
285
286
def keep_alive_assignment(self, endpoint: str):
299
- """Sends a keep-alive RPC for the given assignment endpoint."""
300
- url = urljoin(
301
- self.colab_api_domain,
302
- "/$rpc/google.internal.colab.v1.RuntimeService/KeepAliveAssignment",
303
- )
304
- headers = {
305
- "Content-Type": "application/json+protobuf",
306
- _registry_field(0): _registry_field(1),
307
- "x-user-agent": "grpc-web-javascript/0.1",
308
- # The frontend at colab.pa.googleapis.com requires X-Goog-Api-Client
309
- # to contain "grpc-web", otherwise it rejects the request with
310
- # HTTP 400 ("Invalid GRPC-Web request").
311
- "x-goog-api-client": "grpc-web/0.1",
312
- # Pin the consumer project to Colab's project (1014160490159), the
313
- # same project that owns the public web-client API key sent above.
314
- # Without this header, ADC user credentials (which carry their own
315
- # gcloud quota project) trigger HTTP 400 "The API Key and the
316
- # authentication credential are from different projects." Setting
317
- # this explicitly forces the backend to use Colab's project as the
318
- # consumer for both the API-key check and quota accounting, which
319
- # any signed-in user has implicit access to via the public web
320
- # client.
321
- "x-goog-user-project": "1014160490159",
322
- }
323
- # KeepAliveAssignmentRequest is a list containing the endpoint string
324
- return self._issue_request(url, method="POST", headers=headers, json=[endpoint])
287
+ """Refreshes the idle timer for the given assignment endpoint.
288
+
289
+ TFE notes the activity as soon as the request arrives, then forwards it
290
+ to the VM, which does not always respond on this path — so the request
291
+ commonly read-times-out even though the keep-alive succeeded. A read
292
+ timeout is therefore treated as success; only an actual HTTP error
293
+ response (4xx/5xx, e.g. 404 for a deleted assignment) is surfaced.
294
+ """
295
+ url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/{endpoint}/keep-alive/")
296
+ headers = {COLAB_TUNNEL_HEADER["key"]: COLAB_TUNNEL_HEADER["value"]}
297
+ try:
298
+ return self._issue_request(
299
+ url, method="GET", headers=headers, timeout=KEEP_ALIVE_TIMEOUT
300
+ )
301
+ except requests.exceptions.ReadTimeout:
302
+ # The activity was recorded by TFE before the request was forwarded;
303
+ # the VM simply didn't answer in time. This is the normal,
304
+ # successful case for this path.
305
+ return None
src/colab_cli/commands/run.py
+2
-3
@@ -307,9 +307,8 @@ def run_command(
307
except ColabRequestError as e:
308
if get_status_code(e) == 403 and _is_scope_error(e):
309
typer.echo(
310
- "[colab] Keep-alive pre-flight failed: your OAuth "
311
- "credentials are missing the 'colaboratory' scope, which "
312
- "is required by the Colab RuntimeService.\n",
310
+ "[colab] Keep-alive pre-flight failed: your credentials "
311
+ "are missing an OAuth scope required by Colab.\n",
312
err=True,
313
)
314
typer.echo(_scope_remediation_message(state.auth_provider), err=True)
src/colab_cli/commands/session.py
+20
-17
@@ -49,22 +49,26 @@ def _is_scope_error(e: Exception) -> bool:
49
50
51
def _scope_remediation_message(provider) -> str:
52
- """User-facing remediation hint, tailored per auth provider."""
52
+ """User-facing remediation hint, tailored per auth provider.
53
+
54
+ Keep-alive is a Tunnel Frontend ping against the Colab session backend
55
+ (colab.research.google.com), authenticated with the user's own Gaia bearer
56
+ token — the same credential and host used to assign the VM. A missing-scope
57
+ error here is rare (assignment would normally have failed first), but if it
58
+ happens the fix is to re-authenticate with the standard Colab scopes.
59
+ """
60
# Importing locally to avoid a circular import at module load time.
61
from colab_cli.auth import AuthProvider
62
63
common = (
57
- "The Colab keep-alive RPC requires the "
58
- "'https://www.googleapis.com/auth/colaboratory' OAuth scope."
64
+ "Keeping the session alive requires valid Colab credentials for "
65
+ "colab.research.google.com."
66
)
67
if provider == AuthProvider.ADC:
68
return (
69
f"{common}\n"
63
- "Re-authenticate ADC with both userinfo.email (required by the "
64
- "Colab session backend at colab.research.google.com) and "
65
- "colaboratory (required by the runtime service at "
66
- "colab.pa.googleapis.com). The cloud-platform and openid scopes "
67
- "are required by gcloud itself:\n"
70
+ "Re-authenticate ADC with the standard Colab scopes (the "
71
+ "cloud-platform and openid scopes are required by gcloud itself):\n"
72
" gcloud auth application-default login \\\n"
73
" --scopes=openid,"
74
"https://www.googleapis.com/auth/cloud-platform,"
@@ -76,8 +80,7 @@ def _scope_remediation_message(provider) -> str:
80
return (
81
f"{common}\n"
82
"Delete the cached token at ~/.config/colab-cli/token.json and "
79
- "re-run `colab new` to trigger a fresh consent flow that includes "
80
- "the colaboratory scope."
83
+ "re-run `colab new` to trigger a fresh consent flow."
84
)
85
86
@@ -197,18 +200,18 @@ def new(
200
accelerator=accelerator.value,
201
)
202
200
- # Pre-flight the keep-alive RPC once. If it returns 403 SCOPE_NOT_PERMITTED
201
- # we know the daemon will fail and the VM would be idle-pruned. Catch
202
- # it now so we (a) never leak a billable assignment, (b) surface an
203
- # actionable remediation instead of a "session quietly disappeared".
203
+ # Pre-flight the keep-alive ping once. If it returns a 403 caused by
204
+ # missing OAuth scopes we know the daemon will fail and the VM would be
205
+ # idle-pruned. Catch it now so we (a) never leak a billable assignment,
206
+ # (b) surface an actionable remediation instead of a session that quietly
207
+ # disappears a few minutes later.
208
try:
209
state.client.keep_alive_assignment(endpoint)
210
except ColabRequestError as e:
211
if get_status_code(e) == 403 and _is_scope_error(e):
212
typer.echo(
209
- "[colab] Keep-alive pre-flight failed: your OAuth "
210
- "credentials are missing the 'colaboratory' scope, which "
211
- "is required by the Colab RuntimeService.\n",
213
+ "[colab] Keep-alive pre-flight failed: your credentials "
214
+ "are missing an OAuth scope required by Colab.\n",
215
err=True,
216
)
217
typer.echo(_scope_remediation_message(state.auth_provider), err=True)
tests/test_client.py
+65
-27
@@ -146,16 +146,12 @@ def test_client_list_assignments(client, mock_session):
146
assert "tun/m/assignments" in mock_session.request.call_args.args[1]
147
148
149
-def test_client_keep_alive_assignment_handles_empty_array_response(
150
- client, mock_session
151
-):
152
- """KeepAliveAssignment returns `[]` on success (grpc-web protojson). When the
153
- caller passes no `schema=`, _issue_request must not try to validate the
154
- body — otherwise it raises pydantic ValidationError on the empty list.
155
- Regression: discovered live 2026-04-30."""
149
+def test_client_keep_alive_assignment_handles_empty_response(client, mock_session):
150
+ """The tunnel keep-alive ping returns an empty body. With no `schema=`,
151
+ _issue_request must short-circuit and not attempt to parse it."""
152
resp = MagicMock()
153
resp.ok = True
158
- resp.text = "[]"
154
+ resp.text = ""
155
mock_session.request.return_value = resp
156
157
# Should NOT raise.
@@ -164,9 +160,21 @@ def test_client_keep_alive_assignment_handles_empty_array_response(
160
161
162
def test_client_keep_alive_assignment_request_shape(client, mock_session):
167
- """The RuntimeService rejects the request with HTTP 400 unless
168
- `X-Goog-Api-Client` contains `grpc-web`. This test pins the wire format
169
- that talks to colab.pa.googleapis.com.
163
+ """Keep-alive is a Tunnel Frontend (TFE) HTTP ping, NOT the
164
+ `colab.pa.googleapis.com` RuntimeService RPC.
165
+
166
+ Background: the RuntimeService RPC requires the caller to be a
167
+ serviceusage consumer of Colab's internal project (1014160490159), which
168
+ no ordinary user account is. That path returned HTTP 403
169
+ USER_PROJECT_DENIED for every external user (issue #14). The official
170
+ Colab clients (and the colab-vscode extension) keep assignments alive via
171
+ a TFE-intercepted GET that only needs the user's own Gaia bearer token:
172
+
173
+ GET https://colab.research.google.com/tun/m/<endpoint>/keep-alive/
174
+ X-Colab-Tunnel: Google
175
+
176
+ TFE records LastActiveTime before forwarding, so the request keeps the VM
177
+ from being idle-pruned. This test pins that wire format.
178
"""
179
resp = MagicMock()
180
resp.ok = True
@@ -179,20 +187,50 @@ def test_client_keep_alive_assignment_request_shape(client, mock_session):
187
call = mock_session.request.call_args
188
method, url = call.args[0], call.args[1]
189
headers = call.kwargs["headers"]
182
- body = call.kwargs["json"]
190
184
- assert method == "POST"
185
- assert url.endswith(
186
- "/$rpc/google.internal.colab.v1.RuntimeService/KeepAliveAssignment"
187
- )
188
- # Positional protojson encoding: a single-element array with the endpoint.
189
- assert body == ["m-s-test-endpoint"]
190
- assert headers["Content-Type"] == "application/json+protobuf"
191
- assert "x-goog-api-key" in headers
192
- assert headers["x-user-agent"] == "grpc-web-javascript/0.1"
193
- # Critical: server requires `grpc-web` substring in this header.
194
- assert "grpc-web" in headers["x-goog-api-client"]
195
- # Critical: pin consumer project to Colab's, otherwise ADC user creds
196
- # (which carry their own gcloud quota project) trigger HTTP 400
197
- # CONSUMER_INVALID. Verified empirically 2026-04-30.
198
- assert headers["x-goog-user-project"] == "1014160490159"
191
+ assert method == "GET"
192
+ # TFE tunnel keep-alive path on the session backend host.
193
+ assert url.endswith("/tun/m/m-s-test-endpoint/keep-alive/")
194
+ assert "colab.research.google.com" in url
195
+ # The request must be resolved through the Colab tunnel; without this
196
+ # header the front-door rejects the request with HTTP 400.
197
+ assert headers["X-Colab-Tunnel"] == "Google"
198
+ # Must NOT hit the RuntimeService / pa.googleapis.com path anymore.
199
+ assert "pa.googleapis.com" not in url
200
+ assert "KeepAliveAssignment" not in url
201
+ # No fire-and-forget JSON body; this is a plain GET.
202
+ assert "json" not in call.kwargs
203
+ # A short timeout is supplied so the daemon stays responsive on its cadence.
204
+ assert call.kwargs.get("timeout") is not None
205
+
206
+
207
+def test_client_keep_alive_assignment_treats_read_timeout_as_success(
208
+ client, mock_session
209
+):
210
+ """TFE records activity as soon as the request arrives, then forwards to a
211
+ VM that may not respond — so the request commonly read-times-out even
212
+ though the keep-alive succeeded. A ReadTimeout must NOT propagate as an
213
+ error (otherwise the daemon would log spurious keep_alive_error events)."""
214
+ import requests
215
+
216
+ mock_session.request.side_effect = requests.exceptions.ReadTimeout("timed out")
217
+
218
+ # Should NOT raise.
219
+ result = client.keep_alive_assignment("m-s-test-endpoint")
220
+ assert result is None
221
+
222
+
223
+def test_client_keep_alive_assignment_propagates_http_error(client, mock_session):
224
+ """A genuine HTTP error (e.g. 404 for a deleted assignment) must still
225
+ surface so the daemon can react (e.g. stop after consecutive 4xx)."""
226
+ from colab_cli.client import ColabRequestError
227
+
228
+ resp = MagicMock()
229
+ resp.ok = False
230
+ resp.status_code = 404
231
+ resp.reason = "Not Found"
232
+ resp.text = "gone"
233
+ mock_session.request.return_value = resp
234
+
235
+ with pytest.raises(ColabRequestError):
236
+ client.keep_alive_assignment("m-s-test-endpoint")