From e717770cff5fd88b1bb2c7f4c15a1fd8c8f62037 Mon Sep 17 00:00:00 2001 From: Duc Hieu Ho <150573299+hohieuai@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:48:23 +1100 Subject: [PATCH] fix: enforce --timeout as an idle limit without busy-waiting `--timeout` (default 30 s) is documented as an idle threshold, so that long-running silent tasks can raise it and code that keeps printing is never cut off. Runtime.execute_code passed it to jupyter-kernel-client's execute_interactive, where it is a deadline for the whole run: 0.8.0 and 0.9.0 never enforce it and poll at ~100% CPU once it passes, and 1.0.2 raises TimeoutError even while output flows. For streaming execution without stdin (run, exec, repl), the runtime now waits for kernel messages itself, through the websocket client's execute, get_iopub_msg and get_shell_msg, in blocking slices of up to one second: - Every message from the execution resets the idle clock; messages from other requests are ignored. - After `timeout` seconds without a message it raises TimeoutError, now with a hint to use --timeout. `--timeout 0` (or None) means no limit. - It raises RuntimeError if the connection drops or the kernel dies or restarts, instead of waiting forever. The --timeout help text and the design docs describe the idle semantics. Executions with stdin (auth, drivemount) and the non-streaming path (install, the /content prelude) keep using jupyter-kernel-client. Fixes #183 --- docs/02_execution_and_interactive.md | 4 +- docs/05_run_command.md | 3 +- src/colab_cli/commands/execution.py | 5 +- src/colab_cli/commands/run.py | 5 +- src/colab_cli/runtime.py | 103 +++++++++++--- tests/mock_kernel.py | 119 +++++++++++++++++ tests/test_runtime.py | 34 ++++- tests/test_streaming.py | 192 ++++++++++++++++++--------- 8 files changed, 381 insertions(+), 84 deletions(-) create mode 100644 tests/mock_kernel.py diff --git a/docs/02_execution_and_interactive.md b/docs/02_execution_and_interactive.md index ed1d782a..061cafb8 100644 --- a/docs/02_execution_and_interactive.md +++ b/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 `. - **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`. diff --git a/docs/05_run_command.md b/docs/05_run_command.md index 9a1d4214..f8414616 100644 --- a/docs/05_run_command.md +++ b/docs/05_run_command.md @@ -5,6 +5,7 @@ log: 2026-05-12: Initial design and implementation of `colab run [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: diff --git a/src/colab_cli/commands/execution.py b/src/colab_cli/commands/execution.py index 336a5ffb..4f6ac44a 100644 --- a/src/colab_cli/commands/execution.py +++ b/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]], diff --git a/src/colab_cli/commands/run.py b/src/colab_cli/commands/run.py index 0b222170..73893cb2 100644 --- a/src/colab_cli/commands/run.py +++ b/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]], diff --git a/src/colab_cli/runtime.py b/src/colab_cli/runtime.py index 57e09fde..b2a41463 100644 --- a/src/colab_cli/runtime.py +++ b/src/colab_cli/runtime.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +import queue import time from typing import Any, Callable, Dict, List, Optional @@ -20,6 +21,12 @@ 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: diff --git a/tests/mock_kernel.py b/tests/mock_kernel.py new file mode 100644 index 00000000..a1579702 --- /dev/null +++ b/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 diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 8d439d0a..d50c160a 100644 --- a/tests/test_runtime.py +++ b/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() diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 50e9f4c2..bfe93391 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -12,87 +12,153 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import MagicMock +import pytest +from mock_kernel import ( + MockKernelClient, + connect_runtime, + reply_message, + status_message, + stdout_message, +) + from colab_cli.runtime import ColabRuntime -def test_runtime_execute_code_streaming(): - runtime = ColabRuntime("http://url", "token") - mock_client = MagicMock() +@pytest.fixture +def run(monkeypatch): + """Run execute_code against a MockKernelClient; returns (outputs, streamed).""" - # Mock the return value of execute_interactive (the raw reply) - mock_client.execute_interactive.return_value = {"content": {"status": "ok"}} + def _run(mock_kc, code="work()", timeout=30.0): + runtime = ColabRuntime("http://url", "token") + connect_runtime(runtime, mock_kc, monkeypatch) + streamed = [] + outputs = runtime.execute_code( + code, output_hook=streamed.append, timeout=timeout + ) + return outputs, streamed - # Inject our mock client - runtime._kernel_client = mock_client + return _run - streamed_outputs = [] - def output_hook(o): - streamed_outputs.append(o) +def test_runtime_execute_code_streaming(run): + mock_kc = MockKernelClient( + [ + (0.1, stdout_message("1\n")), + (0.2, stdout_message("2\n")), + (0.3, status_message("idle")), + ] + ) + outputs, streamed = run(mock_kc, code="print(1); print(2)") - code = "print(1); print(2)" + assert [o["text"] for o in outputs] == ["1\n", "2\n"] + assert [o["text"] for o in streamed] == ["1\n", "2\n"] + assert mock_kc.code == "print(1); print(2)" + assert mock_kc.execute_kwargs == {"allow_stdin": False} - # We need to simulate the execution where execute_interactive is called - # and it calls our wrapped_output_hook. - # Since wrapped_output_hook is defined inside execute_code, we have to - # intercept the call to execute_interactive to get a reference to it. - def side_effect(code, output_hook=None, **kwargs): - # Simulate messages arriving - msg1 = { - "header": {"msg_type": "stream"}, - "content": {"name": "stdout", "text": "1\n"}, - } - msg2 = { - "header": {"msg_type": "stream"}, - "content": {"name": "stdout", "text": "2\n"}, - } - if output_hook: - output_hook(msg1) - output_hook(msg2) - return {"content": {"status": "ok"}} +def test_runtime_execute_code_streaming_error_synthesis(run): + reply = reply_message( + "error", ename="RuntimeError", evalue="something went wrong", traceback=["tb"] + ) + mock_kc = MockKernelClient([(0.1, status_message("idle"))], shell=[(0.1, reply)]) + outputs, streamed = run(mock_kc, code="fail()") - mock_client.execute_interactive.side_effect = side_effect + # The synthesized error is added after the run, so it is returned but not streamed. + assert len(outputs) == 1 + assert outputs[0]["output_type"] == "error" + assert outputs[0]["ename"] == "RuntimeError" + assert streamed == [] - outputs = runtime.execute_code(code, output_hook=output_hook) - assert len(outputs) == 2 - assert outputs[0]["text"] == "1\n" - assert outputs[1]["text"] == "2\n" +def test_a_run_that_keeps_printing_is_not_cut_off(run): + """A 10-minute run that prints every 20 s finishes under a 30 s idle limit.""" + lines = [(t, stdout_message(f"step {t}\n")) for t in range(20, 601, 20)] + mock_kc = MockKernelClient([*lines, (601, status_message("idle"))]) + outputs, _ = run(mock_kc, timeout=30.0) - # Verify streaming hook was called - assert len(streamed_outputs) == 2 - assert streamed_outputs[0]["text"] == "1\n" - assert streamed_outputs[1]["text"] == "2\n" + assert "".join(o["text"] for o in outputs).count("step") == 30 -def test_runtime_execute_code_streaming_error_synthesis(): - runtime = ColabRuntime("http://url", "token") - mock_client = MagicMock() +def test_silence_longer_than_the_timeout_raises(run): + mock_kc = MockKernelClient( + [(5, stdout_message("loading\n")), (100, status_message("idle"))] + ) + with pytest.raises(TimeoutError, match="Timeout waiting for output.*--timeout"): + run(mock_kc, timeout=30.0) + # 30 s after the last message, not 30 s after the start. + assert 35 <= mock_kc.now <= 36 - # Simulate an error reply but NO error output message - mock_client.execute_interactive.return_value = { - "content": { - "status": "error", - "ename": "RuntimeError", - "evalue": "something went wrong", - "traceback": ["tb line 1"], - } - } - runtime._kernel_client = mock_client - streamed_outputs = [] - outputs = runtime.execute_code( - "fail()", output_hook=lambda o: streamed_outputs.append(o) +@pytest.mark.parametrize("timeout", [None, 0]) +def test_no_limit_waits_through_any_silence(run, timeout): + mock_kc = MockKernelClient( + [(3600, stdout_message("done\n")), (3600.5, status_message("idle"))] ) + outputs, _ = run(mock_kc, timeout=timeout) + assert outputs[0]["text"] == "done\n" - # Final outputs should include synthesized error - assert len(outputs) == 1 - assert outputs[0]["output_type"] == "error" - assert outputs[0]["ename"] == "RuntimeError" - # Note: synthesized error is added AFTER execute_interactive returns, - # so it won't be in streamed_outputs unless we specifically add logic for it. - # Currently it's only in the returned list. - assert len(streamed_outputs) == 0 +def test_the_reply_may_arrive_after_the_idle_status(run): + mock_kc = MockKernelClient( + [(5, stdout_message("x\n")), (6, status_message("idle"))], + shell=[(26, reply_message())], + ) + outputs, _ = run(mock_kc, timeout=30.0) + assert outputs[0]["text"] == "x\n" + + +def test_a_reply_that_never_arrives_times_out(run): + mock_kc = MockKernelClient([(5, status_message("idle"))], shell=[]) + with pytest.raises(TimeoutError, match="Timeout waiting for reply"): + run(mock_kc, timeout=30.0) + + +def test_messages_from_other_requests_are_ignored(run): + mock_kc = MockKernelClient( + [ + (1, status_message("idle", parent="req-0")), + (2, stdout_message("ours\n")), + (3, status_message("idle")), + ], + shell=[(1, reply_message("error", parent="req-0")), (3, reply_message())], + ) + outputs, _ = run(mock_kc, timeout=30.0) + assert [o["text"] for o in outputs] == ["ours\n"] + + +def test_output_from_other_requests_does_not_reset_the_timeout(run): + others = [(t, stdout_message("other\n", parent="req-0")) for t in range(5, 100, 5)] + mock_kc = MockKernelClient([*others, (100, status_message("idle"))]) + with pytest.raises(TimeoutError): + run(mock_kc, timeout=30.0) + assert 30 <= mock_kc.now <= 31 + + +def test_waiting_never_spins(run): + """Each wait blocks for up to a second, so a long run costs a few calls a second.""" + lines = [(t, stdout_message("tick\n")) for t in range(25, 301, 25)] + mock_kc = MockKernelClient([*lines, (301, status_message("idle"))]) + run(mock_kc, timeout=30.0) + assert 0 < len(mock_kc.waits) < 330 # about one wait per second over 300 s + + +def test_a_lost_connection_stops_the_wait(run): + mock_kc = MockKernelClient( + [(5, stdout_message("x\n")), (100, status_message("idle"))], disconnect_at=10 + ) + with pytest.raises(RuntimeError, match="Connection was lost"): + run(mock_kc, timeout=None) + assert mock_kc.now < 12 + + +@pytest.mark.parametrize("state", ["dead", "restarting"]) +def test_a_kernel_that_dies_or_restarts_ends_the_wait(run, state): + mock_kc = MockKernelClient( + [ + (5, stdout_message("x\n")), + (8, status_message(state, parent=None)), + (100, status_message("idle")), + ] + ) + with pytest.raises(RuntimeError, match=f"Kernel {state}"): + run(mock_kc, timeout=None)