Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions docs/01_session_management.md
Original file line number Diff line number Diff line change
@@ -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/<endpoint>/keep-alive/` with `X-Colab-Tunnel: Google`) on `colab.research.google.com`. The RPC required `serviceusage` consumer access to Colab's internal project `1014160490159`, which ordinary user accounts lack, so every external user hit HTTP 403 `USER_PROJECT_DENIED` and their CLI sessions were idle-pruned within minutes (issue #14). Reproduced live with a third-party account; verified the tunnel ping is accepted by the same bearer-token credential that already works for `assign`. A `ReadTimeout` on the ping is treated as success (TFE records activity before forwarding to the often-non-responding VM). Generalized the pre-flight remediation messaging away from the now-irrelevant `colaboratory`/`pa.googleapis.com` framing, and removed the dead grpc-web client-registry/API-key code.
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.
Expand Down Expand Up @@ -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/<endpoint>/keep-alive/` with the header `X-Colab-Tunnel: Google`, authenticated with the user's own Gaia bearer token (the same credential and host used for `/tun/m/assign`). The Tunnel Frontend (TFE) records `LastActiveTime` before forwarding the request, which refreshes the idle timer. This matches the official `colab-vscode` extension's `sendKeepAlive`. TFE notes the activity on arrival and then forwards to the VM, which often does not answer on this path — so the request commonly read-times-out even though the keep-alive succeeded; a `ReadTimeout` is therefore treated as success, while genuine HTTP errors (e.g. 404 for a deleted assignment) propagate.
Expand All @@ -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
Expand Down Expand Up @@ -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.
26 changes: 26 additions & 0 deletions docs/02_execution_and_interactive.md
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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.
10 changes: 8 additions & 2 deletions docs/03_file_management.md
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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.
- **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.
Loading