colab drivemount: root cause of #113 ("Connection was lost" / mount failed after ~120 s) and a tested fix, plus native Windows console

[arlinamid]

CONTRIBUTING.md says external PRs aren't being accepted, so I'm sharing this here instead of opening one. The fix lives in a branch on my fork, in case it is useful to anyone hitting the same problems or to the maintainers:

• Branch: https://github.com/arlinamid/google-colab-cli/tree/cursor/windows-tty-drive-mount-8411 (commits c19ce50, ef19bd1, f3ef7c8, a0ea62e, 98ad068, 5beff34, 29f8bfb and 0f3ad51, on top of main incl. fix: Windows termios crash and portable Drive auth prompt #135)
• Write-up with all related issues: fix: Windows consoles (PowerShell and cmd) and Drive mount consent arlinamid/google-colab-cli#1

Updates (2026-10-07):

• ef19bd1 fixes colab edit on Windows.
• f3ef7c8 escapes % in the ssh ProxyCommand (a bug on all platforms) and makes the whole test suite pass on Windows.
• a0ea62e stops a piped colab console from leaking terminal replies into the next command (all platforms), and makes colab edit honor a Windows $EDITOR path.
• 98ad068 stops a piped colab console from cutting off its own output on a freshly started runtime (all platforms).
• 5beff34 explains a dropped kernel connection in one line instead of a traceback, and stops a finished command from writing a stopped session back into sessions.json (all platforms).
• 29f8bfb and 0f3ad51 add notebook control to colab exec: --cells, --stop-on-error, a local pre-flight --check/--check-only, and a non-zero exit code when a cell fails (all platforms; this is a new feature, not a fix).

Details are below.

Symptom

On 0.7.4 (WSL Ubuntu, T4 runtime), colab drivemount failed with:

jupyter_kernel_client/wsclient.py:1073 in execute_interactive
RuntimeError: Connection was lost.

This is the same failure as #113: the consent step completes in the browser, but the command never recovers.

Root cause (two separate problems)

1. Kernel-side timeout. drive.mount() defaults to timeout_ms=120000, and that budget includes the time the user spends on the consent screen. A normal click-through can exceed it, so the kernel gives up and the mount fails.
2. The handshake blocks the websocket thread. The dfs_ephemeral colab_request hook ran the whole credential-propagation handshake, including waiting for the user to press Enter, inside the websocket message callback. While it waits, the receive loop is blocked, so a long consent step ends in "Connection was lost" or a missed execute reply.

What the branch changes

• drivemount calls drive.mount(path, timeout_ms=600000), which matches the CLI's existing 10-minute interactive timeout.
• The handshake runs on a separate daemon thread, so the websocket keeps receiving.
• On any failure the CLI sends a colab_reply with an error to the kernel, so drive.mount fails at once instead of waiting out the timeout.
• The consent URL is opened in the system browser (Feature Request: Automatically open Drive authorization URL in system browser #108), and the multi-account HTTP 400 (drivemount fails with "Error propagating: 400" when the authenticated account is not authuser=0 #103) is explained, telling the user to pick the login_hint account, instead of printing Google's raw HTML error page. An optional --authuser flag is only appended to the consent URL when given explicitly.
• Windows: raw console input via msvcrt, virtual-terminal processing enabled for cmd.exe and PowerShell, and fixes for ssh --proxy-mode stdin and quoting.
• colab edit on Windows (ef19bd1): the local copy was kept in an open NamedTemporaryFile, which Windows does not allow to be opened a second time. The download, the hash and the editor all failed with PermissionError. The download error was also swallowed, so without the hash crash the editor would have received an empty file and saving would have overwritten the remote file. The file now lives in a TemporaryDirectory, named after the remote file.
• % in the ssh ProxyCommand (f3ef7c8, all platforms): OpenSSH expands %-tokens (%h, %p, ...) in ProxyCommand before running it. A session name or identity path containing % therefore either failed (unknown key %x) or was silently rewritten, for example a%h became a<host>. % is now written as %%. Checked against OpenSSH 10.5p1 and OpenSSH_for_Windows_9.5p2.
• Correction to the branch's own comments: OpenSSH for Windows (9.5p2) does not run ProxyCommand through cmd.exe. a&echo INJECTED arrived as one literal argument. The command is passed to CreateProcess and split with the MSVCRT rules, which is exactly what subprocess.list2cmdline produces, so the quoting code was right and only the comments were wrong.
• Windows test fixes (f3ef7c8), all test-side:
  • the ProxyCommand round-trip test splits with CommandLineToArgvW on Windows,
  • the REPL tests run inside a prompt_toolkit app session with DummyOutput,
  • the SIGHUP tests expect only the signals the platform has,
  • the --help ordering parser skips wrapped description lines,
  • the TTY-EOF test pins the POSIX branch, plus a new test for the Windows key path.
• Piped colab console (a0ea62e, all platforms):
  • What happened: when it attaches, tmux on the runtime sends terminal queries: DA1 ESC[c, DA2 ESC[>c, XTVERSION ESC[>q and the OSC 10/11 colour queries. The local terminal answers by typing the reply into its own input.
  • Why it leaked: with piped stdin the CLI is not reading that input, so the replies could never reach the remote side. They stayed behind for the next reader; on Windows they appeared in the next command's prompt as ESC[?61;...c ESC[>0;10;1c.
  • Fix: when stdin is not a TTY, the console drops terminal queries from the remote output. A sequence split across websocket messages is held back until it completes. Interactive sessions are unchanged, because there the replies do reach tmux.
• colab edit with a Windows $EDITOR (a0ea62e):
  • Cause: click.edit splits $VISUAL/$EDITOR with POSIX shlex on every platform (click.edit on Windows throws WinError 14001 or 87 pallets/click#3840), so python D:\x\editor.py became D:xeditor.py.
  • Fix: on Windows, when one of those variables is set, colab edit splits it with CommandLineToArgvW, resolves the program with shutil.which (so .cmd launchers such as VS Code's code work), and starts the editor itself. POSIX, and click's own default editor, still go through click.edit.
• Piped colab console on a fresh runtime (98ad068, all platforms):
  • Cause: after stdin EOF the console sent exit\n and closed the websocket itself a fixed 0.5s later (PIPED_EOF_GRACE_SECONDS). On a freshly started runtime the first attach often needs 0.5-1s before the shell answers, so the client closed first and nothing came back.
  • Evidence: I traced each run with timestamps, using a fresh CPU runtime and one attach per run. With the 0.5s grace, only 1 of 3 runs got its output. With a 5s wait, 3 of 3 did. The backend also closed the websocket itself every time, 0.1-1s after exit (tmux [exited], then on_close).
  • Fix: the handler now waits for that server-side close. It only closes from the client side after 30s, as a fallback for a backend that never closes; as with echo cmd | ssh host, a piped command running longer than that is cut off.
• Running command when its session is stopped (5beff34, all platforms):
  • Traceback: a colab stop from another terminal, while colab exec (or repl, run, drivemount) was running against the same session, ended that command with a jupyter_kernel_client traceback (RuntimeError: Connection was lost.). ColabRuntime now raises SessionConnectionLost instead, and the colab entry point prints one line with exit code 1: "Session 'X' was stopped, so the running command ended.", or, if the session is still known locally, "Lost the connection … check it with colab status -s X".
  • Stale session: those commands also saved the session back to sessions.json in their finally blocks, after colab stop had removed it, which left a stopped session behind locally. They now use StateStore.update_if_present, which only writes if the session is still recorded (under the same lock).
  • Missing name: six ColabRuntime constructions did not pass session_name; they now do, so the message names the session.
• colab exec notebook control (29f8bfb, 0f3ad51):
  • Before: colab exec -f nb.ipynb always ran every code cell top to bottom and kept going after a failing cell. It exited 0 either way, so a script could not tell that a notebook had failed. A missing file or an unreadable notebook ended in a traceback.
  • --cells SPEC: runs only the selected code cells, in the given order. SPEC accepts code-cell numbers (from 1, markdown cells not counted), ranges (2-5), #@title values and cell ids. An entry that looks like a number or a range is always read as one; nbformat's random ids are occasionally all digits, so such cells are selected by number. An unknown or ambiguous entry is an error (exit 2) before anything runs.
  • --stop-on-error: stops at the first cell with an error output; the partial outputs are still saved.
  • Exit code: exec exits 1 when the code, or any notebook cell, raised, and lists the failed cells. This is a behaviour change, and it also applies to .py files and stdin.
  • --check / --check-only: a local pre-flight that runs before anything is sent. It validates the notebook, resolves --cells, and syntax-checks each selected cell; IPython !cmd, %magic, x = !cmd and top-level await are accepted, %% cell magics are skipped, and line numbers are kept. If a check fails, nothing runs (exit 2). --check-only prints the plan without needing a session.
  • One-line errors: a missing file or an unreadable notebook prints one line instead of a traceback.

How it was tested

• Linux (WSL Ubuntu, Python 3.12): full test suite, 462 passed and 4 skipped (Windows-only tests). Two cross-process timing tests in test_state.py (test_write_lock_blocks_across_processes, test_readers_are_concurrent_across_processes) failed in one run on a busy machine and passed on rerun; they look timing-sensitive, and these commits do not touch the locking code.
• Windows 11 (Python 3.12): full test suite, 464 passed and 2 skipped (both need termios).
• colab exec notebook options, Linux (WSL), live, CPU runtime: I used a 3-cell notebook whose cell 2 raises.
  • --check on a notebook with a syntax error reported cell 2 (Broken), line 2: expected ':' and ran nothing (exit 2).
  • A plain run executed all three cells and exited 1, with 1 cell(s) raised an error: cell 2 (Two).
  • --stop-on-error stopped after cell 2 (1 cell(s) not run).
  • --cells "Three,1" ran cell 3, then cell 1.
  • A missing file printed File not found (exit 1).
• Session stopped mid-run, Linux (WSL), live, with 5beff34: I started a 2-minute colab exec and stopped its session from another process after 15s. It printed "[colab] Session 'losttest' was stopped, so the running command ended." with exit code 1, and sessions.json stayed empty.
• Piped console on fresh runtimes, Linux (WSL), live, with 98ad068: 3 of 3 runs got their output. In one run the output arrived 1.08s after exit, which the old 0.5s grace would have cut off. The backend closed the socket every time, so the 30s fallback was never used. Before f3ef7c8, 30 tests failed and 1 hung. The new % test cases fail without the fix.
• Piped console, Windows, live: after a piped colab console, I counted the events waiting in the console input buffer (CONIN$, via PeekConsoleInputW). Before a0ea62e there were 41, containing the DA1 and DA2 replies; after it, none.
• colab edit with EDITOR=python D:\tool\...\fake_editor.py (unquoted backslashes), Windows, live: an existing and a new remote file both round-trip. Before a0ea62e the same setting failed with "can't open file 'D:\toolqwen21...'".
• Linux (WSL Ubuntu), live on a CPU runtime: colab new → colab drivemount → Credentials propagated. Resuming mount... → Mounted at /content/drive, exit code 0, confirmed with os.path.ismount('/content/drive') == True → colab stop. This is the same machine and account where 0.7.4 failed with "Connection was lost". With no browser available in WSL, the CLI printed "Could not launch a browser here. Open the link above." and the copy-the-link flow worked.
• Windows 11, PowerShell 7.6.6, Python 3.12, live end-to-end on a CPU runtime: oauth2 login → colab new → colab exec (stdin) → colab console (piped) → colab drivemount → Credentials propagated. Resuming mount... → Mounted at /content/drive, confirmed with os.path.ismount('/content/drive') == True → colab stop. This was a single run with one Google account.
• colab edit, Windows: the two new regression tests (and the two existing edit tests) fail without ef19bd1 and pass with it. Live on a CPU runtime with a scripted editor: editing an existing remote file and creating a new one both round-trip correctly, and no temporary directories are left behind.
• Not tested: the multi-account (drivemount fails with "Error propagating: 400" when the authenticated account is not authuser=0 #103) path, and GPU runtimes with the branch.

Correction

An earlier version of this post listed an issue where the update-check banner ran into the next line ("…settings.json[colab] Mounting…"). That was wrong. The CLI's raw output contains the line breaks (settings.json\r\n\r\n[colab] … on Windows, \n\n on Linux); they were only lost in the terminal view I was reading the output from.

The changes were drafted with AI coding agents: Cursor wrote c19ce50, and Claude Code wrote ef19bd1, f3ef7c8, a0ea62e, 98ad068, 5beff34, 29f8bfb and 0f3ad51. All the tests above were run on my own machines.

[MattBetancourt]

Validating this on native Windows 11 — thank you.

@arlinamid, I independently validated your branch (0f3ad51) end-to-end and wanted to leave the receipts here:

• Suite: 464 passed / 2 skipped on native Windows 11 build 26200 / Python 3.12.10 — exact match to your 464/2 claim. Same commit on Ubuntu 6.17: 462 passed / 4 skipped (your Windows-only CommandLineToArgvW tests), zero failures both sides.
• Live smoke (--auth adc, CPU runtime): new → READY in ~7s, exec returned output from the VM, stop → terminated with zero sessions remaining.
• Merge check: applies onto current main (4025416) with zero conflicts.
• Parity: everything is 1:1 except update --install and shebang, both of which your write-up already scopes out. Confirmed your two known-gaps (resize polling, proxy---rm teardown) as documented, not discovered.

This is the real thing — a complete, tested port, not another partial patch. Thank you for doing the work properly — the write-up (traced timings, ConPTY keystroke tests, per-commit test evidence) made independent validation straightforward.