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

[maxz411]

Symptom

About an hour after colab new, every command that reaches the VM through the tunnel (exec, ls, upload, download, repl, console, ...) fails with HTTP 404 while the runtime is still assigned and running (it keeps executing jobs and consuming compute units). Because 404/401 is treated as "session lost", the CLI then prunes the local session and kills its keep-alive daemon, leaving a VM that can no longer be addressed or stopped from the CLI, only from the web UI. Reported in #106; reproduced today with 0.6.0 and on main.

Cause

The runtime proxy token stored in sessions.json is issued with tokenExpiresInSeconds: 3600 (RuntimeProxyInfo already parses the field, but nothing reads it). Its lifetime is independent of the assignment's, and the CLI never renews it.

Client.list_assignments() returns a freshly issued runtimeProxyInfo (valid for another 3600 s) for every assignment. Writing that token/url into the session's entry restores access; repeating it periodically keeps a runtime reachable for as long as it is assigned.

Fix

The keep-alive daemon, which already runs every 60 s per session, now renews the token before it expires:

• SessionState.token_expires_at records when the token stops being accepted; colab new and colab run populate it from token_expires_in_seconds. A missing value (state written by an older CLI) means "unknown, renew now".
• Each daemon iteration calls _renew_token_if_expiring: with less than TOKEN_RENEWAL_MARGIN (10 min) left it calls list_assignments(), finds its own endpoint and stores the new token, url and expiry. Failures are logged as token_renewal_error and retried on the next iteration; they never stop the daemon or count as keep-alive ping errors. Successful renewals log token_renewed; both events are rendered by colab log.
• StateStore.update(name, **fields) is a locked field-level read-modify-write, so the daemon's write cannot clobber kernel_id/running/last_execution that a concurrently running command writes (a get followed by add from a second process could).
• docs/01_session_management.md and the README keep-alive bullet describe the behaviour.

No new dependencies, no other behaviour change. renew_runtime_token(session_name, endpoint) is a plain function so a reactive path can reuse it.

Known gap, documented in the design doc: a command that holds a SessionState across a renewal (e.g. a long exec) writes its stale copy back when it finishes; the daemon sees the stale expiry and renews again within one iteration (about 60 s). #123's recover-before-prune on 404/401 covers that window; the two changes are complementary (this one keeps the stored token valid for every command, including the file operations noted in the #123 validation, and avoids the failed call altogether, as suggested at the end of #106).

Verification

• Mechanism verified live: a runtime whose stored token had expired became reachable again after copying the fresh runtimeProxyInfo from the assignment listing into the state file, and repeating the renewal every ~45 min kept a runtime usable for hours (well past the one-hour mark at which stock 0.6.0 loses it).
• The renewal was also exercised against a real sessions.json with a fake client: a legacy entry without expiry is renewed at once, other fields survive the write, a fresh token triggers no listing call, and a token inside the margin is renewed.

Tests

uv run pytest tests/ (353 passed; 345 before) and uv run ruff check src tests pass; the touched files are ruff format clean.

New tests in tests/test_keep_alive.py: colab new records the expiry; the daemon renews an expiring token with the listing's values for its own endpoint; leaves a fresh token alone; renews a token of unknown expiry; survives a listing failure (loop continues, no keep_alive_error); does not write anything for an endpoint that is no longer listed. In tests/test_state.py: StateStore.update sets only the given fields and is a no-op for an unknown session.

Related: #106, #123 (reactive recovery), #109 (broader rework that also covers this).

I read CONTRIBUTING.md; I am opening this in case it is useful as a small, self-contained reference for the proactive half of #106. Feel free to close if it does not fit your plans.

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.