diff --git a/AGENTS.md b/AGENTS.md index 6e84b75b..668ff1e2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,7 +5,7 @@ - **Common**: `common.py` centralizes shared `State` (lazy-loading) and session resolution. - **Client**: `ColabClient` handles API interactions (assignment, unassignment). - **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`): - - `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). + - `oauth2`: public `google-auth-oauthlib` `InstalledAppFlow`, token cached in the OS keyring via the Python `keyring` library (`service_name="google-colab-cli"`, `username="oauth_token"`), falling back to `_run_remote_flow` if the keyring is not accessible or locked (keyring failures are echoed to stderr; logs alone go to a file). A pre-keyring `~/.config/colab-cli/token.json` (`LEGACY_TOKEN_PATH`) is migrated into the keyring and deleted only after a successful keyring write, or immediately when the keyring already holds a token. Integration scripts assume no legacy file exists and detect OAuth2 availability via a keyring-only `has_oauth2_token` helper; the migration path is covered by `integration/repro_legacy_token_migration/test.sh` (isolated HOME + file-backed test keyring). Tests that must ignore the host keyring set `PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring`. 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). - `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 retained for forward compatibility and other Colab features; `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`. - **Backend Hosts**: Two distinct backends with different requirements: - `colab.research.google.com` (session backend / `tun/m/...`): accepts the `userinfo.email` scope. Handles assign, unassign, CCU info, and the contents API. @@ -71,7 +71,7 @@ The package version is derived from the git tag via `hatch-vcs` (see `pyproject. 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`. 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 `~/.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. +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. 17. **Suggest the branch-diff review command after committing**: The user reviews changes with `git diff main..` (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 ` (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. 18. **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. diff --git a/docs/01_session_management.md b/docs/01_session_management.md index 5112ffc7..78e8244f 100644 --- a/docs/01_session_management.md +++ b/docs/01_session_management.md @@ -1,5 +1,6 @@ --- log: +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). 2026-10-14: Added configurable HTTP request timeouts and bounded allocation retries. Uncertain allocation POSTs are reconciled by GET using the same notebook hash before resubmission. 2026-09-25: Runtime proxy tokens are now refreshed before expiry. `SessionState` persists `token_expires_at`, and `State.get_session` re-fetches the token via `GET /tun/m/assignments` when it's within 5 minutes of expiring. 2026-09-25: Removed the keep-alive background daemon and related client/CLI ping functions. VM liveness is automatically maintained by the Colab backend based on kernel activity. @@ -122,7 +123,7 @@ The runtime proxy token (`runtimeProxyInfo.token`) authenticates every request t - **Global Flags**: - `-c`, `--client-oauth-config`: Path to the client secrets JSON file (default: `~/.colab-cli-oauth-config.json`). - `--config`: Path to the session state JSON file (default: `~/.config/colab-cli/sessions.json`). -- **Token Storage**: Credentials are persisted to `~/.config/colab-cli/token.json` after the initial flow. +- **Token Storage**: Credentials are persisted to the OS native keyring via the `keyring` library (`service_name="google-colab-cli"`, `username="oauth_token"`) after the initial flow. If the keyring is inaccessible or locked, `_run_remote_flow` is invoked for the current run and a warning is printed to stderr. A legacy `~/.config/colab-cli/token.json` is migrated into the keyring on first use and deleted after a successful keyring write, or immediately if the keyring already holds a token. - Use `requests` for robust HTTP interactions and `pydantic` for schema validation. - Handle authentication headers (likely `Authorization: Bearer ` or cookies). diff --git a/docs/04_automation_and_utility.md b/docs/04_automation_and_utility.md index 1a0e2287..9a544e73 100644 --- a/docs/04_automation_and_utility.md +++ b/docs/04_automation_and_utility.md @@ -1,5 +1,6 @@ --- log: +2026-10-21: Migrated OAuth2 token persistence in `auth.py` from `~/.config/colab-cli/token.json` to the OS native password storage via the `keyring` library (`KEYRING_SERVICE_NAME = "google-colab-cli"`, `KEYRING_USERNAME = "oauth_token"`). When the keyring is empty, locked, or inaccessible, `_run_remote_flow` is invoked; keyring read/write failures print a warning to stderr without failing the command. An existing `~/.config/colab-cli/token.json` (`LEGACY_TOKEN_PATH`) is used when the keyring has no token, copied into the keyring, and deleted only after the keyring write succeeds; it is also deleted when the keyring already holds a token. 2026-09-25: Cleaned up references to keep-alive background daemon and pre-flight pings. VM liveness is automatically maintained by the Colab backend based on active kernel activity. 2026-08-10: Added `colab usage` for account-level compute-unit rate/balance via `GET /tun/m/ccu-info` on the session backend (same bearer token as `colab new`). 2026-06-11: Replaced the `oauth2` provider's `run_local_server()` (localhost redirect) with a remote copy-paste flow (`_run_remote_flow` in `auth.py`). The CLI now prints an authorization URL built with `redirect_uri=https://sdk.cloud.google.com/applicationdefaultauthcode.html` and `token_usage=remote`, then reads the pasted authorization code via `input()` and exchanges it with `flow.fetch_token(code=...)`. This is the same flow `gcloud auth application-default login` uses and works identically in local and remote/headless/container environments, removing the heuristic of whether to auto-open a browser. Confirmed server-side acceptance with a live GET-only check against the bundled cloud-SDK client (`764086051850-...`); the OOB redirect and a non-bundled client id were both verified to be rejected (`OOB flow has been blocked` / `redirect_uri_mismatch`). Unit tests in `tests/test_auth.py` assert no localhost server is started, the redirect URI + `token_usage=remote` are set, and the pasted code is exchanged. @@ -32,7 +33,15 @@ backend, selected via the global `--auth=` flag: `token_usage=remote`) using the registered HTTPS landing page `https://sdk.cloud.google.com/applicationdefaultauthcode.html`; the user signs in, copies the code Google displays, and pastes it back at the - prompt. The refresh token is cached at `~/.config/colab-cli/token.json`. + prompt. The refresh token is cached in the OS native keyring via the + `keyring` library (`service_name="google-colab-cli"`, `username="oauth_token"`), + falling back to running `_run_remote_flow` if the keyring is not accessible + or locked. Keyring read/write failures are printed to stderr, since a + failed write means the user is prompted again on every command. A token + left at the pre-keyring path `~/.config/colab-cli/token.json` is migrated + into the keyring on first use and the file is deleted once the keyring + write succeeds (it is kept if the write fails). If the keyring already + holds a token, the legacy file is deleted. This is the same mechanism `gcloud auth application-default login` uses, and it behaves identically on local, remote, headless, and container hosts (no auto-opened browser, no bound port). We deliberately do **not** @@ -65,9 +74,9 @@ sufficient for this host. How each provider supplies the scope: - **`oauth2`**: `PUBLIC_SCOPES` already includes `colaboratory`, so the - InstalledAppFlow consent screen lists it. Existing cached tokens at - `~/.config/colab-cli/token.json` that were minted before this change must - be deleted to trigger a fresh consent flow. + InstalledAppFlow consent screen lists it. Existing cached tokens in the + keyring that were minted before a scope change must be cleared to trigger a + fresh consent flow. - **`adc`**: `google.auth.default(scopes=PUBLIC_SCOPES)` is called, and for credential subclasses that support `with_scopes` (service accounts, GCE/GKE metadata, impersonated) we re-apply via `creds.with_scopes(...)`. diff --git a/integration/README.md b/integration/README.md index 3e0cce6b..3588f5ec 100644 --- a/integration/README.md +++ b/integration/README.md @@ -16,6 +16,7 @@ End-to-end tests that run against a **live Colab backend** (unlike the mocked un | `repro_variable_persistence/` | Variables persist across `colab exec` calls in the same session. | | `repro_piped_console/` | Fast smoke test (~5s including session creation): `echo cmd \| colab console -s s` runs the command and exits within 30s. Regression test for the 2026-05-07 EOF-handler fix. | | `repro_bundled_oauth/` | Fast smoke test (~5s): verifies that the fallback OAuth configuration is loaded and starts the OAuth flow with the default client ID when local config is missing. | +| `repro_legacy_token_migration/` | Fast smoke test (~5s, no VM): a pre-keyring `~/.config/colab-cli/token.json` is migrated into the keyring and deleted, deleted when the keyring already has a token, and kept (with a stderr warning) when the keyring is inaccessible. Uses a throwaway HOME and file-backed test keyring. | | `repro_ssh/` | Fast smoke test (~5s): `--help` advertises the flags and an unknown session exits. Slow soak test (~95s): Live e2e allocates a CPU VM, runs a real remote command over `colab ssh --proxy-mode` | | `repro_windows_ssh/` | Native Windows live CPU test: Windows OpenSSH launches the generated ProxyCommand, executes a remote marker command, and releases the recorded runtime. Requires cached OAuth2 credentials; exits 77 if prerequisites are missing. | diff --git a/integration/repro_bundled_oauth/test.sh b/integration/repro_bundled_oauth/test.sh index 9b639b19..51fd5b58 100644 --- a/integration/repro_bundled_oauth/test.sh +++ b/integration/repro_bundled_oauth/test.sh @@ -17,26 +17,20 @@ # # This test verifies that the CLI correctly falls back to using the bundled # oauth_config.json containing the default client ID when no local config is present. -# It simulates a fresh user flow by temporarily hiding existing tokens/configs, -# starting the OAuth flow, and asserting that the printed authorization URL -# contains the correct Client ID. +# It simulates a fresh user flow by temporarily hiding an existing config and +# the keyring token, starting the OAuth flow, and asserting that the printed +# authorization URL contains the correct Client ID. # set -e -TOKEN_PATH="$HOME/.config/colab-cli/token.json" CONFIG_PATH="$HOME/.colab-cli-oauth-config.json" BACKUP_SUFFIX=".backup.$(date +%s)" -TOKEN_BACKUP="" CONFIG_BACKUP="" cleanup() { echo "[*] Cleaning up backups..." - if [ -n "$TOKEN_BACKUP" ] && [ -f "$TOKEN_BACKUP" ]; then - mv "$TOKEN_BACKUP" "$TOKEN_PATH" - echo "[*] Restored token.json" - fi if [ -n "$CONFIG_BACKUP" ] && [ -f "$CONFIG_BACKUP" ]; then mv "$CONFIG_BACKUP" "$CONFIG_PATH" echo "[*] Restored .colab-cli-oauth-config.json" @@ -44,13 +38,7 @@ cleanup() { } trap cleanup EXIT -# 1. Back up existing configuration files if they exist -if [ -f "$TOKEN_PATH" ]; then - TOKEN_BACKUP="${TOKEN_PATH}${BACKUP_SUFFIX}" - mv "$TOKEN_PATH" "$TOKEN_BACKUP" - echo "[*] Temporarily backed up token.json to $TOKEN_BACKUP" -fi - +# 1. Back up the existing client config if it exists if [ -f "$CONFIG_PATH" ]; then CONFIG_BACKUP="${CONFIG_PATH}${BACKUP_SUFFIX}" mv "$CONFIG_PATH" "$CONFIG_BACKUP" @@ -63,7 +51,9 @@ echo "[*] Running 'colab --auth=oauth2 sessions' (expecting to trigger browser f # We redirect output to a file so we can inspect it. OUTPUT_LOG=$(mktemp) set +e -PYTHONUNBUFFERED=1 timeout 5 uv run colab --auth=oauth2 sessions > "$OUTPUT_LOG" 2>&1 +# The null keyring backend hides any token cached in the host's OS keyring. +PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring \ + PYTHONUNBUFFERED=1 timeout 5 uv run colab --auth=oauth2 sessions > "$OUTPUT_LOG" 2>&1 EXIT_CODE=$? set -e @@ -76,7 +66,7 @@ echo "----------------------------------------------" EXPECTED_CLIENT_ID="764086051850-6qr4p6gpi6hn506pt8ejuq83di341hur.apps.googleusercontent.com" # The command should have printed the authorization URL -if ! grep -q "Please visit this URL to authorize this application" "$OUTPUT_LOG"; then +if ! grep -q "To authorize colab-cli, visit this URL" "$OUTPUT_LOG"; then echo "[FAILURE] OAuth prompt message not found in output." rm -f "$OUTPUT_LOG" exit 1 diff --git a/integration/repro_legacy_token_migration/test.sh b/integration/repro_legacy_token_migration/test.sh new file mode 100755 index 00000000..7941ab2e --- /dev/null +++ b/integration/repro_legacy_token_migration/test.sh @@ -0,0 +1,136 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Integration Test: Legacy token.json -> keyring migration +# +# Verifies how `colab --auth=oauth2` handles a pre-keyring plaintext token at +# ~/.config/colab-cli/token.json: +# 1. Keyring empty: token is used, copied into the keyring, file deleted. +# 2. Keyring has token: keyring token is kept, file deleted. +# 3. Keyring broken: token is used, file kept, warning printed to stderr. +# +# Runs against a throwaway HOME and a file-backed test keyring, so the real +# token and OS keyring are never touched. The token is fake (unexpired), so the +# backend request after authentication fails; only the local token handling is +# asserted. No VM is allocated. +# + +set -u + +# Keep uv's cache when HOME is overridden below. +export UV_CACHE_DIR="${UV_CACHE_DIR:-$HOME/.cache/uv}" + +TMP_DIR=$(mktemp -d) +trap 'rm -rf "$TMP_DIR"' EXIT + +TEST_HOME="$TMP_DIR/home" +LEGACY="$TEST_HOME/.config/colab-cli/token.json" +KEYRING_FILE="$TMP_DIR/keyring.json" +KEY="google-colab-cli/oauth_token" +mkdir -p "$TMP_DIR/mods" + +cat > "$TMP_DIR/mods/filekeyring.py" <<'EOF' +import json +import os + +from keyring.backend import KeyringBackend + +PATH = os.environ["FILE_KEYRING_PATH"] + + +class Keyring(KeyringBackend): + priority = 1 + + def _load(self): + if not os.path.exists(PATH): + return {} + with open(PATH) as f: + return json.load(f) + + def get_password(self, service, username): + return self._load().get(f"{service}/{username}") + + def set_password(self, service, username, password): + data = self._load() + data[f"{service}/{username}"] = password + with open(PATH, "w") as f: + json.dump(data, f) + + def delete_password(self, service, username): + pass +EOF + +fake_token() { + echo "{\"token\": \"$1\", \"refresh_token\": \"r\", \"client_id\": \"c\", \"client_secret\": \"s\", \"token_uri\": \"https://oauth2.googleapis.com/token\", \"expiry\": \"2099-01-01T00:00:00Z\"}" +} + +keyring_access_token() { + python3 -c "import json, sys; d = json.load(open(sys.argv[1])); print(json.loads(d[sys.argv[2]])['token'])" "$KEYRING_FILE" "$KEY" 2>/dev/null +} + +reset_state() { + rm -rf "$TEST_HOME" "$KEYRING_FILE" + mkdir -p "$(dirname "$LEGACY")" + fake_token "legacy" > "$LEGACY" +} + +# Runs the CLI with the given keyring backend; output goes to $OUT. +run_colab() { + OUT="$TMP_DIR/out.log" + HOME="$TEST_HOME" PYTHONPATH="$TMP_DIR/mods" FILE_KEYRING_PATH="$KEYRING_FILE" \ + PYTHON_KEYRING_BACKEND="$1" PYTHONUNBUFFERED=1 \ + timeout 60 uv run colab --auth=oauth2 sessions < /dev/null > "$OUT" 2>&1 +} + +FAILURES=0 +check() { + if eval "$2"; then + echo " [PASS] $1" + else + echo " [FAIL] $1" + FAILURES=$((FAILURES + 1)) + echo " ---- CLI output ----" + sed 's/^/ | /' "$OUT" + fi +} + +echo "[*] Scenario 1: keyring empty, legacy token.json present" +reset_state +run_colab filekeyring.Keyring +check "no OAuth prompt" '! grep -q "To authorize colab-cli" "$OUT"' +check "legacy token copied into keyring" '[ "$(keyring_access_token)" = "legacy" ]' +check "legacy token.json deleted" '[ ! -f "$LEGACY" ]' + +echo "[*] Scenario 2: keyring already has a token, legacy token.json present" +reset_state +python3 -c "import json, sys; json.dump({sys.argv[2]: sys.argv[3]}, open(sys.argv[1], 'w'))" \ + "$KEYRING_FILE" "$KEY" "$(fake_token keyring)" +run_colab filekeyring.Keyring +check "no OAuth prompt" '! grep -q "To authorize colab-cli" "$OUT"' +check "keyring token unchanged" '[ "$(keyring_access_token)" = "keyring" ]' +check "legacy token.json deleted" '[ ! -f "$LEGACY" ]' + +echo "[*] Scenario 3: keyring inaccessible, legacy token.json present" +reset_state +run_colab keyring.backends.fail.Keyring +check "no OAuth prompt" '! grep -q "To authorize colab-cli" "$OUT"' +check "keyring save warning on stderr" 'grep -q "Could not save the OAuth token to the OS keyring" "$OUT"' +check "legacy token.json kept" '[ -f "$LEGACY" ]' + +if [ "$FAILURES" -ne 0 ]; then + echo "[FAILURE] $FAILURES check(s) failed." + exit 1 +fi +echo "[SUCCESS] Legacy token.json migration behaves as expected." diff --git a/integration/repro_piped_console/test.sh b/integration/repro_piped_console/test.sh index ca16ed37..fe21ca80 100755 --- a/integration/repro_piped_console/test.sh +++ b/integration/repro_piped_console/test.sh @@ -28,8 +28,12 @@ TMP_DIR=$(mktemp -d) SESSION_FILE="$TMP_DIR/sessions.json" SESSION_NAME="test-piped-console" +# OAuth2 is usable if a token is cached in the OS keyring. +has_oauth2_token() { + timeout 10 uv run python -c "import sys, keyring; sys.exit(0 if keyring.get_password('google-colab-cli', 'oauth_token') else 1)" > /dev/null 2>&1 +} # Auth selection (OAuth2 if token cached, else ADC). -if [ -f "$HOME/.config/colab-cli/token.json" ]; then +if has_oauth2_token; then AUTH_FLAGS="--auth=oauth2" elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) diff --git a/integration/repro_run_command/test.sh b/integration/repro_run_command/test.sh index c8d83a74..94d458bf 100644 --- a/integration/repro_run_command/test.sh +++ b/integration/repro_run_command/test.sh @@ -24,7 +24,11 @@ # Don't `set -e` so we can capture failures and clean up explicitly. # ---------- Auth detection ---------------------------------------------------- -if [ -f "$HOME/.config/colab-cli/token.json" ]; then +# OAuth2 is usable if a token is cached in the OS keyring. +has_oauth2_token() { + timeout 10 uv run python -c "import sys, keyring; sys.exit(0 if keyring.get_password('google-colab-cli', 'oauth_token') else 1)" > /dev/null 2>&1 +} +if has_oauth2_token; then AUTH_FLAGS="--auth=oauth2" elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) diff --git a/integration/repro_ssh/test.sh b/integration/repro_ssh/test.sh index 5dbdbf16..e34f2b78 100644 --- a/integration/repro_ssh/test.sh +++ b/integration/repro_ssh/test.sh @@ -65,8 +65,12 @@ fi echo "OFFLINE SMOKE PASSED" # ---------- Part B: live end-to-end (needs auth + a VM) ---------------------- +# OAuth2 is usable if a token is cached in the OS keyring. +has_oauth2_token() { + timeout 10 uv run python -c "import sys, keyring; sys.exit(0 if keyring.get_password('google-colab-cli', 'oauth_token') else 1)" > /dev/null 2>&1 +} # Auth detection (mirrors integration/repro_run_command/test.sh). -if [ -f "$HOME/.config/colab-cli/token.json" ]; then +if has_oauth2_token; then AUTH_FLAGS="--auth=oauth2" elif command -v gcloud >/dev/null && gcloud auth application-default print-access-token >/dev/null 2>&1; then ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) diff --git a/pyproject.toml b/pyproject.toml index cbf86622..06c6d1bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "google-auth-oauthlib>=1.3.0", "html2text>=2024.2.26", "jupyter-kernel-client==0.9.0", + "keyring>=25.7.0", "nbformat>=5.10.4", "packaging>=24.0", "prompt-toolkit>=3.0.52", diff --git a/skills/colab-operator/SKILL.md b/skills/colab-operator/SKILL.md index b5f2a687..c8892d70 100644 --- a/skills/colab-operator/SKILL.md +++ b/skills/colab-operator/SKILL.md @@ -36,7 +36,7 @@ by running `uv tool install google-colab-cli` or `pip install google-colab-cli`. https://www.googleapis.com/auth/colaboratory ``` Why all four: `userinfo.email` (session backend `colab.research.google.com`, else 401), `colaboratory` (Colab features/APIs), `openid`+`cloud-platform` (mandated by gcloud itself; it rejects scope lists missing `cloud-platform`). -- **oauth2 setup**: `colab --auth=oauth2 ` triggers a browser consent flow on first use (token cached at `~/.config/colab-cli/token.json`). Requires a client config at `~/.colab-cli-oauth-config.json` (or `-c PATH`). The browser step means it usually needs a human; prefer ADC for agents. +- **oauth2 setup**: `colab --auth=oauth2 ` triggers a browser consent flow on first use (token cached in the OS keyring via `keyring`). Requires a client config at `~/.colab-cli-oauth-config.json` (or `-c PATH`), falling back to the bundled config. The browser step means it usually needs a human; prefer ADC for agents. - **Verify auth in one shot**: `colab sessions` (read-only, lists server assignments) or `colab whoami` (hidden debug command: prints the active email, scopes, audience, and expiry). When any call 401s or 403s, the cause is almost always missing credentials or scopes — `colab whoami` shows it instantly. - **Do NOT confuse `colab auth` with CLI authentication.** `colab auth` injects *VM-side* GCP credentials into the running kernel (so notebook code can call BigQuery/GCS); it is orthogonal to how the CLI itself authenticates. Never suggest "run `colab auth`" to fix a CLI 401/403 — that's a scope/identity problem fixed via the `gcloud` command above. diff --git a/src/colab_cli/auth.py b/src/colab_cli/auth.py index 3c530219..b95d5513 100644 --- a/src/colab_cli/auth.py +++ b/src/colab_cli/auth.py @@ -17,10 +17,12 @@ import logging from importlib import resources import os +import threading import warnings from typing import Optional import google.auth +import keyring import typer from google.auth.transport import requests from google.auth.transport.requests import Request @@ -51,7 +53,58 @@ class AuthProvider(str, enum.Enum): ] -TOKEN_CONFIG_PATH = os.path.expanduser("~/.config/colab-cli/token.json") +KEYRING_SERVICE_NAME = "google-colab-cli" +KEYRING_USERNAME = "oauth_token" +KEYRING_TIMEOUT_SECONDS = 30.0 + +# Plaintext token location used before tokens moved to the keyring. Read once +# to migrate existing users, then deleted after a successful keyring write. +LEGACY_TOKEN_PATH = os.path.expanduser("~/.config/colab-cli/token.json") + + +def _run_with_timeout(func, *args, timeout: float): + """Run a blocking callable in a daemon thread and raise TimeoutError if it hangs.""" + result = None + error: Optional[BaseException] = None + + def target() -> None: + nonlocal result, error + try: + result = func(*args) + except BaseException as exc: + error = exc + + thread = threading.Thread(target=target, daemon=True) + thread.start() + thread.join(timeout=timeout) + if thread.is_alive(): + # Python's threading.Thread has no API to terminate a thread blocked in + # an OS/D-Bus call; daemon=True ensures it won't block interpreter exit + # and is reaped by the OS when the CLI process finishes. + op_name = getattr(func, "__name__", None) or repr(func) + raise TimeoutError(f"Operation '{op_name}' timed out after {timeout:.1f}s") + if error is not None: + raise error + return result + + +def _get_token_from_keyring() -> Optional[str]: + return _run_with_timeout( + keyring.get_password, + KEYRING_SERVICE_NAME, + KEYRING_USERNAME, + timeout=KEYRING_TIMEOUT_SECONDS, + ) + + +def _save_token_to_keyring(token_json: str) -> None: + _run_with_timeout( + keyring.set_password, + KEYRING_SERVICE_NAME, + KEYRING_USERNAME, + token_json, + timeout=KEYRING_TIMEOUT_SECONDS, + ) # Remote copy-paste OAuth flow. # @@ -118,17 +171,33 @@ def _get_google_auth_credentials(config_path: str) -> Credentials: ) creds = None + needs_save = False - # Ensure config directory exists for the token file - os.makedirs(os.path.dirname(TOKEN_CONFIG_PATH), exist_ok=True) + try: + token_json = _get_token_from_keyring() + if token_json: + creds = Credentials.from_authorized_user_info( + json.loads(token_json), PUBLIC_SCOPES + ) + _remove_legacy_token() + except Exception as e: + logger.warning(f"Failed to load token from keyring: {e}") + # Each command is a separate process, so without this the user would be + # re-prompted on every run with no visible reason. + typer.echo( + f"Warning: Could not read the OAuth token from the OS keyring ({e}).", + err=True, + ) - if os.path.exists(TOKEN_CONFIG_PATH): + if not creds and os.path.exists(LEGACY_TOKEN_PATH): try: creds = Credentials.from_authorized_user_file( - TOKEN_CONFIG_PATH, PUBLIC_SCOPES + LEGACY_TOKEN_PATH, PUBLIC_SCOPES ) + needs_save = True except Exception as e: - logger.warning(f"Failed to load token from {TOKEN_CONFIG_PATH}: {e}") + logger.warning(f"Failed to load legacy token from {LEGACY_TOKEN_PATH}: {e}") + _remove_legacy_token() if not creds or not creds.valid: if creds and creds.expired and creds.refresh_token: @@ -140,17 +209,37 @@ def _get_google_auth_credentials(config_path: str) -> Credentials: if not creds: creds = _run_remote_flow(client_config) + needs_save = True - # Save the credentials for the next run + if needs_save: try: - with open(TOKEN_CONFIG_PATH, "w") as token_file: - token_file.write(creds.to_json()) + _save_token_to_keyring(creds.to_json()) except Exception as e: - logger.error(f"Failed to save token to {TOKEN_CONFIG_PATH}: {e}") + logger.warning(f"Failed to save token to keyring: {e}") + # Each command is a separate process, so without this the user would be + # re-prompted on every run with no visible reason. + typer.echo( + f"Warning: Could not save the OAuth token to the OS keyring ({e}). " + "You will be asked to sign in again on the next command. " + "Unlock or configure your keyring, or use --auth=adc.", + err=True, + ) + else: + _remove_legacy_token() return creds +def _remove_legacy_token() -> None: + """Removes the legacy token file if it exists.""" + if not os.path.exists(LEGACY_TOKEN_PATH): + return + try: + os.remove(LEGACY_TOKEN_PATH) + except OSError as e: + logger.warning(f"Failed to remove legacy token {LEGACY_TOKEN_PATH}: {e}") + + def _get_adc_credentials() -> Credentials: """Retrieves credentials using Google Application Default Credentials. diff --git a/tests/test_auth.py b/tests/test_auth.py index cff95586..83d70c7a 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -14,12 +14,17 @@ from unittest.mock import MagicMock, mock_open, patch +from keyring.errors import KeyringError import pytest from colab_cli.auth import ( + KEYRING_SERVICE_NAME, + KEYRING_USERNAME, + LEGACY_TOKEN_PATH, + PUBLIC_SCOPES, REMOTE_REDIRECT_URI, - TOKEN_CONFIG_PATH, AuthProvider, + _run_with_timeout, get_credentials, ) @@ -27,20 +32,21 @@ @pytest.fixture def mock_deps(mocker): m_exists = mocker.patch("os.path.exists") - m_makedirs = mocker.patch("os.makedirs") + m_keyring = mocker.patch("colab_cli.auth.keyring", create=True) m_creds_cls = mocker.patch("colab_cli.auth.Credentials") m_flow_cls = mocker.patch("colab_cli.auth.InstalledAppFlow") m_request = mocker.patch("colab_cli.auth.Request") m_session = mocker.patch("colab_cli.auth.requests.AuthorizedSession") m_resources = mocker.patch("colab_cli.auth.resources") - # By default, pretend oauth config doesn't exist anywhere + # By default, pretend oauth config doesn't exist anywhere and keyring is empty m_exists.return_value = False + m_keyring.get_password.return_value = None m_resources.files.return_value.joinpath.return_value.is_file.return_value = False return { "exists": m_exists, - "makedirs": m_makedirs, + "keyring": m_keyring, "creds_cls": m_creds_cls, "flow_cls": m_flow_cls, "request": m_request, @@ -58,52 +64,54 @@ def test_get_credentials_no_config(mock_deps): def test_get_credentials_valid_token(mock_deps): - # Setup token exists - def exists_side_effect(path): - return path in ["dummy_config.json", TOKEN_CONFIG_PATH] - - mock_deps["exists"].side_effect = exists_side_effect + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].get_password.return_value = '{"token": "valid"}' - # Valid creds mock_creds = MagicMock() mock_creds.valid = True - mock_deps["creds_cls"].from_authorized_user_file.return_value = mock_creds + mock_deps["creds_cls"].from_authorized_user_info.return_value = mock_creds - # Mock open for config m_open = mock_open(read_data='{"web":{"client_id":"id"}}') with patch("builtins.open", m_open): res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) - mock_deps["creds_cls"].from_authorized_user_file.assert_called_once() + mock_deps["keyring"].get_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME + ) + mock_deps["creds_cls"].from_authorized_user_info.assert_called_once_with( + {"token": "valid"}, PUBLIC_SCOPES + ) + mock_deps["keyring"].set_password.assert_not_called() mock_deps["session"].assert_called_once_with(mock_creds) assert res == mock_deps["session"].return_value def test_get_credentials_expired_token_refresh(mock_deps): - def exists_side_effect(path): - return path in ["dummy_config.json", TOKEN_CONFIG_PATH] - - mock_deps["exists"].side_effect = exists_side_effect + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].get_password.return_value = '{"token": "expired"}' mock_creds = MagicMock() mock_creds.valid = False mock_creds.expired = True mock_creds.refresh_token = "some_token" mock_creds.to_json.return_value = '{"token":"refreshed"}' - mock_deps["creds_cls"].from_authorized_user_file.return_value = mock_creds + mock_deps["creds_cls"].from_authorized_user_info.return_value = mock_creds m_open = mock_open(read_data='{"web":{"client_id":"id"}}') with patch("builtins.open", m_open): res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) mock_creds.refresh.assert_called_once() - m_open.assert_any_call(TOKEN_CONFIG_PATH, "w") + mock_deps["keyring"].set_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME, '{"token":"refreshed"}' + ) assert res == mock_deps["session"].return_value def test_get_credentials_no_token(mock_deps, mocker): """With no cached token, the remote copy-paste flow runs and exchanges code.""" mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].get_password.return_value = None mock_flow = MagicMock() mock_creds_new = MagicMock() @@ -128,6 +136,313 @@ def test_get_credentials_no_token(mock_deps, mocker): assert kwargs.get("token_usage") == "remote" # The pasted code is exchanged for a token. mock_flow.fetch_token.assert_called_once_with(code="pasted-code") + mock_deps["keyring"].set_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME, '{"token":"new"}' + ) + + +def test_get_credentials_keyring_inaccessible(mock_deps, mocker): + """When keyring is inaccessible, invoke _run_remote_flow and return credentials.""" + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].get_password.side_effect = KeyringError( + "No recommended backend was available" + ) + mock_deps["keyring"].set_password.side_effect = KeyringError( + "No recommended backend was available" + ) + + mock_flow = MagicMock() + mock_creds_new = MagicMock() + mock_creds_new.to_json.return_value = '{"token":"new"}' + mock_flow.authorization_url.return_value = ("https://auth.example/url", "state") + mock_flow.credentials = mock_creds_new + mock_deps["flow_cls"].from_client_config.return_value = mock_flow + + mocker.patch("colab_cli.auth.input", create=True, return_value="pasted-code") + + m_open = mock_open(read_data='{"web":{"client_id":"id"}}') + with patch("builtins.open", m_open): + res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + mock_deps["keyring"].get_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME + ) + mock_deps["flow_cls"].from_client_config.assert_called_once() + mock_flow.fetch_token.assert_called_once_with(code="pasted-code") + mock_deps["keyring"].set_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME, '{"token":"new"}' + ) + mock_deps["session"].assert_called_once_with(mock_creds_new) + assert res == mock_deps["session"].return_value + + +def test_get_credentials_corrupt_keyring_data(mock_deps, mocker): + """When keyring contains corrupt JSON, fall back to _run_remote_flow and overwrite.""" + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].get_password.return_value = "not-valid-json" + + mock_flow = MagicMock() + mock_creds_new = MagicMock() + mock_creds_new.to_json.return_value = '{"token":"new"}' + mock_flow.authorization_url.return_value = ("https://auth.example/url", "state") + mock_flow.credentials = mock_creds_new + mock_deps["flow_cls"].from_client_config.return_value = mock_flow + + mocker.patch("colab_cli.auth.input", create=True, return_value="pasted-code") + + m_open = mock_open(read_data='{"web":{"client_id":"id"}}') + with patch("builtins.open", m_open): + res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + mock_flow.fetch_token.assert_called_once_with(code="pasted-code") + mock_deps["keyring"].set_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME, '{"token":"new"}' + ) + assert res == mock_deps["session"].return_value + + +def test_get_credentials_refresh_failure_falls_back_to_remote_flow(mock_deps, mocker): + """When refreshing an expired token fails, fall back to _run_remote_flow.""" + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].get_password.return_value = '{"token": "expired"}' + + mock_creds_expired = MagicMock() + mock_creds_expired.valid = False + mock_creds_expired.expired = True + mock_creds_expired.refresh_token = "bad_refresh_token" + mock_creds_expired.refresh.side_effect = Exception("Token revoked") + mock_deps["creds_cls"].from_authorized_user_info.return_value = mock_creds_expired + + mock_flow = MagicMock() + mock_creds_new = MagicMock() + mock_creds_new.to_json.return_value = '{"token":"new"}' + mock_flow.authorization_url.return_value = ("https://auth.example/url", "state") + mock_flow.credentials = mock_creds_new + mock_deps["flow_cls"].from_client_config.return_value = mock_flow + + mocker.patch("colab_cli.auth.input", create=True, return_value="pasted-code") + + m_open = mock_open(read_data='{"web":{"client_id":"id"}}') + with patch("builtins.open", m_open): + res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + mock_creds_expired.refresh.assert_called_once() + mock_flow.fetch_token.assert_called_once_with(code="pasted-code") + mock_deps["keyring"].set_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME, '{"token":"new"}' + ) + mock_deps["session"].assert_called_once_with(mock_creds_new) + assert res == mock_deps["session"].return_value + + +def test_get_credentials_keyring_timeout_falls_back_to_remote_flow(mock_deps, mocker): + """When a keyring call hangs beyond the timeout, fall back to _run_remote_flow.""" + import threading + + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mocker.patch("colab_cli.auth.KEYRING_TIMEOUT_SECONDS", 0.05) + + hang_event = threading.Event() + mock_deps["keyring"].get_password.side_effect = lambda *a: hang_event.wait(1.0) + mock_deps["keyring"].set_password.side_effect = lambda *a: hang_event.wait(1.0) + + mock_flow = MagicMock() + mock_creds_new = MagicMock() + mock_creds_new.to_json.return_value = '{"token":"new"}' + mock_flow.authorization_url.return_value = ("https://auth.example/url", "state") + mock_flow.credentials = mock_creds_new + mock_deps["flow_cls"].from_client_config.return_value = mock_flow + + mocker.patch("colab_cli.auth.input", create=True, return_value="pasted-code") + + m_open = mock_open(read_data='{"web":{"client_id":"id"}}') + try: + with patch("builtins.open", m_open): + res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + finally: + hang_event.set() + + mock_flow.fetch_token.assert_called_once_with(code="pasted-code") + mock_deps["session"].assert_called_once_with(mock_creds_new) + assert res == mock_deps["session"].return_value + + +def _setup_remote_flow(mock_deps, mocker): + mock_flow = MagicMock() + mock_creds_new = MagicMock() + mock_creds_new.to_json.return_value = '{"token":"new"}' + mock_flow.authorization_url.return_value = ("https://auth.example/url", "state") + mock_flow.credentials = mock_creds_new + mock_deps["flow_cls"].from_client_config.return_value = mock_flow + mocker.patch("colab_cli.auth.input", create=True, return_value="pasted-code") + return mock_flow, mock_creds_new + + +def test_keyring_read_failure_warns_on_stderr(mock_deps, mocker, capsys): + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].get_password.side_effect = KeyringError("locked") + _setup_remote_flow(mock_deps, mocker) + + with patch("builtins.open", mock_open(read_data='{"web":{"client_id":"id"}}')): + get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + err = capsys.readouterr().err + assert "Could not read the OAuth token from the OS keyring" in err + assert "locked" in err + + +def test_keyring_save_failure_warns_on_stderr(mock_deps, mocker, capsys): + mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json" + mock_deps["keyring"].set_password.side_effect = KeyringError("no backend") + _setup_remote_flow(mock_deps, mocker) + + with patch("builtins.open", mock_open(read_data='{"web":{"client_id":"id"}}')): + get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + err = capsys.readouterr().err + assert "Could not save the OAuth token to the OS keyring" in err + assert "no backend" in err + assert "--auth=adc" in err + + +def test_legacy_token_file_migrated_to_keyring(mock_deps, mocker): + """A pre-keyring token.json is used, copied into the keyring, then deleted.""" + mock_deps["exists"].side_effect = lambda path: path in ( + "dummy_config.json", + LEGACY_TOKEN_PATH, + ) + m_remove = mocker.patch("colab_cli.auth.os.remove") + + legacy_creds = MagicMock() + legacy_creds.valid = True + legacy_creds.to_json.return_value = '{"token":"legacy"}' + mock_deps["creds_cls"].from_authorized_user_file.return_value = legacy_creds + + with patch("builtins.open", mock_open(read_data='{"web":{"client_id":"id"}}')): + res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + mock_deps["creds_cls"].from_authorized_user_file.assert_called_once_with( + LEGACY_TOKEN_PATH, PUBLIC_SCOPES + ) + mock_deps["flow_cls"].from_client_config.assert_not_called() + mock_deps["keyring"].set_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME, '{"token":"legacy"}' + ) + m_remove.assert_called_once_with(LEGACY_TOKEN_PATH) + mock_deps["session"].assert_called_once_with(legacy_creds) + assert res == mock_deps["session"].return_value + + +def test_legacy_token_file_kept_when_keyring_save_fails(mock_deps, mocker): + """Never delete the legacy token unless it was persisted to the keyring.""" + mock_deps["exists"].side_effect = lambda path: path in ( + "dummy_config.json", + LEGACY_TOKEN_PATH, + ) + mock_deps["keyring"].get_password.side_effect = KeyringError("locked") + mock_deps["keyring"].set_password.side_effect = KeyringError("locked") + m_remove = mocker.patch("colab_cli.auth.os.remove") + + legacy_creds = MagicMock() + legacy_creds.valid = True + legacy_creds.to_json.return_value = '{"token":"legacy"}' + mock_deps["creds_cls"].from_authorized_user_file.return_value = legacy_creds + + with patch("builtins.open", mock_open(read_data='{"web":{"client_id":"id"}}')): + get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + mock_deps["flow_cls"].from_client_config.assert_not_called() + m_remove.assert_not_called() + mock_deps["session"].assert_called_once_with(legacy_creds) + + +def test_legacy_token_file_deleted_when_keyring_has_token(mock_deps, mocker): + mock_deps["exists"].side_effect = lambda path: path in ( + "dummy_config.json", + LEGACY_TOKEN_PATH, + ) + m_remove = mocker.patch("colab_cli.auth.os.remove") + mock_deps["keyring"].get_password.return_value = '{"token": "valid"}' + mock_creds = MagicMock() + mock_creds.valid = True + mock_deps["creds_cls"].from_authorized_user_info.return_value = mock_creds + + with patch("builtins.open", mock_open(read_data='{"web":{"client_id":"id"}}')): + get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + mock_deps["creds_cls"].from_authorized_user_file.assert_not_called() + mock_deps["keyring"].set_password.assert_not_called() + m_remove.assert_called_once_with(LEGACY_TOKEN_PATH) + mock_deps["session"].assert_called_once_with(mock_creds) + + +def test_corrupt_legacy_token_file_falls_back_to_remote_flow(mock_deps, mocker): + """A corrupt legacy token is deleted right away, then the remote flow runs.""" + files = {"dummy_config.json", LEGACY_TOKEN_PATH} + mock_deps["exists"].side_effect = lambda path: path in files + m_remove = mocker.patch( + "colab_cli.auth.os.remove", side_effect=lambda path: files.discard(path) + ) + mock_deps["creds_cls"].from_authorized_user_file.side_effect = ValueError("bad") + mock_flow, mock_creds_new = _setup_remote_flow(mock_deps, mocker) + + legacy_present_at_prompt = [] + mocker.patch( + "colab_cli.auth.input", + create=True, + side_effect=lambda _: legacy_present_at_prompt.append( + LEGACY_TOKEN_PATH in files + ) + or "pasted-code", + ) + + with patch("builtins.open", mock_open(read_data='{"web":{"client_id":"id"}}')): + get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + assert legacy_present_at_prompt == [False] + m_remove.assert_called_once_with(LEGACY_TOKEN_PATH) + mock_flow.fetch_token.assert_called_once_with(code="pasted-code") + mock_deps["keyring"].set_password.assert_called_once_with( + KEYRING_SERVICE_NAME, KEYRING_USERNAME, '{"token":"new"}' + ) + + +def test_corrupt_legacy_token_file_deleted_even_if_keyring_save_fails( + mock_deps, mocker +): + """The corrupt file holds nothing worth keeping, unlike a valid legacy token.""" + mock_deps["exists"].side_effect = lambda path: path in ( + "dummy_config.json", + LEGACY_TOKEN_PATH, + ) + mock_deps["keyring"].set_password.side_effect = KeyringError("locked") + m_remove = mocker.patch("colab_cli.auth.os.remove") + mock_deps["creds_cls"].from_authorized_user_file.side_effect = ValueError("bad") + _setup_remote_flow(mock_deps, mocker) + + with patch("builtins.open", mock_open(read_data='{"web":{"client_id":"id"}}')): + get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2) + + m_remove.assert_called_once_with(LEGACY_TOKEN_PATH) + + +def test_run_with_timeout_infers_operation_name(): + """TimeoutError message should include the inferred callable name.""" + import threading + + hang_event = threading.Event() + + def custom_slow_operation(): + hang_event.wait(1.0) + + try: + with pytest.raises( + TimeoutError, + match="Operation 'custom_slow_operation' timed out after 0.0s", + ): + _run_with_timeout(custom_slow_operation, timeout=0.01) + finally: + hang_event.set() def test_remote_redirect_is_not_oob(): @@ -143,16 +458,16 @@ def test_get_credentials_fallback_config(mock_deps): m_file.is_file.return_value = True m_file.read_text.return_value = '{"installed":{"client_id":"fallback_id"}}' - # Valid creds in token - mock_deps["exists"].side_effect = lambda path: path == TOKEN_CONFIG_PATH + # Valid creds in keyring + mock_deps["keyring"].get_password.return_value = '{"token": "valid"}' mock_creds = MagicMock() mock_creds.valid = True - mock_deps["creds_cls"].from_authorized_user_file.return_value = mock_creds + mock_deps["creds_cls"].from_authorized_user_info.return_value = mock_creds res = get_credentials("missing_config.json", provider=AuthProvider.OAUTH2) mock_deps["resources"].files.assert_called_once_with("colab_cli") m_file.is_file.assert_called_once() m_file.read_text.assert_called_once() - mock_deps["creds_cls"].from_authorized_user_file.assert_called_once() + mock_deps["creds_cls"].from_authorized_user_info.assert_called_once() assert res == mock_deps["session"].return_value diff --git a/uv.lock b/uv.lock index 75068327..ccd5e622 100644 --- a/uv.lock +++ b/uv.lock @@ -377,6 +377,7 @@ dependencies = [ { name = "google-auth-oauthlib" }, { name = "html2text" }, { name = "jupyter-kernel-client" }, + { name = "keyring" }, { name = "nbformat" }, { name = "packaging" }, { name = "prompt-toolkit" }, @@ -405,6 +406,7 @@ requires-dist = [ { name = "google-auth-oauthlib", specifier = ">=1.3.0" }, { name = "html2text", specifier = ">=2024.2.26" }, { name = "jupyter-kernel-client", git = "https://github.com/googlecolab/jupyter-kernel-client.git" }, + { name = "keyring", specifier = ">=25.7.0" }, { name = "nbformat", specifier = ">=5.10.4" }, { name = "packaging", specifier = ">=24.0" }, { name = "prompt-toolkit", specifier = ">=3.0.52" }, @@ -452,6 +454,48 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "jaraco-classes" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "more-itertools" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/c0/ed4a27bc5571b99e3cff68f8a9fa5b56ff7df1c2251cc715a652ddd26402/jaraco.classes-3.4.0.tar.gz", hash = "sha256:47a024b51d0239c0dd8c8540c6c7f484be3b8fcf0b2d85c13825780d3b3f3acd", size = 11780, upload-time = "2024-03-31T07:27:36.643Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl", hash = "sha256:f662826b6bed8cace05e7ff873ce0f9283b5c924470fe664fff1c2f00f581790", size = 6777, upload-time = "2024-03-31T07:27:34.792Z" }, +] + +[[package]] +name = "jaraco-context" +version = "6.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/50/4763cd07e722bb6285316d390a164bc7e479db9d90daa769f22578f698b4/jaraco_context-6.1.2.tar.gz", hash = "sha256:f1a6c9d391e661cc5b8d39861ff077a7dc24dc23833ccee564b234b81c82dfe3", size = 16801, upload-time = "2026-03-20T22:13:33.922Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/58/bc8954bda5fcda97bd7c19be11b85f91973d67a706ed4a3aec33e7de22db/jaraco_context-6.1.2-py3-none-any.whl", hash = "sha256:bf8150b79a2d5d91ae48629d8b427a8f7ba0e1097dd6202a9059f29a36379535", size = 7871, upload-time = "2026-03-20T22:13:32.808Z" }, +] + +[[package]] +name = "jaraco-functools" +version = "4.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "more-itertools" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6c/1f/c23395957d41ccf27c4e535c3d334c4051e5395b3752057ba4cbaec35c56/jaraco_functools-4.6.0.tar.gz", hash = "sha256:880c577ec9720b3a052d5bc611fb9f2269b3d87902ef42440df443b88e443280", size = 20837, upload-time = "2026-07-14T01:28:02.544Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/36/ecc85bc96c273dc8a11273ed4782272975e6338d4a3e9228621175edf0e3/jaraco_functools-4.6.0-py3-none-any.whl", hash = "sha256:99e3dc0060c5cbe8fcd1cdb36258e2a65ca40f1566b2033b12abb1bb44dd3c30", size = 11677, upload-time = "2026-07-14T01:28:01.59Z" }, +] + +[[package]] +name = "jeepney" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/6f/357efd7602486741aa73ffc0617fb310a29b588ed0fd69c2399acbb85b0c/jeepney-0.9.0.tar.gz", hash = "sha256:cf0e9e845622b81e4a28df94c40345400256ec608d0e55bb8a3feaa9163f5732", size = 106758, upload-time = "2025-02-27T18:51:01.684Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl", hash = "sha256:97e5714520c16fc0a45695e5365a2e11b81ea79bba796e26f9f1d178cb182683", size = 49010, upload-time = "2025-02-27T18:51:00.104Z" }, +] + [[package]] name = "jsonschema" version = "4.26.0" @@ -535,6 +579,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/72/45/cb4671e13fed39f721066ad1a00714d4b607982b8d3e97a25f836198d1df/jupyter_mimetypes-0.2.0-py3-none-any.whl", hash = "sha256:e6dcd989258e3fc944365b656d9173191517e0e393bd878e97ce500e5b388527", size = 16724, upload-time = "2025-08-10T18:18:27.309Z" }, ] +[[package]] +name = "keyring" +version = "25.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jaraco-classes" }, + { name = "jaraco-context" }, + { name = "jaraco-functools" }, + { name = "jeepney", marker = "sys_platform == 'linux'" }, + { name = "pywin32-ctypes", marker = "sys_platform == 'win32'" }, + { name = "secretstorage", marker = "sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/4b/674af6ef2f97d56f0ab5153bf0bfa28ccb6c3ed4d1babf4305449668807b/keyring-25.7.0.tar.gz", hash = "sha256:fe01bd85eb3f8fb3dd0405defdeac9a5b4f6f0439edbb3149577f244a2e8245b", size = 63516, upload-time = "2025-11-16T16:26:09.482Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" }, +] + [[package]] name = "markdown-it-py" version = "4.2.0" @@ -556,6 +617,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "more-itertools" +version = "11.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/52/b8/c63995bebbb9ac9c059350e0fb8c82c2420c85aeca7d34e40aecf3df9281/more_itertools-11.2.1.tar.gz", hash = "sha256:cbf08fd0af284dc69718b9b76a8e6203df624d70209ea511e4219d350c856f63", size = 154133, upload-time = "2026-10-14T18:50:30.407Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/57/3c4b1c6f2e0cc29ba023802458ea3d2a996f87085559e4a5f2192d549532/more_itertools-11.2.1-py3-none-any.whl", hash = "sha256:35a7377edd1dd6608dcb2cdf534ded55ea32d49448875ad042cd3f879fb1ded0", size = 74405, upload-time = "2026-10-14T18:50:29.116Z" }, +] + [[package]] name = "nbformat" version = "5.10.4" @@ -845,6 +915,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "pywin32-ctypes" +version = "0.2.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/85/9f/01a1a99704853cb63f253eea009390c88e7131c67e66a0a02099a8c917cb/pywin32-ctypes-0.2.3.tar.gz", hash = "sha256:d162dc04946d704503b2edc4d55f3dba5c1d539ead017afa00142c38b9885755", size = 29471, upload-time = "2024-08-14T10:15:34.626Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl", hash = "sha256:8a1513379d709975552d202d942d9837758905c8d01eb82b8bcc30918929e7b8", size = 30756, upload-time = "2024-08-14T10:15:33.187Z" }, +] + [[package]] name = "pyzmq" version = "27.1.0" @@ -1078,6 +1157,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2d/c7/c53e8dbff9c9dc4b7928773421ae294a5d28fcb8dcda1a089579d3a7e510/ruff-0.15.17-py3-none-win_arm64.whl", hash = "sha256:f3be1fbb34bcdfd146240d8fb92a709d4c2c8191348580a3c044ec60fa0b4456", size = 11355275, upload-time = "2026-06-11T17:54:43.635Z" }, ] +[[package]] +name = "secretstorage" +version = "3.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "jeepney" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl", hash = "sha256:0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137", size = 15554, upload-time = "2025-11-23T19:02:51.545Z" }, +] + [[package]] name = "shellingham" version = "1.5.4"