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. `colab repl`, which passes no timeout, waits until each cell finishes. Neither busy-loops while waiting. See #183.
 ---
 
 # Design: Execution and Interactive Interaction (`repl`, `exec`, `console`)

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 enforced as 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. The CLI waits for kernel messages itself, in one-second blocking slices, instead of relying on `jupyter-kernel-client`'s `execute_interactive`, whose deadline covered the whole run and which, in 0.9.0, held a CPU core at ~100% once the deadline passed. See #183.
 ---
 
 # Design: `colab run` — Shebang-Compatible One-Shot Execution

src/colab_cli/runtime.py

@@ -13,13 +13,18 @@
 # 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
+
+
 class ColabRuntime:
     def __init__(
         self,
@@ -234,9 +239,12 @@ 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
-            )
+            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)
             # execute_interactive returns the raw reply message
             reply_content = reply["content"] if reply else {"status": "error"}
         else:
@@ -261,6 +269,53 @@ 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`` waits without a limit. The deadline in
+        ``jupyter_kernel_client.execute_interactive`` covers the whole run
+        instead, and past it 0.9.0 polls without waiting, holding a CPU core
+        until the run ends.
+        """
+        client = self.kernel_client._manager.client
+        msg_id = client.execute(code, allow_stdin=False)
+        last_message = time.monotonic()
+
+        def next_message(get_msg):
+            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("Timeout waiting for output")
+                    wait = min(wait, left)
+                try:
+                    msg = get_msg(timeout=wait)
+                except queue.Empty:
+                    continue
+                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_hook(msg)
+            if (
+                msg["header"]["msg_type"] == "status"
+                and msg["content"].get("execution_state") == "idle"
+            ):
+                return next_message(client.get_shell_msg)
+
     def stop(self, shutdown_kernel: bool = False):
         if self._kernel_client:
             try:

tests/mock_kernel.py

@@ -0,0 +1,102 @@
+# 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 queue
+import threading
+import time
+import types
+from unittest.mock import MagicMock
+
+from colab_cli import runtime as runtime_module
+
+REQUEST = "req-1"
+
+
+def _msg(msg_type, content, parent=REQUEST):
+    return {
+        "header": {"msg_type": msg_type},
+        "parent_header": {"msg_id": parent},
+        "content": content,
+    }
+
+
+def _stdout(text, parent=REQUEST):
+    return _msg("stream", {"name": "stdout", "text": text}, parent)
+
+
+def _idle(parent=REQUEST):
+    return _msg("status", {"execution_state": "idle"}, parent)
+
+
+def _reply(status="ok", **content):
+    return _msg("execute_reply", {"status": status, **content})
+
+
+class MockKernelClient:
+    """The websocket client's public API, on a simulated clock.
+
+    `iopub` and `shell` are lists of (seconds after the request, message). 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.
+    """
+
+    def __init__(self, iopub, shell=None):
+        self.now = 0.0
+        self.iopub = list(iopub)
+        self.shell = list(
+            shell if shell is not None else [(self._last(iopub), _reply())]
+        )
+        self.waits = []
+        self.connection_ready = threading.Event()
+        self.connection_ready.set()
+
+    @staticmethod
+    def _last(messages):
+        return messages[-1][0] if messages else 0.0
+
+    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, (
+            "must wait in bounded, non-zero slices"
+        )
+        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
+        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)
+
+
+def connect_runtime(runtime, mock_kc, monkeypatch):
+    """Make `runtime` talk to `mock_kc`, with the runtime's clock set to its clock."""
+    clock = types.SimpleNamespace(monotonic=lambda: mock_kc.now, sleep=time.sleep)
+    monkeypatch.setattr(runtime_module, "time", clock)
+    runtime._kernel_client = MagicMock()
+    runtime._kernel_client._manager.client = mock_kc

tests/test_runtime.py

@@ -14,6 +14,8 @@
 
 from unittest.mock import MagicMock, patch
 
+import pytest
+
 import jupyter_kernel_client
 
 from colab_cli.runtime import ColabRuntime
@@ -160,18 +162,22 @@ def test_colab_runtime_execute_code_with_timeout():
     assert kwargs.get("timeout") == 600
 
 
-def test_colab_runtime_execute_interactive_with_timeout():
+def test_colab_runtime_execute_interactive_with_timeout(monkeypatch):
     """timeout must also be plumbed through the execute_interactive branch
-    (used when an output_hook is supplied)."""
-    runtime = ColabRuntime("http://url", "token123")
-    mock_kc = MagicMock()
-    runtime._kernel_client = mock_kc
+    (used when an output_hook is supplied). There it is an idle limit: the
+    longest silence allowed between messages from the execution."""
+    from mock_kernel import MockKernelClient, _idle, _stdout, connect_runtime
 
-    mock_kc.execute_interactive.return_value = {"content": {"status": "ok"}}
+    runtime = ColabRuntime("http://url", "token123")
+    mock_kc = MockKernelClient([(599, _stdout("1\n")), (599.5, _idle())])
+    connect_runtime(runtime, mock_kc, monkeypatch)
     runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600)
 
-    _, kwargs = mock_kc.execute_interactive.call_args
-    assert kwargs.get("timeout") == 600
+    runtime = ColabRuntime("http://url", "token123")
+    mock_kc = MockKernelClient([(601, _stdout("1\n")), (601.5, _idle())])
+    connect_runtime(runtime, mock_kc, monkeypatch)
+    with pytest.raises(TimeoutError):
+        runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600)
 
 
 def test_colab_runtime_stop():