Mirrored from https://github.com/247TaxiDriver/QOSami-colab-cli.git
Synced Nov 6, 02:00 UTC
fix: stop spurious 'drivemount failed' on OAuth waits over 10s (#11)
The drivefs handshake blocks the kernel on input_request while the user OAuths in their browser. jupyter_kernel_client's default 10s wall-clock timeout fires during this silent stretch, raising TimeoutError mid-flow even though the mount has actually succeeded server-side. Add an optional timeout= parameter to ColabRuntime.execute_code that is forwarded to both execute() and execute_interactive(). Pass timeout=600 (10 min) from the human-in-the-loop subcommands -- colab drivemount and colab auth -- via a shared INTERACTIVE_AUTOMATION_TIMEOUT_SEC constant. Non-interactive paths keep the upstream default since their continuous iopub traffic never trips the practical inactivity ceiling.
Tyler committed
May 12, 2026 at 17:41 UTC
510115b0c3fff627e749144d84bad5fb8f748c57
5 files changed
+117
-4
docs/04_automation_and_utility.md
+6
@@ -2,6 +2,7 @@
2
log:
3
2026-05-07: Added a developer-only `colab whoami` subcommand (hidden from `colab --help`). Mints an access token via the same `auth.get_credentials(...)` path the rest of the CLI uses (honoring the global `--auth=...` flag), refreshes the credentials, then queries `https://oauth2.googleapis.com/tokeninfo` to print the email, scopes, audience, and expiry of whatever the CLI is about to send. Built specifically to short-circuit the "why is my call to colab.pa.googleapis.com 403-ing" debugging loop — the answer is almost always "missing scope" or "wrong identity", both of which `whoami` makes immediately visible. Hidden via `app.command(hidden=True)`; reachable via `colab whoami` or `colab whoami --help`. Suppressed from the daily-update banner check (added to `_AUTO_UPDATE_SUPPRESSED` in `cli.py`) so the banner doesn't obscure the auth output.
4
2026-05-11: Removed the local-file update source (`update_file_path` setting and `_fetch_local` helper); `colab update` now consults PyPI only. Switched the default `update_url` to the canonical PyPI JSON API (`https://pypi.org/pypi/google-colab-cli/json`), which already exposes the `info.version` schema the auto-update subsystem expects. Re-added `colab update --install` as a public self-install path that runs `pip install -U google-colab-cli` against the current `sys.executable`; Linux-only (other platforms exit non-zero with an explanatory message), and a silent no-op when the cached `latest_version` is already at or below the current install.
5
+2026-05-12: Added an optional `timeout=` parameter to `ColabRuntime.execute_code` that flows through to both the `execute()` and `execute_interactive()` branches. `colab auth` and `colab drivemount` now pass `timeout=600` (10 min) via a shared `INTERACTIVE_AUTOMATION_TIMEOUT_SEC` constant in `commands/automation.py`. Background: `jupyter_kernel_client` defaults to a 10s wall-clock timeout that is consumed even when the kernel is idle waiting on `input_request`. With the drivefs hook intercepting that request and prompting the user to OAuth in their browser, any user that takes >10s to click through (essentially everyone) hit `TimeoutError` and saw "drivemount failed" even though the mount had actually succeeded server-side. The fix is scoped narrowly to the two human-in-the-loop subcommands; non-interactive paths (`colab exec`, `colab run`, `colab install`, `colab repl --pipe`, `colab console --pipe`) keep the upstream default since they receive continuous iopub traffic that resets the practical inactivity ceiling.
6
---
7
8
# Design: Automation and Utility (`auth`, `install`, `log`, `pay`, `version`, `update`, `whoami`)
@@ -114,6 +115,11 @@ remediation guidance) rather than silently after ~1 minute via the daemon.
115
(`/tun/m/credentials-propagation/`), prompts the user with the Google OAuth
116
consent URL if needed, and dispatches the required `colab_reply` message to
117
the `stdin` channel to unlock the kernel thread.
118
+- **Timeout**: The kernel is silent (no iopub traffic) the entire time the
119
+ user is OAuthing in their browser. To avoid the upstream 10s
120
+ `jupyter_kernel_client` default raising `TimeoutError` mid-flow, this
121
+ subcommand passes `timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC` (600s) to
122
+ `ColabRuntime.execute_code`. Same applies to `colab auth`.
123
124
### 4. Logging and Notebook Capture (`colab log`)
125
src/colab_cli/commands/automation.py
+31
-4
@@ -26,8 +26,22 @@ from colab_cli.auth import get_credentials
26
from colab_cli.utils import get_status_code
27
28
29
+# Default execute() timeout for human-in-the-loop automations (auth /
30
+# drivemount). The kernel goes silent while the user completes a browser
31
+# OAuth flow, which can routinely take 30s+; the upstream 10s default
32
+# raises ``TimeoutError`` mid-flow even though the mount actually succeeds.
33
+# 10 minutes is long enough for any realistic interactive auth ceremony
34
+# without leaving CI hangs unbounded.
35
+INTERACTIVE_AUTOMATION_TIMEOUT_SEC = 600
36
+
37
+
38
def run_automation(
30
- name: str, op: str, code: str, allow_stdin: bool = False, path: str = None
39
+ name: str,
40
+ op: str,
41
+ code: str,
42
+ allow_stdin: bool = False,
43
+ path: str = None,
44
+ timeout: Optional[float] = None,
45
):
46
from colab_cli.common import state
47
@@ -130,7 +144,7 @@ def run_automation(
144
else:
145
state.history.log_event(name, "automation", {"op": op, "code": code})
146
133
- outputs = runtime.execute_code(code, allow_stdin=allow_stdin)
147
+ outputs = runtime.execute_code(code, allow_stdin=allow_stdin, timeout=timeout)
148
state.history.log_event(
149
name, "automation_result", {"op": op, "outputs": outputs}
150
)
@@ -166,7 +180,13 @@ def auth(
180
name = state.resolve_session(session)
181
code = "import os\nos.environ['USE_AUTH_EPHEM'] = '0'\nfrom google.colab import auth\nauth.authenticate_user()"
182
typer.echo(f"[colab] Starting Google Auth flow on {name}...")
169
- run_automation(name, "auth", code, allow_stdin=True)
183
+ run_automation(
184
+ name,
185
+ "auth",
186
+ code,
187
+ allow_stdin=True,
188
+ timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC,
189
+ )
190
191
192
def drivemount(
@@ -181,7 +201,14 @@ def drivemount(
201
name = state.resolve_session(session)
202
code = f"from google.colab import drive\ndrive.mount('{path}')"
203
typer.echo(f"[colab] Mounting Google Drive to '{path}' on {name}...")
184
- run_automation(name, "drivemount", code, allow_stdin=True, path=path)
204
+ run_automation(
205
+ name,
206
+ "drivemount",
207
+ code,
208
+ allow_stdin=True,
209
+ path=path,
210
+ timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC,
211
+ )
212
213
214
def install(
src/colab_cli/runtime.py
+13
@@ -161,8 +161,21 @@ class ColabRuntime:
161
allow_stdin: bool = False,
162
stdin_hook: Any = None,
163
output_hook: Optional[Callable[[Dict[str, Any]], None]] = None,
164
+ timeout: Optional[float] = None,
165
) -> List[Dict[str, Any]]:
166
+ # ``jupyter_kernel_client`` defaults ``timeout`` to ``REQUEST_TIMEOUT``
167
+ # (10 seconds) on both ``execute`` and ``execute_interactive``. That
168
+ # value is a wall-clock budget that shrinks every time the poll loop
169
+ # iterates -- as long as iopub/stdin events arrive back-to-back the
170
+ # call survives, but a single >10s quiet stretch (e.g. a kernel
171
+ # blocked on ``input_request`` while the user OAuths in the browser)
172
+ # will raise ``TimeoutError`` even though the underlying execution is
173
+ # still healthy. Callers that know they need a longer ceiling can
174
+ # pass ``timeout=`` here; otherwise we forward whatever the upstream
175
+ # default is (currently 10s).
176
kwargs = {"allow_stdin": allow_stdin}
177
+ if timeout is not None:
178
+ kwargs["timeout"] = timeout
179
180
# Wrap stdin_hook to log inputs
181
original_stdin_hook = stdin_hook
tests/test_automation.py
+25
@@ -98,3 +98,28 @@ def test_cli_drivemount(mock_state, mock_runtime_class, mock_session):
98
99
assert "drive.mount('/foo/bar')" in called_code
100
assert mock_runtime.colab_request_hook is not None
101
+ # Drivemount waits for the user to OAuth in their browser; the kernel
102
+ # goes silent during that wait and the default 10s execute() timeout
103
+ # would raise TimeoutError mid-flow. Insist on a generous timeout
104
+ # (>= 5 minutes) being forwarded to runtime.execute_code.
105
+ _, kwargs = mock_runtime.execute_code.call_args
106
+ assert kwargs.get("timeout") is not None and kwargs["timeout"] >= 300
107
+
108
+
109
+@patch("colab_cli.commands.automation.ColabRuntime")
110
+@patch("colab_cli.common.state")
111
+def test_cli_auth_uses_long_timeout(mock_state, mock_runtime_class, mock_session):
112
+ """`colab auth` walks the user through a paste-the-code flow that
113
+ routinely takes >10s, so it must pass a generous timeout to
114
+ runtime.execute_code or the call will TimeoutError mid-flow."""
115
+ mock_state.store.get.return_value = mock_session
116
+ mock_state.resolve_session.return_value = "test-session"
117
+
118
+ mock_runtime = mock_runtime_class.return_value
119
+ mock_runtime.execute_code.return_value = [{"text": "Authenticated"}]
120
+
121
+ result = runner.invoke(app, ["auth", "-s", "test-session"])
122
+ assert result.exit_code == 0
123
+
124
+ _, kwargs = mock_runtime.execute_code.call_args
125
+ assert kwargs.get("timeout") is not None and kwargs["timeout"] >= 300
tests/test_runtime.py
+42
@@ -77,6 +77,48 @@ def test_colab_runtime_execute_code():
77
}
78
79
80
+def test_colab_runtime_execute_code_default_no_timeout():
81
+ """By default, execute_code should NOT pass a timeout (relies on jupyter
82
+ kernel client default), preserving existing behavior for fast / streaming
83
+ workloads."""
84
+ runtime = ColabRuntime("http://url", "token123")
85
+ mock_kc = MagicMock()
86
+ runtime._kernel_client = mock_kc
87
+
88
+ mock_kc.execute.return_value = {"outputs": []}
89
+ runtime.execute_code("print(1)")
90
+
91
+ _, kwargs = mock_kc.execute.call_args
92
+ assert "timeout" not in kwargs
93
+
94
+
95
+def test_colab_runtime_execute_code_with_timeout():
96
+ """When a timeout is supplied, it must be forwarded to kernel_client.execute."""
97
+ runtime = ColabRuntime("http://url", "token123")
98
+ mock_kc = MagicMock()
99
+ runtime._kernel_client = mock_kc
100
+
101
+ mock_kc.execute.return_value = {"outputs": []}
102
+ runtime.execute_code("print(1)", timeout=600)
103
+
104
+ _, kwargs = mock_kc.execute.call_args
105
+ assert kwargs.get("timeout") == 600
106
+
107
+
108
+def test_colab_runtime_execute_interactive_with_timeout():
109
+ """timeout must also be plumbed through the execute_interactive branch
110
+ (used when an output_hook is supplied)."""
111
+ runtime = ColabRuntime("http://url", "token123")
112
+ mock_kc = MagicMock()
113
+ runtime._kernel_client = mock_kc
114
+
115
+ mock_kc.execute_interactive.return_value = {"content": {"status": "ok"}}
116
+ runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600)
117
+
118
+ _, kwargs = mock_kc.execute_interactive.call_args
119
+ assert kwargs.get("timeout") == 600
120
+
121
+
122
def test_colab_runtime_stop():
123
runtime = ColabRuntime("http://url", "token123")
124
mock_kc = MagicMock()