main
sh 164 lines 7.83 KB
Raw
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."