| 1 | #!/bin/bash |
| 2 | # Copyright 2026 Google LLC |
| 3 | # |
| 4 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | # you may not use this file except in compliance with the License. |
| 6 | # You may obtain a copy of the License at |
| 7 | # |
| 8 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | # |
| 10 | # Unless required by applicable law or agreed to in writing, software |
| 11 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | # See the License for the specific language governing permissions and |
| 14 | # limitations under the License. |
| 15 | # |
| 16 | # Integration Test: Keep-Alive Daemon Soak (OAuth Scope Regression Guard) |
| 17 | # |
| 18 | # Background: |
| 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`. |
| 38 | # 2. Waits 90 seconds — long enough for the daemon to hit at least one |
| 39 | # ping iteration *after* the pre-flight (the loop sleeps 60s between |
| 40 | # pings). |
| 41 | # 3. Reads the structured history via `colab log` and asserts NO |
| 42 | # `KEEP: error` events were recorded. Any error event means a daemon |
| 43 | # ping was rejected by the server, which is the regression. |
| 44 | # 4. Verifies the daemon process is still alive. |
| 45 | # 5. Cleans up via `colab stop`. |
| 46 | # |
| 47 | # Cost: ~95 seconds of real wall-clock time + one short-lived Colab CPU |
| 48 | # assignment. |
| 49 | |
| 50 | set -e |
| 51 | |
| 52 | # Use a uniquely-named session per run so we don't trip on stale history |
| 53 | # from previous runs (history files are keyed by session name and live at |
| 54 | # ~/.config/colab-cli/history/<name>.jsonl). |
| 55 | SESSION_NAME="repro-keep-alive-scope-$(date +%s)" |
| 56 | |
| 57 | TMP_DIR=$(mktemp -d) |
| 58 | SESSION_FILE="$TMP_DIR/sessions.json" |
| 59 | |
| 60 | # This test soaks the keep-alive daemon for 90s, so it is meaningful only on |
| 61 | # auth providers that actually spawn a daemon. We require either |
| 62 | # OAuth2 or properly-scoped ADC. |
| 63 | if [ -f "$HOME/.config/colab-cli/token.json" ]; then |
| 64 | AUTH_FLAGS="--auth=oauth2" |
| 65 | elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then |
| 66 | ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) |
| 67 | ADC_SCOPES=$(curl -s "https://www.googleapis.com/oauth2/v3/tokeninfo?access_token=$ADC_TOKEN" | python3 -c "import json,sys; print(json.load(sys.stdin).get('scope',''))" 2>/dev/null) |
| 68 | if echo "$ADC_SCOPES" | grep -q "colaboratory" && echo "$ADC_SCOPES" | grep -q "userinfo.email"; then |
| 69 | AUTH_FLAGS="--auth=adc" |
| 70 | else |
| 71 | echo "Skipping: ADC token lacks required scopes (need both" |
| 72 | echo " userinfo.email and colaboratory). Re-issue with:" |
| 73 | echo " gcloud auth application-default login \\" |
| 74 | echo " --scopes=openid,\\" |
| 75 | echo " https://www.googleapis.com/auth/cloud-platform,\\" |
| 76 | echo " https://www.googleapis.com/auth/userinfo.email,\\" |
| 77 | echo " https://www.googleapis.com/auth/colaboratory" |
| 78 | exit 0 # environment-not-applicable |
| 79 | fi |
| 80 | else |
| 81 | echo "Skipping: this test requires --auth=oauth2 or properly-scoped ADC" |
| 82 | echo " Bootstrap options:" |
| 83 | echo " - OAuth2: 'uv run colab --auth=oauth2 sessions' (browser consent)" |
| 84 | echo " - ADC: gcloud auth application-default login \\" |
| 85 | echo " --scopes=openid,\\" |
| 86 | echo " https://www.googleapis.com/auth/cloud-platform,\\" |
| 87 | echo " https://www.googleapis.com/auth/userinfo.email,\\" |
| 88 | echo " https://www.googleapis.com/auth/colaboratory" |
| 89 | exit 0 # environment-not-applicable |
| 90 | fi |
| 91 | |
| 92 | cleanup() { |
| 93 | echo "[*] Cleaning up..." |
| 94 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" 2>/dev/null || true |
| 95 | rm -rf "$TMP_DIR" |
| 96 | # Best-effort: scrub the history file so the test is idempotent. |
| 97 | rm -f "$HOME/.config/colab-cli/history/${SESSION_NAME}.jsonl" |
| 98 | } |
| 99 | trap cleanup EXIT |
| 100 | |
| 101 | echo "[*] Creating session '$SESSION_NAME' (REAL API CALL) using $AUTH_FLAGS..." |
| 102 | # Note: `colab new` now performs a synchronous keep-alive pre-flight. If the |
| 103 | # OAuth scope is missing, this command itself will fail fast with an |
| 104 | # actionable remediation message — so step (1) of the regression already |
| 105 | # fires here. |
| 106 | if ! uv run colab $AUTH_FLAGS --config "$SESSION_FILE" new -s "$SESSION_NAME"; then |
| 107 | echo "[FAILURE] 'colab new' failed. If this is a SCOPE_NOT_PERMITTED error," |
| 108 | echo " the colaboratory scope is missing from your auth provider." |
| 109 | echo " For ADC: gcloud auth application-default login \\" |
| 110 | echo " --scopes=openid,\\" |
| 111 | echo " https://www.googleapis.com/auth/cloud-platform,\\" |
| 112 | echo " https://www.googleapis.com/auth/userinfo.email,\\" |
| 113 | echo " https://www.googleapis.com/auth/colaboratory" |
| 114 | exit 1 |
| 115 | fi |
| 116 | |
| 117 | # Sanity-check the session was persisted with a daemon PID. |
| 118 | PID=$(grep -A 15 "$SESSION_NAME" "$SESSION_FILE" | grep "keep_alive_pid" | awk '{print $2}' | tr -d ',') |
| 119 | if [ -z "$PID" ] || [ "$PID" == "null" ]; then |
| 120 | echo "[FAILURE] No keep_alive_pid recorded for session." |
| 121 | cat "$SESSION_FILE" |
| 122 | exit 1 |
| 123 | fi |
| 124 | echo "[*] Keep-alive daemon PID: $PID" |
| 125 | |
| 126 | # Soak: wait long enough for at least one daemon-driven ping (loop sleeps |
| 127 | # 60s) to land *after* the pre-flight that `colab new` did. 90s gives us a |
| 128 | # comfortable margin. |
| 129 | echo "[*] Soaking for 90s to let the daemon perform at least one ping..." |
| 130 | sleep 90 |
| 131 | |
| 132 | # The daemon must still be alive. |
| 133 | if ! ps -p $PID > /dev/null; then |
| 134 | echo "[FAILURE] Keep-alive daemon (pid=$PID) died during soak." |
| 135 | echo " History dump:" |
| 136 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" log -s "$SESSION_NAME" || true |
| 137 | exit 1 |
| 138 | fi |
| 139 | echo "[*] Daemon still alive after 90s." |
| 140 | |
| 141 | # The structured history must NOT contain any keep_alive_error events. Any |
| 142 | # error here means a server-side rejection (auth, headers, payload) — the |
| 143 | # exact class of bug this test guards against. |
| 144 | LOG_OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" log -s "$SESSION_NAME") |
| 145 | echo "----- colab log output -----" |
| 146 | echo "$LOG_OUTPUT" |
| 147 | echo "----------------------------" |
| 148 | |
| 149 | if echo "$LOG_OUTPUT" | grep -q "KEEP: error"; then |
| 150 | echo "[FAILURE] keep_alive_error events recorded during soak." |
| 151 | echo " This indicates the daemon's pings are being rejected." |
| 152 | echo " Common causes:" |
| 153 | echo " - Missing 'colaboratory' OAuth scope (403 SCOPE_NOT_PERMITTED)" |
| 154 | echo " - Missing X-Goog-Api-Client: grpc-web header (400 Invalid GRPC-Web)" |
| 155 | exit 1 |
| 156 | fi |
| 157 | |
| 158 | # Positive assertion: we expect at least one KEEP: started event. |
| 159 | if ! echo "$LOG_OUTPUT" | grep -q "KEEP: started"; then |
| 160 | echo "[FAILURE] No KEEP: started event recorded — daemon never ran?" |
| 161 | exit 1 |
| 162 | fi |
| 163 | |
| 164 | echo "[SUCCESS] Keep-alive daemon survived 90s soak with zero error events." |