[hohieuai]

Fixes #183.

Summary

--timeout now works the way its design describes. It is an idle threshold: an execution fails only after that many seconds without output, so code that keeps printing is never cut off. While it waits, the local client no longer busy-loops. The flag, its 30-second default, and the TimeoutError message are unchanged.

Background

#38 added --timeout so users can "increase the idle execution threshold" for tasks that produce no output, and the docs describe it as letting long-running silent tasks run "without being prematurely killed". Runtime.execute_code, however, passed the value to jupyter-kernel-client's execute_interactive, which treats it as a deadline for the whole run. The measurements in #183 show the result:

• With the pinned 0.9.0: past the deadline, the client polls without waiting and holds a CPU core at about 100% until the run ends. The deadline never fires.
• With 1.0.2, which fixes that loop (Fix interactive execution timeout spin datalayer/jupyter-kernel-client#41): the run fails with TimeoutError after 30 seconds, even while output is flowing.

Changes

• runtime.py: for streaming execution without stdin (colab run, colab exec, colab repl), execute_code now waits for kernel messages itself, through the websocket client's public API (execute, get_iopub_msg, get_shell_msg).
  • It waits in blocking slices of up to one second.
  • Each message from this execution resets the idle clock. Messages from other requests don't count.
  • With no message for timeout seconds, it raises the same TimeoutError("Timeout waiting for output"). None waits without a limit.
  • It raises RuntimeError("Connection was lost.") if the websocket drops, which the old loop only noticed once a message arrived.
• Unchanged paths: executions that use stdin (the auth and Drive automations, with their explicit 120-second timeout) and the non-streaming execute path keep using the library as before.
• Tests:
  • The existing tests keep their names and what they assert: test_colab_runtime_execute_interactive_with_timeout and the two tests in test_streaming.py. They now drive MockKernelClient (tests/mock_kernel.py), a mock of the websocket client that replays scripted kernel messages on a simulated clock, instead of mocking execute_interactive, because the streaming path no longer calls it. For example, the timeout test still checks that timeout=600 reaches that branch. A 599-second silence completes, and a 601-second silence raises TimeoutError.
  • New tests in test_streaming.py cover:
    • a 10-minute run that prints every 20 seconds and is not cut off;
    • a timeout 30 seconds after the last message, not after the start;
    • no limit with None;
    • waits that are always bounded and non-zero;
    • messages from other requests, which don't count as output;
    • a lost connection.
• docs/: dated log entries in both design docs.

Testing

• uv run pytest: 366 passed, 19 skipped. Ruff check and format pass on the changed files.

• I ran the branch on real CPU runtimes with jupyter-kernel-client 0.8.0 (Google's fork, via tool.uv.sources), 0.9.0 (the PyPI pin), and 1.0.2:

  Case	--timeout	Result
  colab run, prints every 5 s for 60 s (0.8.0, 0.9.0, 1.0.2)	default (30)	12/12 steps, exit 0. The client used 0.0–0.1% CPU at 40 s (before: 98–99% with 0.9.0, and a failure after 6 steps with 1.0.2)
  colab run, silent for 45 s	default (30)	TimeoutError: Timeout waiting for output, exit 1, session stopped
  colab run, silent for 45 s	60	Completes, exit 0
  colab exec on an existing session, prints every 5 s for 60 s	default (30)	Completes, exit 0, 0.0% CPU, 0.28 s of CPU time in 71 s
  colab exec, silent for 45 s	default (30)	TimeoutError: Timeout waiting for output, exit 1
  Ctrl-C 40 s into a silent colab run	600	Exits 1.8 s after SIGINT, exit code 130, session stopped

Notes

• Default --timeout 30 caps the whole run: every longer colab run busy-loops a CPU core today, and fails once jupyter-kernel-client is upgraded #183 first suggested removing the default timeout. This PR keeps the default and the flag, and makes them behave as feat: Expose --timeout flag for exec and run commands #38 and the docs describe instead.
• With this change, it is safe to raise the jupyter-kernel-client pin to >=1.0.2 later. That is left out of this PR because it also touches tool.uv.sources.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[hohieuai]

Superseded by #185, which contains the same fix as a single commit with an updated description.