fix: enforce --timeout as an idle limit without busy-waiting

docs/02_execution_and_interactive.md

@@ -5,6 +5,8 @@ log:
 2026-05-07: Fixed `print_kitty` (used by `colab exec --output-image` and any image-producing exec) to no-op when `sys.stdout.isatty()` is false. The Kitty Graphics Protocol escape sequence is meaningless when stdout is a file or pipe and was visually corrupting captured output (a multi-KB base64 PNG blob would land in log files, grep targets, or showboat captures). Image bytes are still saved to disk via `handle_image`'s file-write path; only the inline-render attempt is suppressed.
 
 2026-06-04: Bumped the default `--timeout` for `colab exec` from 10s to 30s (and the matching `colab run` default) so brief silent tasks are less likely to hit a premature `TimeoutError`. Explicit `--timeout` overrides are unaffected.
+
+2026-10-24: `--timeout` for `colab exec` (and `colab run`) is now an idle limit: execution fails only after that many seconds without output, so code that keeps printing is never cut off, and `--timeout 0` means no limit. On 0.8.0/0.9.0 the old deadline was never enforced, so code that is silent for longer than `--timeout` now fails where it used to complete. `colab repl`, which passes no timeout, waits until each cell finishes, or until the kernel dies or restarts. None of them busy-loop while waiting. `colab install`, `colab auth`, and `colab drivemount` are unchanged. See #183.
 ---
 
 # Design: Execution and Interactive Interaction (`repl`, `exec`, `console`)
@@ -28,7 +30,7 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel
     - If file path is local: Read content, send as code.
     - If file path is remote: Execute `!python <path>`.
 - **Multi-Modal Output**: Handle `display_data` messages (e.g., `image/png`, `text/html`). For the CLI, we'll save images to temporary files and print their paths, or if the terminal supports it (e.g., iTerm2), inline them.
-- **Timeout Configuration**: Exposes a `--timeout` flag (default 30s) to allow long-running silent tasks (like model compilation or data downloading) to execute without being prematurely killed.
+- **Timeout Configuration**: Exposes a `--timeout` flag (default 30s): the longest silence allowed before execution fails. Raise it so long-running silent tasks (like model compilation or data downloading) are not prematurely killed; `0` means no limit.
 
 ### 3. Console (`colab console`)
 - **Implementation**: Connects directly to the backend terminal endpoint (`/colab/tty`) via WebSockets using `websocket-client`.

docs/05_run_command.md

@@ -5,6 +5,7 @@ log:
 2026-05-12: Initial design and implementation of `colab run <script.py> [args...]`. Combines `colab new` + `colab exec` + `colab stop` into a single fire-and-forget invocation so a Python file can use `#!/usr/bin/env -S colab run` as a shebang line and execute on a freshly-allocated Colab VM. Adds `--keep` (skip auto-stop), `--gpu` / `--tpu` (passthrough to session creation), `-s/--session` (name the ephemeral session), and propagates the script's exit status (non-zero on any uncaught exception in the kernel). The script's `sys.argv` is re-set inside the kernel to mirror native `python script.py arg1 arg2` semantics, and `__name__` is set to `"__main__"`.
 2026-05-12: Native CPython exit-code semantics for `sys.exit()` / `raise SystemExit(...)` from the script body. The Colab kernel reports a `SystemExit` as `output_type=='error'`, which under the previous logic would have (a) printed the IPython traceback (`An exception has occurred, use %tb...`) and (b) flagged the run as a failure regardless of the integer exit code. Now: `sys.exit()` / `sys.exit(0)` exit 0 silently; `sys.exit(N)` exits N; `sys.exit('msg')` exits 1 (matching CPython). The IPython "To exit: use 'exit', 'quit', or Ctrl-D." UserWarning is filtered via the prelude. Encoded after running `examples/gpu_hello.py` end-to-end and seeing the noisy `SystemExit: 0` traceback at the end of an otherwise-successful GPU run.
 2026-06-04: Bumped the default value of the `--timeout` flag from 10.0s to 30.0s so short-but-silent tasks aren't prematurely killed out of the box. Mirrors the same change for `colab exec`.
+2026-10-24: `--timeout` is now an idle limit: the run fails only after that many seconds without output, so a script that keeps printing can run as long as it needs, and `--timeout 0` means no limit. The CLI waits for kernel messages itself, in blocking slices of up to one second, instead of using `jupyter-kernel-client`'s `execute_interactive`. That deadline covered the whole run: 1.0.2 enforces it, failing runs after 30s even while they print, and 0.8.0/0.9.0 never enforce it, polling at ~100% CPU once it passes. As a result, on 0.8.0/0.9.0 a script that is silent for longer than `--timeout` now fails (and, without `--keep`, its session is stopped) where it used to complete. See #183.
 ---
 
 # Design: `colab run` — Shebang-Compatible One-Shot Execution
@@ -33,7 +34,7 @@ colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...
 | `--tpu` | str | None | Same set as `colab new --tpu` (v5e1, v6e1). |
 | `--high-mem` | bool | False | Same as `colab new --high-mem` — request high-RAM when supported. |
 | `--keep` | bool | False | Do **not** stop the session after the script finishes. |
-| `--timeout` | float | 30.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. |
+| `--timeout` | float | 30.0 | Seconds to wait without output before failing, so a hung or silent task doesn't block forever. Raise it for long silent tasks; `0` means no limit. |
 
 ### Shebang usage
 With `--keep` and `--gpu` baked into the shebang line, an entire one-file workload becomes:

src/colab_cli/commands/execution.py

@@ -149,7 +149,10 @@ def exec_command(
     ] = None,
     timeout: Annotated[
         Optional[float],
-        typer.Option("--timeout", help="Timeout in seconds for code execution"),
+        typer.Option(
+            "--timeout",
+            help="Seconds to wait without output before failing (0 for no limit)",
+        ),
     ] = 30.0,
     env: Annotated[
         Optional[List[str]],

src/colab_cli/commands/run.py

@@ -258,7 +258,10 @@ def run_command(
     ] = False,
     timeout: Annotated[
         Optional[float],
-        typer.Option("--timeout", help="Timeout in seconds for code execution"),
+        typer.Option(
+            "--timeout",
+            help="Seconds to wait without output before failing (0 for no limit)",
+        ),
     ] = 30.0,
     env: Annotated[
         Optional[List[str]],

src/colab_cli/runtime.py

@@ -13,13 +13,20 @@
 # limitations under the License.
 
 import logging
+import queue
 import time
 from typing import Any, Callable, Dict, List, Optional
 
 import jupyter_kernel_client
 import requests
 
 
+# The longest single wait for a kernel message, so a lost connection is noticed.
+_WAIT_SLICE_SEC = 1.0
+# Kernel states that end any execution in progress.
+_KERNEL_GONE = frozenset({"dead", "restarting"})
+
+
 class ColabRuntime:
     def __init__(
         self,
@@ -185,16 +192,17 @@ def execute_code(
         output_hook: Optional[Callable[[Dict[str, Any]], None]] = None,
         timeout: Optional[float] = None,
     ) -> List[Dict[str, Any]]:
-        # ``jupyter_kernel_client`` defaults ``timeout`` to ``REQUEST_TIMEOUT``
-        # (10 seconds) on both ``execute`` and ``execute_interactive``. That
-        # value is a wall-clock budget that shrinks every time the poll loop
-        # iterates -- as long as iopub/stdin events arrive back-to-back the
-        # call survives, but a single >10s quiet stretch (e.g. a kernel
-        # blocked on ``input_request`` while the user OAuths in the browser)
-        # will raise ``TimeoutError`` even though the underlying execution is
-        # still healthy. Callers that know they need a longer ceiling can
-        # pass ``timeout=`` here; otherwise we forward whatever the upstream
-        # default is (currently 10s).
+        # ``timeout`` means one of two things, depending on the path:
+        # - Streaming without stdin (``output_hook`` set, ``allow_stdin``
+        #   False; used by run, exec and repl): ``_execute_streaming`` waits
+        #   for messages itself, and ``timeout`` is an idle limit. ``None``
+        #   or a value <= 0 means no limit.
+        # - Otherwise the call goes to ``jupyter_kernel_client``, where
+        #   ``timeout`` is a deadline for the whole call, defaulting to
+        #   ``REQUEST_TIMEOUT`` (10 s). 0.8.0 and 0.9.0 never enforce it and
+        #   poll without waiting once it passes; 1.0.2 raises
+        #   ``TimeoutError``. Callers that need a longer ceiling, such as
+        #   the interactive automations, pass ``timeout=`` explicitly.
         kwargs = {"allow_stdin": allow_stdin}
         if timeout is not None:
             kwargs["timeout"] = timeout
@@ -218,7 +226,7 @@ def wrapped_stdin_hook(prompt):
             kwargs["stdin_hook"] = wrapped_stdin_hook
 
         if output_hook:
-            # If we have an output hook, we use execute_interactive and manage buffering ourselves
+            # With an output hook, we stream messages and manage buffering ourselves
             outputs = []
 
             def wrapped_output_hook(msg):
@@ -234,10 +242,13 @@ def wrapped_output_hook(msg):
                         if idx < len(outputs):
                             output_hook(outputs[idx])
 
-            reply = self.kernel_client.execute_interactive(
-                code, output_hook=wrapped_output_hook, **kwargs
-            )
-            # execute_interactive returns the raw reply message
+            if allow_stdin:
+                reply = self.kernel_client.execute_interactive(
+                    code, output_hook=wrapped_output_hook, **kwargs
+                )
+            else:
+                reply = self._execute_streaming(code, wrapped_output_hook, timeout)
+            # Both paths return the raw execute_reply message
             reply_content = reply["content"] if reply else {"status": "error"}
         else:
             reply = self.kernel_client.execute(code, **kwargs)
@@ -261,6 +272,68 @@ def wrapped_output_hook(msg):
 
         return outputs
 
+    def _execute_streaming(
+        self,
+        code: str,
+        output_hook: Callable[[Dict[str, Any]], None],
+        timeout: Optional[float],
+    ) -> Dict[str, Any]:
+        """Run ``code``, pass each output message to ``output_hook``, return the reply.
+
+        ``timeout`` is an idle limit: how long to wait, in seconds, for the next
+        message from this execution. A run that keeps printing can take as long
+        as it needs, and ``None`` or a value <= 0 waits without a limit. This
+        replaces ``jupyter_kernel_client``'s ``execute_interactive``, whose
+        deadline covers the whole run and which, in 0.8.0 and 0.9.0, polls
+        without waiting once the deadline passes.
+
+        Like ``_apply_ws_hook``, this uses the websocket client behind the
+        kernel client. It skips that library's per-client execution lock,
+        because a ColabRuntime runs one execution at a time.
+        """
+        if timeout is not None and timeout <= 0:
+            timeout = None
+        client = self.kernel_client._manager.client
+        msg_id = client.execute(code, allow_stdin=False)
+        last_message = time.monotonic()
+
+        def next_message(get_msg, waiting_for):
+            nonlocal last_message
+            while True:
+                if not client.connection_ready.is_set():
+                    raise RuntimeError("Connection was lost.")
+                wait = _WAIT_SLICE_SEC
+                if timeout is not None:
+                    left = last_message + timeout - time.monotonic()
+                    if left <= 0:
+                        raise TimeoutError(
+                            f"Timeout waiting for {waiting_for}: no message from "
+                            f"the kernel for {timeout:g}s. Use --timeout to allow "
+                            "longer silences."
+                        )
+                    wait = min(wait, left)
+                try:
+                    msg = get_msg(timeout=wait)
+                except queue.Empty:
+                    continue
+                state = msg["content"].get("execution_state")
+                if msg["header"]["msg_type"] == "status" and state in _KERNEL_GONE:
+                    # The kernel announces these to every client, without a
+                    # parent, and our execution can no longer finish.
+                    raise RuntimeError(f"Kernel {state} before the code finished.")
+                if msg["parent_header"].get("msg_id") == msg_id:
+                    last_message = time.monotonic()
+                    return msg
+
+        while True:
+            msg = next_message(client.get_iopub_msg, "output")
+            output_hook(msg)
+            if (
+                msg["header"]["msg_type"] == "status"
+                and msg["content"].get("execution_state") == "idle"
+            ):
+                return next_message(client.get_shell_msg, "reply")
+
     def stop(self, shutdown_kernel: bool = False):
         if self._kernel_client:
             try:

tests/mock_kernel.py

@@ -0,0 +1,119 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A mock of jupyter-kernel-client's websocket client for streaming tests.
+
+MockKernelClient replays a script of kernel messages on a simulated clock, so
+tests can cover long runs and long silences without waiting.
+"""
+
+import itertools
+import queue
+import threading
+import time
+from unittest.mock import MagicMock
+
+from colab_cli import runtime as runtime_module
+
+REQUEST = "req-1"
+_msg_ids = itertools.count()
+
+
+def kernel_message(msg_type, content, parent=REQUEST):
+    """A Jupyter message as the websocket client delivers it."""
+    return {
+        "msg_type": msg_type,
+        "header": {"msg_type": msg_type, "msg_id": f"msg-{next(_msg_ids)}"},
+        "parent_header": {"msg_id": parent} if parent else {},
+        "content": content,
+    }
+
+
+def stdout_message(text, parent=REQUEST):
+    return kernel_message("stream", {"name": "stdout", "text": text}, parent)
+
+
+def status_message(state, parent=REQUEST):
+    return kernel_message("status", {"execution_state": state}, parent)
+
+
+def reply_message(status="ok", parent=REQUEST, **content):
+    return kernel_message("execute_reply", {"status": status, **content}, parent)
+
+
+class MockKernelClient:
+    """The websocket client's methods that ColabRuntime uses, on a simulated clock.
+
+    ``iopub`` and ``shell`` are lists of (seconds after the request, message).
+    By default the shell channel replies when the last iopub message arrives. A
+    get_*_msg call that would wait past the next message advances the clock by
+    ``timeout`` and raises queue.Empty, as a quiet kernel would. The connection
+    drops when the clock reaches ``disconnect_at``.
+    """
+
+    def __init__(self, iopub, shell=None, disconnect_at=None):
+        self.now = 0.0
+        self.iopub = list(iopub)
+        last = iopub[-1][0] if iopub else 0.0
+        self.shell = list(shell if shell is not None else [(last, reply_message())])
+        self.disconnect_at = disconnect_at
+        self.waits = []
+        self.connection_ready = threading.Event()
+        self.connection_ready.set()
+
+    def execute(self, code, **kwargs):
+        self.code, self.execute_kwargs = code, kwargs
+        return REQUEST
+
+    def _get(self, messages, timeout):
+        self.waits.append(timeout)
+        assert timeout is not None and timeout > 0, "waits must be bounded"
+        if messages and messages[0][0] <= self.now + timeout:
+            at, msg = messages.pop(0)
+            self.now = max(self.now, at)
+            return msg
+        self.now += timeout
+        if self.disconnect_at is not None and self.now >= self.disconnect_at:
+            self.connection_ready.clear()
+        raise queue.Empty
+
+    def get_iopub_msg(self, timeout=None):
+        return self._get(self.iopub, timeout)
+
+    def get_shell_msg(self, timeout=None):
+        return self._get(self.shell, timeout)
+
+
+class _SimulatedTime:
+    """The ``time`` module, with ``monotonic`` read from a MockKernelClient."""
+
+    def __init__(self, mock_kc):
+        self._mock_kc = mock_kc
+
+    def monotonic(self):
+        return self._mock_kc.now
+
+    def __getattr__(self, name):
+        return getattr(time, name)
+
+
+def connect_runtime(runtime, mock_kc, monkeypatch):
+    """Make ``runtime`` use ``mock_kc``, and read its clock from ``mock_kc``.
+
+    Only ``runtime``'s view of ``time.monotonic`` changes; everything else in
+    ``time`` stays real.
+    """
+    monkeypatch.setattr(runtime_module, "time", _SimulatedTime(mock_kc))
+    runtime._kernel_client = MagicMock()
+    runtime._kernel_client._manager.client = mock_kc

tests/test_runtime.py

@@ -15,6 +15,13 @@
 from unittest.mock import MagicMock, patch
 
 import jupyter_kernel_client
+import pytest
+from mock_kernel import (
+    MockKernelClient,
+    connect_runtime,
+    status_message,
+    stdout_message,
+)
 
 from colab_cli.runtime import ColabRuntime
 
@@ -162,18 +169,41 @@ def test_colab_runtime_execute_code_with_timeout():
 
 def test_colab_runtime_execute_interactive_with_timeout():
     """timeout must also be plumbed through the execute_interactive branch
-    (used when an output_hook is supplied)."""
+    (used when an output_hook is supplied with allow_stdin)."""
     runtime = ColabRuntime("http://url", "token123")
     mock_kc = MagicMock()
     runtime._kernel_client = mock_kc
 
     mock_kc.execute_interactive.return_value = {"content": {"status": "ok"}}
-    runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600)
+    runtime.execute_code(
+        "print(1)", output_hook=lambda o: None, allow_stdin=True, timeout=600
+    )
 
     _, kwargs = mock_kc.execute_interactive.call_args
     assert kwargs.get("timeout") == 600
 
 
+def test_colab_runtime_streaming_timeout_is_an_idle_limit(monkeypatch):
+    """Without stdin, an output_hook streams through _execute_streaming, where
+    timeout is the longest silence allowed between messages."""
+    for silence, finishes in ((599, True), (601, False)):
+        runtime = ColabRuntime("http://url", "token123")
+        mock_kc = MockKernelClient(
+            [(silence, stdout_message("1\n")), (silence + 0.5, status_message("idle"))]
+        )
+        connect_runtime(runtime, mock_kc, monkeypatch)
+        if finishes:
+            outputs = runtime.execute_code(
+                "print(1)", output_hook=lambda o: None, timeout=600
+            )
+            assert outputs[0]["text"] == "1\n"
+        else:
+            with pytest.raises(TimeoutError):
+                runtime.execute_code(
+                    "print(1)", output_hook=lambda o: None, timeout=600
+                )
+
+
 def test_colab_runtime_stop():
     runtime = ColabRuntime("http://url", "token123")
     mock_kc = MagicMock()