diff --git a/docs/01_session_management.md b/docs/01_session_management.md index 62793fe..e8498db 100644 --- a/docs/01_session_management.md +++ b/docs/01_session_management.md @@ -1,5 +1,6 @@ --- log: +2026-08-13: Fixed issue #106 by reconciling saved runtime-proxy credentials with every `/tun/m/assignments` response. Explicit and implicit session resolution now adopt the fresh token/URL returned by the control plane; 401/404 runtime handshakes refresh and retry once; local bindings are pruned only after the server confirms their endpoint is absent. State writes are endpoint-guarded and field-level so concurrent commands cannot overwrite refreshed credentials, revive a removed session, or mutate a same-name replacement. 2026-08-09: Added `--high-mem` to `colab new`, `colab run`, and `colab ssh` (auto-create). Assign requests now send `shape=hm` when high-RAM is requested; `colab sessions` and `colab status` display machine shape. 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//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. 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. @@ -76,9 +77,16 @@ The CLI maps user flags to these backend parameters: ### 4. Session Listing (`colab sessions`) - **API**: `GET https://colab.research.google.com/tun/m/assignments` (based on `colab-agent` implementation). -- **Function**: Lists all active VM assignments for the user. This is useful for synchronizing local state with actual backend sessions. +- **Function**: Lists all active VM assignments for the user and reconciles local bindings by endpoint. Every response carries a fresh `RuntimeProxyInfo.token` and URL; locally tracked sessions adopt those values while preserving kernel/session IDs, keep-alive PID, and execution metadata. An endpoint absent from a successful response is pruned. If the lookup fails, local state is preserved because the assignment's absence was not confirmed. -### 5. Keep-Alive Protocol +### 5. Runtime-Proxy Credential Refresh + +- Runtime-proxy tokens expire independently of the VM assignment (issue #106). Both explicit `-s NAME` resolution and unique-session resolution refresh from `/tun/m/assignments` before opening a runtime connection. +- If a runtime connection still returns a proxy-auth 401/404, the operation performs one fresh assignments lookup and retries once, but only when the returned token or URL changed. A second failure is surfaced without an unbounded retry loop. +- `prune_session()` never treats a runtime 401/404 alone as proof that the VM is gone. It removes a binding and stops its keep-alive daemon only after the control plane confirms the exact endpoint is absent. +- State updates merge selected metadata fields into the latest stored object and require the expected endpoint. This prevents a long-running command's `finally` block from restoring an expired token, resurrecting a removed binding, or touching a newly-created session that reused the same name. + +### 6. Keep-Alive Protocol To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 minutes), the CLI implements a background keep-alive mechanism. - **Daemon Process**: Since the CLI is a fire-and-forget tool, `colab new` spawns a detached background process running a hidden `keep-alive` command. - **Tunnel ping**: Every 60 seconds, the daemon issues `GET https://colab.research.google.com/tun/m//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. @@ -93,7 +101,6 @@ To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 mi - **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. ## TODO / Future Work -- **Backend Sync**: Implement a way to reconcile the local `sessions.json` with the output of `colab sessions`. - **Resource Usage**: Add real-time resource usage (CPU/RAM/GPU) to the `status` output by executing a diagnostic snippet on the VM. ## Implementation Details @@ -128,3 +135,5 @@ TDD is mandatory for all session management features. - **Test Case (cross-process exclusion)**: Hold the write lock from a separate process and confirm the store's in-process write blocks until release. - **Test Case (concurrent readers)**: Hold a read lock from a separate process and confirm the store can still complete a read concurrently. - **Test Case (multi-thread regression)**: Two `StateStore` instances writing from different threads must serialize without raising `RuntimeError` (guards the `is_singleton=False` choice). +- **Test Case (token refresh)**: Verify named and implicit resolution adopt fresh token/URL values, runtime-proxy failures retry once, failed assignments lookups preserve bindings, and server-confirmed missing endpoints are pruned. +- **Test Case (stale-writer safety)**: Verify field-level updates preserve refreshed credentials and cannot revive a removed session or modify a same-name replacement endpoint. diff --git a/docs/02_execution_and_interactive.md b/docs/02_execution_and_interactive.md index ed1d782..11b035a 100644 --- a/docs/02_execution_and_interactive.md +++ b/docs/02_execution_and_interactive.md @@ -1,5 +1,10 @@ --- log: +2026-08-15: Fixed long-running Console output freezes caused by the raw `/colab/tty` flow-control protocol. The runtime requests an application-level acknowledgement about every 100 KB and pauses its PTY after six unacknowledged chunks; Console now acknowledges each request after flushing the corresponding terminal output. Protocol ping/pong remains transport-only liveness and is not treated as PTY progress. Added a CPU-only live regression that streams beyond the former 600 KB stall threshold without touching pre-existing assignments. +2026-08-14: Prevented Console reconnects from appearing frozen after a concurrent `colab stop` or local-proxy failure. Reconnect now treats a removed local binding as conclusive before calling the control plane, bounds its assignments refresh to ten seconds, continues with the last runtime token after transient refresh errors, and makes raw-mode Ctrl-C cancellation explicit. The CPU-only live regression covers both proxy-failure fallback and cross-process binding removal without touching pre-existing assignments. +2026-08-14: Tightened Console reconnect semantics after review. Shell `exit`/`logout`/Ctrl-D input now suppresses reconnect only when the peer closes within two seconds, so leaving a nested shell or foreground program cannot permanently disable later recovery. A connection that remains healthy for 30 seconds resets the 1/2/5/10/30-second backoff, and every independent loss reports its reason. The CPU-only live regressions now use pipeline failure propagation and recover the exact test endpoint during EXIT cleanup, closing the assignment-leak window between successful creation and endpoint capture. +2026-08-13: Made `colab console` resilient to long-lived proxy disconnects. Console now sends protocol ping/pong heartbeats, reports connection loss and retry progress, refreshes runtime-proxy credentials before reconnecting to the same endpoint, and reuses one stoppable stdin forwarding thread across attempts. Interactive sessions retry with 1/2/5/10/30-second backoff (then every 30 seconds); piped input is never replayed. Normal shell exits remain terminal, while an endpoint deletion or same-name replacement stops reconnecting safely. Added bounded loopback and resource-retention regressions after an unbounded zero-delay test loop caused pytest output capture to exhaust local memory. +2026-08-13: Fixed issue #106 for `exec`, `repl`, `console`, and `restart-kernel`: runtime startup proactively uses refreshed proxy credentials and retries one proxy-auth failure after a control-plane refresh. Terminal failures no longer unconditionally prune local state, and endpoint-guarded metadata cleanup prevents stale `finally` blocks from restoring expired tokens or reviving removed sessions. 2026-05-07: Fixed `colab console` piped-stdin handling. Previously a piped invocation (e.g. `echo 'cmd' | colab console -s s`) sent the command and then hung indefinitely because the previous EOF handler emitted a bare `\x04` (Ctrl-D), which the remote `tmux`-wrapped bash treats as a literal character rather than a session terminator. The new handler sends `exit\n` (which bash actually exits on) and then closes the websocket from the client side after a short grace period (`PIPED_EOF_GRACE_SECONDS = 0.5s`) so any tail output (bash `logout`, tmux `[exited]`) makes it back to the user. TTY mode is unchanged: real-terminal EOF is left to the remote shell. Verified live: `echo 'echo HELLO' | colab console -s s` now exits in ~1.2s instead of hanging. 2026-05-07: Fixed `print_kitty` (used by `colab exec --output-image` and any image-producing exec) to no-op when `sys.stdout.isatty()` is false. The Kitty Graphics Protocol escape sequence is meaningless when stdout is a file or pipe and was visually corrupting captured output (a multi-KB base64 PNG blob would land in log files, grep targets, or showboat captures). Image bytes are still saved to disk via `handle_image`'s file-write path; only the inline-render attempt is suppressed. @@ -34,7 +39,19 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel - **Implementation**: Connects directly to the backend terminal endpoint (`/colab/tty`) via WebSockets using `websocket-client`. - **Interactive**: Bypasses the Jupyter kernel entirely to provide a raw, PTY-backed bash session on the Colab VM. - **Terminal Management**: Configures `sys.stdin` to raw mode using `termios` and `tty`, passing single characters to the socket and writing raw ANSI escape sequences directly to `sys.stdout.buffer`. Hooks into `SIGWINCH` to communicate local terminal dimensions (`cols`/`rows`) to the remote bash environment so output rendering works perfectly during resizing. +- **Liveness detection**: `websocket-client` sends a protocol ping every 20 seconds and requires a pong within 10 seconds. This both keeps an otherwise-idle proxy path active and turns a silently dead path into a detectable disconnect. +- **PTY output flow control**: The runtime marks roughly every 100 KB of terminal output with `"ack": true` and pauses the PTY after six unacknowledged chunks. Console writes and flushes each marked chunk before replying `{"ack": true}` on the same WebSocket. Protocol pong frames prove only transport liveness and do not satisfy this application-level acknowledgement. +- **Interactive reconnects**: An abnormal close after a successful handshake is reported on stderr, then retried after 1, 2, 5, 10, and 30 seconds, followed by 30-second retries until the user presses Ctrl-C or the assignment is confirmed gone. A connection healthy for 30 seconds resets this backoff so an unrelated later outage starts again at one second. Each retry first checks that the original local binding still exists, then refreshes its assignment credentials with a ten-second HTTP timeout; a transient refresh failure falls back to the last token, while a removed binding or same-name replacement stops safely. One stdin-forwarding thread is reused across every attempt and is stopped when Console exits, so reconnects cannot accumulate competing terminal readers. +- **Close semantics**: Close codes 1000 (normal) and 1001 (going away) end Console without reconnecting. A locally recognized `exit`, `logout`, or Ctrl-D request suppresses an abnormal-close reconnect for two seconds; this bounded hint covers immediate shell termination without permanently disabling recovery when the input only leaves a nested shell or foreground program. A 401/404 initial-handshake error is returned to the shared one-time credential-refresh path rather than entering the unbounded transport reconnect loop. +- **Status visibility**: Connection loss, every retry delay/attempt, successful reconnection, raw-mode Ctrl-C cancellation, and final closure are printed as concise `[colab]` messages on stderr, separate from the remote terminal byte stream. - **Piped stdin**: Detected via `sys.stdin.isatty()`. When piped, the input characters are forwarded one at a time to the remote pty, and on EOF the client sends `exit\n` and then closes the websocket itself after `PIPED_EOF_GRACE_SECONDS` (0.5s) so the user's shell goodbye text drains back. The remote `/colab/tty` endpoint wraps bash in tmux, which intercepts a bare `\x04` as a literal character — that is why we send `exit\n` rather than Ctrl-D. +- **Piped disconnects**: Piped input is never reconnected or replayed because the CLI cannot know which bytes the remote shell already consumed. An abnormal close returns a non-zero exit with a concise error instead. + +### 4. Expired Runtime-Proxy Credentials + +- Session resolution adopts the latest runtime-proxy token and URL returned by `/tun/m/assignments` before Jupyter or terminal connection startup. +- A proxy-auth 401/404 triggers one refresh-and-retry. A repeated failure is reported without deleting the local binding unless the assignments endpoint independently confirms the VM endpoint is gone. +- Kernel/session ID callbacks and `running`/`last_execution` cleanup use endpoint-guarded field updates, so a stale command cannot overwrite a token refreshed by another invocation or recreate a deleted session. This specifically prevents the former Console `finally` resurrection path. ## Implementation Details - **Kernel Management**: `ColabRuntime` (from `colab-agent`) already handles message signing and message types. @@ -54,4 +71,13 @@ TDD is mandatory for all execution features. - **Test Case**: Verify large piped inputs are handled without buffer overflow or truncation. - **Test Case**: `colab console` with piped stdin sends `exit\n` and calls `ws.close()` on EOF (regression: previously sent `\x04` only and hung). - **Test Case**: `colab console` in TTY mode does not synthesize an exit on EOF (the user owns the session lifecycle). +- **Test Case**: An abnormal interactive close prints visible state changes, refreshes credentials for the same endpoint, and reconnects using one stdin reader; deletion and same-name endpoint replacement stop retries. +- **Test Case**: Protocol ping/pong settings are enabled, normal closes do not reconnect, and initial 401/404 handshake failures return to the bounded shared token-refresh path. +- **Test Case**: A terminal message marked `"ack": true` is flushed before Console replies with `{"ack": true}`; unmarked messages do not emit acknowledgements. A live isolated CPU Console streams beyond 600 KB and still receives a trailing sentinel. +- **Test Case**: A real loopback WebSocket closes once abnormally and then normally, proving that the refreshed token is used on the second handshake. The fake peer has a hard timeout and propagates server-thread exceptions. `integration/repro_console_reconnect/test.sh` repeats the fault injection against an isolated live CPU assignment, injects one control-plane `ProxyError`, removes a copied binding from another process, and verifies fallback/local-stop status without changing the pre-existing endpoint snapshot. +- **Test Case**: Test-only reconnect limits prevent zero-delay retry fixtures from spinning forever. A 500-attempt object-retention regression verifies that old WebSocket attempts are collectable, and an idle-stdin regression verifies that the forwarding thread stops without waiting for another keystroke. +- **Test Case**: Piped disconnects fail once without retrying or replaying input. +- **Test Case**: Shell-exit intent expires after two seconds, short-lived reconnect failures continue their backoff, and a connection healthy for 30 seconds resets the next outage to attempt one. - **Test Case**: `print_kitty` is a no-op when `sys.stdout.isatty()` is false (regression: previously emitted ANSI/base64 into pipes and files). +- **Test Case**: Runtime-proxy 401/404 startup failures refresh and retry once without unconditional pruning. +- **Test Case**: Console and execution cleanup merge metadata into the latest endpoint-matching state and never revive a removed session. diff --git a/docs/03_file_management.md b/docs/03_file_management.md index be44c9e..c194d92 100644 --- a/docs/03_file_management.md +++ b/docs/03_file_management.md @@ -1,3 +1,8 @@ +--- +log: +2026-08-13: Fixed issue #106 for all Contents API commands. File operations now use refreshed runtime-proxy credentials and retry once after an empty 401/404 tunnel response; a non-empty Jupyter Contents 404 remains a normal missing-file error. +--- + # Design: File Management (`ls`, `rm`, `upload`, `download`, `edit`) ## Overview @@ -44,7 +49,7 @@ File management on the Colab VM will be implemented using the Jupyter Contents A ## Implementation Details - **Base URL**: The backend URL obtained during session assignment. - **Proxy Token**: The `colab-runtime-proxy-token` is required for each request. -- **Error Handling**: Handle 404 (not found) and 403 (unauthorized). +- **Error Handling**: An empty tunnel-level 401/404 is classified as an expired/invalid runtime-proxy credential and retried once after refreshing from `/tun/m/assignments`. A non-empty Jupyter Contents 404 remains `FileNotFoundError`; missing paths are never mistaken for token expiry. - **Large Files**: The Contents API might have limitations for very large files. If so, we'll implement a fallback via the kernel (streaming chunks). ## Testing Strategy @@ -59,4 +64,5 @@ TDD is mandatory for all file management features. ### 2. Error Cases - **Test Case**: Verify 404 responses are correctly caught and presented as a "File not found" error to the user. -- **Test Case**: Verify correct handling of large file uploads exceeding API limits via kernel streaming. \ No newline at end of file +- **Test Case**: Verify correct handling of large file uploads exceeding API limits via kernel streaming. +- **Test Case**: Verify empty 401/404 responses request refreshed credentials and retry once, while JSON/non-empty 404 responses remain missing-file errors. diff --git a/docs/04_automation_and_utility.md b/docs/04_automation_and_utility.md index 0a3d69a..cbad451 100644 --- a/docs/04_automation_and_utility.md +++ b/docs/04_automation_and_utility.md @@ -1,5 +1,6 @@ --- log: +2026-08-13: Fixed issue #106 for VM-side automation. `auth`, `drivemount`, and `install` start kernels through the shared runtime-proxy refresh/retry path; `install -r` also uploads its requirements file with refreshed Contents API credentials. Metadata cleanup uses endpoint-guarded field updates. 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. 2026-06-01: Enabled `colab update --install` self-update on macOS in addition to Linux. Refactored platform check logic to keep the implementation DRY and updated both tests and documentation. Also, on these platforms, an additional message is shown recommending `colab update --install` to upgrade in place, positioned above the standard `pip`/`uv` installation command. 2026-05-29: Added default OAuth2 client config (`oauth_config.json`) as a bundled package resource and restored fallback loading logic in `get_credentials()`. The CLI now falls back to using these default credentials when no explicit local config is found. Added `integration/repro_bundled_oauth` integration test. @@ -121,6 +122,7 @@ remediation guidance) rather than silently after ~1 minute via the daemon. subprocess.check_call([sys.executable, "-m", "pip", "install", "..."])` - **Requirements File**: Upload `requirements.txt` if provided with `-r` and then run `pip install -r`. +- **Runtime credentials**: Kernel startup and the optional requirements-file upload both use the shared runtime-proxy refresh path. A proxy-auth failure refreshes from the assignments endpoint and retries once. ### 3. Drive Mounting (`colab drivemount`) diff --git a/docs/05_run_command.md b/docs/05_run_command.md index 24df987..7b8f156 100644 --- a/docs/05_run_command.md +++ b/docs/05_run_command.md @@ -1,5 +1,6 @@ --- log: +2026-08-13: Fixed issue #106 during one-shot runtime startup. `colab run` adopts refreshed runtime-proxy credentials and retries one proxy-auth failure; endpoint-guarded cleanup preserves refreshed tokens under `--keep`. Startup failures before a runtime object exists now retain the original error and still release the newly allocated assignment. 2026-08-09: Added `--high-mem` flag (passthrough to session creation; sends `shape=hm` on assign when supported). 2026-05-12: Initial design and implementation of `colab run [args...]`. Combines `colab new` + `colab exec` + `colab stop` into a single fire-and-forget invocation so a Python file can use `#!/usr/bin/env -S colab run` as a shebang line and execute on a freshly-allocated Colab VM. Adds `--keep` (skip auto-stop), `--gpu` / `--tpu` (passthrough to session creation), `-s/--session` (name the ephemeral session), and propagates the script's exit status (non-zero on any uncaught exception in the kernel). The script's `sys.argv` is re-set inside the kernel to mirror native `python script.py arg1 arg2` semantics, and `__name__` is set to `"__main__"`. 2026-05-12: Native CPython exit-code semantics for `sys.exit()` / `raise SystemExit(...)` from the script body. The Colab kernel reports a `SystemExit` as `output_type=='error'`, which under the previous logic would have (a) printed the IPython traceback (`An exception has occurred, use %tb...`) and (b) flagged the run as a failure regardless of the integer exit code. Now: `sys.exit()` / `sys.exit(0)` exit 0 silently; `sys.exit(N)` exits N; `sys.exit('msg')` exits 1 (matching CPython). The IPython "To exit: use 'exit', 'quit', or Ctrl-D." UserWarning is filtered via the prelude. Encoded after running `examples/gpu_hello.py` end-to-end and seeing the noisy `SystemExit: 0` traceback at the end of an otherwise-successful GPU run. @@ -57,6 +58,7 @@ print(torch.cuda.get_device_name(0)) __name__ = '__main__' ``` Then executes the script body in the same kernel cell so any `if __name__ == "__main__":` guard fires. + Runtime startup uses the shared proxy-credential refresh path and retries one proxy-auth failure with the latest token/URL from `/tun/m/assignments`. 3. **Detect failure**: If the kernel returns any output of `output_type == "error"` (uncaught exception, syntax error, etc.) the CLI exits non-zero. 4. **Tear down**: In a `finally` block, unless `--keep` was passed, the CLI: - Sends `runtime.stop(shutdown_kernel=True)` (best-effort). @@ -66,6 +68,7 @@ print(torch.cuda.get_device_name(0)) - Logs `session_terminated` with `reason="run_completed"` (or `"run_failed"`). If `--keep` is set, the session remains visible in `colab sessions` and `colab status` and can be reused with `colab exec -s `, `colab repl -s `, etc., until the user runs `colab stop` (or the keep-alive daemon hits its 24h cap). +Metadata cleanup is an endpoint-guarded field update, so it cannot overwrite credentials refreshed during startup. If startup fails before a runtime object is created, teardown still unassigns the fresh VM without masking the original exception. ## AGENTS.md Constraints Honoured - **Item 7 (no background threads)**: The keep-alive daemon is the existing detached process from `colab new`; this command introduces no new threads. diff --git a/docs/06_ssh_access.md b/docs/06_ssh_access.md index dd8636d..642082b 100644 --- a/docs/06_ssh_access.md +++ b/docs/06_ssh_access.md @@ -1,5 +1,6 @@ --- log: +2026-08-13: Fixed issue #106 for SSH proxy mode. Existing-session resolution proactively refreshes the runtime-proxy token; a structured HTTP 401 handshake failure refreshes and retries once, while HTTP 404 retains its distinct “SSH endpoint is not exposed” meaning. 2026-08-09: Added `--high-mem` passthrough when `colab ssh` auto-creates a runtime (forwards to `colab new --high-mem`). 2026-07-17: Initial design and implementation of `colab ssh` — client side of SSH-over-WebSocket runtime access. Adds three modes (interactive shell, `-s SESSION`, and `--proxy-mode` OpenSSH ProxyCommand bridge), `--identity/-i` key selection, and per-HTTP-status handshake error messages. Server side is out of scope for this repo; the subcommand is a no-op against runtimes that do not expose the `/colab/ssh` endpoint (surfaces an actionable HTTP 404 message). 2026-07-22: Bare `colab ssh` now auto-creates a runtime (via `colab new`) when you have no active session, with `--gpu/--tpu` passthrough and `--rm` to stop an auto-created runtime on exit. Fixed two client bugs: the dead 403 branch (feature-off returns 404, not 403) and the RSA guidance (all `ssh-rsa` keys are server-rejected, so `id_rsa` is no longer auto-scanned and the 400 message no longer advertises `rsa-sha2`). Added `tests/test_ssh_wire_contract.py` (real loopback-server wire assertions) and `tests/test_ssh_autocreate.py`. @@ -64,7 +65,10 @@ remote command, so to also land in `/content` add `RequestTTY yes` and and sends the resolved public key verbatim in the `X-Colab-Ssh-Pubkey` header (no transformation -- the bytes the user controls are exactly what the server receives). Only `ssh-ed25519` / `ecdsa-sha2-nistp{256,384,521}` keys are - accepted. + accepted. Session resolution first adopts the latest proxy token. If the + handshake returns HTTP 401, proxy mode refreshes and retries once; HTTP 404 + is not retried because it normally means that this runtime was created + without the SSH endpoint. 3. **Interactive shell**: Spawns the system `ssh` binary with the CLI re-invoked as its own `ProxyCommand` (`python -m colab_cli.cli ssh --proxy-mode`), so the WebSocket bridge and the interactive shell share one code path. It forces a @@ -89,7 +93,7 @@ remote command, so to also land in `/content` add `RequestTTY yes` and | Status | Meaning surfaced to the user | | --- | --- | | 400 | Bad/unsupported/missing pubkey, with remediation (`ssh-keygen -t ed25519`) | - | 401 | Token invalid/expired — try `colab new` | + | 401 | Token invalid/expired — refresh and retry once, then report authentication failure | | 403 | Forbidden — token lacks permission for this action (feature-off returns 404, not 403) | | 404 | SSH not exposed on this runtime — SSH is baked in at creation, so run `colab new` | | 429 | Another `colab ssh` is already connected — disconnect first | diff --git a/integration/README.md b/integration/README.md index 2b79610..22687e3 100644 --- a/integration/README.md +++ b/integration/README.md @@ -16,6 +16,9 @@ End-to-end tests that run against a **live Colab backend** (unlike the mocked un | `repro_keep_alive_scope/` | Slow soak test (~95s): runs the daemon long enough for one ping past the pre-flight, asserts no `keep_alive_error` events. | | `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_console_reconnect/` | CPU-only live fault injection: verifies fallback after a transient control-plane proxy failure, prompt exit after another process removes the local binding, and preservation of pre-existing assignments. | +| `repro_console_flow_control/` | CPU-only live regression: streams more than 600 KB through `/colab/tty`, verifies application-level acknowledgements prevent the remote PTY from pausing, and preserves pre-existing assignments. | +| `repro_runtime_token_refresh/` | CPU-only regression for issue #106: corrupts the saved runtime-proxy token and verifies `ls`, `exec`, and piped `console` self-heal without disturbing pre-existing assignments. | | `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_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` | diff --git a/integration/repro_console_flow_control/test.sh b/integration/repro_console_flow_control/test.sh new file mode 100755 index 0000000..34c1aa6 --- /dev/null +++ b/integration/repro_console_flow_control/test.sh @@ -0,0 +1,254 @@ +#!/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 +# +# https://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. + +# CPU-only live regression for the raw /colab/tty application-level flow +# control. The runtime requests an acknowledgement roughly every 100 KB and +# pauses its PTY after six missed acknowledgements, so a short Console smoke +# test cannot expose the failure. + +set -euo pipefail + +TMP_DIR=$(mktemp -d) +SESSION_FILE="$TMP_DIR/sessions.json" +SERVER_SESSION_FILE="$TMP_DIR/server-snapshot.json" +OUTPUT_FILE="$TMP_DIR/console.out" +SESSION_NAME="console-flow-control-$PPID-$$" +TEST_ENDPOINT="" + +if [ -f "$HOME/.config/colab-cli/token.json" ]; then + AUTH_PROVIDER="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) + 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) + if echo "$ADC_SCOPES" | grep -q "userinfo.email"; then + AUTH_PROVIDER="adc" + else + echo "Error: ADC token lacks the userinfo.email scope." >&2 + exit 1 + fi +else + echo "Error: No usable auth provider found." >&2 + exit 1 +fi +AUTH_FLAGS="--auth=$AUTH_PROVIDER" + +server_endpoints() { + uv run colab $AUTH_FLAGS --config "$SERVER_SESSION_FILE" sessions 2>/dev/null | sed -n 's/^\[[^]]*\] \([^ ]*\) .*/\1/p' | sort +} + +BEFORE_ENDPOINTS=$(server_endpoints) + +cleanup() { + exit_code=$? + cleanup_endpoint="$TEST_ENDPOINT" + if [ -z "$cleanup_endpoint" ] && [ -s "$SESSION_FILE" ]; then + cleanup_endpoint=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' 2>/dev/null || true +import json +import sys + +print(json.load(open(sys.argv[1])).get(sys.argv[2], {}).get("endpoint", "")) +PY + ) + fi + + uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" >/dev/null 2>&1 || true + if [ -n "$cleanup_endpoint" ]; then + uv run python - "$AUTH_PROVIDER" "$cleanup_endpoint" <<'PY' >/dev/null 2>&1 || true +import sys +from colab_cli.auth import AuthProvider +from colab_cli.common import state + +state.auth_provider = AuthProvider(sys.argv[1]) +state.client.unassign(sys.argv[2]) +PY + fi + if [ "$exit_code" -ne 0 ] && [ -s "$OUTPUT_FILE" ]; then + echo "[FAILURE] Last 1024 bytes of Console output:" >&2 + tail -c 1024 "$OUTPUT_FILE" >&2 + fi + rm -rf "$TMP_DIR" + trap - EXIT + exit "$exit_code" +} +trap cleanup EXIT + +echo "[*] Creating isolated CPU session $SESSION_NAME..." +# Deliberately omit --gpu/--tpu. This test must not consume or reconnect to an +# accelerator assignment. +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" new -s "$SESSION_NAME" + +TEST_ENDPOINT=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' +import json +import sys + +print(json.load(open(sys.argv[1]))[sys.argv[2]]["endpoint"]) +PY +) +TEST_ACCELERATOR=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' +import json +import sys + +print(json.load(open(sys.argv[1]))[sys.argv[2]]["accelerator"]) +PY +) +if [ "$TEST_ACCELERATOR" != "NONE" ]; then + echo "[FAILURE] Refusing to probe non-CPU session: $TEST_ACCELERATOR" >&2 + exit 1 +fi + +echo "[*] Streaming beyond the PTY pause threshold on $TEST_ENDPOINT..." +timeout 120 uv run python - "$AUTH_PROVIDER" "$SESSION_FILE" "$SESSION_NAME" >"$OUTPUT_FILE" 2>&1 <<'PY' +import json +import os +import pty +import sys +import threading + +import colab_cli.console as console +from colab_cli.auth import AuthProvider +from colab_cli.common import State + +auth_provider, config_path, session_name = sys.argv[1:] +state = State() +state.auth_provider = AuthProvider(auth_provider) +state.config_path = config_path +session = state.store.get(session_name) +if session is None: + raise SystemExit("Isolated CPU session is missing") +if str(session.accelerator) not in ("AcceleratorType.NONE", "NONE"): + raise SystemExit(f"Refusing to probe non-CPU session: {session.accelerator}") + +real_websocket_app = console.websocket.WebSocketApp +real_forwarder = console._ConsoleInputForwarder +opened = threading.Event() +sentinel_received = threading.Event() +ack_requests = 0 +active_ws = None +active_forwarder = None +output_tail = "" + + +class ObservedForwarder(real_forwarder): + def __init__(self, *args, **kwargs): + global active_forwarder + super().__init__(*args, **kwargs) + active_forwarder = self + + +def websocket_app(**kwargs): + original_open = kwargs["on_open"] + original_message = kwargs["on_message"] + + def observed_open(ws): + global active_ws + original_open(ws) + active_ws = ws + opened.set() + + def observed_message(ws, message): + global ack_requests, output_tail + saw_sentinel = False + try: + payload = json.loads(message) + except (TypeError, json.JSONDecodeError): + pass + else: + if payload.get("ack") is True: + ack_requests += 1 + output_tail = (output_tail + payload.get("data", ""))[-64:] + if "CONSOLE_FLOW_CONTROL_OK" in output_tail: + saw_sentinel = True + original_message(ws, message) + if saw_sentinel: + sentinel_received.set() + + kwargs["on_open"] = observed_open + kwargs["on_message"] = observed_message + return real_websocket_app(**kwargs) + + +master_fd, slave_fd = pty.openpty() +original_stdin = sys.stdin +stdin_stream = os.fdopen(slave_fd, encoding="utf-8", buffering=1) +sys.stdin = stdin_stream +console.websocket.WebSocketApp = websocket_app +console._ConsoleInputForwarder = ObservedForwarder + + +def drive_shell(): + if not opened.wait(60): + os.write(master_fd, b"\x03") + return + command = ( + b"python3 -c \"import sys,time; " + b"[(sys.stdout.write('F' * 1024 + '\\\\n'), sys.stdout.flush(), " + b"time.sleep(0.02)) for _ in range(800)]; " + b"print('CONSOLE_FLOW_' + 'CONTROL_OK')\"\n" + ) + os.write(master_fd, command) + if not sentinel_received.wait(90): + os.write(master_fd, b"\x03") + return + if active_forwarder is not None: + active_forwarder.user_requested_close = True + if active_ws is not None: + active_ws.close() + + +threading.Thread(target=drive_shell, daemon=True).start() +try: + console.connect_console( + session, + refresh_session=lambda expected: state.refresh_session( + session_name, expected_session=expected, timeout=10 + ), + retry_delays=(1,), + _max_reconnect_attempts=1, + ) +finally: + console.websocket.WebSocketApp = real_websocket_app + console._ConsoleInputForwarder = real_forwarder + sys.stdin = original_stdin + stdin_stream.close() + os.close(master_fd) + +if ack_requests < 6: + raise SystemExit(f"Expected at least 6 PTY ACK requests, got {ack_requests}") +print(f"CONSOLE_FLOW_CONTROL_ACK_REQUESTS={ack_requests}") +PY + +grep -a -q "CONSOLE_FLOW_CONTROL_OK" "$OUTPUT_FILE" +grep -a -E -q "CONSOLE_FLOW_CONTROL_ACK_REQUESTS=([6-9]|[1-9][0-9]+)" "$OUTPUT_FILE" +OUTPUT_BYTES=$(wc -c < "$OUTPUT_FILE") +if [ "$OUTPUT_BYTES" -lt 650000 ]; then + echo "[FAILURE] Console returned only $OUTPUT_BYTES bytes." >&2 + exit 1 +fi + +echo "[*] Stopping isolated CPU session..." +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" + +AFTER_ENDPOINTS=$(server_endpoints) +if [ "$AFTER_ENDPOINTS" != "$BEFORE_ENDPOINTS" ]; then + echo "[FAILURE] Pre-existing assignments changed during the test." >&2 + echo "Before:" >&2 + echo "$BEFORE_ENDPOINTS" >&2 + echo "After:" >&2 + echo "$AFTER_ENDPOINTS" >&2 + exit 1 +fi +TEST_ENDPOINT="" + +echo "[SUCCESS] Console acknowledged PTY flow control beyond 600 KB on CPU." diff --git a/integration/repro_console_reconnect/test.sh b/integration/repro_console_reconnect/test.sh new file mode 100755 index 0000000..2e7de1e --- /dev/null +++ b/integration/repro_console_reconnect/test.sh @@ -0,0 +1,338 @@ +#!/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 +# +# https://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. + +# CPU-only live regression for interactive Console reconnects. The probe +# establishes genuine /colab/tty WebSockets and exercises two recovery paths: +# a transient control-plane proxy failure falls back to the last token, and a +# binding removed by another process stops without another HTTP lookup. Existing +# assignments are only snapshotted and must be unchanged after the isolated test +# assignment is removed. + +set -euo pipefail + +TMP_DIR=$(mktemp -d) +SESSION_FILE="$TMP_DIR/sessions.json" +SERVER_SESSION_FILE="$TMP_DIR/server-snapshot.json" +OUTPUT_FILE="$TMP_DIR/console.out" +STOP_SESSION_FILE="$TMP_DIR/stop-sessions.json" +STOP_OUTPUT_FILE="$TMP_DIR/console-stop.out" +SESSION_NAME="console-reconnect-$PPID-$$" +TEST_ENDPOINT="" + +if [ -f "$HOME/.config/colab-cli/token.json" ]; then + AUTH_PROVIDER="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) + 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) + if echo "$ADC_SCOPES" | grep -q "userinfo.email"; then + AUTH_PROVIDER="adc" + else + echo "Error: ADC token lacks the userinfo.email scope." >&2 + exit 1 + fi +else + echo "Error: No usable auth provider found." >&2 + exit 1 +fi +AUTH_FLAGS="--auth=$AUTH_PROVIDER" + +server_endpoints() { + uv run colab $AUTH_FLAGS --config "$SERVER_SESSION_FILE" sessions 2>/dev/null | sed -n 's/^\[[^]]*\] \([^ ]*\) .*/\1/p' | sort +} + +BEFORE_ENDPOINTS=$(server_endpoints) + +cleanup() { + exit_code=$? + cleanup_endpoint="$TEST_ENDPOINT" + if [ -z "$cleanup_endpoint" ] && [ -s "$SESSION_FILE" ]; then + cleanup_endpoint=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' 2>/dev/null || true +import json +import sys + +print(json.load(open(sys.argv[1])).get(sys.argv[2], {}).get("endpoint", "")) +PY + ) + fi + + # `new` may have succeeded even if the first endpoint read failed. Always + # stop by the isolated name, then unassign the exact recovered endpoint. + uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" >/dev/null 2>&1 || true + if [ -n "$cleanup_endpoint" ]; then + uv run python - "$AUTH_PROVIDER" "$cleanup_endpoint" <<'PY' >/dev/null 2>&1 || true +import sys +from colab_cli.auth import AuthProvider +from colab_cli.common import state + +state.auth_provider = AuthProvider(sys.argv[1]) +state.client.unassign(sys.argv[2]) +PY + fi + if [ "$exit_code" -ne 0 ] && [ -s "$OUTPUT_FILE" ]; then + echo "[FAILURE] Captured Console probe output:" >&2 + sed 's/^/ /' "$OUTPUT_FILE" >&2 + fi + if [ "$exit_code" -ne 0 ] && [ -s "$STOP_OUTPUT_FILE" ]; then + echo "[FAILURE] Captured local-stop probe output:" >&2 + sed 's/^/ /' "$STOP_OUTPUT_FILE" >&2 + fi + rm -rf "$TMP_DIR" + trap - EXIT + exit "$exit_code" +} +trap cleanup EXIT + +echo "[*] Creating isolated CPU session $SESSION_NAME..." +# Deliberately omit --gpu/--tpu. This test must never consume an accelerator or +# interact with a pre-existing accelerator assignment. +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" new -s "$SESSION_NAME" + +TEST_ENDPOINT=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' +import json +import sys + +session = json.load(open(sys.argv[1]))[sys.argv[2]] +print(session["endpoint"]) +PY +) +TEST_ACCELERATOR=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' +import json +import sys + +print(json.load(open(sys.argv[1]))[sys.argv[2]]["accelerator"]) +PY +) +if [ "$TEST_ACCELERATOR" != "NONE" ]; then + echo "[FAILURE] Refusing to probe non-CPU session: $TEST_ACCELERATOR" >&2 + exit 1 +fi + +echo "[*] Injecting one transient refresh failure on CPU endpoint $TEST_ENDPOINT..." +timeout 120 uv run python - "$AUTH_PROVIDER" "$SESSION_FILE" "$SESSION_NAME" >"$OUTPUT_FILE" 2>&1 <<'PY' +import os +import pty +import sys +import threading + +import colab_cli.console as console +from colab_cli.auth import AuthProvider +from colab_cli.common import State +from requests.exceptions import ProxyError + +auth_provider, config_path, session_name = sys.argv[1:] +state = State() +state.auth_provider = AuthProvider(auth_provider) +state.config_path = config_path +session = state.store.get(session_name) +if session is None: + raise SystemExit("Isolated CPU session is missing") +if str(session.accelerator) not in ("AcceleratorType.NONE", "NONE"): + raise SystemExit(f"Refusing to probe non-CPU session: {session.accelerator}") + +real_websocket_app = console.websocket.WebSocketApp +created = 0 +second_open = threading.Event() +forced_disconnect = threading.Event() +refresh_failures = 0 + + +def websocket_app(**kwargs): + global created + created += 1 + attempt_number = created + original_open = kwargs["on_open"] + + def injected_open(ws): + original_open(ws) + if attempt_number == 1: + def disconnect(): + forced_disconnect.set() + ws.sock.shutdown() + + threading.Timer(1.0, disconnect).start() + elif attempt_number == 2: + second_open.set() + + kwargs["on_open"] = injected_open + return real_websocket_app(**kwargs) + + +master_fd, slave_fd = pty.openpty() +original_stdin = sys.stdin +stdin_stream = os.fdopen(slave_fd, encoding="utf-8", buffering=1) +sys.stdin = stdin_stream +console.websocket.WebSocketApp = websocket_app + + +def drive_reconnected_shell(): + if not second_open.wait(90): + os.write(master_fd, b"\x03") + return + os.write(master_fd, b"printf 'COLAB_CONSOLE_RECONNECT_OK\\n'\nexit\n") + + +def refresh_session(expected): + global refresh_failures + if refresh_failures == 0: + refresh_failures += 1 + raise ProxyError("injected transient control-plane proxy failure") + return state.refresh_session( + session_name, expected_session=expected, timeout=10 + ) + + +threading.Thread(target=drive_reconnected_shell, daemon=True).start() +try: + console.connect_console( + session, + refresh_session=refresh_session, + retry_delays=(1,), + _max_reconnect_attempts=2, + ) +finally: + console.websocket.WebSocketApp = real_websocket_app + sys.stdin = original_stdin + stdin_stream.close() + os.close(master_fd) + +if not forced_disconnect.is_set(): + raise SystemExit("Fault injection did not run") +if not second_open.is_set(): + raise SystemExit("Console did not reconnect") +print( + f"INTEGRATION_OK attempts={created} refresh_failures={refresh_failures} " + f"endpoint={session.endpoint}" +) +PY + +grep -a -q "Console connection lost" "$OUTPUT_FILE" +grep -a -q "Reconnecting in 1s (attempt 1" "$OUTPUT_FILE" +grep -a -q "could not refresh credentials" "$OUTPUT_FILE" +grep -a -q "Console reconnected (attempt 1, endpoint $TEST_ENDPOINT)" "$OUTPUT_FILE" +grep -a -q "COLAB_CONSOLE_RECONNECT_OK" "$OUTPUT_FILE" +grep -a -q "INTEGRATION_OK attempts=2 refresh_failures=1 endpoint=$TEST_ENDPOINT" "$OUTPUT_FILE" + +echo "[*] Simulating a completed concurrent stop in an isolated binding copy..." +cp "$SESSION_FILE" "$STOP_SESSION_FILE" +timeout 60 uv run python - "$AUTH_PROVIDER" "$STOP_SESSION_FILE" "$SESSION_NAME" >"$STOP_OUTPUT_FILE" 2>&1 <<'PY' +import os +import pty +import subprocess +import sys +import threading + +import colab_cli.console as console +from colab_cli.auth import AuthProvider +from colab_cli.commands.execution import _refresh_console_session +from colab_cli.common import State + +auth_provider, config_path, session_name = sys.argv[1:] +state = State() +state.auth_provider = AuthProvider(auth_provider) +state.config_path = config_path +session = state.store.get(session_name) +if session is None: + raise SystemExit("Isolated CPU session is missing") +if str(session.accelerator) not in ("AcceleratorType.NONE", "NONE"): + raise SystemExit(f"Refusing to probe non-CPU session: {session.accelerator}") + +real_websocket_app = console.websocket.WebSocketApp +binding_removed = threading.Event() +refresh_calls = 0 + + +def websocket_app(**kwargs): + original_open = kwargs["on_open"] + + def injected_open(ws): + original_open(ws) + + def remove_binding_and_disconnect(): + code = ( + "import sys; from colab_cli.state import StateStore; " + "removed = StateStore(sys.argv[1]).remove_if_endpoint(" + "sys.argv[2], sys.argv[3]); raise SystemExit(removed is None)" + ) + subprocess.run( + [ + sys.executable, + "-c", + code, + config_path, + session_name, + session.endpoint, + ], + check=True, + ) + binding_removed.set() + ws.sock.shutdown() + + threading.Timer(1.0, remove_binding_and_disconnect).start() + + kwargs["on_open"] = injected_open + return real_websocket_app(**kwargs) + + +def refresh_session(expected): + global refresh_calls + refresh_calls += 1 + return _refresh_console_session(state, session_name, expected) + + +master_fd, slave_fd = pty.openpty() +original_stdin = sys.stdin +stdin_stream = os.fdopen(slave_fd, encoding="utf-8", buffering=1) +sys.stdin = stdin_stream +console.websocket.WebSocketApp = websocket_app +try: + console.connect_console( + session, + refresh_session=refresh_session, + retry_delays=(1,), + _max_reconnect_attempts=1, + ) +finally: + console.websocket.WebSocketApp = real_websocket_app + sys.stdin = original_stdin + stdin_stream.close() + os.close(master_fd) + +if not binding_removed.is_set(): + raise SystemExit("Concurrent binding removal did not run") +if refresh_calls != 1: + raise SystemExit(f"Expected one local refresh check, got {refresh_calls}") +if state._client is not None: + raise SystemExit("Local stop unexpectedly initialized the HTTP client") +print(f"LOCAL_STOP_OK refresh_calls={refresh_calls} endpoint={session.endpoint}") +PY + +grep -a -q "Session '$SESSION_NAME' is no longer active" "$STOP_OUTPUT_FILE" +grep -a -q "LOCAL_STOP_OK refresh_calls=1 endpoint=$TEST_ENDPOINT" "$STOP_OUTPUT_FILE" + +echo "[*] Stopping isolated CPU session..." +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" + +AFTER_ENDPOINTS=$(server_endpoints) +if [ "$AFTER_ENDPOINTS" != "$BEFORE_ENDPOINTS" ]; then + echo "[FAILURE] Pre-existing assignments changed during the test." >&2 + echo "Before:" >&2 + echo "$BEFORE_ENDPOINTS" >&2 + echo "After:" >&2 + echo "$AFTER_ENDPOINTS" >&2 + exit 1 +fi +TEST_ENDPOINT="" + +echo "[SUCCESS] Console visibly reconnected to the same CPU endpoint." diff --git a/integration/repro_runtime_token_refresh/test.sh b/integration/repro_runtime_token_refresh/test.sh new file mode 100644 index 0000000..605a718 --- /dev/null +++ b/integration/repro_runtime_token_refresh/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 +# +# https://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. + +# CPU-only end-to-end regression for issue #106. Every command starts with a +# deliberately invalid saved runtime-proxy token; session resolution must adopt +# the fresh token returned by /tun/m/assignments before touching the VM. + +set -euo pipefail + +TMP_DIR=$(mktemp -d) +SESSION_FILE="$TMP_DIR/sessions.json" +SERVER_SESSION_FILE="$TMP_DIR/server-snapshot.json" +SESSION_NAME="token-refresh-$PPID-$$" +TEST_ENDPOINT="" + +if [ -f "$HOME/.config/colab-cli/token.json" ]; then + AUTH_PROVIDER="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) + 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) + if echo "$ADC_SCOPES" | grep -q "userinfo.email"; then + AUTH_PROVIDER="adc" + else + echo "Error: ADC token lacks the userinfo.email scope." >&2 + exit 1 + fi +else + echo "Error: No usable auth provider found." >&2 + exit 1 +fi +AUTH_FLAGS="--auth=$AUTH_PROVIDER" + +server_endpoints() { + # Isolate the listing from the user's normal sessions.json. The command + # reads server assignments but cannot adopt, rename, or rewrite any of the + # user's pre-existing local bindings. + uv run colab $AUTH_FLAGS --config "$SERVER_SESSION_FILE" sessions 2>/dev/null | sed -n 's/^\[[^]]*\] \([^ ]*\) .*/\1/p' | sort +} + +BEFORE_ENDPOINTS=$(server_endpoints) + +cleanup() { + cleanup_endpoint="$TEST_ENDPOINT" + if [ -z "$cleanup_endpoint" ] && [ -s "$SESSION_FILE" ]; then + cleanup_endpoint=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' 2>/dev/null || true +import json +import sys + +print(json.load(open(sys.argv[1])).get(sys.argv[2], {}).get("endpoint", "")) +PY + ) + fi + + # Always try the isolated local binding: `new` may have succeeded even if + # the immediately-following endpoint read failed under `set -e`. + uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" >/dev/null 2>&1 || true + if [ -n "$cleanup_endpoint" ]; then + # Exact, idempotent fallback. Do not gate cleanup on another assignments + # listing: that lookup may be the failure that triggered this trap. + uv run python - "$AUTH_PROVIDER" "$cleanup_endpoint" <<'PY' >/dev/null 2>&1 || true +import sys +from colab_cli.auth import AuthProvider +from colab_cli.common import state + +state.auth_provider = AuthProvider(sys.argv[1]) +state.client.unassign(sys.argv[2]) +PY + fi + rm -rf "$TMP_DIR" +} +trap cleanup EXIT + +echo "[*] Creating isolated CPU session $SESSION_NAME..." +# Intentionally omit --gpu/--tpu: this regression must never consume an +# accelerator allocation. +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" new -s "$SESSION_NAME" +TEST_ENDPOINT=$(uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' +import json +import sys + +print(json.load(open(sys.argv[1]))[sys.argv[2]]["endpoint"]) +PY +) + +expire_saved_token() { + uv run python - "$SESSION_FILE" "$SESSION_NAME" <<'PY' +import json +import sys + +path, name = sys.argv[1:] +with open(path) as f: + data = json.load(f) +data[name]["token"] = "deliberately-expired-runtime-proxy-token" +with open(path, "w") as f: + json.dump(data, f, indent=2) +PY +} + +expire_saved_token +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" ls -s "$SESSION_NAME" content >/dev/null + +expire_saved_token +EXEC_OUT=$(echo 'print("TOKEN-REFRESH-EXEC-OK")' | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" exec -s "$SESSION_NAME") +echo "$EXEC_OUT" | grep -q "TOKEN-REFRESH-EXEC-OK" + +expire_saved_token +CONSOLE_OUT="$TMP_DIR/console.out" +timeout 30 bash -c "echo 'echo TOKEN-REFRESH-CONSOLE-OK' | uv run colab $AUTH_FLAGS --config '$SESSION_FILE' console -s '$SESSION_NAME'" >"$CONSOLE_OUT" 2>&1 +grep -a -q "TOKEN-REFRESH-CONSOLE-OK" "$CONSOLE_OUT" + +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" + +AFTER_ENDPOINTS=$(server_endpoints) +if [ "$AFTER_ENDPOINTS" != "$BEFORE_ENDPOINTS" ]; then + echo "[FAILURE] Pre-existing assignments changed during the test." >&2 + echo "Before:" >&2 + echo "$BEFORE_ENDPOINTS" >&2 + echo "After:" >&2 + echo "$AFTER_ENDPOINTS" >&2 + exit 1 +fi +TEST_ENDPOINT="" + +echo "[SUCCESS] Runtime token refresh works for ls, exec, and console on CPU." diff --git a/src/colab_cli/client.py b/src/colab_cli/client.py index b3380e1..44774bb 100644 --- a/src/colab_cli/client.py +++ b/src/colab_cli/client.py @@ -241,9 +241,14 @@ def _issue_request( return return TypeAdapter(schema).validate_python(json.loads(body)) - def list_assignments(self) -> List[ListedAssignment]: + def list_assignments( + self, *, timeout: Optional[float] = None + ) -> List[ListedAssignment]: url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/assignments") - assignments = self._issue_request(url, schema=ListedAssignments) + request_kwargs = {"timeout": timeout} if timeout is not None else {} + assignments = self._issue_request( + url, schema=ListedAssignments, **request_kwargs + ) return assignments.assignments def unassign(self, endpoint: str): @@ -261,9 +266,7 @@ def assign( accelerator: Optional[Accelerator] = None, shape: Optional[Shape] = None, ) -> Union[PostAssignmentResponse, Assignment]: - assignment = self._get_assignment( - notebook_hash, variant, accelerator, shape - ) + assignment = self._get_assignment(notebook_hash, variant, accelerator, shape) if isinstance(assignment, Assignment): return assignment diff --git a/src/colab_cli/commands/automation.py b/src/colab_cli/commands/automation.py index 18c02a5..09b44d9 100644 --- a/src/colab_cli/commands/automation.py +++ b/src/colab_cli/commands/automation.py @@ -38,7 +38,6 @@ INTERACTIVE_AUTOMATION_TIMEOUT_SEC = 600 - def run_automation( name: str, op: str, @@ -49,8 +48,18 @@ def run_automation( ): from colab_cli.common import state - s = state.store.get(name) - runtime = ColabRuntime(s.url, s.token, session_name=s.name, history=state.history) + def start_runtime(s): + runtime = ColabRuntime( + s.url, s.token, session_name=s.name, history=state.history + ) + try: + _ = runtime.kernel_client + except Exception: + runtime.stop() + raise + return runtime, s + + runtime, s = state.run_with_runtime_proxy_retry(name, start_runtime) def drivefs_hook(deserialize_msg, wsclient): content = deserialize_msg.get("content", {}) @@ -133,13 +142,17 @@ def drivefs_hook(deserialize_msg, wsclient): runtime.colab_request_hook = drivefs_hook try: - s.running = f"automation({op})" s.last_execution = ( f"automation:{op}", None, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), ) - state.store.add(s) + state.store.update_fields( + name, + s.endpoint, + running=f"automation({op})", + last_execution=s.last_execution, + ) if op == "drivemount": state.history.log_event( @@ -169,8 +182,7 @@ def drivefs_hook(deserialize_msg, wsclient): else: sys.stderr.write(f"{ename}: {evalue}\n") finally: - s.running = None - state.store.add(s) + state.store.update_fields(name, s.endpoint, running=None) runtime.stop() @@ -240,9 +252,11 @@ def install( if not os.path.isfile(requirement): typer.echo(f"[colab] Requirements file '{requirement}' not found locally.") raise typer.Exit(1) - contents = ContentsClient(state.store.get(name)) remote_path = f"content/{os.path.basename(requirement)}" - contents.upload(requirement, remote_path) + state.run_with_runtime_proxy_retry( + name, + lambda s: ContentsClient(s).upload(requirement, remote_path), + ) commands.extend(["-r", f"/{remote_path}"]) if packages: commands.extend(packages) diff --git a/src/colab_cli/commands/execution.py b/src/colab_cli/commands/execution.py index 0ae23d1..81a3ada 100644 --- a/src/colab_cli/commands/execution.py +++ b/src/colab_cli/commands/execution.py @@ -24,12 +24,14 @@ from typing import List, Optional from typing_extensions import Annotated +from colab_cli.console import ConsoleConnectionError, connect_console from colab_cli.runtime import ColabRuntime -from colab_cli.utils import handle_image, is_terminal_error, render_display_data -from colab_cli.console import connect_console +from colab_cli.utils import handle_image, is_runtime_proxy_error, render_display_data _console = Console() +CONSOLE_REFRESH_TIMEOUT_SECONDS = 10 + TITLE_REGEX = re.compile(r"^\s*#\s*@title\s+(.*)", re.MULTILINE) ENV_KEY_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") @@ -38,6 +40,23 @@ def is_stdin_tty(): return sys.stdin.isatty() +def _refresh_console_session(state, name, expected): + """Refreshes one Console binding without reviving a locally stopped VM.""" + current = state.store.get(name) + if current is None: + # `colab stop` removes the endpoint only after unassign succeeds, so a + # missing local binding is already conclusive and needs no HTTP lookup. + return None + if current.endpoint != expected.endpoint: + # Let Console's endpoint guard report and reject the replacement. + return current + return state.refresh_session( + name, + expected_session=current, + timeout=CONSOLE_REFRESH_TIMEOUT_SECONDS, + ) + + def _parse_env_vars(env: Optional[List[str]]) -> dict[str, str]: """Parse repeatable --env KEY=VALUE entries into an ordered mapping.""" env_vars = {} @@ -72,6 +91,52 @@ def _build_env_prelude(env_vars: dict[str, str]) -> str: return "\n".join(lines) + "\n" +def _start_runtime(state, name, session): + """Starts a runtime, allowing the state layer to retry with fresh creds.""" + endpoint = session.endpoint + + def on_started(kernel_id): + state.store.update_fields(name, endpoint, kernel_id=kernel_id) + + def on_session_started(session_id): + state.store.update_fields(name, endpoint, session_id=session_id) + + runtime = ColabRuntime( + session.url, + session.token, + kernel_id=session.kernel_id, + session_id=session.session_id, + on_kernel_started=on_started, + on_session_started=on_session_started, + ) + try: + runtime.execute_code( + "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" + ) + except Exception: + runtime.stop() + raise + return runtime, session + + +def _connect_runtime(state, name): + return state.run_with_runtime_proxy_retry( + name, lambda session: _start_runtime(state, name, session) + ) + + +def _raise_runtime_connection_error(state, name, error): + if is_runtime_proxy_error(error): + if state.store.get(name) is None: + typer.echo(f"[colab] Session '{name}' is no longer active.") + else: + typer.echo( + f"[colab] Session '{name}' rejected refreshed runtime credentials." + ) + raise typer.Exit(1) + raise error + + def save_output(outputs, cell): if cell is None: return @@ -199,40 +264,14 @@ def exec_command( if not any(b["code"].strip() for b in code_blocks): raise typer.Exit(0) - def on_started(kid): - s.kernel_id = kid - state.store.add(s) - - def on_sess_started(sid): - s.session_id = sid - state.store.add(s) - - runtime = ColabRuntime( - s.url, - s.token, - kernel_id=s.kernel_id, - session_id=s.session_id, - on_kernel_started=on_started, - on_session_started=on_sess_started, - ) try: - # Ensure we are in /content which is the standard Colab working directory - runtime.execute_code( - "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" - ) + runtime, s = _connect_runtime(state, name) except Exception as e: - if is_terminal_error(e): - typer.echo( - f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." - ) - state.prune_session(name) - raise typer.Exit(1) - raise e + _raise_runtime_connection_error(state, name, e) try: is_nb = file and file.endswith(".ipynb") - s.running = f"exec({file or 'stdin'})" - state.store.add(s) + state.store.update_fields(name, s.endpoint, running=f"exec({file or 'stdin'})") for i, block in enumerate(code_blocks): code = _build_env_prelude(env_vars) + block["code"] @@ -256,7 +295,7 @@ def on_sess_started(sid): identifier, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), ) - state.store.add(s) + state.store.update_fields(name, s.endpoint, last_execution=s.last_execution) outputs = runtime.execute_code( code, @@ -276,8 +315,7 @@ def on_sess_started(sid): }, ) finally: - s.running = None - state.store.add(s) + state.store.update_fields(name, s.endpoint, running=None) runtime.stop() if file and file.endswith(".ipynb"): output_file = os.path.splitext(file)[0] + "_output.ipynb" @@ -303,35 +341,10 @@ def repl( typer.echo(f"[colab] Session '{name}' not found.") raise typer.Exit(1) - def on_started(kid): - s.kernel_id = kid - state.store.add(s) - - def on_sess_started(sid): - s.session_id = sid - state.store.add(s) - - runtime = ColabRuntime( - s.url, - s.token, - kernel_id=s.kernel_id, - session_id=s.session_id, - on_kernel_started=on_started, - on_session_started=on_sess_started, - ) try: - # Ensure we are in /content which is the standard Colab working directory - runtime.execute_code( - "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" - ) + runtime, s = _connect_runtime(state, name) except Exception as e: - if is_terminal_error(e): - typer.echo( - f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." - ) - state.prune_session(name) - raise typer.Exit(1) - raise e + _raise_runtime_connection_error(state, name, e) if not is_stdin_tty(): code = sys.stdin.read() @@ -343,8 +356,12 @@ def on_sess_started(sid): None, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), ) - s.running = "repl(stdin)" - state.store.add(s) + state.store.update_fields( + name, + s.endpoint, + last_execution=s.last_execution, + running="repl(stdin)", + ) try: outputs = runtime.execute_code( code, output_hook=lambda o: display_output(o, output_image) @@ -353,14 +370,12 @@ def on_sess_started(sid): name, "execution", {"code": code, "outputs": outputs, "source": "piped"} ) finally: - s.running = None - state.store.add(s) + state.store.update_fields(name, s.endpoint, running=None) runtime.stop() else: from colab_cli.repl import ColabREPL - s.running = "repl" - state.store.add(s) + state.store.update_fields(name, s.endpoint, running="repl") try: repl_inst = ColabREPL( runtime, @@ -371,8 +386,7 @@ def on_sess_started(sid): state.history.log_event(name, "repl_started", {}) repl_inst.run() finally: - s.running = None - state.store.add(s) + state.store.update_fields(name, s.endpoint, running=None) def console( @@ -389,21 +403,24 @@ def console( typer.echo(f"[colab] Session '{name}' not found.") raise typer.Exit(1) state.history.log_event(s.name, "console_started", {}) - s.running = "console" - state.store.add(s) + state.store.update_fields(name, s.endpoint, running="console") try: - connect_console(s) + state.run_with_runtime_proxy_retry( + name, + lambda current: connect_console( + current, + refresh_session=lambda expected: _refresh_console_session( + state, name, expected + ), + ), + ) + except ConsoleConnectionError as e: + typer.echo(f"[colab] Console disconnected: {e}", err=True) + raise typer.Exit(1) from e except Exception as e: - if is_terminal_error(e): - typer.echo( - f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." - ) - state.prune_session(name) - raise typer.Exit(1) - raise e + _raise_runtime_connection_error(state, name, e) finally: - s.running = None - state.store.add(s) + state.store.update_fields(name, s.endpoint, running=None) def register(app: typer.Typer): diff --git a/src/colab_cli/commands/files.py b/src/colab_cli/commands/files.py index 494de1e..fd1738f 100644 --- a/src/colab_cli/commands/files.py +++ b/src/colab_cli/commands/files.py @@ -33,13 +33,13 @@ def ls( from colab_cli.common import state name = state.resolve_session(session) - s = state.store.get(name) - if not s: + if not state.store.get(name): typer.echo(f"[colab] Session '{name}' not found.") raise typer.Exit(1) - contents = ContentsClient(s) try: - data = contents.list_dir(path) + data = state.run_with_runtime_proxy_retry( + name, lambda s: ContentsClient(s).list_dir(path) + ) state.history.log_event(name, "file_operation", {"op": "ls", "path": path}) if data.get("type") == "directory": items = data.get("content", []) @@ -65,13 +65,11 @@ def rm( from colab_cli.common import state name = state.resolve_session(session) - s = state.store.get(name) - if not s: + if not state.store.get(name): typer.echo(f"[colab] Session '{name}' not found.") raise typer.Exit(1) - contents = ContentsClient(s) try: - contents.rm(path) + state.run_with_runtime_proxy_retry(name, lambda s: ContentsClient(s).rm(path)) state.history.log_event(name, "file_operation", {"op": "rm", "path": path}) typer.echo(f"[colab] Deleted {path}") except Exception as e: @@ -90,16 +88,16 @@ def upload( from colab_cli.common import state name = state.resolve_session(session) - s = state.store.get(name) - if not s: + if not state.store.get(name): typer.echo(f"[colab] Session '{name}' not found.") raise typer.Exit(1) if not os.path.isfile(local_path): typer.echo(f"[colab] Local file '{local_path}' not found.") raise typer.Exit(1) - contents = ContentsClient(s) try: - contents.upload(local_path, remote_path) + state.run_with_runtime_proxy_retry( + name, lambda s: ContentsClient(s).upload(local_path, remote_path) + ) state.history.log_event( name, "file_operation", @@ -126,13 +124,13 @@ def download( from colab_cli.common import state name = state.resolve_session(session) - s = state.store.get(name) - if not s: + if not state.store.get(name): typer.echo(f"[colab] Session '{name}' not found.") raise typer.Exit(1) - contents = ContentsClient(s) try: - contents.download(remote_path, local_path) + state.run_with_runtime_proxy_retry( + name, lambda s: ContentsClient(s).download(remote_path, local_path) + ) state.history.log_event( name, "file_operation", @@ -154,13 +152,10 @@ def edit( from colab_cli.common import state name = state.resolve_session(session) - s = state.store.get(name) - if not s: + if not state.store.get(name): typer.echo(f"[colab] Session '{name}' not found.") raise typer.Exit(1) - contents = ContentsClient(s) - def get_file_hash(path): if not os.path.exists(path): return None @@ -173,7 +168,9 @@ def get_file_hash(path): local_path = tf.name try: - contents.download(remote_path, local_path) + state.run_with_runtime_proxy_retry( + name, lambda s: ContentsClient(s).download(remote_path, local_path) + ) except Exception: # If download fails, assume file doesn't exist and start empty pass @@ -185,7 +182,9 @@ def get_file_hash(path): hash_after = get_file_hash(local_path) if hash_after != hash_before: - contents.upload(local_path, remote_path) + state.run_with_runtime_proxy_retry( + name, lambda s: ContentsClient(s).upload(local_path, remote_path) + ) state.history.log_event( name, "file_operation", diff --git a/src/colab_cli/commands/run.py b/src/colab_cli/commands/run.py index 550854a..968fd36 100644 --- a/src/colab_cli/commands/run.py +++ b/src/colab_cli/commands/run.py @@ -45,7 +45,10 @@ PostAssignmentResponse, Shape, ) -from colab_cli.commands.execution import _build_env_prelude, _parse_env_vars +from colab_cli.commands.execution import ( + _build_env_prelude, + _parse_env_vars, +) from colab_cli.commands.session import ( _is_scope_error, _scope_remediation_message, @@ -54,7 +57,35 @@ ) from colab_cli.runtime import ColabRuntime from colab_cli.state import SessionState -from colab_cli.utils import get_status_code, is_terminal_error +from colab_cli.utils import get_status_code + + +def _start_run_runtime(state, name, session): + """Starts the one-shot runtime using this module's client seam.""" + endpoint = session.endpoint + + def on_started(kernel_id): + state.store.update_fields(name, endpoint, kernel_id=kernel_id) + + def on_session_started(session_id): + state.store.update_fields(name, endpoint, session_id=session_id) + + runtime = ColabRuntime( + session.url, + session.token, + kernel_id=session.kernel_id, + session_id=session.session_id, + on_kernel_started=on_started, + on_session_started=on_session_started, + ) + try: + runtime.execute_code( + "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" + ) + except Exception: + runtime.stop() + raise + return runtime, session def _build_script_payload( @@ -297,9 +328,7 @@ def run_command( raise typer.Exit(2) name = session or f"run-{uuid.uuid4().hex[:6]}" - variant, accelerator, shape = resolve_runtime_options( - gpu, tpu, high_mem=high_mem - ) + variant, accelerator, shape = resolve_runtime_options(gpu, tpu, high_mem=high_mem) if high_mem and accelerator in HIGH_MEM_ONLY_ACCELERATORS: typer.echo( @@ -396,40 +425,14 @@ def run_command( # ----- Execute the script ------------------------------------------------- exit_code = 0 cleanup_reason = "run_completed" - - def on_started(kid): - s.kernel_id = kid - state.store.add(s) - - def on_sess_started(sid): - s.session_id = sid - state.store.add(s) - - runtime = ColabRuntime( - s.url, - s.token, - kernel_id=s.kernel_id, - session_id=s.session_id, - on_kernel_started=on_started, - on_session_started=on_sess_started, - ) + runtime = None try: - # Same /content prelude as `colab exec` for consistency. - try: - runtime.execute_code( - "import os; os.makedirs('/content', exist_ok=True); " - "os.chdir('/content')" - ) - except Exception as e: - if is_terminal_error(e): - typer.echo( - f"[colab] Session '{name}' appears to be lost (404/401).", - err=True, - ) - state.prune_session(name) - raise typer.Exit(1) - raise + runtime, s = state.run_with_runtime_proxy_retry( + name, + lambda current: _start_run_runtime(state, name, current), + initial_session=s, + ) payload = _build_script_payload(script, script_args, env_vars) s.running = f"run({os.path.basename(script)})" @@ -438,7 +441,12 @@ def on_sess_started(sid): None, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), ) - state.store.add(s) + state.store.update_fields( + name, + s.endpoint, + running=s.running, + last_execution=s.last_execution, + ) try: outputs = runtime.execute_code( @@ -461,12 +469,13 @@ def on_sess_started(sid): ) finally: s.running = None - state.store.add(s) + state.store.update_fields(name, s.endpoint, running=None) # Best-effort runtime close (keeps remote kernel alive for --keep). - try: - runtime.stop() - except Exception: - pass + if runtime is not None: + try: + runtime.stop() + except Exception: + pass if not keep: _teardown(name, s, reason=cleanup_reason) @@ -504,7 +513,7 @@ def _teardown(name: str, s: SessionState, *, reason: str) -> None: pass try: - state.store.remove(name) + state.store.remove_if_endpoint(name, s.endpoint) except Exception: pass diff --git a/src/colab_cli/commands/session.py b/src/colab_cli/commands/session.py index 537262d..28004da 100644 --- a/src/colab_cli/commands/session.py +++ b/src/colab_cli/commands/session.py @@ -299,29 +299,30 @@ def restart_kernel( from colab_cli.common import state name = state.resolve_session(session) - s = state.store.get(name) - def on_started(kid): - s.kernel_id = kid - state.store.add(s) - - def on_sess_started(sid): - s.session_id = sid - state.store.add(s) - - runtime = ColabRuntime( - s.url, - s.token, - kernel_id=s.kernel_id, - session_id=s.session_id, - on_kernel_started=on_started, - on_session_started=on_sess_started, - ) + def restart(s): + endpoint = s.endpoint - try: - runtime.restart() - finally: - runtime.stop() + def on_started(kernel_id): + state.store.update_fields(name, endpoint, kernel_id=kernel_id) + + def on_session_started(session_id): + state.store.update_fields(name, endpoint, session_id=session_id) + + runtime = ColabRuntime( + s.url, + s.token, + kernel_id=s.kernel_id, + session_id=s.session_id, + on_kernel_started=on_started, + on_session_started=on_session_started, + ) + try: + return runtime.restart() + finally: + runtime.stop() + + state.run_with_runtime_proxy_retry(name, restart) def sessions_command(): @@ -421,7 +422,7 @@ def stop( pass state.client.unassign(s.endpoint) - state.store.remove(name) + state.store.remove_if_endpoint(name, s.endpoint) state.history.log_event(name, "session_terminated", {"reason": "user_requested"}) typer.echo("[colab] Session terminated.") diff --git a/src/colab_cli/commands/ssh.py b/src/colab_cli/commands/ssh.py index 86226e4..ec41c9d 100644 --- a/src/colab_cli/commands/ssh.py +++ b/src/colab_cli/commands/ssh.py @@ -48,6 +48,7 @@ import uuid from colab_cli.state import SessionState +from colab_cli.utils import RuntimeProxyError import typer from typing_extensions import Annotated import websocket @@ -93,9 +94,7 @@ def _pubkey_from_identity(identity: str) -> str: raise typer.Exit(code=2) pubkey = res.stdout.strip() if not pubkey: - typer.echo( - f"[colab] ssh-keygen produced no key for {identity}.", err=True - ) + typer.echo(f"[colab] ssh-keygen produced no key for {identity}.", err=True) raise typer.Exit(code=2) return pubkey @@ -319,6 +318,8 @@ def _connect_websocket(url: str, pubkey: str) -> websocket.WebSocket: body = getattr(e, "resp_body", b"") or b"" if isinstance(body, str): body = body.encode("utf-8", errors="replace") + if status == 401: + raise RuntimeProxyError(status, body) from e msg = _explain_handshake_failure(status, body) typer.echo(f"[colab] {msg}", err=True) raise typer.Exit(code=1) @@ -483,9 +484,7 @@ def _select_proxy_session( """ if session and not _session_exists(session): with contextlib.redirect_stdout(sys.stderr): - return _auto_create_session( - gpu, tpu, name=session, high_mem=high_mem - ), True + return _auto_create_session(gpu, tpu, name=session, high_mem=high_mem), True return _resolve_session(session), False @@ -554,9 +553,7 @@ def _on_signal(signum, frame): pass # e.g. not running in the main thread -def _run_proxy_bridge( - s: SessionState, identity: Optional[str], rm: bool -) -> int: +def _run_proxy_bridge(s: SessionState, identity: Optional[str], rm: bool) -> int: """Runs the ``--proxy-mode`` WebSocket-stdio bridge, honoring ``--rm``. Args: @@ -583,7 +580,18 @@ def _do_rm() -> None: _install_rm_signal_handlers(_do_rm) pubkey = _resolve_pubkey(identity) - ws = _connect_websocket(_build_ws_url(s), pubkey) + from colab_cli.common import state + + try: + ws = state.run_with_runtime_proxy_retry( + s.name, + lambda current: _connect_websocket(_build_ws_url(current), pubkey), + initial_session=s, + ) + except RuntimeProxyError as error: + message = _explain_handshake_failure(error.status_code, error.response_body) + typer.echo(f"[colab] {message}", err=True) + raise typer.Exit(code=1) from error try: return _bridge_proxy_mode(ws) finally: diff --git a/src/colab_cli/common.py b/src/colab_cli/common.py index 2fea125..22dad5d 100644 --- a/src/colab_cli/common.py +++ b/src/colab_cli/common.py @@ -17,14 +17,18 @@ import signal import sys import time -from typing import Optional +from typing import Callable, Optional, TypeVar import typer from colab_cli.auth import AuthProvider, get_credentials from colab_cli.client import Client, Prod from colab_cli.history import HistoryLogger -from colab_cli.state import StateStore, SettingsStore +from colab_cli.state import SessionState, StateStore, SettingsStore +from colab_cli.utils import is_runtime_proxy_error + + +T = TypeVar("T") class State: @@ -68,20 +72,120 @@ def client(self): self._client = Client(Prod(), creds) return self._client - def prune_session(self, name: str): - """Removes a session from local state and kills its keep-alive process.""" - s = self.store.get(name) - if s and s.keep_alive_pid: + def _remove_session(self, name: str, endpoint: str) -> bool: + """Removes a binding only if it still refers to ``endpoint``.""" + s = self.store.remove_if_endpoint(name, endpoint) + if s is None: + return False + if s.keep_alive_pid: kill_process(s.keep_alive_pid) - self.store.remove(name) if self._sessions and name in self._sessions: del self._sessions[name] self.history.log_event(name, "session_terminated", {"reason": "pruned"}) + return True - def sync_sessions(self): - if self._sessions is not None: - return self._sessions, self.client.list_assignments() + def _refresh_session_from_assignments( + self, + name: str, + assignments, + expected_session: Optional[SessionState] = None, + ): + """Reconciles one local binding with a known server-side snapshot.""" + s = expected_session or self.store.get(name) + if s is None: + return None + assignment = next( + ( + candidate + for candidate in assignments + if candidate.endpoint == s.endpoint + ), + None, + ) + if assignment is None: + if self._remove_session(name, s.endpoint): + return None + # Another process replaced the name after this server snapshot + # began. The snapshot says nothing about that new endpoint, so + # preserve it and let a later resolution reconcile it. + current = self.store.get(name) + if current is not None and self._sessions is not None: + self._sessions[name] = current + return current + + rpi = assignment.runtime_proxy_info + refreshed = self.store.update_fields( + name, s.endpoint, token=rpi.token, url=rpi.url + ) + if refreshed is None: + # A same-name replacement won the endpoint guard. Do not apply + # credentials from the old endpoint to it. + refreshed = self.store.get(name) + if refreshed is not None and self._sessions is not None: + self._sessions[name] = refreshed + return refreshed + + def refresh_session( + self, + name: str, + expected_session: Optional[SessionState] = None, + timeout: Optional[float] = None, + ): + """Fetches and adopts the current runtime-proxy credentials.""" + s = expected_session or self.store.get(name) + if s is None: + return None + request_kwargs = {"timeout": timeout} if timeout is not None else {} + return self._refresh_session_from_assignments( + name, + self.client.list_assignments(**request_kwargs), + expected_session=s, + ) + + def prune_session(self, name: str) -> bool: + """Prunes only after the control plane confirms the assignment is gone. + If the check is inconclusive, preserve the binding: deleting local + access to a potentially live, billable VM is the worse failure mode. + """ + s = self.store.get(name) + if s is None: + return False + try: + refreshed = self.refresh_session(name, expected_session=s) + except Exception: + return False + return refreshed is None + + def run_with_runtime_proxy_retry( + self, + name: str, + operation: Callable[[SessionState], T], + initial_session: Optional[SessionState] = None, + ) -> T: + """Runs an operation and retries once with fresh proxy credentials.""" + s = initial_session or self.store.get(name) + if s is None: + raise RuntimeError(f"Session '{name}' not found") + try: + return operation(s) + except Exception as error: + if not is_runtime_proxy_error(error): + raise + old_credentials = (s.token, s.url) + try: + refreshed = self.refresh_session(name, expected_session=s) + except Exception: + raise error + if ( + refreshed is None + or refreshed.endpoint != s.endpoint + or (refreshed.token, refreshed.url) == old_credentials + ): + raise error + return operation(refreshed) + + def sync_sessions(self): # Check local store first. If it's empty, we don't necessarily need to hit the backend # unless we are specifically looking for server-side assignments (e.g. 'colab sessions'). local_sessions = self.store.list() @@ -97,13 +201,15 @@ def sync_sessions(self): return self._sessions, assignments assignments = self.client.list_assignments() - active_endpoints = {a.endpoint for a in assignments} - self._sessions = local_sessions pruned = 0 for name, s in list(self._sessions.items()): - if s.endpoint not in active_endpoints: - self.prune_session(name) + if ( + self._refresh_session_from_assignments( + name, assignments, expected_session=s + ) + is None + ): pruned += 1 if pruned > 0: @@ -113,6 +219,14 @@ def sync_sessions(self): def resolve_session(self, session_name: Optional[str]) -> str: if session_name: + s = self.store.get(session_name) + if s is not None: + try: + self.refresh_session(session_name, expected_session=s) + except Exception: + # The runtime path may still be reachable when the control + # plane is temporarily unavailable. Preserve and try it. + pass return session_name # Check local store first to avoid hitting the backend (and triggering auth) if we don't have to diff --git a/src/colab_cli/console.py b/src/colab_cli/console.py index 0766e1d..4ee92c1 100644 --- a/src/colab_cli/console.py +++ b/src/colab_cli/console.py @@ -12,58 +12,82 @@ # See the License for the specific language governing permissions and # limitations under the License. +import codecs +import itertools import json import logging import os +import select import signal +import struct import sys import termios import threading import time import tty +from collections.abc import Callable, Iterable, Iterator +from dataclasses import dataclass +from typing import Optional from urllib.parse import urlparse import websocket from colab_cli.state import SessionState +from colab_cli.utils import is_runtime_proxy_error logger = logging.getLogger(__name__) -# Global flag to stop the read thread when the websocket closes +# These globals and module-level callbacks are retained for compatibility with +# callers that exercise the individual callback functions. ``connect_console`` +# uses connection-local callbacks so multiple Console processes cannot share +# lifecycle state. _is_running = False _last_error = None -# When stdin is piped and reaches EOF, we send "exit\n" to the remote shell and -# then wait this many seconds for any remaining output (the shell's goodbye, -# tmux teardown messages, etc.) to flush before closing the websocket from the -# client side. Empirically 0.5s is enough for the typical /colab/tty backend -# wrapped in tmux + bash; bumping it just delays exit, lowering it risks -# truncating tail output. +# When stdin is piped and reaches EOF, let the remote shell flush its goodbye +# before the client closes the websocket. PIPED_EOF_GRACE_SECONDS = 0.5 +# Protocol pings keep an otherwise idle TTY tunnel visible to HTTP proxies and +# bound how long a silently dead connection can look healthy. +CONSOLE_PING_INTERVAL_SECONDS = 20 +CONSOLE_PING_TIMEOUT_SECONDS = 10 + +# Retry quickly through brief proxy churn, then settle into a low-frequency +# retry cadence until the user cancels or Colab confirms that the VM is gone. +CONSOLE_RETRY_DELAYS_SECONDS = (1, 2, 5, 10, 30) + +# Shell-exit input is only a hint: ``exit`` may leave a nested shell and +# Ctrl-D may close a foreground program without closing the Console transport. +# Suppress reconnect only when the peer closes immediately after that input. +CONSOLE_SHELL_EXIT_INTENT_SECONDS = 2 + +# Treat an established socket that survives this long as a recovered +# connection. A later outage starts a fresh backoff sequence instead of +# inheriting the retry delay from an unrelated earlier outage. +CONSOLE_STABLE_CONNECTION_SECONDS = 30 + +_NORMAL_CLOSE_CODES = (1000, 1001) + + +class ConsoleConnectionError(RuntimeError): + """The Console transport failed before piped input completed.""" + def on_message(ws, message): - """Callback for when a message is received from the server.""" - try: - data = json.loads(message) - if "data" in data: - # The backend sends raw ANSI escape sequences and string content. - # We write it directly to stdout buffer to avoid python print() formatting. - sys.stdout.buffer.write(data["data"].encode("utf-8")) - sys.stdout.buffer.flush() - except Exception as e: - logger.debug(f"Error parsing message: {e}") + """Compatibility callback for writing remote terminal output.""" + _write_terminal_message(ws, message) def on_error(ws, error): - """Callback for when a websocket error occurs.""" + """Compatibility callback for recording a websocket error.""" global _last_error _last_error = error - logger.error(f"WebSocket Error: {error}") + logger.error("WebSocket Error: %s", error) def on_close(ws, close_status_code, close_msg): - """Callback for when the websocket is closed.""" + """Compatibility callback for recording websocket closure.""" global _is_running _is_running = False @@ -75,33 +99,22 @@ def send_terminal_size(ws): payload = json.dumps({"cols": size.columns, "rows": size.lines}) ws.send(payload) except Exception as e: - logger.debug(f"Failed to send terminal size: {e}") + logger.debug("Failed to send terminal size: %s", e) def on_open(ws): - """Callback for when the websocket connection is opened.""" + """Compatibility callback that opens one legacy stdin forwarding thread.""" global _is_running _is_running = True - - # Send initial terminal size send_terminal_size(ws) - # Setup the background thread to read from stdin def read_stdin(): is_tty = sys.stdin.isatty() while _is_running: try: - # Read a single character (or escape sequence byte) char = sys.stdin.read(1) if not char: if not is_tty: - # Piped input has reached EOF. The remote /colab/tty - # endpoint wraps bash in tmux which intercepts \x04 - # (Ctrl-D) as a literal character, so it never exits. - # Instead send "exit\n" so bash voluntarily terminates, - # wait a short grace period for the shell's goodbye - # output to drain back to us, then close the websocket - # ourselves to guarantee the client unblocks. try: ws.send(json.dumps({"data": "exit\n"})) except Exception: @@ -116,57 +129,474 @@ def read_stdin(): except Exception: break - thread = threading.Thread(target=read_stdin, daemon=True) - thread.start() + threading.Thread(target=read_stdin, daemon=True).start() -def connect_console(session: SessionState): - """ - Connects to the Colab TTY endpoint and sets up a raw terminal session. - """ - global _is_running, _last_error - _last_error = None +def _write_terminal_message(ws, message) -> None: + """Writes terminal output and acknowledges the server's PTY flow control.""" + try: + data = json.loads(message) + if "data" in data: + sys.stdout.buffer.write(data["data"].encode("utf-8")) + sys.stdout.buffer.flush() + if data.get("ack") is True: + ws.send(json.dumps({"ack": True})) + except Exception as e: + logger.debug("Error handling Console message: %s", e) + + +def _status(message: str) -> None: + """Writes connection state outside the remote terminal byte stream.""" + print(f"\r\n[colab] {message}", file=sys.stderr, flush=True) + + +def _retry_delays(delays: Iterable[float]) -> Iterator[float]: + configured = tuple(delays) + if not configured: + configured = (30,) + yield from configured + yield from itertools.repeat(configured[-1]) + + +@dataclass +class _Attempt: + opened: bool = False + error: Optional[object] = None + received_close_frame: bool = False + close_code: Optional[int] = None + close_reason: str = "" + opened_at: Optional[float] = None + duration: float = 0.0 + + @property + def abnormal(self) -> bool: + # websocket-client 1.9 invokes on_error with the received ABNF close + # frame before on_close. A normal close code must win over that + # compatibility quirk or `exit` would spuriously reconnect. + if self.close_code in _NORMAL_CLOSE_CODES: + return False + if self.error is not None: + return True + if self.received_close_frame: + return True + if self.close_code is None: + # A mocked/no-op WebSocketApp has no callbacks. A real peer close + # invokes on_close and supplies either a code or an error callback. + return False + return True + + def description(self) -> str: + if self.error is not None: + return str(self.error) + if self.close_code is not None: + reason = f" ({self.close_reason})" if self.close_reason else "" + return f"WebSocket closed with code {self.close_code}{reason}" + return "WebSocket connection closed" + + +class _ConsoleInputForwarder: + """Owns the command's single stdin reader across websocket reconnects.""" + + def __init__(self, is_tty: bool): + self.is_tty = is_tty + self.stop_event = threading.Event() + self._user_requested_close = False + self._shell_close_intent_deadline: Optional[float] = None + self._active_event = threading.Event() + self._lock = threading.Lock() + self._ws = None + self._started = False + self._thread: Optional[threading.Thread] = None + self._line: list[str] = [] + self._stdin = sys.stdin + try: + self._stdin_fd: Optional[int] = self._stdin.fileno() + except (AttributeError, OSError, TypeError, ValueError): + self._stdin_fd = None + + def start(self) -> None: + if self._started: + return + self._started = True + self._thread = threading.Thread(target=self._read_stdin, daemon=True) + self._thread.start() + + def attach(self, ws) -> None: + with self._lock: + self._ws = ws + self._active_event.set() + + def detach(self, ws) -> None: + with self._lock: + if self._ws is ws: + self._ws = None + self._active_event.clear() + + def close_active(self) -> None: + with self._lock: + ws = self._ws + if ws is not None: + try: + ws.close() + except Exception: + pass + + def stop(self) -> None: + self.stop_event.set() + self._active_event.set() + + def join(self, timeout: Optional[float] = None) -> bool: + if self._thread is None: + return True + self._thread.join(timeout) + return not self._thread.is_alive() + + def _current_ws(self): + with self._lock: + return self._ws - # Construct the WebSocket URL from the base URL + def _wait_for_first_connection(self) -> bool: + while not self.stop_event.is_set(): + if self._active_event.wait(0.1): + return True + return False + + @property + def user_requested_close(self) -> bool: + """Whether this close should terminate rather than reconnect.""" + if self._user_requested_close: + return True + deadline = self._shell_close_intent_deadline + if deadline is None: + return False + if time.monotonic() <= deadline: + return True + self._shell_close_intent_deadline = None + return False + + @user_requested_close.setter + def user_requested_close(self, value: bool) -> None: + # Retain the attribute-style API for permanent cancellation paths and + # compatibility with existing callers/tests. Shell command detection + # uses the bounded intent helper below instead. + self._user_requested_close = value + if not value: + self._shell_close_intent_deadline = None + + def _mark_shell_close_intent(self) -> None: + self._shell_close_intent_deadline = ( + time.monotonic() + CONSOLE_SHELL_EXIT_INTENT_SECONDS + ) + + def _track_exit_request(self, char: str) -> None: + if char in ("\r", "\n"): + command = "".join(self._line).strip() + if command in ("exit", "logout") or command.startswith("exit "): + self._mark_shell_close_intent() + self._line.clear() + elif char in ("\x7f", "\b"): + if self._line: + self._line.pop() + elif char == "\x04": + self._mark_shell_close_intent() + elif char.isprintable(): + self._line.append(char) + + def _read_stdin(self) -> None: + # Piped bytes must not be consumed until the first socket is ready. + if not self.is_tty and not self._wait_for_first_connection(): + return + + if self._stdin_fd is None: + self._read_text_stdin() + return + + encoding = getattr(self._stdin, "encoding", None) or "utf-8" + errors = getattr(self._stdin, "errors", None) or "strict" + decoder = codecs.getincrementaldecoder(encoding)(errors=errors) + while not self.stop_event.is_set(): + try: + readable, _, _ = select.select([self._stdin_fd], [], [], 0.1) + if not readable: + continue + chunk = os.read(self._stdin_fd, 4096) + except (OSError, TypeError, ValueError): + return + + if not chunk: + tail = decoder.decode(b"", final=True) + for char in tail: + self._forward_char(char) + self._handle_eof() + return + + for char in decoder.decode(chunk): + self._forward_char(char) + + def _read_text_stdin(self) -> None: + """Fallback for synthetic streams without a selectable file descriptor.""" + while not self.stop_event.is_set(): + try: + char = self._stdin.read(1) + except Exception: + return + if not char: + self._handle_eof() + return + self._forward_char(char) + + def _handle_eof(self) -> None: + if self.is_tty: + return + self.user_requested_close = True + ws = self._current_ws() + if ws is not None: + try: + ws.send(json.dumps({"data": "exit\n"})) + except Exception: + pass + time.sleep(PIPED_EOF_GRACE_SECONDS) + try: + ws.close() + except Exception: + pass + self.stop_event.set() + + def _forward_char(self, char: str) -> None: + if self.stop_event.is_set(): + return + ws = self._current_ws() + if ws is None: + # Raw mode suppresses SIGINT. During a reconnect delay, make + # Ctrl-C an explicit request to stop retrying. + if self.is_tty and char == "\x03": + self.user_requested_close = True + self.stop_event.set() + _status("Console reconnect cancelled by user.") + return + + if self.is_tty: + self._track_exit_request(char) + try: + ws.send(json.dumps({"data": char})) + except Exception: + # The websocket callback/loop owns reconnect decisions. + pass + + +def _build_ws_url(session: SessionState) -> str: parsed = urlparse(session.url) ws_scheme = "wss" if parsed.scheme == "https" else "ws" - ws_url = f"{ws_scheme}://{parsed.netloc}/colab/tty?colab-runtime-proxy-token={session.token}" + return ( + f"{ws_scheme}://{parsed.netloc}/colab/tty" + f"?colab-runtime-proxy-token={session.token}" + ) + +def connect_console( + session: SessionState, + *, + refresh_session: Optional[Callable[[SessionState], Optional[SessionState]]] = None, + retry_delays: Iterable[float] = CONSOLE_RETRY_DELAYS_SECONDS, + _max_reconnect_attempts: Optional[int] = None, +) -> None: + """Connects to the Colab TTY and reconnects an interrupted TTY session. + + Reconnection never creates, stops, or replaces a runtime. The optional + refresh callback must return credentials for the same endpoint, ``None`` + when the control plane confirms it is gone, or raise when the lookup is + inconclusive. + """ is_tty = sys.stdin.isatty() fd = sys.stdin.fileno() if is_tty else None old_settings = termios.tcgetattr(fd) if is_tty else None - - ws = websocket.WebSocketApp( - url=ws_url, - on_open=on_open, - on_message=on_message, - on_error=on_error, - on_close=on_close, - ) + old_sigwinch = signal.getsignal(signal.SIGWINCH) if is_tty else None + forwarder = _ConsoleInputForwarder(is_tty) + current = session + original_endpoint = session.endpoint + delays = _retry_delays(retry_delays) + reconnect_attempt = 0 + total_reconnect_attempts = 0 + active_ws = {"ws": None} + pid = os.getpid() def handle_sigwinch(signum, frame): - """Handle window resize events.""" - if _is_running: + ws = active_ws["ws"] + if ws is not None: send_terminal_size(ws) try: if is_tty: tty.setraw(fd, termios.TCSANOW) signal.signal(signal.SIGWINCH, handle_sigwinch) + while not forwarder.stop_event.is_set(): + attempt = _Attempt() + ws_url = _build_ws_url(current) + + def attempt_open(ws): + attempt.opened = True + attempt.opened_at = time.monotonic() + active_ws["ws"] = ws + forwarder.attach(ws) + forwarder.start() + send_terminal_size(ws) + logger.info( + "Console connected pid=%s endpoint=%s reconnect_attempt=%s", + pid, + original_endpoint, + reconnect_attempt, + ) + if reconnect_attempt: + _status( + f"Console reconnected (attempt {reconnect_attempt}, " + f"endpoint {original_endpoint})." + ) + + def attempt_message(ws, message): + _write_terminal_message(ws, message) + + def attempt_error(ws, error): + # websocket-client 1.9 passes a received ABNF close frame to + # on_error, then invokes on_close with a lost/None status. Parse + # that frame here so normal 1000 closes do not reconnect while + # abnormal/empty closes do. + if ( + isinstance(error, websocket.ABNF) + and error.opcode == websocket.ABNF.OPCODE_CLOSE + ): + attempt.received_close_frame = True + if len(error.data) >= 2: + attempt.close_code = struct.unpack("!H", error.data[:2])[0] + attempt.close_reason = error.data[2:].decode( + "utf-8", errors="replace" + ) + else: + attempt.error = error + logger.debug( + "Console WebSocket error pid=%s endpoint=%s " + "reconnect_attempt=%s error=%s", + pid, + original_endpoint, + reconnect_attempt, + error, + ) + + def attempt_close(ws, close_status_code, close_msg): + attempt.received_close_frame = True + if close_status_code is not None: + attempt.close_code = close_status_code + attempt.close_reason = close_msg or "" + active_ws["ws"] = None + forwarder.detach(ws) + attempt.duration = ( + time.monotonic() - attempt.opened_at + if attempt.opened_at is not None + else 0.0 + ) + logger.info( + "Console closed pid=%s endpoint=%s reconnect_attempt=%s " + "code=%s reason=%r duration=%.1fs", + pid, + original_endpoint, + reconnect_attempt, + close_status_code, + close_msg or "", + attempt.duration, + ) - # This is a blocking call until the connection is closed - ws.run_forever() + ws = websocket.WebSocketApp( + url=ws_url, + on_open=attempt_open, + on_message=attempt_message, + on_error=attempt_error, + on_close=attempt_close, + ) + active_ws["ws"] = ws + try: + ws.run_forever( + ping_interval=CONSOLE_PING_INTERVAL_SECONDS, + ping_timeout=CONSOLE_PING_TIMEOUT_SECONDS, + ) + finally: + active_ws["ws"] = None + forwarder.detach(ws) + + if forwarder.user_requested_close or not attempt.abnormal: + break + + if attempt.opened and attempt.duration >= CONSOLE_STABLE_CONNECTION_SECONDS: + reconnect_attempt = 0 + delays = _retry_delays(retry_delays) + + if not attempt.opened and isinstance(attempt.error, Exception): + if is_runtime_proxy_error(attempt.error): + # Let State's bounded startup retry refresh an expired + # runtime-proxy token. Generic Console reconnects are for + # established sockets, not authentication failures. + raise attempt.error + + if not is_tty: + raise ConsoleConnectionError(attempt.description()) + + if ( + _max_reconnect_attempts is not None + and total_reconnect_attempts >= _max_reconnect_attempts + ): + raise ConsoleConnectionError( + "Console reconnect limit reached after " + f"{total_reconnect_attempts} attempt(s): " + f"{attempt.description()}" + ) + + _status(f"Console connection lost: {attempt.description()}.") + + reconnect_attempt += 1 + total_reconnect_attempts += 1 + delay = next(delays) + _status( + f"Reconnecting in {delay:g}s (attempt {reconnect_attempt}; " + "press Ctrl-C to stop)..." + ) + if forwarder.stop_event.wait(delay): + break - if _last_error: - # Re-raise or wrap terminal errors - err_msg = str(_last_error) - if "404" in err_msg or "401" in err_msg: - # We raise a standard exception that the caller can recognize - raise RuntimeError(f"Connection failed: {err_msg}") + if refresh_session is not None: + try: + refreshed = refresh_session(current) + except Exception as error: + logger.debug( + "Console credential refresh failed pid=%s endpoint=%s " + "reconnect_attempt=%s error=%s", + pid, + original_endpoint, + reconnect_attempt, + error, + ) + _status( + "Warning: could not refresh credentials; retrying with " + f"the last known token ({error})." + ) + else: + if refreshed is None: + _status( + f"Session '{session.name}' is no longer active; " + "stopping Console reconnects." + ) + break + if refreshed.endpoint != original_endpoint: + _status( + f"Session '{session.name}' now refers to a different " + "runtime; refusing to reconnect to it." + ) + break + current = refreshed + except KeyboardInterrupt: + forwarder.user_requested_close = True finally: + forwarder.stop() + forwarder.close_active() + forwarder.join(timeout=0.2) if is_tty: - # Always ensure the terminal is restored to its original state termios.tcsetattr(fd, termios.TCSANOW, old_settings) - # Restore the default signal handler for resize - signal.signal(signal.SIGWINCH, signal.SIG_DFL) - print("\r\nConnection closed.") + signal.signal(signal.SIGWINCH, old_sigwinch) + _status("Console connection closed.") diff --git a/src/colab_cli/contents.py b/src/colab_cli/contents.py index 72d6d06..c8d0ba9 100644 --- a/src/colab_cli/contents.py +++ b/src/colab_cli/contents.py @@ -18,7 +18,7 @@ import requests from colab_cli.state import SessionState -from colab_cli.utils import get_status_code +from colab_cli.utils import RuntimeProxyError, get_status_code class ContentsClient: @@ -39,7 +39,14 @@ def _request( response = requests.request(method, url, params=req_params, json=json_data) - if get_status_code(response) == 404: + status_code = get_status_code(response) + if status_code in (401, 404) and not response.content: + # The Tunnel Frontend returns an empty 404 for an expired runtime + # proxy token. A real Contents API missing-path response has a JSON + # body, so preserve FileNotFoundError for that case. + raise RuntimeProxyError(status_code) + + if status_code == 404: raise FileNotFoundError(f"File or directory not found: {path}") response.raise_for_status() diff --git a/src/colab_cli/repl.py b/src/colab_cli/repl.py index 3d36ca3..b5efabe 100644 --- a/src/colab_cli/repl.py +++ b/src/colab_cli/repl.py @@ -28,7 +28,6 @@ from colab_cli.utils import handle_image, render_display_data - class ColabREPL: def __init__( self, @@ -120,7 +119,11 @@ def execute(self, code: str): None, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), ) - state.store.add(s) + state.store.update_fields( + s.name, + s.endpoint, + last_execution=s.last_execution, + ) try: outputs = self.runtime.execute_code( diff --git a/src/colab_cli/state.py b/src/colab_cli/state.py index e404984..aa91649 100644 --- a/src/colab_cli/state.py +++ b/src/colab_cli/state.py @@ -16,7 +16,7 @@ import json import os from datetime import datetime -from typing import Dict, Optional, Tuple, Iterator, IO +from typing import Any, Dict, Optional, Tuple, Iterator, IO import filelock from pydantic import BaseModel @@ -136,6 +136,28 @@ def add(self, state: SessionState): sessions[state.name] = state self._save_raw(f, sessions) + def update_fields( + self, name: str, endpoint: str, **changes: Any + ) -> Optional[SessionState]: + """Updates selected fields without overwriting concurrent changes. + + Runtime-proxy credentials are refreshed by separate short-lived CLI + processes. Long-running commands must therefore merge their metadata + (``running``, kernel IDs, execution timestamps) into the latest stored + object instead of writing an old in-memory ``SessionState`` back over + a fresh token. The endpoint guard also prevents a late writer from + mutating or reviving a same-name replacement session. + """ + with self._lock_exclusive() as f: + sessions = self._load_raw(f) + current = sessions.get(name) + if current is None or current.endpoint != endpoint: + return None + updated = current.model_copy(update=changes) + sessions[name] = updated + self._save_raw(f, sessions) + return updated + def get(self, name: str) -> Optional[SessionState]: with self._lock_shared() as f: if f is None: @@ -150,6 +172,17 @@ def remove(self, name: str): del sessions[name] self._save_raw(f, sessions) + def remove_if_endpoint(self, name: str, endpoint: str) -> Optional[SessionState]: + """Removes and returns a session only when its endpoint still matches.""" + with self._lock_exclusive() as f: + sessions = self._load_raw(f) + current = sessions.get(name) + if current is None or current.endpoint != endpoint: + return None + del sessions[name] + self._save_raw(f, sessions) + return current + def list(self) -> Dict[str, SessionState]: with self._lock_shared() as f: if f is None: diff --git a/src/colab_cli/utils.py b/src/colab_cli/utils.py index 8f9791b..4073ad7 100644 --- a/src/colab_cli/utils.py +++ b/src/colab_cli/utils.py @@ -24,6 +24,15 @@ from rich.text import Text +class RuntimeProxyError(Exception): + """The runtime tunnel rejected an expired or invalid proxy token.""" + + def __init__(self, status_code: int, response_body: bytes = b""): + super().__init__(f"Runtime proxy rejected credentials (HTTP {status_code})") + self.status_code = status_code + self.response_body = response_body + + def get_status_code(e: Exception) -> Optional[int]: """Safely extracts status code from various exception types.""" if hasattr(e, "response") and e.response is not None: @@ -34,16 +43,32 @@ def get_status_code(e: Exception) -> Optional[int]: return None -def is_terminal_error(e: Exception) -> bool: - """Checks if an exception indicates a lost session (404/401).""" - code = get_status_code(e) - if code in (404, 401): +def is_runtime_proxy_error(e: Exception) -> bool: + """Returns whether a connection failed due to proxy authentication. + + Prefer structured status codes. A narrow text fallback is retained for + jupyter-kernel-client and websocket-client versions that only expose the + HTTP handshake status in their exception message. + """ + if isinstance(e, RuntimeProxyError): return True - # Some exceptions from jupyter-kernel-client might wrap the real one or be different - err_msg = str(e) - if "404" in err_msg or "401" in err_msg: + if isinstance(e, FileNotFoundError): + return False + code = get_status_code(e) + if code in (401, 404): return True - return False + message = str(e).lower() + return any( + marker in message + for marker in ( + "handshake status 401", + "handshake status 404", + "http 401", + "http 404", + "401 unauthorized", + "404 not found", + ) + ) def print_kitty(image_bytes: bytes): diff --git a/tests/conftest.py b/tests/conftest.py index 9481ef5..4223aaa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,6 +26,14 @@ def mock_common_state(mocker): mock_state.client = MagicMock() mock_state.history = MagicMock() + # By default execute retry-wrapped operations immediately with the session + # returned by the mocked store. Individual tests can override this seam. + mock_state.run_with_runtime_proxy_retry.side_effect = ( + lambda name, operation, initial_session=None: operation( + initial_session or mock_state.store.get(name) + ) + ) + # Default behavior for sync_sessions mock_state.sync_sessions.return_value = ({}, []) diff --git a/tests/test_automation.py b/tests/test_automation.py index 009fb59..23de2f4 100644 --- a/tests/test_automation.py +++ b/tests/test_automation.py @@ -36,6 +36,9 @@ def mock_session(): def test_cli_auth(mock_state, mock_runtime_class, mock_session): mock_state.store.get.return_value = mock_session mock_state.resolve_session.return_value = "test-session" + mock_state.run_with_runtime_proxy_retry.side_effect = lambda name, operation: ( + operation(mock_session) + ) mock_runtime = mock_runtime_class.return_value mock_runtime.execute_code.return_value = [{"text": "Success"}] @@ -46,7 +49,12 @@ def test_cli_auth(mock_state, mock_runtime_class, mock_session): assert mock_session.last_execution[0] == "automation:auth" assert mock_session.last_execution[1] is None assert mock_session.last_execution[2] is not None - mock_state.store.add.assert_called_with(mock_session) + mock_state.store.update_fields.assert_any_call( + "test-session", + "e1", + running="automation(auth)", + last_execution=mock_session.last_execution, + ) # Verify ColabRuntime was invoked with the correct code mock_runtime.execute_code.assert_called_once() @@ -61,6 +69,9 @@ def test_cli_auth(mock_state, mock_runtime_class, mock_session): def test_cli_install(mock_state, mock_runtime_class, mock_session): mock_state.store.get.return_value = mock_session mock_state.resolve_session.return_value = "test-session" + mock_state.run_with_runtime_proxy_retry.side_effect = lambda name, operation: ( + operation(mock_session) + ) mock_runtime = mock_runtime_class.return_value mock_runtime.execute_code.return_value = [{"text": "Installed"}] @@ -69,7 +80,12 @@ def test_cli_install(mock_state, mock_runtime_class, mock_session): assert result.exit_code == 0 assert mock_session.last_execution[0] == "automation:install" assert mock_session.last_execution[2] is not None - mock_state.store.add.assert_called_with(mock_session) + mock_state.store.update_fields.assert_any_call( + "test-session", + "e1", + running="automation(install)", + last_execution=mock_session.last_execution, + ) mock_runtime.execute_code.assert_called_once() called_code = mock_runtime.execute_code.call_args[0][0] @@ -85,6 +101,9 @@ def test_cli_install(mock_state, mock_runtime_class, mock_session): def test_cli_drivemount(mock_state, mock_runtime_class, mock_session): mock_state.store.get.return_value = mock_session mock_state.resolve_session.return_value = "test-session" + mock_state.run_with_runtime_proxy_retry.side_effect = lambda name, operation: ( + operation(mock_session) + ) mock_runtime = mock_runtime_class.return_value mock_runtime.execute_code.return_value = [{"text": "Mounted"}] @@ -114,6 +133,9 @@ def test_cli_auth_uses_long_timeout(mock_state, mock_runtime_class, mock_session runtime.execute_code or the call will TimeoutError mid-flow.""" mock_state.store.get.return_value = mock_session mock_state.resolve_session.return_value = "test-session" + mock_state.run_with_runtime_proxy_retry.side_effect = lambda name, operation: ( + operation(mock_session) + ) mock_runtime = mock_runtime_class.return_value mock_runtime.execute_code.return_value = [{"text": "Authenticated"}] @@ -123,3 +145,38 @@ def test_cli_auth_uses_long_timeout(mock_state, mock_runtime_class, mock_session _, kwargs = mock_runtime.execute_code.call_args assert kwargs.get("timeout") is not None and kwargs["timeout"] >= 300 + + +@patch("colab_cli.commands.automation.ContentsClient") +@patch("colab_cli.commands.automation.ColabRuntime") +@patch("colab_cli.common.state") +def test_cli_install_requirement_upload_uses_refreshed_session( + mock_state, mock_runtime_class, mock_contents_class, tmp_path +): + stale = SessionState( + name="test-session", + token="expired-token", + url="https://old.url", + endpoint="e1", + ) + fresh = stale.model_copy( + update={"token": "fresh-token", "url": "https://fresh.url"} + ) + requirement = tmp_path / "requirements.txt" + requirement.write_text("pandas\n") + mock_state.store.get.return_value = stale + mock_state.resolve_session.return_value = "test-session" + mock_state.run_with_runtime_proxy_retry.side_effect = lambda name, operation: ( + operation(fresh) + ) + mock_runtime_class.return_value.execute_code.return_value = [] + + result = runner.invoke( + app, ["install", "-s", "test-session", "-r", str(requirement)] + ) + + assert result.exit_code == 0, result.output + mock_contents_class.assert_called_once_with(fresh) + mock_contents_class.return_value.upload.assert_called_once_with( + str(requirement), "content/requirements.txt" + ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 3741d7a..a38ec8d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -13,7 +13,7 @@ # limitations under the License. import time -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from typer.testing import CliRunner @@ -261,7 +261,7 @@ def test_cli_session_resolution(mock_store, mock_common_state): result = runner.invoke(app, ["stop"]) assert result.exit_code == 0 - mock_store.remove.assert_called_with("unique-session") + mock_store.remove_if_endpoint.assert_called_with("unique-session", "e1") def test_cli_stop(mock_client, mock_store, mock_common_state): @@ -278,7 +278,7 @@ def test_cli_stop(mock_client, mock_store, mock_common_state): assert result.exit_code == 0 mock_client.unassign.assert_called_with("e1") - mock_store.remove.assert_called_with("s1") + mock_store.remove_if_endpoint.assert_called_with("s1", "e1") def test_cli_sessions_prune(mock_common_state): @@ -412,7 +412,72 @@ def test_cli_console(mock_store, mock_common_state): with patch("colab_cli.commands.execution.connect_console") as mock_connect: result = runner.invoke(app, ["console", "-s", "s1"]) assert result.exit_code == 0 - mock_connect.assert_called_once_with(mock_session_state) + mock_connect.assert_called_once() + assert mock_connect.call_args.args == (mock_session_state,) + refresh_session = mock_connect.call_args.kwargs["refresh_session"] + refresh_session(mock_session_state) + mock_common_state.refresh_session.assert_called_once_with( + "s1", expected_session=mock_session_state, timeout=10 + ) + + +def test_cli_console_reconnect_stops_without_http_after_local_stop( + mock_store, mock_common_state +): + """A concurrent `colab stop` is already conclusive local evidence.""" + mock_session_state = MagicMock() + mock_session_state.name = "s1" + mock_session_state.endpoint = "endpoint-1" + mock_store.get.return_value = mock_session_state + mock_common_state.resolve_session.return_value = "s1" + + with patch("colab_cli.commands.execution.connect_console") as mock_connect: + result = runner.invoke(app, ["console", "-s", "s1"]) + + assert result.exit_code == 0 + refresh_session = mock_connect.call_args.kwargs["refresh_session"] + mock_store.get.return_value = None + + assert refresh_session(mock_session_state) is None + mock_common_state.refresh_session.assert_not_called() + + +def test_cli_console_auth_failure_does_not_prune_or_revive( + mock_store, mock_common_state +): + from colab_cli.utils import RuntimeProxyError + + mock_session_state = MagicMock() + mock_session_state.name = "s1" + mock_session_state.endpoint = "endpoint-1" + mock_store.get.return_value = mock_session_state + mock_common_state.resolve_session.return_value = "s1" + mock_common_state.run_with_runtime_proxy_retry.side_effect = RuntimeProxyError(401) + + result = runner.invoke(app, ["console", "-s", "s1"]) + + assert result.exit_code == 1 + mock_common_state.prune_session.assert_not_called() + mock_store.add.assert_not_called() + mock_store.update_fields.assert_any_call("s1", "endpoint-1", running=None) + + +def test_cli_console_piped_disconnect_is_reported(mock_store, mock_common_state): + from colab_cli.console import ConsoleConnectionError + + mock_session_state = MagicMock() + mock_session_state.name = "s1" + mock_session_state.endpoint = "endpoint-1" + mock_store.get.return_value = mock_session_state + mock_common_state.resolve_session.return_value = "s1" + mock_common_state.run_with_runtime_proxy_retry.side_effect = ConsoleConnectionError( + "proxy link lost" + ) + + result = runner.invoke(app, ["console", "-s", "s1"]) + + assert result.exit_code == 1 + assert "Console disconnected: proxy link lost" in result.output @patch("colab_cli.commands.files.ContentsClient") @@ -437,6 +502,28 @@ def test_cli_ls(mock_contents_class, mock_store, mock_common_state): assert "b_file" in result.output +@patch("colab_cli.commands.files.ContentsClient") +def test_cli_ls_uses_refreshed_session( + mock_contents_class, mock_store, mock_common_state +): + stale = MagicMock() + fresh = MagicMock() + mock_store.get.return_value = stale + mock_common_state.resolve_session.return_value = "s1" + mock_common_state.run_with_runtime_proxy_retry.side_effect = ( + lambda name, operation: operation(fresh) + ) + mock_contents_class.return_value.list_dir.return_value = { + "type": "directory", + "content": [], + } + + result = runner.invoke(app, ["ls", "-s", "s1", "content"]) + + assert result.exit_code == 0, result.output + mock_contents_class.assert_called_once_with(fresh) + + @patch("colab_cli.commands.files.ContentsClient") def test_cli_rm(mock_contents_class, mock_store, mock_common_state): mock_session_state = MagicMock() @@ -450,6 +537,36 @@ def test_cli_rm(mock_contents_class, mock_store, mock_common_state): assert "Deleted content/file.txt" in result.output +def test_restart_kernel_uses_runtime_proxy_retry(mock_store, mock_common_state): + stale = MagicMock() + stale.endpoint = "endpoint-1" + stale.url = "https://old" + stale.token = "expired" + fresh = MagicMock() + fresh.endpoint = "endpoint-1" + fresh.url = "https://fresh" + fresh.token = "fresh" + mock_store.get.return_value = stale + mock_common_state.resolve_session.return_value = "s1" + mock_common_state.run_with_runtime_proxy_retry.side_effect = ( + lambda name, operation: operation(fresh) + ) + + with patch("colab_cli.commands.session.ColabRuntime") as runtime_class: + result = runner.invoke(app, ["restart-kernel", "-s", "s1"]) + + assert result.exit_code == 0, result.output + runtime_class.assert_called_once_with( + "https://fresh", + "fresh", + kernel_id=fresh.kernel_id, + session_id=fresh.session_id, + on_kernel_started=ANY, + on_session_started=ANY, + ) + runtime_class.return_value.restart.assert_called_once_with() + + @patch("colab_cli.commands.files.os.path.isfile") @patch("colab_cli.commands.files.ContentsClient") def test_cli_upload(mock_contents_class, mock_isfile, mock_store, mock_common_state): diff --git a/tests/test_client.py b/tests/test_client.py index d222675..161d65b 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -155,6 +155,18 @@ def test_client_list_assignments(client, mock_session): assert "tun/m/assignments" in mock_session.request.call_args.args[1] +def test_client_list_assignments_forwards_request_timeout(client, mock_session): + """Interactive reconnects can bound a stalled control-plane lookup.""" + resp = MagicMock() + resp.ok = True + resp.text = ")]}'\n" + json.dumps({"assignments": []}) + mock_session.request.return_value = resp + + assert client.list_assignments(timeout=10) == [] + + assert mock_session.request.call_args.kwargs["timeout"] == 10 + + def test_client_keep_alive_assignment_handles_empty_response(client, mock_session): """The tunnel keep-alive ping returns an empty body. With no `schema=`, _issue_request must short-circuit and not attempt to parse it.""" diff --git a/tests/test_console.py b/tests/test_console.py index 310e98e..444d509 100644 --- a/tests/test_console.py +++ b/tests/test_console.py @@ -12,13 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. +import base64 +import gc +import hashlib import json import os +import socket import sys import termios -from unittest.mock import MagicMock, patch - -from colab_cli.console import connect_console, on_message, on_open +import threading +import time +import weakref +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock, call, patch + +import websocket + +from colab_cli.console import ( + ConsoleConnectionError, + _ConsoleInputForwarder, + connect_console, + on_message, + on_open, +) from colab_cli.state import SessionState import pytest @@ -69,6 +85,11 @@ def test_console_initialization( mock_ws_app.assert_called_once() assert mock_ws_app.call_args[1]["url"] == expected_url + # Long-lived interactive consoles must actively detect dead proxy links. + mock_ws_instance.run_forever.assert_called_once_with( + ping_interval=20, ping_timeout=10 + ) + # 2. Verify raw mode setup and teardown mock_tcgetattr.assert_called_once_with(sys.stdin.fileno()) mock_setraw.assert_called_once_with(sys.stdin.fileno(), termios.TCSANOW) @@ -106,6 +127,536 @@ def test_console_piped_input( mock_tcsetattr.assert_not_called() +def _websocket_attempt(opened=True, error=None, close_code=1000, close_reason=""): + """Returns a WebSocketApp mock that emits one deterministic lifecycle.""" + ws = MagicMock() + + def run_forever(**kwargs): + if opened: + ws._on_open(ws) + if error is not None: + ws._on_error(ws, error) + ws._on_close(ws, close_code, close_reason) + + ws.run_forever.side_effect = run_forever + return ws + + +def _connect_as_tty(*args, **kwargs): + """Runs ``connect_console`` with terminal syscalls isolated from pytest.""" + with ( + patch("colab_cli.console.sys.stdin.fileno", return_value=0), + patch("colab_cli.console.termios.tcgetattr", return_value=["attrs"]), + patch("colab_cli.console.termios.tcsetattr"), + patch("colab_cli.console.tty.setraw"), + patch("colab_cli.console.signal.getsignal", return_value=MagicMock()), + patch("colab_cli.console.signal.signal"), + ): + connect_console(*args, **kwargs) + + +@patch("colab_cli.console.threading.Thread") +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_reconnects_with_fresh_credentials_and_visible_status( + _mock_isatty, mock_ws_app, mock_thread, mock_session, capsys +): + """An abnormal post-connect close refreshes credentials and reconnects.""" + first = _websocket_attempt( + error=websocket.WebSocketConnectionClosedException("proxy link lost"), + close_code=1006, + close_reason="abnormal closure", + ) + second = _websocket_attempt(close_code=1000, close_reason="shell exited") + attempts = iter([first, second]) + + def make_ws(**kwargs): + ws = next(attempts) + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + return ws + + mock_ws_app.side_effect = make_ws + refreshed = SessionState( + name=mock_session.name, + token="fresh-token", + url="https://fresh-runtime.example.test", + endpoint=mock_session.endpoint, + ) + refresh = MagicMock(return_value=refreshed) + + _connect_as_tty( + mock_session, + refresh_session=refresh, + retry_delays=(0,), + _max_reconnect_attempts=1, + ) + + refresh.assert_called_once_with(mock_session) + assert mock_ws_app.call_count == 2 + assert "fresh-token" in mock_ws_app.call_args.kwargs["url"] + # Reconnects reuse the one stdin reader instead of racing for terminal input. + assert mock_thread.call_count == 1 + stderr = capsys.readouterr().err + assert "Console connection lost" in stderr + assert "Reconnecting in 0s (attempt 1" in stderr + assert "Console reconnected" in stderr + + +@patch("colab_cli.console.threading.Thread") +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_stops_reconnecting_when_endpoint_is_confirmed_gone( + _mock_isatty, mock_ws_app, _mock_thread, mock_session, capsys +): + ws = _websocket_attempt( + error=websocket.WebSocketConnectionClosedException("proxy link lost"), + close_code=1006, + ) + + def make_ws(**kwargs): + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + return ws + + mock_ws_app.side_effect = make_ws + refresh = MagicMock(return_value=None) + + _connect_as_tty( + mock_session, + refresh_session=refresh, + retry_delays=(0,), + _max_reconnect_attempts=1, + ) + + assert mock_ws_app.call_count == 1 + assert "no longer active" in capsys.readouterr().err + + +@patch("colab_cli.console.threading.Thread") +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_transient_refresh_failure_preserves_binding_and_retries( + _mock_isatty, mock_ws_app, _mock_thread, mock_session, capsys +): + first = _websocket_attempt( + error=websocket.WebSocketConnectionClosedException("proxy link lost"), + close_code=1006, + ) + second = _websocket_attempt(close_code=1000) + attempts = iter([first, second]) + + def make_ws(**kwargs): + ws = next(attempts) + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + return ws + + mock_ws_app.side_effect = make_ws + refresh = MagicMock(side_effect=OSError("control plane unavailable")) + + _connect_as_tty( + mock_session, + refresh_session=refresh, + retry_delays=(0,), + _max_reconnect_attempts=1, + ) + + assert mock_ws_app.call_count == 2 + assert "could not refresh credentials" in capsys.readouterr().err + + +@patch("colab_cli.console.threading.Thread") +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_refuses_same_name_replacement_endpoint( + _mock_isatty, mock_ws_app, _mock_thread, mock_session, capsys +): + ws = _websocket_attempt( + error=websocket.WebSocketConnectionClosedException("proxy link lost"), + close_code=1006, + ) + + def make_ws(**kwargs): + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + return ws + + mock_ws_app.side_effect = make_ws + replacement = SessionState( + name=mock_session.name, + token="replacement-token", + url="https://replacement.example.test", + endpoint="replacement-endpoint", + ) + + _connect_as_tty( + mock_session, + refresh_session=MagicMock(return_value=replacement), + retry_delays=(0,), + _max_reconnect_attempts=1, + ) + + assert mock_ws_app.call_count == 1 + assert "different runtime" in capsys.readouterr().err + + +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_user_exit_does_not_reconnect(_mock_isatty, mock_ws_app, mock_session): + ws = _websocket_attempt( + error=websocket.WebSocketProtocolException("empty close frame"), + close_code=None, + ) + + forwarder = MagicMock() + forwarder.stop_event.is_set.return_value = False + forwarder.user_requested_close = False + + def make_ws(**kwargs): + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + original_open = ws._on_open + + def opened(current_ws): + original_open(current_ws) + # Models the stdin worker recognizing `exit\n` before peer close. + forwarder.user_requested_close = True + + ws._on_open = opened + return ws + + mock_ws_app.side_effect = make_ws + with patch("colab_cli.console._ConsoleInputForwarder", return_value=forwarder): + _connect_as_tty(mock_session, retry_delays=(0,), _max_reconnect_attempts=1) + + assert mock_ws_app.call_count == 1 + + +@pytest.mark.parametrize("command", ["exit\n", "logout\n", "\x04"]) +def test_console_shell_exit_intent_expires(command): + """Nested-shell exit input must not disable reconnect for the process lifetime.""" + forwarder = _ConsoleInputForwarder(is_tty=True) + + with patch("colab_cli.console.time.monotonic", return_value=10.0): + for char in command: + forwarder._track_exit_request(char) + + with patch("colab_cli.console.time.monotonic", return_value=11.9): + assert forwarder.user_requested_close is True + with patch("colab_cli.console.time.monotonic", return_value=12.1): + assert forwarder.user_requested_close is False + + +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_stable_connection_resets_retry_backoff( + _mock_isatty, mock_ws_app, mock_session, capsys +): + """A later independent outage starts a fresh visible retry sequence.""" + attempts = iter( + [ + _websocket_attempt(close_code=1006, close_reason="first loss"), + _websocket_attempt(close_code=1006, close_reason="short recovery"), + _websocket_attempt(close_code=1006, close_reason="later loss"), + _websocket_attempt(close_code=1000, close_reason="shell exited"), + ] + ) + + def make_ws(**kwargs): + ws = next(attempts) + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + return ws + + mock_ws_app.side_effect = make_ws + forwarder = MagicMock() + forwarder.stop_event.is_set.return_value = False + forwarder.stop_event.wait.return_value = False + forwarder.user_requested_close = False + + with ( + patch("colab_cli.console._ConsoleInputForwarder", return_value=forwarder), + # Each attempt records an open and close time. The third connection is + # healthy for 31 seconds, so its later loss starts a new retry series. + patch( + "colab_cli.console.time.monotonic", + side_effect=[0, 1, 2, 3, 4, 35, 36, 37], + ), + ): + _connect_as_tty( + mock_session, + refresh_session=MagicMock(return_value=mock_session), + retry_delays=(1, 2), + _max_reconnect_attempts=3, + ) + + assert forwarder.stop_event.wait.call_args_list == [call(1), call(2), call(1)] + stderr = capsys.readouterr().err + assert stderr.count("Console connection lost") == 3 + assert stderr.count("Reconnecting in 1s (attempt 1") == 2 + assert "Reconnecting in 2s (attempt 2" in stderr + + +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=False) +def test_console_piped_disconnect_is_not_retried_or_replayed( + _mock_isatty, mock_ws_app, mock_session +): + ws = _websocket_attempt( + error=websocket.WebSocketConnectionClosedException("proxy link lost"), + close_code=1006, + ) + + def make_ws(**kwargs): + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + return ws + + mock_ws_app.side_effect = make_ws + + with patch("colab_cli.console._ConsoleInputForwarder.start"): + with pytest.raises(ConsoleConnectionError, match="proxy link lost"): + connect_console(mock_session, retry_delays=(0,), _max_reconnect_attempts=1) + + assert mock_ws_app.call_count == 1 + + +@patch("colab_cli.console.websocket.WebSocketApp") +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_initial_proxy_auth_error_is_returned_to_outer_refresh( + _mock_isatty, mock_ws_app, mock_session +): + """A failed initial handshake uses State's bounded credential retry.""" + error = RuntimeError("Handshake status 401 Unauthorized") + ws = _websocket_attempt(opened=False, error=error, close_code=None) + + def make_ws(**kwargs): + ws._on_open = kwargs["on_open"] + ws._on_error = kwargs["on_error"] + ws._on_close = kwargs["on_close"] + return ws + + mock_ws_app.side_effect = make_ws + refresh = MagicMock() + + with pytest.raises(RuntimeError, match="Handshake status 401"): + _connect_as_tty( + mock_session, + refresh_session=refresh, + retry_delays=(0,), + _max_reconnect_attempts=1, + ) + + assert mock_ws_app.call_count == 1 + refresh.assert_not_called() + + +@patch("colab_cli.console.sys.stdin.isatty", return_value=True) +def test_console_retry_limit_bounds_failures_and_releases_attempts( + _mock_isatty, mock_session +): + """A failed reconnect test cannot spin forever or retain every socket.""" + sockets = [] + + class DisconnectingWebSocket: + def __init__(self, **kwargs): + self.on_open = kwargs["on_open"] + self.on_error = kwargs["on_error"] + self.on_close = kwargs["on_close"] + sockets.append(weakref.ref(self)) + + def run_forever(self, **_kwargs): + self.on_open(self) + self.on_error( + self, + websocket.WebSocketConnectionClosedException("proxy link lost"), + ) + self.on_close(self, 1006, "abnormal closure") + + def send(self, _payload): + pass + + def close(self): + pass + + with ( + patch("colab_cli.console.websocket.WebSocketApp", DisconnectingWebSocket), + patch("colab_cli.console.threading.Thread"), + patch("colab_cli.console._status"), + patch("colab_cli.console.send_terminal_size"), + patch("colab_cli.console.logger.debug"), + patch("colab_cli.console.logger.info"), + pytest.raises(ConsoleConnectionError, match="reconnect limit"), + ): + _connect_as_tty( + mock_session, + refresh_session=lambda _session: mock_session, + retry_delays=(0,), + _max_reconnect_attempts=500, + ) + + assert len(sockets) == 501 + gc.collect() + assert not [ref for ref in sockets if ref() is not None] + + +def test_console_input_forwarder_thread_stops_while_stdin_is_idle(): + """Stopping Console wakes its one stdin thread without waiting for input.""" + read_fd, write_fd = os.pipe() + with ( + os.fdopen(read_fd) as reader, + os.fdopen(write_fd, "w") as _writer, + patch("colab_cli.console.sys.stdin", reader), + ): + forwarder = _ConsoleInputForwarder(is_tty=True) + forwarder.start() + deadline = time.monotonic() + 1 + while not forwarder._thread.is_alive() and time.monotonic() < deadline: + time.sleep(0.01) + + forwarder.stop() + + assert forwarder.join(timeout=1) + + +def test_console_input_forwarder_drains_buffered_input_without_waiting_for_eof(): + """Text already buffered in stdin is forwarded while the pipe stays open.""" + read_fd, write_fd = os.pipe() + with ( + os.fdopen(read_fd) as reader, + os.fdopen(write_fd, "w") as writer, + patch("colab_cli.console.sys.stdin", reader), + ): + writer.write("abc") + writer.flush() + ws = MagicMock() + forwarder = _ConsoleInputForwarder(is_tty=True) + forwarder.attach(ws) + forwarder.start() + + deadline = time.monotonic() + 1 + while ws.send.call_count < 3 and time.monotonic() < deadline: + time.sleep(0.01) + forwarder.stop() + assert forwarder.join(timeout=1) + + assert [json.loads(call.args[0]) for call in ws.send.call_args_list] == [ + {"data": "a"}, + {"data": "b"}, + {"data": "c"}, + ] + + +def test_console_ctrl_c_during_reconnect_is_visible(capsys): + """Raw-mode Ctrl-C cancellation must not look like another silent freeze.""" + forwarder = _ConsoleInputForwarder(is_tty=True) + + forwarder._forward_char("\x03") + + assert forwarder.user_requested_close is True + assert forwarder.stop_event.is_set() + assert "Console reconnect cancelled by user" in capsys.readouterr().err + + +def test_console_real_loopback_reconnect_uses_refreshed_token( + mock_session, monkeypatch, capsys +): + """Exercise WebSocketApp's real close/reconnect path over loopback.""" + requests = [] + ready = threading.Event() + port = [] + + def serve_two_connections(): + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + listener.listen(2) + listener.settimeout(5) + port.append(listener.getsockname()[1]) + ready.set() + for close_code in (1011, 1000): + conn, _ = listener.accept() + with conn: + conn.settimeout(5) + request = b"" + while b"\r\n\r\n" not in request: + request += conn.recv(4096) + requests.append(request.split(b"\r\n", 1)[0].decode()) + key = next( + line.split(b":", 1)[1].strip() + for line in request.split(b"\r\n") + if line.lower().startswith(b"sec-websocket-key:") + ) + accept = base64.b64encode( + hashlib.sha1( + key + b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11" + ).digest() + ) + conn.sendall( + b"HTTP/1.1 101 Switching Protocols\r\n" + b"Upgrade: websocket\r\n" + b"Connection: Upgrade\r\n" + b"Sec-WebSocket-Accept: " + accept + b"\r\n\r\n" + ) + conn.sendall(b"\x88\x02" + close_code.to_bytes(2, "big")) + + with ThreadPoolExecutor(max_workers=1) as executor: + server = executor.submit(serve_two_connections) + assert ready.wait(2) + + for key in ("NO_PROXY", "no_proxy"): + current = os.environ.get(key, "") + monkeypatch.setenv( + key, f"127.0.0.1,localhost{',' + current if current else ''}" + ) + + initial = SessionState( + name=mock_session.name, + token="old-token", + url=f"http://127.0.0.1:{port[0]}", + endpoint=mock_session.endpoint, + ) + refreshed = SessionState( + name=mock_session.name, + token="fresh-token", + url=f"http://127.0.0.1:{port[0]}", + endpoint=mock_session.endpoint, + ) + + with ( + patch("colab_cli.console.sys.stdin.isatty", return_value=True), + patch("colab_cli.console._ConsoleInputForwarder.start"), + patch("colab_cli.console.sys.stdin.fileno", return_value=0), + patch("colab_cli.console.termios.tcgetattr", return_value=["attrs"]), + patch("colab_cli.console.termios.tcsetattr"), + patch("colab_cli.console.tty.setraw"), + patch("colab_cli.console.signal.getsignal", return_value=MagicMock()), + patch("colab_cli.console.signal.signal"), + ): + connect_console( + initial, + refresh_session=MagicMock(return_value=refreshed), + retry_delays=(0,), + _max_reconnect_attempts=1, + ) + + # Propagate errors from the fake peer instead of leaving Console in a + # zero-delay retry loop while pytest captures output without bounds. + server.result(timeout=2) + + assert len(requests) == 2 + assert "colab-runtime-proxy-token=old-token" in requests[0] + assert "colab-runtime-proxy-token=fresh-token" in requests[1] + assert "Console reconnected" in capsys.readouterr().err + + @patch("colab_cli.console.os.get_terminal_size") def test_on_open_sends_terminal_size(mock_get_term_size): mock_ws = MagicMock() @@ -131,6 +682,26 @@ def test_on_message_writes_to_stdout(mock_flush, mock_write): # Verify that the data is written exactly as received mock_write.assert_called_once_with(test_data.encode("utf-8")) mock_flush.assert_called_once() + mock_ws.send.assert_not_called() + + +@patch("colab_cli.console.sys.stdout.buffer.write") +@patch("colab_cli.console.sys.stdout.buffer.flush") +def test_on_message_acknowledges_pty_flow_control(mock_flush, mock_write): + """The raw /colab/tty peer pauses after six unacknowledged output chunks.""" + mock_ws = MagicMock() + message_json = json.dumps({"data": "flow-controlled output", "ack": True}) + events = [] + mock_write.side_effect = lambda _data: events.append("write") + mock_flush.side_effect = lambda: events.append("flush") + mock_ws.send.side_effect = lambda _payload: events.append("ack") + + on_message(mock_ws, message_json) + + mock_write.assert_called_once_with(b"flow-controlled output") + mock_flush.assert_called_once() + mock_ws.send.assert_called_once_with(json.dumps({"ack": True})) + assert events == ["write", "flush", "ack"] @patch("colab_cli.console.os.get_terminal_size") diff --git a/tests/test_contents.py b/tests/test_contents.py index 4c910c8..b888679 100644 --- a/tests/test_contents.py +++ b/tests/test_contents.py @@ -16,7 +16,7 @@ from unittest.mock import MagicMock, patch import pytest -from colab_cli.contents import ContentsClient +from colab_cli.contents import ContentsClient, RuntimeProxyError from requests import Response from colab_cli.state import SessionState @@ -63,6 +63,34 @@ def test_list_dir(mock_request, client): assert len(res["content"]) == 2 +@pytest.mark.parametrize("status", [401, 404]) +@patch("colab_cli.contents.requests.request") +def test_empty_auth_failure_is_runtime_proxy_error_not_missing_file( + mock_request, client, status +): + """Expired proxy tokens return an empty 404 at the tunnel frontend.""" + mock_resp = MagicMock(spec=Response) + mock_resp.status_code = status + mock_resp.content = b"" + mock_request.return_value = mock_resp + + with pytest.raises(RuntimeProxyError) as exc_info: + client.list_dir("content") + + assert exc_info.value.status_code == status + + +@patch("colab_cli.contents.requests.request") +def test_nonempty_404_remains_file_not_found(mock_request, client): + mock_resp = MagicMock(spec=Response) + mock_resp.status_code = 404 + mock_resp.content = b'{"message":"not found"}' + mock_request.return_value = mock_resp + + with pytest.raises(FileNotFoundError): + client.list_dir("content/missing.txt") + + @patch("colab_cli.contents.requests.request") def test_rm_file(mock_request, client): mock_resp = MagicMock(spec=Response) diff --git a/tests/test_exec.py b/tests/test_exec.py index 203405c..d198068 100644 --- a/tests/test_exec.py +++ b/tests/test_exec.py @@ -54,7 +54,9 @@ def test_cli_exec_file(mock_store, mock_runtime_class, mock_common_state, tmp_pa assert mock_session.last_execution[0] == str(script) assert mock_session.last_execution[1] is None assert mock_session.last_execution[2] is not None - mock_store.add.assert_called_with(mock_session) + mock_store.update_fields.assert_any_call( + "s1", mock_session.endpoint, last_execution=mock_session.last_execution + ) mock_runtime.execute_code.assert_any_call( "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" ) @@ -81,7 +83,9 @@ def test_cli_exec_stdin(mock_store, mock_runtime_class, mock_common_state): assert mock_session.last_execution[0] == "stdin" assert mock_session.last_execution[1] is None assert mock_session.last_execution[2] is not None - mock_store.add.assert_called_with(mock_session) + mock_store.update_fields.assert_any_call( + "s1", mock_session.endpoint, last_execution=mock_session.last_execution + ) mock_runtime.execute_code.assert_any_call( "print(42)", output_hook=ANY, timeout=30.0 ) @@ -258,7 +262,7 @@ def test_cli_exec_empty_code(mock_runtime_class, mock_store, mock_common_state): assert result.exit_code == 0 -def test_cli_exec_lost_session_prunes( +def test_cli_exec_runtime_proxy_failure_does_not_prune( mock_runtime_class, mock_store, mock_common_state ): mock_session = MagicMock() @@ -272,8 +276,8 @@ def test_cli_exec_lost_session_prunes( result = runner.invoke(app, ["exec", "-s", "lost-sess"], input="print(1)") assert result.exit_code == 1 - assert "appears to be lost" in result.output - mock_common_state.prune_session.assert_called_once_with("lost-sess") + assert "rejected refreshed runtime credentials" in result.output + mock_common_state.prune_session.assert_not_called() def test_cli_exec_timeout(mock_store, mock_runtime_class, mock_common_state, tmp_path): diff --git a/tests/test_integration_scripts.py b/tests/test_integration_scripts.py new file mode 100644 index 0000000..fc26eda --- /dev/null +++ b/tests/test_integration_scripts.py @@ -0,0 +1,144 @@ +# 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 +# +# https://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. + +import os +from pathlib import Path +import subprocess + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +LIVE_SCRIPTS = ( + REPO_ROOT / "integration/repro_runtime_token_refresh/test.sh", + REPO_ROOT / "integration/repro_console_reconnect/test.sh", + REPO_ROOT / "integration/repro_console_flow_control/test.sh", +) + + +def _write_fake_uv(tmp_path: Path) -> tuple[Path, Path]: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + log_path = tmp_path / "uv.log" + fake_uv = bin_dir / "uv" + fake_uv.write_text( + """#!/bin/bash +set -eu +printf '%s\\n' "$*" >> "$FAKE_UV_LOG" + +if [ "$1" = run ] && [ "$2" = colab ]; then + shift 2 + config="" + previous="" + for arg in "$@"; do + if [ "$previous" = --config ]; then + config="$arg" + fi + previous="$arg" + done + + case " $* " in + *" sessions "*) + if [ "$FAKE_UV_MODE" = sessions-fail ]; then + exit 7 + fi + exit 0 + ;; + *" new "*) + printf 'new\\n' >> "$FAKE_UV_LOG" + mkdir -p "$(dirname "$config")" + printf '{"fake-session":{"endpoint":"test-endpoint","accelerator":"NONE"}}\\n' > "$config" + exit 0 + ;; + *" stop "*) + printf 'stop\\n' >> "$FAKE_UV_LOG" + exit 0 + ;; + esac +fi + +if [ "$1" = run ] && [ "$2" = python ]; then + count_file="$FAKE_UV_PYTHON_COUNT" + count=0 + if [ -f "$count_file" ]; then + count=$(cat "$count_file") + fi + count=$((count + 1)) + printf '%s' "$count" > "$count_file" + if [ "$FAKE_UV_MODE" = endpoint-fail ] && [ "$count" -eq 1 ]; then + exit 8 + fi + if [ "$#" -ge 5 ] && [ "$4" = oauth2 ]; then + printf 'unassign %s\\n' "$5" >> "$FAKE_UV_LOG" + exit 0 + fi + printf 'test-endpoint\\n' + exit 0 +fi + +exit 99 +""" + ) + fake_uv.chmod(0o755) + return bin_dir, log_path + + +def _run_script( + script: Path, tmp_path: Path, mode: str +) -> tuple[subprocess.CompletedProcess, str]: + bin_dir, log_path = _write_fake_uv(tmp_path) + fake_home = tmp_path / "home" + token = fake_home / ".config/colab-cli/token.json" + token.parent.mkdir(parents=True) + token.write_text("{}") + env = os.environ.copy() + env.update( + { + "PATH": f"{bin_dir}:{env['PATH']}", + "HOME": str(fake_home), + "FAKE_UV_LOG": str(log_path), + "FAKE_UV_MODE": mode, + "FAKE_UV_PYTHON_COUNT": str(tmp_path / "python-count"), + } + ) + result = subprocess.run( + ["bash", str(script)], + cwd=REPO_ROOT, + env=env, + text=True, + capture_output=True, + timeout=10, + check=False, + ) + return result, log_path.read_text() if log_path.exists() else "" + + +@pytest.mark.parametrize("script", LIVE_SCRIPTS, ids=lambda path: path.parent.name) +def test_live_script_cleans_assignment_when_endpoint_read_fails(script, tmp_path): + result, log = _run_script(script, tmp_path, "endpoint-fail") + + assert result.returncode != 0 + assert "new" in log + assert "stop" in log + assert "unassign test-endpoint" in log + + +@pytest.mark.parametrize("script", LIVE_SCRIPTS, ids=lambda path: path.parent.name) +def test_live_script_stops_before_allocating_when_assignment_snapshot_fails( + script, tmp_path +): + result, log = _run_script(script, tmp_path, "sessions-fail") + + assert result.returncode != 0 + assert "new" not in log diff --git a/tests/test_keep_alive.py b/tests/test_keep_alive.py index ed023d7..01db223 100644 --- a/tests/test_keep_alive.py +++ b/tests/test_keep_alive.py @@ -310,6 +310,7 @@ def test_sync_sessions_handles_lost_vm(mock_kill, mock_common_state): mock_common_state.store.list.return_value = {"lost-sess": lost_session} # Ensure store.get returns the session too mock_common_state.store.get.return_value = lost_session + mock_common_state.store.remove_if_endpoint.return_value = lost_session from colab_cli.common import State diff --git a/tests/test_repl.py b/tests/test_repl.py index 36dbcdb..2d8b35d 100644 --- a/tests/test_repl.py +++ b/tests/test_repl.py @@ -83,7 +83,11 @@ def mock_execute_code(code, output_hook=None, **kwargs): assert mock_session.last_execution[0] == "REPL" assert mock_session.last_execution[1] is None assert mock_session.last_execution[2] is not None - mock_store.add.assert_called_with(mock_session) + mock_store.update_fields.assert_called_with( + mock_session.name, + mock_session.endpoint, + last_execution=mock_session.last_execution, + ) def test_cli_repl_interactive( @@ -127,7 +131,12 @@ def test_cli_repl_piped(mock_runtime_class, mock_store, mock_common_state): assert result.exit_code == 0 assert mock_session.last_execution[0] == "stdin" assert mock_session.last_execution[2] is not None - mock_store.add.assert_called_with(mock_session) + mock_store.update_fields.assert_any_call( + "s1", + mock_session.endpoint, + last_execution=mock_session.last_execution, + running="repl(stdin)", + ) mock_runtime.execute_code.assert_any_call("print(1)", output_hook=ANY) diff --git a/tests/test_resolution_logic.py b/tests/test_resolution_logic.py index 4690140..6398fae 100644 --- a/tests/test_resolution_logic.py +++ b/tests/test_resolution_logic.py @@ -17,6 +17,8 @@ import typer from colab_cli.auth import AuthProvider from colab_cli.common import State +from colab_cli.state import SessionState, StateStore +from colab_cli.utils import RuntimeProxyError def test_resolve_session_no_local_sessions(): @@ -39,6 +41,7 @@ def test_resolve_session_with_local_but_none_on_server(): mock_session = MagicMock() mock_session.endpoint = "e1" state._store.list.return_value = {"s1": mock_session} + state._store.get.return_value = mock_session # But server says no assignments state._client = MagicMock() @@ -55,7 +58,7 @@ def test_resolve_session_with_local_but_none_on_server(): "[colab] Error: No active sessions found. Create one with 'colab new'." ) - state._store.remove.assert_called_with("s1") + state._store.remove_if_endpoint.assert_called_with("s1", "e1") def test_sync_sessions_avoids_client_if_no_local(): @@ -117,3 +120,251 @@ def test_state_client_auth_provider_adc(): _ = state.client args, kwargs = mock_get_creds.call_args assert kwargs["provider"] is AuthProvider.ADC + + +def _listed_assignment(endpoint="e1", token="fresh-token", url="https://fresh"): + assignment = MagicMock() + assignment.endpoint = endpoint + assignment.runtime_proxy_info.token = token + assignment.runtime_proxy_info.url = url + return assignment + + +def test_resolve_named_session_refreshes_runtime_proxy(tmp_path): + """Explicit ``-s NAME`` must not bypass runtime-proxy token refresh.""" + store = StateStore(str(tmp_path / "sessions.json")) + store.add( + SessionState( + name="s1", + token="expired-token", + url="https://old", + endpoint="e1", + keep_alive_pid=123, + kernel_id="kernel-1", + ) + ) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.return_value = [_listed_assignment()] + + assert state.resolve_session("s1") == "s1" + + refreshed = store.get("s1") + assert refreshed.token == "fresh-token" + assert refreshed.url == "https://fresh" + assert refreshed.keep_alive_pid == 123 + assert refreshed.kernel_id == "kernel-1" + + +def test_refresh_session_forwards_request_timeout(tmp_path): + store = StateStore(str(tmp_path / "sessions.json")) + original = SessionState(name="s1", token="old", url="old", endpoint="e1") + store.add(original) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.return_value = [_listed_assignment()] + + state.refresh_session("s1", expected_session=original, timeout=10) + + state._client.list_assignments.assert_called_once_with(timeout=10) + + +def test_resolve_named_session_removes_binding_only_when_server_confirms_missing( + tmp_path, +): + store = StateStore(str(tmp_path / "sessions.json")) + store.add(SessionState(name="s1", token="token", url="url", endpoint="e1")) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.return_value = [] + state._history = MagicMock() + + assert state.resolve_session("s1") == "s1" + assert store.get("s1") is None + + +def test_resolve_named_session_keeps_binding_when_refresh_fails(tmp_path): + store = StateStore(str(tmp_path / "sessions.json")) + original = SessionState(name="s1", token="token", url="url", endpoint="e1") + store.add(original) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.side_effect = RuntimeError("control plane down") + + assert state.resolve_session("s1") == "s1" + assert store.get("s1") == original + + +def test_refresh_does_not_remove_same_name_replacement_created_during_lookup( + tmp_path, +): + """A server snapshot for e1 must never delete a concurrent e2 binding.""" + store = StateStore(str(tmp_path / "sessions.json")) + store.add(SessionState(name="s1", token="old", url="old", endpoint="e1")) + replacement = SessionState( + name="s1", token="replacement", url="replacement", endpoint="e2" + ) + state = State() + state._store = store + state._client = MagicMock() + state._history = MagicMock() + + def list_then_replace(): + # Model a control-plane response captured before another process + # replaces the local binding, but delivered after that replacement. + store.add(replacement) + return [] + + state._client.list_assignments.side_effect = list_then_replace + + assert state.refresh_session("s1") == replacement + assert store.get("s1") == replacement + state._history.log_event.assert_not_called() + + +def test_sync_does_not_count_concurrent_same_name_replacement_as_pruned( + tmp_path, +): + store = StateStore(str(tmp_path / "sessions.json")) + store.add(SessionState(name="s1", token="old", url="old", endpoint="e1")) + replacement = SessionState( + name="s1", token="replacement", url="replacement", endpoint="e2" + ) + state = State() + state._store = store + state._client = MagicMock() + state._history = MagicMock() + + def list_then_replace(): + store.add(replacement) + return [] + + state._client.list_assignments.side_effect = list_then_replace + + with patch("typer.echo") as echo: + sessions, assignments = state.sync_sessions() + + assert assignments == [] + assert sessions == {"s1": replacement} + assert store.get("s1") == replacement + echo.assert_not_called() + + +def test_prune_session_keeps_binding_when_server_check_fails(tmp_path): + """An inconclusive control-plane check must never delete local state.""" + store = StateStore(str(tmp_path / "sessions.json")) + store.add(SessionState(name="s1", token="t", url="u", endpoint="e1")) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.side_effect = RuntimeError("network down") + + assert state.prune_session("s1") is False + assert store.get("s1") is not None + + +def test_runtime_proxy_error_refreshes_and_retries_once(tmp_path): + store = StateStore(str(tmp_path / "sessions.json")) + store.add(SessionState(name="s1", token="expired-token", url="old", endpoint="e1")) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.return_value = [_listed_assignment()] + seen_tokens = [] + + def operation(session): + seen_tokens.append(session.token) + if len(seen_tokens) == 1: + raise RuntimeError("Handshake status 404 Not Found") + return "connected" + + assert state.run_with_runtime_proxy_retry("s1", operation) == "connected" + assert seen_tokens == ["expired-token", "fresh-token"] + assert state._client.list_assignments.call_count == 1 + + +def test_runtime_proxy_retry_is_bounded_to_one_retry(tmp_path): + store = StateStore(str(tmp_path / "sessions.json")) + store.add(SessionState(name="s1", token="expired-token", url="old", endpoint="e1")) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.return_value = [_listed_assignment()] + operation = MagicMock(side_effect=RuntimeError("HTTP 404 Not Found")) + + with pytest.raises(RuntimeError, match="404"): + state.run_with_runtime_proxy_retry("s1", operation) + + assert operation.call_count == 2 + assert state._client.list_assignments.call_count == 1 + + +def test_runtime_proxy_retry_never_switches_to_same_name_replacement(tmp_path): + """A retry for endpoint e1 must not run against a concurrent endpoint e2.""" + store = StateStore(str(tmp_path / "sessions.json")) + original = SessionState(name="s1", token="expired-token", url="old", endpoint="e1") + replacement = SessionState( + name="s1", token="replacement", url="replacement", endpoint="e2" + ) + store.add(original) + state = State() + state._store = store + state._client = MagicMock() + + def list_then_replace(): + store.add(replacement) + return [] + + state._client.list_assignments.side_effect = list_then_replace + operation = MagicMock(side_effect=RuntimeProxyError(401)) + + with pytest.raises(RuntimeProxyError): + state.run_with_runtime_proxy_retry("s1", operation) + + operation.assert_called_once_with(original) + assert store.get("s1") == replacement + + +def test_runtime_proxy_retry_never_prunes_replacement_created_after_failure( + tmp_path, +): + store = StateStore(str(tmp_path / "sessions.json")) + original = SessionState(name="s1", token="expired-token", url="old", endpoint="e1") + replacement = SessionState( + name="s1", token="replacement", url="replacement", endpoint="e2" + ) + store.add(original) + state = State() + state._store = store + state._client = MagicMock() + state._client.list_assignments.return_value = [] + state._history = MagicMock() + + def fail_after_replacement(session): + store.add(replacement) + raise RuntimeProxyError(401) + + with pytest.raises(RuntimeProxyError): + state.run_with_runtime_proxy_retry("s1", fail_after_replacement) + + assert store.get("s1") == replacement + state._history.log_event.assert_not_called() + + +def test_runtime_proxy_retry_does_not_retry_unrelated_404_text(tmp_path): + store = StateStore(str(tmp_path / "sessions.json")) + store.add(SessionState(name="s1", token="token", url="url", endpoint="e1")) + state = State() + state._store = store + state._client = MagicMock() + operation = MagicMock(side_effect=FileNotFoundError("report-404.txt")) + + with pytest.raises(FileNotFoundError, match="report-404"): + state.run_with_runtime_proxy_retry("s1", operation) + + operation.assert_called_once() + state._client.list_assignments.assert_not_called() diff --git a/tests/test_run.py b/tests/test_run.py index a482180..a6ed255 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -27,6 +27,7 @@ PostAssignmentResponse, Variant, ) +from colab_cli.utils import RuntimeProxyError runner = CliRunner() @@ -170,7 +171,54 @@ def test_run_keep_skips_unassign( assert result.exit_code == 0, result.output mock_client.assign.assert_called_once() mock_client.unassign.assert_not_called() - mock_store.remove.assert_not_called() + mock_store.remove_if_endpoint.assert_not_called() + + +def test_run_keep_finally_preserves_refreshed_runtime_proxy_credentials( + mock_client, + mock_store, + mock_runtime_class, + mock_spawn_keep_alive, + assign_response, + script_path, + mock_common_state, +): + """Late metadata cleanup must merge instead of restoring the old token.""" + mock_client.assign.return_value = assign_response + mock_runtime_class.return_value.execute_code.return_value = [] + persisted = {} + + def add(session): + persisted[session.name] = session.model_copy(deep=True) + + def get(name): + return persisted.get(name) + + def update_fields(name, endpoint, **changes): + current = persisted.get(name) + if current is None or current.endpoint != endpoint: + return None + persisted[name] = current.model_copy(update=changes) + return persisted[name] + + def refresh_then_run(name, operation, initial_session=None): + fresh = persisted[name].model_copy( + update={"token": "fresh-token", "url": "https://fresh.url"} + ) + persisted[name] = fresh + return operation(fresh) + + mock_store.add.side_effect = add + mock_store.get.side_effect = get + mock_store.update_fields.side_effect = update_fields + mock_common_state.run_with_runtime_proxy_retry.side_effect = refresh_then_run + + result = runner.invoke(app, ["run", "--keep", str(script_path)]) + + assert result.exit_code == 0, result.output + assert persisted[next(iter(persisted))].token == "fresh-token" + assert persisted[next(iter(persisted))].url == "https://fresh.url" + assert persisted[next(iter(persisted))].running is None # --------------------------------------------------------------------------- @@ -403,6 +451,30 @@ def test_run_unassign_called_on_exception_during_execute( mock_client.unassign.assert_called_once_with("ep-123") +def test_run_startup_failure_preserves_error_and_releases_assignment( + mock_client, + mock_store, + mock_runtime_class, + mock_spawn_keep_alive, + assign_response, + script_path, + mock_common_state, +): + """A failure before ``runtime`` is assigned must not be masked by finally.""" + mock_client.assign.return_value = assign_response + persisted = {} + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s) + mock_store.get.side_effect = lambda name: persisted.get("s") + startup_error = RuntimeProxyError(401) + mock_common_state.run_with_runtime_proxy_retry.side_effect = startup_error + + result = runner.invoke(app, ["run", str(script_path)]) + + assert result.exit_code != 0 + assert result.exception is startup_error + mock_client.unassign.assert_called_once_with("ep-123") + + # --------------------------------------------------------------------------- # Accelerator passthrough # --------------------------------------------------------------------------- diff --git a/tests/test_ssh.py b/tests/test_ssh.py index 9d8ca80..daf99c2 100644 --- a/tests/test_ssh.py +++ b/tests/test_ssh.py @@ -34,6 +34,8 @@ from typer.testing import CliRunner import websocket +from colab_cli.utils import RuntimeProxyError + runner = CliRunner() @@ -143,9 +145,7 @@ def test_resolve_pubkey_default_scan_key_order( if present: content = f"ssh-key-content-for-{present}\n" (ssh_dir / present).write_text(content) - monkeypatch.setattr( - "os.path.expanduser", lambda p: p.replace("~", str(fake_home)) - ) + monkeypatch.setattr("os.path.expanduser", lambda p: p.replace("~", str(fake_home))) if expect_found: assert ssh_module._resolve_pubkey(None) == content.strip() else: @@ -206,6 +206,22 @@ def test_connect_websocket_network_failure_exits_1(mocker, capsys, exc): assert "WebSocket connection failed" in capsys.readouterr().err +def _bad_status(status, body=b""): + error = websocket.WebSocketBadStatusException(f"Handshake status {status}", status) + error.status_code = status + error.resp_body = body + return error + + +def test_connect_websocket_401_is_retryable_runtime_proxy_error(mocker): + mocker.patch.object(websocket.WebSocket, "connect", side_effect=_bad_status(401)) + + with pytest.raises(RuntimeProxyError) as exc_info: + ssh_module._connect_websocket("wss://host/colab/ssh?x=1", "pk") + + assert exc_info.value.status_code == 401 + + # --- proxy-mode byte bridge (ws <-> stdout) --------------------------------- @@ -343,18 +359,14 @@ def test_ssh_proxy_mode_calls_websocket(mock_common_state, mocker): connect = mocker.patch.object( ssh_module, "_connect_websocket", return_value=fake_ws ) - bridge = mocker.patch.object( - ssh_module, "_bridge_proxy_mode", return_value=0 - ) + bridge = mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) ssh_subprocess = mocker.patch.object(ssh_module, "_run_interactive_ssh") result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) assert result.exit_code == 0 connect.assert_called_once() args, _ = connect.call_args - assert args[0].startswith( - "wss://abc-foo.colab.googleusercontent.com/colab/ssh" - ) + assert args[0].startswith("wss://abc-foo.colab.googleusercontent.com/colab/ssh") assert args[1] == fake_pub bridge.assert_called_once_with(fake_ws) ssh_subprocess.assert_not_called() @@ -401,9 +413,7 @@ def fake_connect(url, pubkey): captured["url"] = url return MagicMock() - mocker.patch.object( - ssh_module, "_connect_websocket", side_effect=fake_connect - ) + mocker.patch.object(ssh_module, "_connect_websocket", side_effect=fake_connect) mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) @@ -430,9 +440,7 @@ def test_ssh_handshake_400_emits_actionable_message( ssh_module, "_resolve_pubkey", return_value="ssh-rsa AAAAfake u@h" ) - err = websocket.WebSocketBadStatusException( - "Handshake status 400 Bad Request", 400 - ) + err = websocket.WebSocketBadStatusException("Handshake status 400 Bad Request", 400) err.status_code = 400 err.resp_body = resp_body mocker.patch.object(websocket.WebSocket, "connect", side_effect=err) diff --git a/tests/test_ssh_lifecycle.py b/tests/test_ssh_lifecycle.py index 3143a3b..d47b1d4 100644 --- a/tests/test_ssh_lifecycle.py +++ b/tests/test_ssh_lifecycle.py @@ -39,6 +39,9 @@ import pytest import typer from typer.testing import CliRunner +import websocket + +from colab_cli.utils import RuntimeProxyError runner = CliRunner() @@ -60,9 +63,7 @@ def _make_session( def _patch_proxy(mocker): """Stub the proxy-mode I/O seams (pubkey, connect, bridge).""" mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") - mocker.patch.object( - ssh_module, "_connect_websocket", return_value=MagicMock() - ) + mocker.patch.object(ssh_module, "_connect_websocket", return_value=MagicMock()) mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) @@ -77,16 +78,12 @@ def test_proxy_mode_rm_stops_even_if_bridge_raises(mock_common_state, mocker): mocker.patch("signal.signal") stop = mocker.patch("colab_cli.commands.session.stop") mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") - mocker.patch.object( - ssh_module, "_connect_websocket", return_value=MagicMock() - ) + mocker.patch.object(ssh_module, "_connect_websocket", return_value=MagicMock()) mocker.patch.object( ssh_module, "_bridge_proxy_mode", side_effect=RuntimeError("ws died") ) - result = runner.invoke( - app, ["ssh", "--proxy-mode", "-s", "colab-ephem", "--rm"] - ) + result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "colab-ephem", "--rm"]) assert result.exit_code != 0 stop.assert_called_once_with(session="colab-ephem") @@ -138,6 +135,64 @@ def test_proxy_mode_exit_code_propagates(mock_common_state, mocker, code): assert result.exit_code == code +def test_proxy_mode_401_retries_once_with_refreshed_session(mock_common_state, mocker): + stale = _make_session(token="expired") + fresh = _make_session(token="fresh") + mock_common_state.store.get.return_value = stale + mock_common_state.resolve_session.return_value = "s1" + mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") + connect = mocker.patch.object( + ssh_module, + "_connect_websocket", + side_effect=[RuntimeProxyError(401), MagicMock()], + ) + mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) + + def retry(name, operation, initial_session=None): + with pytest.raises(RuntimeProxyError): + operation(initial_session) + return operation(fresh) + + mock_common_state.run_with_runtime_proxy_retry.side_effect = retry + + result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) + + assert result.exit_code == 0, result.output + assert connect.call_count == 2 + assert "expired" in connect.call_args_list[0].args[0] + assert "fresh" in connect.call_args_list[1].args[0] + + +def test_proxy_mode_repeated_401_has_actionable_error(mock_common_state, mocker): + sess = _make_session() + mock_common_state.store.get.return_value = sess + mock_common_state.resolve_session.return_value = "s1" + mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") + mock_common_state.run_with_runtime_proxy_retry.side_effect = RuntimeProxyError(401) + + result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) + + assert result.exit_code == 1 + assert "Authentication failed (HTTP 401)" in result.stderr + + +def test_proxy_mode_404_keeps_endpoint_unavailable_message(mock_common_state, mocker): + sess = _make_session() + mock_common_state.store.get.return_value = sess + mock_common_state.resolve_session.return_value = "s1" + mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") + error = websocket.WebSocketBadStatusException("Handshake status 404", 404) + error.status_code = 404 + error.resp_body = b"" + connect = mocker.patch.object(websocket.WebSocket, "connect", side_effect=error) + + result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) + + assert result.exit_code == 1 + assert "Endpoint not found (HTTP 404)" in result.stderr + connect.assert_called_once() + + # --- C. --proxy-mode keeps stdout clean (byte-stream integrity) -------------- @@ -162,9 +217,7 @@ def test_proxy_bridge_routes_rm_output_to_stderr(mocker, capsys): sess = _make_session("colab") mocker.patch("signal.signal") mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") - mocker.patch.object( - ssh_module, "_connect_websocket", return_value=MagicMock() - ) + mocker.patch.object(ssh_module, "_connect_websocket", return_value=MagicMock()) mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) def stop_echo(session=None): @@ -247,21 +300,15 @@ def test_proxy_mode_rm_teardown_idempotent(mock_common_state, mocker): mocker.patch("os._exit") # keep the handler from killing the test process stop = mocker.patch("colab_cli.commands.session.stop") mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") - mocker.patch.object( - ssh_module, "_connect_websocket", return_value=MagicMock() - ) + mocker.patch.object(ssh_module, "_connect_websocket", return_value=MagicMock()) def bridge_then_hup(ws): handlers[_signal.SIGHUP](_signal.SIGHUP, None) # OpenSSH HUPs us return 0 - mocker.patch.object( - ssh_module, "_bridge_proxy_mode", side_effect=bridge_then_hup - ) + mocker.patch.object(ssh_module, "_bridge_proxy_mode", side_effect=bridge_then_hup) - result = runner.invoke( - app, ["ssh", "--proxy-mode", "-s", "colab-ephem", "--rm"] - ) + result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "colab-ephem", "--rm"]) assert result.exit_code == 0 stop.assert_called_once_with(session="colab-ephem") diff --git a/tests/test_state.py b/tests/test_state.py index ccc1b5b..972d8ae 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -68,6 +68,44 @@ def test_state_store_remove(temp_config): assert new_store.get("to-be-removed") is None +def test_state_store_update_fields_preserves_runtime_proxy(temp_config): + """Metadata writers must not overwrite a concurrently refreshed token.""" + store = StateStore(temp_config) + store.add(SessionState(name="s", token="fresh", url="fresh-url", endpoint="e1")) + + updated = store.update_fields("s", "e1", running=None, kernel_id="kernel-1") + + assert updated.token == "fresh" + assert updated.url == "fresh-url" + assert updated.kernel_id == "kernel-1" + + +def test_state_store_update_fields_does_not_revive_removed_session(temp_config): + store = StateStore(temp_config) + store.add(SessionState(name="s", token="old", url="old-url", endpoint="e1")) + store.remove("s") + + assert store.update_fields("s", "e1", running=None) is None + assert store.get("s") is None + + +def test_state_store_update_fields_rejects_same_name_new_endpoint(temp_config): + """A stale command must not overwrite a same-name replacement session.""" + store = StateStore(temp_config) + store.add(SessionState(name="s", token="new", url="new-url", endpoint="e2")) + + assert store.update_fields("s", "e1", running=None) is None + assert store.get("s").endpoint == "e2" + + +def test_state_store_remove_if_endpoint_rejects_same_name_new_endpoint(temp_config): + store = StateStore(temp_config) + store.add(SessionState(name="s", token="new", url="new-url", endpoint="e2")) + + assert store.remove_if_endpoint("s", "e1") is None + assert store.get("s").endpoint == "e2" + + def test_state_store_list(temp_config): store = StateStore(temp_config) s1 = SessionState(name="s1", token="t1", url="u1", endpoint="e1")