[Bug]: Live sessions are pruned locally after 60 minutes — proxy token expiry is misclassified as "session lost"

[willame]

[Bug]: Live sessions are pruned locally after 60 minutes — proxy token expiry is misclassified as "session lost"

Summary

RuntimeProxyInfo.token has a finite TTL (the server reports it in tokenExpiresInSeconds), but the CLI never reads that field and never refreshes the token. Once it expires, the next exec / run gets a 401/404 from the proxy, is_terminal_error() returns True, and state.prune_session() deletes the local binding for a runtime that is still alive on the server.

The user loses access to a running VM — including any work in progress — while the assignment is still listed by list_assignments(). The server actually returns a fresh token on every list_assignments() call; the CLI just discards it.

Observed reproducibly at ~60 minute intervals on long-running jobs.

Version

• google-colab-cli 0.6.0 (current PyPI release)
• Also present on main as of 2026-08-10 (checked below)

Root cause

1 — The token has a TTL that is parsed but never used.

src/colab_cli/client.py:

class RuntimeProxyInfo(BaseModel):
    token: str
    token_expires_in_seconds: int = Field(..., alias="tokenExpiresInSeconds")
    url: str

token_expires_in_seconds has no other reference anywhere in the codebase.

2 — Expiry is indistinguishable from session loss.

src/colab_cli/utils.py:37-45:

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):
        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:
        return True
    return False

An expired-but-valid session and a genuinely gone session produce the same 401/404. The substring branch widens this further — any exception whose text happens to contain 401 or 404 is classified as terminal.

3 — That classification is trusted unconditionally to delete state.

Four call sites act on it, and none re-confirm against the server:

• src/colab_cli/commands/execution.py — 3 sites (exec, and two others in the same file)
• src/colab_cli/commands/run.py — 1 site

each doing:

if is_terminal_error(e):
    ...
    state.prune_session(name)

prune_session() (src/colab_cli/common.py:71) kills the keep-alive process and removes the store entry with no server check.

Note the contrast with sync_sessions() in the same file, which does check active_endpoints from list_assignments() before pruning. The terminal-error path skips that check entirely.

Steps to reproduce

1. colab new (any accelerator)
2. Start a job that runs longer than the token TTL (~60 min), polling with colab exec periodically
3. At roughly the 60-minute mark, an exec fails with 401/404 and the CLI prints that the session was pruned
4. colab sessions still lists the assignment — the VM is alive and the job is still running, but the local binding is gone

Measured on two separate long runs; the failure recurred once per TTL period (~60 min) rather than at random.

Suggested fix

Confirm against the server before deleting, and adopt the fresh token the server already returns. Guarding prune_session() itself covers all four call sites from a single place, so the "should this binding be deleted" decision lives in exactly one location:

def _assignment_gone(self, name: str) -> bool:
    s = self.store.get(name)
    if not s:
        return True
    try:
        assignments = self.client.list_assignments()
    except Exception:
        # Can't confirm -> keep the binding. Deleting on an inconclusive
        # check is the strictly worse failure mode.
        return False
    for a in assignments:
        if a.endpoint == s.endpoint:
            rpi = a.runtime_proxy_info
            if rpi.token != s.token or rpi.url != s.url:
                s.token, s.url = rpi.token, rpi.url
                self.store.add(s)
            return False
    return True

def prune_session(self, name: str):
    if not self._assignment_gone(name):
        return
    ...  # existing body unchanged

With this in place the failing exec still errors, but the binding survives and the next call succeeds with the refreshed token — self-healing within one poll interval.

A proactive refresh keyed on token_expires_in_seconds (renew at, say, 80% of TTL) would avoid the failed call altogether, but the guard above is the minimal change that stops data loss.

Result

Running the patched version locally: a 141-minute session hit the condition twice (at ~49 min and ~102 min). Both times the binding was retained, the token was swapped, and the job continued to completion.

Happy to open a PR if this approach looks right.