Skip to content

fix: refresh runtime credentials and stabilize console sessions - #109

Open
Dts0 wants to merge 1 commit into
googlecolab:mainfrom
Dts0:fix-issue-106-runtime-token-refresh
Open

fix: refresh runtime credentials and stabilize console sessions#109
Dts0 wants to merge 1 commit into
googlecolab:mainfrom
Dts0:fix-issue-106-runtime-token-refresh

Conversation

@Dts0

@Dts0 Dts0 commented Aug 13, 2026 •
edited
Loading

Copy link
Copy Markdown

Summary

Fixes #89.
Fixes #106.

This PR improves the reliability of long-running Colab CLI sessions, particularly around expired runtime-
proxy credentials and interactive Console connections.

Runtime-proxy token refresh (#106)

  • Refreshes runtime-proxy tokens and URLs from /tun/m/assignments during session resolution.
  • Retries runtime-proxy authentication failures once with refreshed credentials.
  • Preserves local session bindings when credential refresh is inconclusive.
  • Removes a local session only after confirming that its exact server-side endpoint is gone.
  • Uses endpoint-guarded, field-level state updates so stale commands cannot overwrite refreshed credentials,
    revive deleted sessions, or modify same-name replacement sessions.
  • Applies credential recovery consistently across execution, file operations, automation, run, restart- kernel, Console, and SSH.
  • Distinguishes tunnel-level authentication failures from normal Contents API “file not found” responses.
  • Keeps SSH 404 behavior unchanged because it normally indicates that the runtime does not expose the SSH
    endpoint.

Console reliability (#89)

  • Adds WebSocket ping/pong liveness checks and visible reconnect status messages.
  • Automatically reconnects dropped interactive Console connections with bounded backoff and refreshed
    credentials.
  • Reuses a single stoppable stdin-forwarding thread across reconnects, preventing thread and object
    accumulation.
  • Preserves normal shell-exit, piped-input, local-stop, and same-name replacement semantics.
  • Bounds control-plane refresh calls so a failed proxy or stopped session cannot leave Console appearing
    frozen.
  • Implements the /colab/tty application-level flow-control protocol: Console now acknowledges marked
    output after flushing it locally, preventing the remote PTY from pausing permanently after approximately 600
    KB of output.

Testing

  • Full unit test suite: 396 passed.
  • Ruff passed for all changed Python files.
  • Added CPU-only end-to-end regressions that:
  • corrupt saved runtime-proxy credentials and verify automatic recovery for ls, exec, and piped
    console;
  • force Console transport failures and verify visible reconnection, credential-refresh fallback, and local
    session removal handling;
  • stream beyond the previous 600 KB PTY pause threshold and verify that output continues to a trailing
    sentinel.
  • All live tests use isolated session state without GPU or TPU flags, clean up their exact test endpoints,
    and verify that pre-existing assignments remain unchanged.

@Dts0
Dts0 marked this pull request as draft August 13, 2026 15:27
@Dts0
Dts0 marked this pull request as ready for review August 15, 2026 01:31
@Dts0 Dts0 changed the title fix: refresh expired runtime proxy tokens fix: refresh runtime credentials and stabilize console sessions Aug 15, 2026
@sinadavinc-hash

sinadavinc-hash commented Aug 26, 2026

Copy link
Copy Markdown

Summary

Fixes #89.
Fixes #106.

This PR improves the reliability of long-running Colab CLI sessions, particularly around expired runtime-
proxy credentials and interactive Console connections.

Runtime-proxy token refresh (#106)

  • Refreshes runtime-proxy tokens and URLs from /tun/m/assignments during session resolution.
  • Retries runtime-proxy authentication failures once with refreshed credentials.
  • Preserves local session bindings when credential refresh is inconclusive.
  • Removes a local session only after confirming that its exact server-side endpoint is gone.
  • Uses endpoint-guarded, field-level state updates so stale commands cannot overwrite refreshed credentials,
    revive deleted sessions, or modify same-name replacement sessions.
  • Applies credential recovery consistently across execution, file operations, automation, run, restart- kernel, Console, and SSH.
  • Distinguishes tunnel-level authentication failures from normal Contents API “file not found” responses.
  • Keeps SSH 404 behavior unchanged because it normally indicates that the runtime does not expose the SSH
    endpoint.

Console reliability (#89)

  • Adds WebSocket ping/pong liveness checks and visible reconnect status messages.
  • Automatically reconnects dropped interactive Console connections with bounded backoff and refreshed
    credentials.
  • Reuses a single stoppable stdin-forwarding thread across reconnects, preventing thread and object
    accumulation.
  • Preserves normal shell-exit, piped-input, local-stop, and same-name replacement semantics.
  • Bounds control-plane refresh calls so a failed proxy or stopped session cannot leave Console appearing
    frozen.
  • Implements the /colab/tty application-level flow-control protocol: Console now acknowledges marked
    output after flushing it locally, preventing the remote PTY from pausing permanently after approximately 600
    KB of output.

Testing

  • Full unit test suite: 396 passed.
  • Ruff passed for all changed Python files.
  • Added CPU-only end-to-end regressions that:
  • corrupt saved runtime-proxy credentials and verify automatic recovery for ls, exec, and piped
    console;
  • force Console transport failures and verify visible reconnection, credential-refresh fallback, and local
    session removal handling;
  • stream beyond the previous 600 KB PTY pause threshold and verify that output continues to a trailing
    sentinel.
  • All live tests use isolated session state without GPU or TPU flags, clean up their exact test endpoints,
    and verify that pre-existing assignments remain unchanged.

sinadavinc-hash
sinadavinc-hash reviewed Aug 26, 2026
View reviewed changes

@sinadavinc-hash sinadavinc-hash left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Proxy tkn

sinadavinc-hash
sinadavinc-hash reviewed Aug 26, 2026
View reviewed changes

@sinadavinc-hash sinadavinc-hash left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Next p

@seantangth seantangth mentioned this pull request Sep 18, 2026
Closed
@maxz411 maxz411 mentioned this pull request Sep 20, 2026
Closed
hcsolakoglu added a commit to hcsolakoglu/google-colab-cli that referenced this pull request Sep 24, 2026
@hcsolakoglu
fix: distinguish proxy 401/empty-404 from genuine contents 404
af6fa1c 
The tunnel frontend returns an empty 401/404 for an expired runtime
proxy token, while a genuine Contents API 'not found' carries a JSON
error body. Mapping every 404 to FileNotFoundError let `colab edit`
treat proxy failures as 'start empty' and overwrite the remote file on
save. Now empty-body 401/404 raises RuntimeProxyError (mirrors the
approach of upstream googlecolab#109), which propagates and aborts the edit before
the editor opens; only a JSON-body 404 stays FileNotFoundError.

Tests: empty-404/401 -> RuntimeProxyError, JSON-404 -> FileNotFoundError,
edit aborts on RuntimeProxyError without opening the editor or uploading.
@hcsolakoglu hcsolakoglu mentioned this pull request Sep 24, 2026
Closed
@EvanWiederspan EvanWiederspan mentioned this pull request Sep 25, 2026
Merged
@EvanWiederspan

EvanWiederspan commented Sep 25, 2026

Copy link
Copy Markdown
Contributor

Thanks for the contribution! We submitted the token-refresh change separately as #149. We can take a look at the console changes still if you remove any of the refresh-token based changes you added

@Dts0

Dts0 commented Sep 26, 2026

Copy link
Copy Markdown
Author

Thanks for the contribution! We submitted the token-refresh change separately as #149. We can take a look at the console changes still if you remove any of the refresh-token based changes you added

Ok, I will do it later.

@Dts0 HZHXDLKJGFYXGS_63699
fix: make colab console resilient to dropped connections
d7160eb 
`colab console` previously died on the first proxy hiccup or on a long
output burst. Keep the raw /colab/tty session usable:

- Reconnect an abnormally closed interactive socket with 1/2/5/10/30s
  backoff (then every 30s) until the user cancels. Connection loss, every
  retry, and successful reconnection are reported on stderr instead of
  leaving the terminal apparently frozen.
- Refresh runtime-proxy credentials before each reconnect, bounded to 10s,
  and never switch to a same-name replacement. A binding removed by a
  concurrent `colab stop` is treated as conclusive before any HTTP lookup.
- Honour /colab/tty application-level flow control: write and flush output
  marked with `"ack": true` before replying on the same socket, so the
  remote PTY stops pausing after six unacknowledged chunks (~100 KB each).
- Keep one stdin-forwarding thread across attempts and never replay piped
  input; normal close codes and a recent shell-exit request end the session.

Session-binding writes on the Console path go through endpoint-guarded
`StateStore.update_fields` / `remove_if_endpoint`, so a stale writer cannot
resurrect a removed binding or clobber a token refreshed by another
invocation. `State.resolve_session` stays offline: refreshing remains the
job of `State.get_session` (also bounded to 10s), which avoids pruning a
just-created binding when the assignments snapshot lags.

Adds CPU-only live regressions for reconnection, >600 KB output flow
control, and runtime-proxy token expiry, plus unit coverage for the
endpoint-guarded store primitives.

Verified: 384 unit tests, ruff on src/tests/integration, and the CPU-only
live integration scenarios (no assignments left behind).
@Dts0
Dts0 force-pushed the fix-issue-106-runtime-token-refresh branch from 5c3c786 to d7160eb Compare September 27, 2026 06:27
@Dts0

Dts0 commented Sep 27, 2026

Copy link
Copy Markdown
Author

Thanks for the contribution! We submitted the token-refresh change separately as #149. We can take a look at the console changes still if you remove any of the refresh-token based changes you added

Ok, I will do it later.

Done. I’ve removed the duplicate changes, resolved the conflicts, squashed the remaining commits into one, and retested the changes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer
@sinadavinc-hash sinadavinc-hash sinadavinc-hash left review comments
Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

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

3 participants

@Dts0 @sinadavinc-hash @EvanWiederspan