google-colab-cli 0.7.4__tar.gz → 0.7.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (89) hide show
  1. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/PKG-INFO +9 -3
  2. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/README.md +7 -1
  3. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/01_session_management.md +26 -0
  4. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/pyproject.toml +1 -1
  5. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/client.py +45 -17
  6. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/automation.py +12 -2
  7. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/ssh.py +0 -4
  8. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/console.py +34 -9
  9. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_client.py +162 -5
  10. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_console.py +40 -1
  11. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_ssh.py +44 -0
  12. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/uv.lock +7 -7
  13. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/.githooks/pre-commit +0 -0
  14. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/.gitignore +0 -0
  15. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/.pre-commit-config.yaml +0 -0
  16. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/.python-version +0 -0
  17. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/AGENTS.md +0 -0
  18. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/CHANGELOG.md +0 -0
  19. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/CONTRIBUTING.md +0 -0
  20. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/LICENSE +0 -0
  21. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/cloudbuild.yaml +0 -0
  22. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/02_execution_and_interactive.md +0 -0
  23. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/03_file_management.md +0 -0
  24. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/04_automation_and_utility.md +0 -0
  25. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/05_run_command.md +0 -0
  26. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/06_ssh_access.md +0 -0
  27. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/3042ab12-2026-05-07.png +0 -0
  28. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/demo.webm +0 -0
  29. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/docs/demos.md +0 -0
  30. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/examples/finetune_run.py +0 -0
  31. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/integration/README.md +0 -0
  32. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/integration/repro_bundled_oauth/test.sh +0 -0
  33. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/integration/repro_piped_console/test.sh +0 -0
  34. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/integration/repro_plot_redirection/test.sh +0 -0
  35. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/integration/repro_run_command/test.sh +0 -0
  36. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/integration/repro_ssh/test.sh +0 -0
  37. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/integration/repro_variable_persistence/test.sh +0 -0
  38. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/skills/colab-operator/SKILL.md +0 -0
  39. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/auth.py +0 -0
  40. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/auto_update.py +0 -0
  41. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/cli.py +0 -0
  42. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/__init__.py +0 -0
  43. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/execution.py +0 -0
  44. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/files.py +0 -0
  45. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/run.py +0 -0
  46. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/session.py +0 -0
  47. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/usage.py +0 -0
  48. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/commands/utility.py +0 -0
  49. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/common.py +0 -0
  50. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/consumption.py +0 -0
  51. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/contents.py +0 -0
  52. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/converter.py +0 -0
  53. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/history.py +0 -0
  54. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/oauth_config.json +0 -0
  55. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/repl.py +0 -0
  56. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/runtime.py +0 -0
  57. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/state.py +0 -0
  58. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/src/colab_cli/utils.py +0 -0
  59. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/conftest.py +0 -0
  60. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_auth.py +0 -0
  61. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_auth_adc.py +0 -0
  62. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_automation.py +0 -0
  63. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_cli.py +0 -0
  64. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_cli_log.py +0 -0
  65. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_common.py +0 -0
  66. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_consumption.py +0 -0
  67. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_consumption_client.py +0 -0
  68. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_contents.py +0 -0
  69. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_exec.py +0 -0
  70. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_history.py +0 -0
  71. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_ipynb_exec.py +0 -0
  72. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_log_export.py +0 -0
  73. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_pay.py +0 -0
  74. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_readme.py +0 -0
  75. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_repl.py +0 -0
  76. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_resolution_logic.py +0 -0
  77. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_run.py +0 -0
  78. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_runtime.py +0 -0
  79. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_ssh_autocreate.py +0 -0
  80. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_ssh_lifecycle.py +0 -0
  81. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_ssh_wire_contract.py +0 -0
  82. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_ssh_workdir.py +0 -0
  83. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_state.py +0 -0
  84. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_streaming.py +0 -0
  85. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_update.py +0 -0
  86. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_url.py +0 -0
  87. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_utils.py +0 -0
  88. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_version.py +0 -0
  89. {google_colab_cli-0.7.4 → google_colab_cli-0.7.5}/tests/test_whoami.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: google-colab-cli
3
- Version: 0.7.4
3
+ Version: 0.7.5
4
4
  Summary: CLI for interacting with Colab.
5
5
  Project-URL: Homepage, https://github.com/googlecolab/google-colab-cli
6
6
  Project-URL: Repository, https://github.com/googlecolab/google-colab-cli
@@ -31,7 +31,7 @@ Requires-Dist: requests>=2.32.5
31
31
  Requires-Dist: rich>=14.3.3
32
32
  Requires-Dist: typer>=0.24.1
33
33
  Requires-Dist: typing-extensions>=4.0
34
- Requires-Dist: websocket-client>=1.0
34
+ Requires-Dist: websocket-client>=1.6
35
35
  Description-Content-Type: text/markdown
36
36
 
37
37
  # Colab CLI
@@ -84,7 +84,7 @@ Run a CPU-based VM runtime, execute some code, and clean up:
84
84
  colab new
85
85
 
86
86
  # 2. Execute code from stdin
87
- echo "print('Hello from Google Colab!')" | colab exec
87
+ echo 'print("Hello from Google Colab!")' | colab exec
88
88
 
89
89
  # 3. Stop and release the VM resource
90
90
  colab stop
@@ -226,3 +226,9 @@ To view interactive walkthroughs of eleven real-world automated scenarios, check
226
226
  ## Contributing
227
227
 
228
228
  Feedback and contributions are welcome! Please read [`CONTRIBUTING.md`](./CONTRIBUTING.md) for details.
229
+
230
+ ---
231
+
232
+ ## Legal & Privacy
233
+
234
+ See Colab's [Terms of Service](https://research.google.com/colaboratory/tos_v6.html) and the [Google Privacy Policy](https://policies.google.com/privacy), which apply to usage of this CLI.
@@ -48,7 +48,7 @@ Run a CPU-based VM runtime, execute some code, and clean up:
48
48
  colab new
49
49
 
50
50
  # 2. Execute code from stdin
51
- echo "print('Hello from Google Colab!')" | colab exec
51
+ echo 'print("Hello from Google Colab!")' | colab exec
52
52
 
53
53
  # 3. Stop and release the VM resource
54
54
  colab stop
@@ -190,3 +190,9 @@ To view interactive walkthroughs of eleven real-world automated scenarios, check
190
190
  ## Contributing
191
191
 
192
192
  Feedback and contributions are welcome! Please read [`CONTRIBUTING.md`](./CONTRIBUTING.md) for details.
193
+
194
+ ---
195
+
196
+ ## Legal & Privacy
197
+
198
+ See Colab's [Terms of Service](https://research.google.com/colaboratory/tos_v6.html) and the [Google Privacy Policy](https://policies.google.com/privacy), which apply to usage of this CLI.
@@ -1,5 +1,6 @@
1
1
  ---
2
2
  log:
3
+ 2026-10-14: Added configurable HTTP request timeouts and bounded allocation retries. Uncertain allocation POSTs are reconciled by GET using the same notebook hash before resubmission.
3
4
  2026-09-25: Runtime proxy tokens are now refreshed before expiry. `SessionState` persists `token_expires_at`, and `State.get_session` re-fetches the token via `GET /tun/m/assignments` when it's within 5 minutes of expiring.
4
5
  2026-09-25: Removed the keep-alive background daemon and related client/CLI ping functions. VM liveness is automatically maintained by the Colab backend based on kernel activity.
5
6
  2026-08-10: Added `colab usage` for account-level compute-unit rate/balance via `GET /tun/m/ccu-info` on the session backend (same bearer token as `colab new`).
@@ -53,6 +54,31 @@ The CLI maps user flags to these backend parameters:
53
54
 
54
55
  ## Approach
55
56
 
57
+ ### HTTP Timeout and Allocation Recovery
58
+
59
+ `COLAB_CLI_REQUEST_TIMEOUT` sets the HTTP request timeout in seconds. If unset,
60
+ the session's existing default applies (AuthorizedSession defaults to 120 seconds).
61
+ An explicit `timeout=` passed by a caller takes precedence. This does not change
62
+ kernel execution timeouts, Colab runtime lifetime, or benchmark job timeouts.
63
+
64
+ Allocation POSTs retry `ReadTimeout` and HTTP 502/503/504 at most
65
+ `COLAB_CLI_ASSIGN_RETRIES` times (default 3; 0 disables retries). Between attempts,
66
+ the client waits `COLAB_CLI_RETRY_BACKOFF * 2**attempt` seconds (default 10, 20, 40).
67
+ It then queries the same notebook hash: an existing allocation is returned directly;
68
+ otherwise the refreshed XSRF token is used for the next POST. HTTP 412 still raises
69
+ `TooManyAssignmentsError`; other errors are not retried. Failed initial/recovery
70
+ GETs and exhausted POST retries propagate to the caller. There is no unlimited wait.
71
+
72
+ ```bash
73
+ COLAB_CLI_REQUEST_TIMEOUT=300 COLAB_CLI_ASSIGN_RETRIES=3 COLAB_CLI_RETRY_BACKOFF=10 \
74
+ colab new --session kernel-benchmark-v6e --tpu v6e1
75
+ ```
76
+
77
+ Regression tests cover timeout precedence, same-request recovery, refreshed
78
+ tokens, retry limits, exponential backoff and permanent-error handling. A live
79
+ allocation/recovery test is also required before declaring the behavior verified
80
+ end to end; mocked timeouts alone do not establish backend compatibility.
81
+
56
82
  ### 1. New Session (`colab new`)
57
83
  - **API**: `GET https://colab.sandbox.google.com/tun/m/assign` (based on HAR).
58
84
  - **Parameters**:
@@ -38,7 +38,7 @@ dependencies = [
38
38
  "rich>=14.3.3",
39
39
  "typer>=0.24.1",
40
40
  "typing-extensions>=4.0",
41
- "websocket-client>=1.0",
41
+ "websocket-client>=1.6",
42
42
  ]
43
43
 
44
44
  [project.scripts]
@@ -13,14 +13,19 @@
13
13
  # limitations under the License.
14
14
 
15
15
  import abc
16
+ import json
17
+ import logging
18
+ import os
19
+ import time
20
+ import uuid
16
21
  from dataclasses import dataclass
17
22
  from datetime import datetime, timedelta, timezone
18
23
  from enum import Enum
19
- import json
20
- import logging
21
24
  from typing import Dict, List, Optional, Union
22
25
  from urllib.parse import urljoin, urlparse
23
- import uuid
26
+
27
+ import requests
28
+ from pydantic import BaseModel, Field, TypeAdapter
24
29
 
25
30
  from colab_cli.consumption import (
26
31
  CcuInfo,
@@ -28,8 +33,6 @@ from colab_cli.consumption import (
28
33
  consumption_user_info_from_tunnel,
29
34
  )
30
35
  from colab_cli.utils import get_status_code
31
- from pydantic import BaseModel, Field, TypeAdapter
32
- import requests
33
36
 
34
37
  # Standard Colab Headers
35
38
  ACCEPT_JSON_HEADER = {"key": "Accept", "value": "application/json"}
@@ -214,6 +217,12 @@ class Client:
214
217
  "value"
215
218
  ]
216
219
 
220
+ # Slow allocations may outlive AuthorizedSession's default read timeout.
221
+ if "timeout" not in kwargs:
222
+ timeout = os.getenv("COLAB_CLI_REQUEST_TIMEOUT")
223
+ if timeout:
224
+ kwargs["timeout"] = float(timeout)
225
+
217
226
  self.logger.debug(f"Request: {method} {endpoint}")
218
227
  self.logger.debug(f"Params: {params}")
219
228
 
@@ -269,20 +278,40 @@ class Client:
269
278
  accelerator: Optional[Accelerator] = None,
270
279
  shape: Optional[Shape] = None,
271
280
  ) -> Union[PostAssignmentResponse, Assignment]:
272
- assignment = self._get_assignment(
273
- notebook_hash, variant, accelerator, shape
274
- )
281
+ assignment = self._get_assignment(notebook_hash, variant, accelerator, shape)
275
282
  if isinstance(assignment, Assignment):
276
283
  return assignment
277
284
 
278
- try:
279
- res = self._post_assignment(
280
- notebook_hash, assignment.token, variant, accelerator, shape
281
- )
282
- except ColabRequestError as e:
283
- if get_status_code(e) == 412:
284
- raise TooManyAssignmentsError(str(e))
285
- raise e
285
+ retries = max(0, int(os.getenv("COLAB_CLI_ASSIGN_RETRIES", "3")))
286
+ backoff = max(0.0, float(os.getenv("COLAB_CLI_RETRY_BACKOFF", "10")))
287
+ for attempt in range(retries + 1):
288
+ try:
289
+ res = self._post_assignment(
290
+ notebook_hash, assignment.token, variant, accelerator, shape
291
+ )
292
+ break
293
+ except (ColabRequestError, requests.exceptions.ReadTimeout) as e:
294
+ status = (
295
+ get_status_code(e) if isinstance(e, ColabRequestError) else None
296
+ )
297
+ if status == 412:
298
+ raise TooManyAssignmentsError(str(e))
299
+ retryable = isinstance(
300
+ e, requests.exceptions.ReadTimeout
301
+ ) or status in (502, 503, 504)
302
+ if not retryable or attempt >= retries:
303
+ raise
304
+ delay = backoff * (2**attempt)
305
+ self.logger.warning(
306
+ "Assignment request failed; retrying in %.0fs", delay
307
+ )
308
+ time.sleep(delay)
309
+ # Reconcile an uncertain POST before submitting the same request again.
310
+ assignment = self._get_assignment(
311
+ notebook_hash, variant, accelerator, shape
312
+ )
313
+ if isinstance(assignment, Assignment):
314
+ return assignment
286
315
 
287
316
  return res
288
317
 
@@ -329,4 +358,3 @@ class Client:
329
358
  return self._issue_request(
330
359
  url, method="POST", headers=headers, schema=PostAssignmentResponse
331
360
  )
332
-
@@ -102,8 +102,18 @@ def run_automation(
102
102
  state.history.log_event(s.name, "drive_auth_needed", {"uri": uri})
103
103
  sys.stdout.write("Press Enter after you have granted access... ")
104
104
  sys.stdout.flush()
105
- with open("/dev/tty") as tty:
106
- tty.readline()
105
+ try:
106
+ with open("/dev/tty") as tty:
107
+ tty.readline()
108
+ except (FileNotFoundError, OSError, IOError):
109
+ # Windows has no /dev/tty; fall back to stdin.
110
+ try:
111
+ sys.stdin.readline()
112
+ except Exception:
113
+ try:
114
+ input()
115
+ except Exception:
116
+ pass
107
117
 
108
118
  typer.echo("[colab] Authorizing VM...")
109
119
  params["dryrun"] = "false"
@@ -37,7 +37,6 @@ the ``X-Colab-Ssh-Pubkey`` header. RSA keys are rejected by the server, so
37
37
  import contextlib
38
38
  import os
39
39
  from pathlib import Path
40
- import select
41
40
  import shlex
42
41
  import signal
43
42
  import subprocess
@@ -360,9 +359,6 @@ def _bridge_proxy_mode(ws: websocket.WebSocket) -> int:
360
359
  def stdin_to_ws():
361
360
  try:
362
361
  while True:
363
- ready, _, _ = select.select([stdin_fd], [], [], None)
364
- if not ready:
365
- continue
366
362
  data = os.read(stdin_fd, 8192)
367
363
  if not data:
368
364
  break
@@ -17,10 +17,17 @@ import logging
17
17
  import os
18
18
  import signal
19
19
  import sys
20
- import termios
21
20
  import threading
22
21
  import time
23
- import tty
22
+
23
+ try:
24
+ import termios
25
+ import tty
26
+ except ImportError:
27
+ # Windows has neither termios nor tty. connect_console() below degrades
28
+ # gracefully to line-buffered input when these are unavailable.
29
+ termios = None # type: ignore
30
+ tty = None # type: ignore
24
31
  from urllib.parse import urlparse
25
32
 
26
33
  import websocket
@@ -133,8 +140,17 @@ def connect_console(session: SessionState):
133
140
  ws_url = f"{ws_scheme}://{parsed.netloc}/colab/tty?colab-runtime-proxy-token={session.token}"
134
141
 
135
142
  is_tty = sys.stdin.isatty()
136
- fd = sys.stdin.fileno() if is_tty else None
137
- old_settings = termios.tcgetattr(fd) if is_tty else None
143
+ fd = None
144
+ old_settings = None
145
+ can_raw = termios is not None and tty is not None and is_tty
146
+ if can_raw:
147
+ try:
148
+ fd = sys.stdin.fileno()
149
+ old_settings = termios.tcgetattr(fd)
150
+ except Exception:
151
+ fd = None
152
+ old_settings = None
153
+ can_raw = False
138
154
 
139
155
  ws = websocket.WebSocketApp(
140
156
  url=ws_url,
@@ -149,10 +165,12 @@ def connect_console(session: SessionState):
149
165
  if _is_running:
150
166
  send_terminal_size(ws)
151
167
 
168
+ sigwinch = getattr(signal, "SIGWINCH", None)
152
169
  try:
153
- if is_tty:
170
+ if can_raw:
154
171
  tty.setraw(fd, termios.TCSANOW)
155
- signal.signal(signal.SIGWINCH, handle_sigwinch)
172
+ if sigwinch is not None:
173
+ signal.signal(sigwinch, handle_sigwinch)
156
174
 
157
175
  # This is a blocking call until the connection is closed
158
176
  ws.run_forever()
@@ -164,9 +182,16 @@ def connect_console(session: SessionState):
164
182
  # We raise a standard exception that the caller can recognize
165
183
  raise RuntimeError(f"Connection failed: {err_msg}")
166
184
  finally:
167
- if is_tty:
185
+ if can_raw:
168
186
  # Always ensure the terminal is restored to its original state
169
- termios.tcsetattr(fd, termios.TCSANOW, old_settings)
187
+ try:
188
+ termios.tcsetattr(fd, termios.TCSANOW, old_settings)
189
+ except Exception:
190
+ pass
170
191
  # Restore the default signal handler for resize
171
- signal.signal(signal.SIGWINCH, signal.SIG_DFL)
192
+ if sigwinch is not None:
193
+ try:
194
+ signal.signal(sigwinch, signal.SIG_DFL)
195
+ except Exception:
196
+ pass
172
197
  print("\r\nConnection closed.")
@@ -12,16 +12,21 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- import uuid
16
15
  import json
17
- import pytest
16
+ import uuid
18
17
  from unittest.mock import MagicMock
18
+
19
+ import pytest
20
+ import requests
21
+
19
22
  from colab_cli.client import (
23
+ Accelerator,
24
+ Assignment,
20
25
  Client,
21
- Prod,
26
+ ColabRequestError,
27
+ GetAssignmentResponse,
22
28
  PostAssignmentResponse,
23
- Assignment,
24
- Accelerator,
29
+ Prod,
25
30
  Shape,
26
31
  TooManyAssignmentsError,
27
32
  Variant,
@@ -39,6 +44,158 @@ def client(mock_session):
39
44
  return Client(Prod(), mock_session)
40
45
 
41
46
 
47
+ def test_request_timeout_override_and_explicit_precedence(
48
+ client, mock_session, monkeypatch
49
+ ):
50
+ monkeypatch.setenv("COLAB_CLI_REQUEST_TIMEOUT", "300")
51
+ mock_session.request.return_value.text = ""
52
+ client._issue_request("https://colab.research.google.com/test")
53
+ assert mock_session.request.call_args.kwargs["timeout"] == 300
54
+ client._issue_request("https://colab.research.google.com/test", timeout=12)
55
+ assert mock_session.request.call_args.kwargs["timeout"] == 12
56
+
57
+
58
+ def test_request_default_timeout_is_preserved(client, mock_session, monkeypatch):
59
+ monkeypatch.delenv("COLAB_CLI_REQUEST_TIMEOUT", raising=False)
60
+ mock_session.request.return_value.text = ""
61
+ client._issue_request("https://colab.research.google.com/test")
62
+ assert "timeout" not in mock_session.request.call_args.kwargs
63
+
64
+
65
+ def pending_assignment(token="initial"):
66
+ return GetAssignmentResponse(acc="NONE", nbh="same", token=token, variant="DEFAULT")
67
+
68
+
69
+ def allocated_assignment():
70
+ return Assignment(
71
+ endpoint="recovered",
72
+ runtimeProxyInfo={
73
+ "token": "proxy",
74
+ "tokenExpiresInSeconds": 3600,
75
+ "url": "http://backend",
76
+ },
77
+ )
78
+
79
+
80
+ @pytest.mark.parametrize("status", [None, 502, 503, 504])
81
+ def test_assign_recovers_before_another_post(client, mocker, monkeypatch, status):
82
+ monkeypatch.setenv("COLAB_CLI_ASSIGN_RETRIES", "3")
83
+ monkeypatch.setenv("COLAB_CLI_RETRY_BACKOFF", "10")
84
+ recovered = allocated_assignment()
85
+ get = mocker.patch.object(
86
+ client, "_get_assignment", side_effect=[pending_assignment(), recovered]
87
+ )
88
+ error = (
89
+ requests.exceptions.ReadTimeout("allocation pending")
90
+ if status is None
91
+ else ColabRequestError(
92
+ "transient", request=None, response=MagicMock(status_code=status)
93
+ )
94
+ )
95
+ post = mocker.patch.object(client, "_post_assignment", side_effect=error)
96
+ sleep = mocker.patch("time.sleep")
97
+ notebook = uuid.uuid4()
98
+ assert client.assign(notebook) is recovered
99
+ assert get.call_count == 2
100
+ assert all(call.args[0] == notebook for call in get.call_args_list)
101
+ post.assert_called_once()
102
+ sleep.assert_called_once_with(10)
103
+
104
+
105
+ def test_assign_refreshes_token_and_uses_bounded_backoff(client, mocker, monkeypatch):
106
+ monkeypatch.setenv("COLAB_CLI_ASSIGN_RETRIES", "3")
107
+ monkeypatch.setenv("COLAB_CLI_RETRY_BACKOFF", "10")
108
+ mocker.patch.object(
109
+ client,
110
+ "_get_assignment",
111
+ side_effect=[pending_assignment(str(i)) for i in range(4)],
112
+ )
113
+ post = mocker.patch.object(
114
+ client,
115
+ "_post_assignment",
116
+ side_effect=requests.exceptions.ReadTimeout("pending"),
117
+ )
118
+ sleep = mocker.patch("time.sleep")
119
+ notebook = uuid.uuid4()
120
+ with pytest.raises(requests.exceptions.ReadTimeout):
121
+ client.assign(notebook)
122
+ assert [call.args[:2] for call in post.call_args_list] == [
123
+ (notebook, str(i)) for i in range(4)
124
+ ]
125
+ assert [call.args[0] for call in sleep.call_args_list] == [10, 20, 40]
126
+
127
+
128
+ @pytest.mark.parametrize("status", [400, 401, 403])
129
+ def test_assign_does_not_retry_nontransient_errors(client, mocker, status):
130
+ get = mocker.patch.object(
131
+ client, "_get_assignment", return_value=pending_assignment()
132
+ )
133
+ post = mocker.patch.object(
134
+ client,
135
+ "_post_assignment",
136
+ side_effect=ColabRequestError(
137
+ "permanent", request=None, response=MagicMock(status_code=status)
138
+ ),
139
+ )
140
+ sleep = mocker.patch("time.sleep")
141
+ with pytest.raises(ColabRequestError):
142
+ client.assign(uuid.uuid4())
143
+ get.assert_called_once()
144
+ post.assert_called_once()
145
+ sleep.assert_not_called()
146
+
147
+
148
+ def test_assign_succeeds_after_recovery_refreshes_token(client, mocker, monkeypatch):
149
+ monkeypatch.setenv("COLAB_CLI_ASSIGN_RETRIES", "1")
150
+ monkeypatch.setenv("COLAB_CLI_RETRY_BACKOFF", "0")
151
+ get = mocker.patch.object(
152
+ client,
153
+ "_get_assignment",
154
+ side_effect=[pending_assignment("first"), pending_assignment("fresh")],
155
+ )
156
+ result = PostAssignmentResponse(
157
+ accelerator="NONE",
158
+ endpoint="new",
159
+ variant=0,
160
+ runtimeProxyInfo={
161
+ "token": "proxy",
162
+ "tokenExpiresInSeconds": 3600,
163
+ "url": "http://backend",
164
+ },
165
+ )
166
+ post = mocker.patch.object(
167
+ client,
168
+ "_post_assignment",
169
+ side_effect=[requests.exceptions.ReadTimeout("pending"), result],
170
+ )
171
+ mocker.patch("time.sleep")
172
+ notebook = uuid.uuid4()
173
+ assert client.assign(notebook) is result
174
+ assert get.call_count == 2
175
+ assert [call.args[:2] for call in post.call_args_list] == [
176
+ (notebook, "first"),
177
+ (notebook, "fresh"),
178
+ ]
179
+
180
+
181
+ def test_assign_retries_can_be_disabled(client, mocker, monkeypatch):
182
+ monkeypatch.setenv("COLAB_CLI_ASSIGN_RETRIES", "0")
183
+ get = mocker.patch.object(
184
+ client, "_get_assignment", return_value=pending_assignment()
185
+ )
186
+ post = mocker.patch.object(
187
+ client,
188
+ "_post_assignment",
189
+ side_effect=requests.exceptions.ReadTimeout("pending"),
190
+ )
191
+ sleep = mocker.patch("time.sleep")
192
+ with pytest.raises(requests.exceptions.ReadTimeout):
193
+ client.assign(uuid.uuid4())
194
+ get.assert_called_once()
195
+ post.assert_called_once()
196
+ sleep.assert_not_called()
197
+
198
+
42
199
  def test_client_assign_new(client, mock_session):
43
200
  # Mock _get_assignment (GET)
44
201
  get_resp = MagicMock()
@@ -15,13 +15,22 @@
15
15
  import json
16
16
  import os
17
17
  import sys
18
- import termios
19
18
  from unittest.mock import MagicMock, patch
20
19
 
20
+ try:
21
+ import termios
22
+ except ImportError:
23
+ # Windows has no termios; tests requiring raw-terminal mode are skipped.
24
+ termios = None # type: ignore
25
+
21
26
  from colab_cli.console import connect_console, on_message, on_open
22
27
  from colab_cli.state import SessionState
23
28
  import pytest
24
29
 
30
+ needs_termios = pytest.mark.skipif(
31
+ termios is None, reason="requires termios (POSIX only)"
32
+ )
33
+
25
34
 
26
35
  @pytest.fixture
27
36
  def mock_session():
@@ -33,6 +42,7 @@ def mock_session():
33
42
  )
34
43
 
35
44
 
45
+ @needs_termios
36
46
  @patch("colab_cli.console.websocket.WebSocketApp")
37
47
  @patch("colab_cli.console.tty.setraw")
38
48
  @patch("colab_cli.console.termios.tcgetattr")
@@ -79,6 +89,7 @@ def test_console_initialization(
79
89
  )
80
90
 
81
91
 
92
+ @needs_termios
82
93
  @patch("colab_cli.console.websocket.WebSocketApp")
83
94
  @patch("colab_cli.console.tty.setraw")
84
95
  @patch("colab_cli.console.termios.tcgetattr")
@@ -220,3 +231,31 @@ def test_read_stdin_eof_tty_does_not_close_ws(
220
231
  sent_payloads = [json.loads(c.args[0]) for c in mock_ws.send.call_args_list]
221
232
  assert {"data": "exit\n"} not in sent_payloads
222
233
  mock_ws.close.assert_not_called()
234
+
235
+
236
+ @patch("colab_cli.console.websocket.WebSocketApp")
237
+ @patch("colab_cli.console.sys.stdin.isatty")
238
+ def test_console_no_termios_degrades_gracefully(
239
+ mock_isatty, mock_ws_app, mock_session
240
+ ):
241
+ """Windows has no termios/tty/SIGWINCH. connect_console must not crash.
242
+
243
+ Regression test for `ModuleNotFoundError: No module named 'termios'`
244
+ which broke every `colab` command on Windows via
245
+ cli -> execution -> console imports.
246
+ """
247
+ import colab_cli.console as console_mod
248
+
249
+ mock_isatty.return_value = True
250
+ mock_ws_instance = MagicMock()
251
+ mock_ws_app.return_value = mock_ws_instance
252
+ mock_ws_instance.run_forever.return_value = None
253
+
254
+ with (
255
+ patch.object(console_mod, "termios", None),
256
+ patch.object(console_mod, "tty", None),
257
+ patch("colab_cli.console.threading.Thread"),
258
+ ):
259
+ connect_console(mock_session)
260
+
261
+ mock_ws_instance.run_forever.assert_called_once()
@@ -25,6 +25,7 @@ import io
25
25
  import shlex
26
26
  import subprocess
27
27
  import sys
28
+ import threading
28
29
  from unittest.mock import MagicMock
29
30
 
30
31
  from colab_cli.cli import app
@@ -235,6 +236,49 @@ def test_bridge_proxy_mode_pumps_ws_to_stdout(mocker):
235
236
  ws.close.assert_called()
236
237
 
237
238
 
239
+ def test_bridge_proxy_mode_pumps_stdin_to_ws(mocker):
240
+ """os.read() chunks leave as binary frames; EOF ends the pump and closes
241
+ the socket.
242
+
243
+ Regression test: the pump used select() on the stdin fd before reading.
244
+ Windows select() only accepts sockets, so there the OSError killed the
245
+ pump thread and the remote side saw no input at all; a plain blocking
246
+ read works on every platform.
247
+ """
248
+ reads = iter([b"ping ", b"hello", b""])
249
+ stdin_eof = threading.Event()
250
+
251
+ def fake_read(fd, size):
252
+ chunk = next(reads)
253
+ if not chunk:
254
+ stdin_eof.set()
255
+ return chunk
256
+
257
+ mocker.patch("sys.stdin")
258
+ sys.stdin.buffer.fileno.return_value = 0
259
+ mocker.patch.object(ssh_module.os, "read", side_effect=fake_read)
260
+ fake_stdout = MagicMock()
261
+ fake_stdout.buffer = io.BytesIO()
262
+ mocker.patch("sys.stdout", fake_stdout)
263
+
264
+ # Hold the ws->stdout loop until the pump has drained stdin, so the
265
+ # send_binary assertions below can't race the reader thread.
266
+ def recv_data(control_frame):
267
+ stdin_eof.wait(timeout=5)
268
+ return (websocket.ABNF.OPCODE_CLOSE, b"")
269
+
270
+ ws = MagicMock()
271
+ ws.recv_data.side_effect = recv_data
272
+
273
+ rc = ssh_module._bridge_proxy_mode(ws)
274
+ assert rc == 0
275
+ assert [call.args[0] for call in ws.send_binary.call_args_list] == [
276
+ b"ping ",
277
+ b"hello",
278
+ ]
279
+ ws.close.assert_called()
280
+
281
+
238
282
  # --- ProxyCommand shell quoting ---------------------------------------------
239
283
 
240
284
 
@@ -414,7 +414,7 @@ requires-dist = [
414
414
  { name = "rich", specifier = ">=14.3.3" },
415
415
  { name = "typer", specifier = ">=0.24.1" },
416
416
  { name = "typing-extensions", specifier = ">=4.0" },
417
- { name = "websocket-client", specifier = ">=1.0" },
417
+ { name = "websocket-client", specifier = ">=1.6" },
418
418
  ]
419
419
 
420
420
  [package.metadata.requires-dev]
@@ -573,11 +573,11 @@ wheels = [
573
573
 
574
574
  [[package]]
575
575
  name = "oauthlib"
576
- version = "3.3.1"
576
+ version = "4.0.0"
577
577
  source = { registry = "https://pypi.org/simple" }
578
- sdist = { url = "https://files.pythonhosted.org/packages/0b/5f/19930f824ffeb0ad4372da4812c50edbd1434f678c90c2733e1188edfc63/oauthlib-3.3.1.tar.gz", hash = "sha256:0f0f8aa759826a193cf66c12ea1af1637f87b9b4622d46e866952bb022e538c9", size = 185918, upload-time = "2025-06-19T22:48:08.269Z" }
578
+ sdist = { url = "https://files.pythonhosted.org/packages/7a/d8/a1bcc8ba112a627f8ffbdc212a78ce18d3ac07e91a5ca65d27918eee25a1/oauthlib-4.0.0.tar.gz", hash = "sha256:efb274799819440f95b4ab3b818869f1ce9ae26c5beacba0201d1a1b76b54f86", size = 187232, upload-time = "2026-09-28T06:01:18.77Z" }
579
579
  wheels = [
580
- { url = "https://files.pythonhosted.org/packages/be/9c/92789c596b8df838baa98fa71844d84283302f7604ed565dafe5a6b5041a/oauthlib-3.3.1-py3-none-any.whl", hash = "sha256:88119c938d2b8fb88561af5f6ee0eec8cc8d552b7bb1f712743136eb7523b7a1", size = 160065, upload-time = "2025-06-19T22:48:06.508Z" },
580
+ { url = "https://files.pythonhosted.org/packages/d9/f4/78229a1066068ca14fc60fb26cf7381cabe4382261392b90e5f9552722d4/oauthlib-4.0.0-py3-none-any.whl", hash = "sha256:624c28c13a0a59cabf9747dfa52af63be3e512a7f2714df16e91b5b3a145e6cd", size = 159715, upload-time = "2026-09-28T06:01:17.008Z" },
581
581
  ]
582
582
 
583
583
  [[package]]
@@ -1160,11 +1160,11 @@ wheels = [
1160
1160
 
1161
1161
  [[package]]
1162
1162
  name = "urllib3"
1163
- version = "2.7.0"
1163
+ version = "2.8.0"
1164
1164
  source = { registry = "https://pypi.org/simple" }
1165
- sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
1165
+ sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" }
1166
1166
  wheels = [
1167
- { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
1167
+ { url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" },
1168
1168
  ]
1169
1169
 
1170
1170
  [[package]]