fix: renew the runtime proxy token before it expires (sessions lost after one hour)

README.md

@@ -19,7 +19,7 @@ Designed to support seamless developer productivity, headless automation, and AI
 * **Instant VM Provisioning:** Spin up CPU, GPU (T4, L4, G4, H100, A100), or TPU (v5e1, v6e1) runtimes in seconds.
 * **Robust Code Execution:** Run local Python scripts, Jupyter Notebooks (`.ipynb`), or piped `stdin` code; launch interactive REPLs or raw TTY console shells.
 * **Ephemeral Job Runner (`colab run`):** Provision a fresh VM, execute a local script with forwarded arguments, retrieve output files, and automatically tear down the runtime in a single command.
-* **Automatic Keep-Alive:** Built-in background daemon automatically prevents idle VM termination, keeping resource allocations active without requiring open browser tabs.
+* **Automatic Keep-Alive:** Built-in background daemon automatically prevents idle VM termination, keeping resource allocations active without requiring open browser tabs. It also renews the session's runtime token before it expires, so sessions stay reachable beyond the token's one-hour lifetime.
 * **Seamless Workspace Automation:** Mount Google Drive, authenticate Google Cloud Platform (GCP) credentials, and install dependencies with high-performance `uv` package management.
 * **State & Log Archival:** Inspect local session states or export interactive history logs to standard Jupyter Notebooks, Markdown, or structured JSONL.
 

docs/01_session_management.md

@@ -1,5 +1,6 @@
 ---
 log:
+2026-09-19: The keep-alive daemon now renews the runtime proxy token before it expires. The token stored at `colab new` is issued for `tokenExpiresInSeconds` (3600 s) and was never refreshed, so about an hour after creation every command through the tunnel failed with 404 and the CLI pruned the local session, although the VM stayed assigned and kept running (issue #106). `SessionState` gained `token_expires_at`; each daemon iteration checks it and, with less than 10 minutes left (or the expiry unknown), copies the fresh token from `/tun/m/assignments` into the store via the new locked field-level `StateStore.update`.
 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.
@@ -84,6 +85,7 @@ To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 mi
 - **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.
   - **Why not the RuntimeService RPC**: The previous implementation called `google.internal.colab.v1.RuntimeService/KeepAliveAssignment` at `colab.pa.googleapis.com` with `X-Goog-User-Project: 1014160490159`. That path requires the caller to be a `serviceusage` consumer of Colab's internal project `1014160490159`, which no ordinary user account is — so it returned HTTP 403 `USER_PROJECT_DENIED` for every external user, causing CLI sessions to be idle-pruned within minutes (issue #14). Dropping the header instead produced HTTP 400 `CONSUMER_INVALID` (public API-key project ≠ bearer-token quota project). The browser only succeeds because it rides the user's `google.com` cookie through an internal cookie-proxy (`colab.clients6.google.com`), which a headless bearer-token client cannot use. The TFE tunnel ping needs no project entitlement and works for any account that can assign a VM.
 - **Pre-flight (`colab new`, OAuth2/ADC only)**: Immediately after a successful `assign`, the CLI invokes `keep_alive_assignment` once synchronously. If the response is 403 with a `SCOPE_NOT_PERMITTED` body, it unassigns the new VM (to avoid leaking a billable assignment) and prints a per-provider remediation message before exiting non-zero. Other errors are tolerated — the daemon will retry and surface them via the structured event log. (Because keep-alive now uses the same backend/credential as `assign`, a scope failure at this stage is rare — assignment would normally have failed first.)
+- **Runtime token renewal**: The runtime proxy token stored at creation (`token`/`url` in `sessions.json`, used by every command that reaches the VM through the tunnel: `exec`, `repl`, `console`, `ls`/`upload`/`download`, `ssh`) is issued for `tokenExpiresInSeconds` (3600 s), independently of the assignment's lifetime. Once it expires, tunnel requests fail with 404 while the VM stays assigned and keeps running (issue #106). `colab new` and `colab run` record the expiry as `token_expires_at`; each daemon iteration checks it and, with less than `TOKEN_RENEWAL_MARGIN` (10 minutes) left or the expiry unknown (state written by an older CLI), fetches `/tun/m/assignments` (which returns a freshly issued token for every assignment) and writes `token`, `url` and `token_expires_at` through `StateStore.update`, a locked field-level write that cannot clobber `kernel_id`/`running` written concurrently by a command. The daemon logs `token_renewed` / `token_renewal_error` (rendered by `colab log`); a failed renewal is retried on the next iteration. Known gap: a command that holds a `SessionState` across a renewal (e.g. a long `exec`) writes its stale copy back when it finishes; the daemon then sees the stale expiry and renews again within one iteration (about 60 s).
 - **Structured logging**: The daemon emits `keep_alive_started` (with `pid`, `endpoint`), one `keep_alive_error` per failed iteration (with `status_code`, `error_type`, truncated `error`, `response_body`, `iteration`, `consecutive_4xx`), and `keep_alive_stopped` (with `reason`, `iterations`, `duration_seconds`, optional `last_error`, optional `expected_endpoint`/`actual_endpoint`). All three are rendered specially by `colab log` so users get diagnostic context without parsing JSONL by hand.
 - **Termination**:
     - **Explicit**: `colab stop` terminates the daemon using its stored PID.

src/colab_cli/commands/run.py

@@ -52,6 +52,7 @@
     _scope_remediation_message,
     resolve_runtime_options,
     spawn_keep_alive,
+    token_expiry,
 )
 from colab_cli.runtime import ColabRuntime
 from colab_cli.state import SessionState
@@ -349,12 +350,15 @@ def run_command(
         )
         url = res.runtime_proxy_info.url if hasattr(res, "runtime_proxy_info") else ""
         endpoint = res.endpoint
+    info = getattr(res, "runtime_proxy_info", None)
+    token_expires_at = token_expiry(info) if info else None
 
     s = SessionState(
         name=name,
         token=token,
         url=url,
         endpoint=endpoint,
+        token_expires_at=token_expires_at,
         variant=variant.value,
         accelerator=accelerator.value,
         machine_shape=(

src/colab_cli/commands/session.py

@@ -26,6 +26,7 @@
     ColabRequestError,
     HIGH_MEM_ONLY_ACCELERATORS,
     PostAssignmentResponse,
+    RuntimeProxyInfo,
     Shape,
     TooManyAssignmentsError,
     Variant,
@@ -147,6 +148,11 @@ def resolve_runtime_options(
     return variant, accelerator, shape
 
 
+def token_expiry(info: RuntimeProxyInfo) -> float:
+    """Unix time at which the runtime proxy token in `info` stops being accepted."""
+    return time.time() + info.token_expires_in_seconds
+
+
 def new(
     session: Annotated[
         Optional[str], typer.Option("-s", "--session", help="Session name")
@@ -239,6 +245,8 @@ def new(
         )
         url = res.runtime_proxy_info.url if hasattr(res, "runtime_proxy_info") else ""
         endpoint = res.endpoint
+    info = getattr(res, "runtime_proxy_info", None)
+    token_expires_at = token_expiry(info) if info else None
 
     # Importing locally to avoid a top-level circular import via auth.
 
@@ -247,6 +255,7 @@ def new(
         token=token,
         url=url,
         endpoint=endpoint,
+        token_expires_at=token_expires_at,
         variant=variant.value,
         accelerator=accelerator.value,
         machine_shape=(
@@ -481,6 +490,62 @@ def spawn_keep_alive(
     return p.pid
 
 
+# Renew the runtime proxy token once it has less than this many seconds left.
+TOKEN_RENEWAL_MARGIN = 10 * 60
+
+
+def renew_runtime_token(session_name: str, endpoint: str) -> bool:
+    """Stores a fresh runtime proxy token for the session.
+
+    The token issued when a runtime is assigned expires after about an hour
+    (`RuntimeProxyInfo.token_expires_in_seconds`), independently of the
+    assignment: afterwards every request through the tunnel fails with 404
+    while the VM stays assigned and keeps running. The assignment listing
+    returns a freshly issued token for every assignment, so copying it into
+    the store keeps the session reachable. Returns False when the endpoint
+    is no longer assigned.
+    """
+    from colab_cli.common import state
+
+    for a in state.client.list_assignments():
+        if a.endpoint == endpoint:
+            info = a.runtime_proxy_info
+            state.store.update(
+                session_name,
+                token=info.token,
+                url=info.url,
+                token_expires_at=token_expiry(info),
+            )
+            return True
+    return False
+
+
+def _renew_token_if_expiring(s: SessionState) -> None:
+    """Renews the token when it is about to expire or its expiry is unknown."""
+    from colab_cli.common import state
+
+    if (
+        s.token_expires_at is not None
+        and s.token_expires_at - time.time() >= TOKEN_RENEWAL_MARGIN
+    ):
+        return
+    try:
+        renewed = renew_runtime_token(s.name, s.endpoint)
+    except Exception as e:
+        state.history.log_event(
+            s.name,
+            "token_renewal_error",
+            {"error_type": type(e).__name__, "error": str(e)[:500]},
+        )
+        return
+    if renewed:
+        state.history.log_event(s.name, "token_renewed", {"endpoint": s.endpoint})
+    else:
+        state.history.log_event(
+            s.name, "token_renewal_error", {"error": "assignment no longer listed"}
+        )
+
+
 def keep_alive(
     endpoint: Annotated[str, typer.Argument(help="Endpoint ID")],
     session_name: Annotated[str, typer.Argument(help="Session name")],
@@ -549,6 +614,8 @@ def keep_alive(
                 # For other errors (network), we retry and don't count as 4xx
                 pass
 
+        _renew_token_if_expiring(s)
+
         time.sleep(60)
 
     payload: Dict[str, Any] = {

src/colab_cli/commands/utility.py

@@ -236,6 +236,13 @@ def log(
                         f"actual={event.get('actual_endpoint')}"
                     )
                 typer.echo(msg)
+            elif etype == "token_renewed":
+                typer.echo(f"[{ts}] KEEP: token renewed")
+            elif etype == "token_renewal_error":
+                typer.echo(
+                    f"[{ts}] KEEP: token renewal error "
+                    f"msg={event.get('error', '')[:120]}"
+                )
             else:
                 typer.echo(f"[{ts}] EVENT: {etype}")
 

src/colab_cli/state.py

@@ -35,6 +35,9 @@ class SessionState(BaseModel):
     last_execution: Optional[Tuple[str, Optional[str], str]] = None
     running: Optional[str] = None
     keep_alive_pid: Optional[int] = None
+    # Unix time at which `token` stops being accepted by the runtime proxy.
+    # None means unknown (state written by an older CLI): renew right away.
+    token_expires_at: Optional[float] = None
 
 
 class Settings(BaseModel):
@@ -136,6 +139,18 @@ def add(self, state: SessionState):
             sessions[state.name] = state
             self._save_raw(f, sessions)
 
+    def update(self, name: str, **fields):
+        """Sets fields on a stored session in one locked read-modify-write.
+
+        Unlike `get` followed by `add`, this cannot overwrite fields another
+        process wrote in between. A no-op if the session does not exist.
+        """
+        with self._lock_exclusive() as f:
+            sessions = self._load_raw(f)
+            if name in sessions:
+                sessions[name] = sessions[name].model_copy(update=fields)
+                self._save_raw(f, sessions)
+
     def get(self, name: str) -> Optional[SessionState]:
         with self._lock_shared() as f:
             if f is None:

tests/test_keep_alive.py

@@ -421,3 +421,121 @@ def test_keep_alive_logs_error_events_and_last_error(mock_common_state):
     assert payload["reason"] == "consecutive_4xx_errors"
     assert payload["last_error"]["status_code"] == 404
     assert payload["last_error"]["error_type"] == "ColabRequestError"
+
+
+# --- Runtime proxy token renewal (issue #106) -------------------------------
+#
+# The token stored at `colab new` expires after `token_expires_in_seconds`
+# (3600 s) although the VM stays assigned; the daemon renews it from the
+# assignment listing before that happens.
+
+
+def _assignment(endpoint, token="t2", url="u2", expires_in=3600):
+    return MagicMock(
+        endpoint=endpoint,
+        runtime_proxy_info=MagicMock(
+            token=token, url=url, token_expires_in_seconds=expires_in
+        ),
+    )
+
+
+def _one_keep_alive_iteration(mock_common_state, session, now=1000.0):
+    """Runs the daemon loop once for `session` with the clock frozen at `now`."""
+    mock_common_state.store.get.return_value = session
+    with (
+        patch("time.time", return_value=now),
+        patch("time.sleep", side_effect=InterruptedError),
+    ):
+        with pytest.raises(InterruptedError):
+            keep_alive(session.endpoint, session.name)
+
+
+@patch("colab_cli.commands.session.spawn_keep_alive")
+def test_new_records_token_expiry(mock_spawn, mock_common_state):
+    """`colab new` stores when the proxy token expires so the daemon can renew
+    it in time."""
+    mock_common_state.client.assign.return_value = _assignment(
+        "e1", token="t1", url="u1", expires_in=3600
+    )
+    mock_spawn.return_value = 9999
+
+    with patch("time.time", return_value=1000.0):
+        new(session="test-sess")
+
+    saved = mock_common_state.store.add.call_args.args[0]
+    assert saved.token_expires_at == 1000.0 + 3600
+
+
+def test_keep_alive_renews_expiring_token(mock_common_state):
+    """With less than TOKEN_RENEWAL_MARGIN left, the daemon copies the fresh
+    token, url and expiry of its own endpoint from the assignment listing."""
+    session = SessionState(
+        name="test", token="old", url="u1", endpoint="e1", token_expires_at=1000.0 + 300
+    )
+    mock_common_state.client.list_assignments.return_value = [
+        _assignment("other"),
+        _assignment("e1", token="new", url="u2"),
+    ]
+
+    _one_keep_alive_iteration(mock_common_state, session)
+
+    mock_common_state.store.update.assert_called_once_with(
+        "test", token="new", url="u2", token_expires_at=1000.0 + 3600
+    )
+    events = [c.args[1] for c in mock_common_state.history.log_event.call_args_list]
+    assert "token_renewed" in events
+
+
+def test_keep_alive_leaves_fresh_token_alone(mock_common_state):
+    session = SessionState(
+        name="test", token="t", url="u", endpoint="e1", token_expires_at=1000.0 + 3000
+    )
+
+    _one_keep_alive_iteration(mock_common_state, session)
+
+    mock_common_state.client.list_assignments.assert_not_called()
+    mock_common_state.store.update.assert_not_called()
+
+
+def test_keep_alive_renews_token_of_unknown_expiry(mock_common_state):
+    """State written by an older CLI has no expiry: renew right away."""
+    session = SessionState(name="test", token="t", url="u", endpoint="e1")
+    mock_common_state.client.list_assignments.return_value = [_assignment("e1")]
+
+    _one_keep_alive_iteration(mock_common_state, session)
+
+    mock_common_state.store.update.assert_called_once()
+
+
+def test_keep_alive_survives_renewal_failure(mock_common_state):
+    """A failed listing is logged and retried next iteration; it neither stops
+    the daemon nor counts as a keep-alive ping error."""
+    session = SessionState(
+        name="test", token="t", url="u", endpoint="e1", token_expires_at=1000.0 + 300
+    )
+    mock_common_state.client.list_assignments.side_effect = RuntimeError("boom")
+
+    # Reaching time.sleep proves the loop went on after the failure.
+    _one_keep_alive_iteration(mock_common_state, session)
+
+    mock_common_state.store.update.assert_not_called()
+    events = {
+        c.args[1]: c.args[2] for c in mock_common_state.history.log_event.call_args_list
+    }
+    assert events["token_renewal_error"]["error"] == "boom"
+    assert "keep_alive_error" not in events
+
+
+def test_keep_alive_does_not_renew_unlisted_assignment(mock_common_state):
+    """An endpoint missing from the listing is gone; nothing is written and the
+    existing 4xx handling of the ping decides when the daemon stops."""
+    session = SessionState(
+        name="test", token="t", url="u", endpoint="e1", token_expires_at=1000.0 + 300
+    )
+    mock_common_state.client.list_assignments.return_value = [_assignment("other")]
+
+    _one_keep_alive_iteration(mock_common_state, session)
+
+    mock_common_state.store.update.assert_not_called()
+    events = [c.args[1] for c in mock_common_state.history.log_event.call_args_list]
+    assert "token_renewal_error" in events

tests/test_state.py

@@ -325,3 +325,26 @@ def test_state_store_multiprocess_concurrency(temp_config):
     assert p1.exitcode == 0
     assert p2.exitcode == 0
     assert len(StateStore(temp_config).list()) == 80
+
+
+def test_state_store_update_sets_only_given_fields(temp_config):
+    """`update` is a locked read-modify-write of the given fields, so the
+    keep-alive daemon can renew the token without clobbering fields a command
+    wrote in between (unlike `get` followed by `add`)."""
+    store = StateStore(temp_config)
+    store.add(
+        SessionState(name="s", token="old", url="u1", endpoint="e", kernel_id="k1")
+    )
+
+    store.update("s", token="new", token_expires_at=123.5)
+
+    loaded = StateStore(temp_config).get("s")
+    assert loaded.token == "new"
+    assert loaded.token_expires_at == 123.5
+    assert loaded.kernel_id == "k1"
+
+
+def test_state_store_update_missing_session_is_noop(temp_config):
+    store = StateStore(temp_config)
+    store.update("nope", token="new")
+    assert store.list() == {}