diff --git a/packages/coding-agent/.changes/kernel-defer-event-loop-imports.md b/packages/coding-agent/.changes/kernel-defer-event-loop-imports.md new file mode 100644 index 0000000000..b69589d24f --- /dev/null +++ b/packages/coding-agent/.changes/kernel-defer-event-loop-imports.md @@ -0,0 +1 @@ +- Reduced Python kernel startup time by deferring the event-loop import stack (asyncio plus the shell tool's heavy stdlib imports) until after the ready event, with no protocol or behavior changes. diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 34e7700bf6..82b9e93b29 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -2,17 +2,11 @@ from __future__ import annotations -import asyncio -import atexit import functools import json import os -import secrets -import selectors -import shutil import signal import socket -import struct import subprocess import sys import threading @@ -20,16 +14,17 @@ from collections import deque from collections.abc import Callable, Generator from dataclasses import dataclass -from datetime import datetime, timezone from typing import Any, cast from . import _winjob -_IS_POSIX = os.name == "posix" +# Boot-lean imports: asyncio, secrets, shutil, datetime, selectors, struct, +# fcntl/termios, and atexit load on first use below so `import rlm` (and with +# it the kernel's pre-ready startup path) stays small. asyncio is bound onto +# this module's globals by BashHandle.__init__ before any code path here can +# touch it; every other user imports inside the function that needs it. -if _IS_POSIX: - import fcntl - import termios +_IS_POSIX = os.name == "posix" _HEAD_CAP = 512 * 1024 _TAIL_CAP = 3 * 512 * 1024 @@ -236,6 +231,12 @@ class BashHandle: """ def __init__(self, command: str) -> None: + # Every asyncio use in this module runs on a handle path (bash() is the + # only constructor), so bind the module global here, before + # _schedule_background_completion_notice or any await can run. + global asyncio + import asyncio + self.command = command completion_context = _current_cell_completion_context() self._creating_cell_finished = completion_context[0] if completion_context else None @@ -271,6 +272,8 @@ def __init__(self, command: str) -> None: self._completion_marker: bytes | None = None status_write = -1 if _IS_POSIX: + import secrets + # Full-duplex status channel: the child end rides in as stdin (fd 0) # and the script remaps it to _STATUS_FD before swapping in /dev/null # (dash rejects multi-digit fds in redirections at parse time). The @@ -421,6 +424,8 @@ def _force_kill(self) -> None: _signal_group(self._pid, signal.SIGKILL) def _pump(self) -> None: + import selectors + stdout = self._proc.stdout assert stdout is not None if not _IS_POSIX: @@ -573,6 +578,8 @@ def _reap_group(self) -> bool: def _read_status(self) -> int | None: if self._status_read < 0: return None + import selectors + try: # DefaultSelector (kqueue/epoll) instead of select(): select() rejects # fds >= FD_SETSIZE (1024) even when the process fd limit is higher. @@ -619,6 +626,9 @@ def _pipe_pending(self) -> bool: # quiescence heuristic (best-effort parity). if not _IS_POSIX or self._eof.is_set(): return False + import fcntl + import struct + import termios stdout = self._proc.stdout if stdout is None: return False @@ -679,6 +689,7 @@ def _schedule_background_completion_notice(self) -> None: except RuntimeError: return from . import repl + import secrets activity = {"id": secrets.token_hex(16), "pid": self._pid, "active": True} # Publish synchronously before bash() returns and the creating cell can end. @@ -959,6 +970,8 @@ def bash(command: str) -> BashHandle: def _shell() -> str: + import shutil + # Read per call so env changes made in the REPL apply to later commands. override = os.environ.get("PRIME_AGENT_BASH_SHELL") if override: @@ -985,6 +998,8 @@ def _with_prefix(command: str) -> str: def _fence_printf() -> str: + import shutil + # `\command -p printf` defeats alias expansion but not a user-defined shell # function named `command`, which would swallow both fence frames and leave # the await hanging until the shell dies (wedged behind background jobs). A @@ -1130,6 +1145,8 @@ def _record_journal(pid: int, active: bool) -> bool: # Returns False only when the journal is configured but enrollment failed; # active-record callers must then fail closed. Active records always carry # a processStartId so host reaping stays identity-verified. + from datetime import datetime, timezone + path = os.environ.get("PRIME_AGENT_INTERNAL_ORPHAN_PROCESS_JOURNAL") owner = os.environ.get("PRIME_AGENT_KERNEL_OWNER_PID") if not path or not owner: @@ -1197,6 +1214,8 @@ def _kill_live_handles() -> None: def _install_shutdown_hook() -> None: global _hook_installed + import atexit + with _hook_lock: if _hook_installed: return diff --git a/prime-agent-runtime/src/rlm/repl.py b/prime-agent-runtime/src/rlm/repl.py index 413b2a83ac..8f888c1c53 100644 --- a/prime-agent-runtime/src/rlm/repl.py +++ b/prime-agent-runtime/src/rlm/repl.py @@ -8,7 +8,6 @@ from __future__ import annotations import ast -import asyncio import codecs import contextvars import ctypes @@ -20,7 +19,6 @@ import platform import signal import sys -import tempfile import threading import time import traceback @@ -49,6 +47,8 @@ class _CellExecution: def __init__(self) -> None: + import asyncio + self.finished = asyncio.Event() self.owner: asyncio.Task[Any] | None = None @@ -119,6 +119,8 @@ def current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] def active_cell_task() -> asyncio.Task[Any] | None: """The cell body task executing right now, or None between cells (global state, not the cell contextvar — detached tasks keep stale context copies).""" + import asyncio + with _interrupt_lock: task = _active["task"] return task if isinstance(task, asyncio.Task) and not task.done() else None @@ -331,6 +333,10 @@ def _consume_task_exception(task: asyncio.Task[Any]) -> None: def _sigint_handler(signum: int, frame: types.FrameType | None) -> None: + # asyncio loads by the time any task can be active (main() imports it), so + # this is a cached sys.modules hit even inside the signal handler. + import asyncio + global _handoff_interrupted task = _active["task"] # No lock (the main thread may hold it): the rid equality revalidates the @@ -529,6 +535,8 @@ async def _run_codes(codes: list[types.CodeType], ns: dict[str, Any]) -> Any: async def _run_guarded(task: asyncio.Task[Any], rid: str) -> tuple[str, Any, dict[str, Any] | None]: """Await a request task; returns (status, value, error event or None).""" + import asyncio + with _interrupt_lock: _active["interrupted"] = False _active["rid"] = rid @@ -644,6 +652,7 @@ def _snapshot_state( committed: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: import datetime + import tempfile try: import dill @@ -851,6 +860,8 @@ def _restore_state( async def _handle_state(req: dict[str, Any], ns: dict[str, Any]) -> None: """Run snapshot/restore as an interruptible task and reply in the done event.""" + import asyncio + rid = req["id"] committed: list[dict[str, Any]] = [] @@ -1173,15 +1184,25 @@ def main() -> None: user_module.__dict__["__builtins__"] = __builtins__ sys.modules["__main__"] = user_module + _send({"event": "ready", "protocol": PROTOCOL_VERSION, "python": platform.python_version()}) + + # The event-loop stack (asyncio plus its ssl, concurrent.futures, and + # logging imports) is the heaviest part of this module's boot chain; load + # it after the ready event so kernel startup stays lean. The loop, reader + # thread, and serve task all come up here before the host's first request + # can be served, and every function that references asyncio runs only + # after this point. + import asyncio _loop = asyncio.new_event_loop() asyncio.set_event_loop(_loop) queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() - signal.signal(signal.SIGINT, _sigint_handler) threading.Thread(target=_read_requests, args=(stdin_fd, queue), daemon=True).start() - _send({"event": "ready", "protocol": PROTOCOL_VERSION, "python": platform.python_version()}) - _serve_task = _loop.create_task(_serve(queue, user_module.__dict__)) + # _sigint_handler has no task to target before serving starts, so installing + # it earlier would silently swallow a Ctrl-C during this boot window; the + # default handler must stay in charge until the loop and serve task exist. + signal.signal(signal.SIGINT, _sigint_handler) # A KeyboardInterrupt escaping a cell or background task stops # run_until_complete; the interrupt is already recorded, so resume serving. while not _serve_task.done(): diff --git a/prime-agent-runtime/test/test_bash.py b/prime-agent-runtime/test/test_bash.py index f11c930b2e..6da47bc48b 100644 --- a/prime-agent-runtime/test/test_bash.py +++ b/prime-agent-runtime/test/test_bash.py @@ -4,6 +4,8 @@ import json import os import resource +import secrets +import shutil import signal import socket import subprocess @@ -56,6 +58,22 @@ async def test_await_returns_result(self): awaited = await handle self.assertEqual(handle.poll(), awaited) + def test_handle_construction_binds_asyncio_without_an_event_loop(self): + # rlm.bash defers its asyncio import and binds it on the first handle + # construction, so a fresh interpreter with no event loop reaps one. + code = ( + "import rlm, sys, time\n" + "assert 'asyncio' not in sys.modules\n" + "handle = rlm.BashHandle('exit 7')\n" + "for _ in range(250):\n" + " if handle.poll() is not None:\n" + " break\n" + " time.sleep(0.02)\n" + "assert handle.poll() is not None, 'handle never completed'\n" + "assert handle.poll().exit_code == 7" + ) + subprocess.run([sys.executable, "-c", code], check=True, timeout=30) + def test_construction_cleanup_uses_windows_signal_without_sigkill(self): failure = RuntimeError("task construction failed") loop = mock.Mock() @@ -599,8 +617,9 @@ async def test_windows_without_bash_raises_teaching_error(self): # Windows must raise without consulting PATH: a which() hit would be # the same repo-controlled-PATH hole the host-side resolution closed. with mock.patch.object(bash_module, "_IS_POSIX", False): + # rlm.bash imports shutil lazily, so patch the stdlib module itself. with mock.patch.object( - bash_module.shutil, "which", return_value=r"C:\evil\bash.exe" + shutil, "which", return_value=r"C:\evil\bash.exe" ) as which: with self.assertRaisesRegex(RuntimeError, "PRIME_AGENT_BASH_SHELL"): bash_module._shell() @@ -649,7 +668,8 @@ async def test_sentinel_like_output_and_echoed_wrapper_do_not_truncate(self): "if [ -r /proc/$$/cmdline ]; then cat /proc/$$/cmdline; fi\n" "printf '\\nafter-sentinel-lookalike\\n'" ) - with mock.patch.object(bash_module.secrets, "token_hex", return_value=token): + # rlm.bash imports secrets lazily, so patch the stdlib module itself. + with mock.patch.object(secrets, "token_hex", return_value=token): result = await asyncio.wait_for(bash(command), timeout=5) actual_marker = ( bash_module._COMPLETION_PREFIX + token.encode() + bash_module._COMPLETION_SUFFIX diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 27b3b45fd8..0c19e19f46 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -137,6 +137,33 @@ def test_ready_handshake_and_startup_time(self): print(f"\n[startup] spawn -> ready: {self.ready_ms:.0f} ms") self.assertLess(self.ready_ms, 500) + def test_import_rlm_defers_the_event_loop_stack(self): + # `import rlm` is the pre-ready boot path: the event loop stack must load + # after the ready event, not during package import, and serving must keep + # it resident. + env = {**os.environ, "PYTHONPATH": SRC + os.pathsep + os.environ.get("PYTHONPATH", "")} + code = "import rlm, sys; assert 'asyncio' not in sys.modules and 'secrets' not in sys.modules" + subprocess.run([sys.executable, "-c", code], env=env, check=True, timeout=30) + events = self.repl.execute("serving", "import sys\n'asyncio' in sys.modules") + self.assertEqual(one(events, "result")["text"], "True") + + def test_sigint_during_the_deferred_boot_stays_fatal(self): + # A fake `asyncio` parks the kernel inside the post-ready deferred + # import. The fifo open below returns only once the kernel is parked in + # that window, where an early _sigint_handler install swallowed the + # Ctrl-C, so only the default handler may be in charge there. + with tempfile.TemporaryDirectory() as tmp: + park = os.path.join(tmp, "deferred-boot-park") + os.mkfifo(park) + with open(os.path.join(tmp, "asyncio.py"), "w") as fake_asyncio: + fake_asyncio.write(f"import os\nos.read(os.open({park!r}, os.O_RDONLY), 1)\n") + repl = ReplProcess(env={"PYTHONPATH": tmp + os.pathsep + SRC}) + self.addCleanup(repl.close) + self.assertEqual(repl.ready()[0]["event"], "ready") + with open(park, "wb"): + os.kill(repl.proc.pid, signal.SIGINT) + self.assertNotEqual(repl.proc.wait(timeout=10), 0) + def test_result_echo(self): events = self.repl.execute("a", "1+1") self.assertEqual(one(events, "result")["text"], "2") @@ -379,19 +406,16 @@ def test_interrupt_without_pthread_kill_cancels_awaited_cell(self): self.assertEqual(error["ename"], "KeyboardInterrupt") self.assertEqual(one(events, "done")["status"], "error") - def test_stdout_buffer_write_works_and_surfaces_as_null(self): - # Libraries write bytes via sys.stdout.buffer; the tagged writer must - # expose a working buffer whose bytes surface (null-attributed) before done. + def test_stdout_buffer_write_surfaces_as_null_and_rejects_int(self): + # Libraries write bytes via sys.stdout.buffer: the tagged writer exposes a + # working buffer whose bytes surface (null-attributed) before done, and it + # raises TypeError for ints (bytes(5) would emit five NULs). events = self.repl.execute( "bufw", "import sys\nsys.stdout.buffer.write(b'buffer-bytes\\n')\nsys.stdout.buffer.flush()" ) - self.assertEqual(one(events, "done")["status"], "ok") buffered = next(e for e in events if e.get("event") == "stdout" and "buffer-bytes" in e["text"]) self.assertIsNone(buffered["id"]) self.assertLess(events.index(buffered), events.index(one(events, "done"))) - - def test_stdout_buffer_write_rejects_int(self): - # A real stdout.buffer raises TypeError for ints; bytes(5) would emit five NULs. events = self.repl.execute("bufint", "import sys\nsys.stdout.buffer.write(5)") self.assertEqual(one(events, "error")["ename"], "TypeError") self.assertEqual(one(events, "done")["status"], "error") @@ -1148,7 +1172,7 @@ def test_malformed_request_line(self): self.assertEqual(one(events, "result")["text"], "'alive'") def test_list_names(self): - self.repl.execute("ln1", "alpha = 1\ndef helper(n):\n return n\n_hidden = 2\nrlm = object()") + self.repl.execute("ln1", "alpha = 1\ndef helper(n):\n return n\n_hidden = 2\nrlm = object()\nglobals()[1] = 2") self.repl.send({"type": "list_names", "id": "ln2"}) done = one(self.repl.until_done("ln2"), "done") self.assertEqual(done["status"], "ok") @@ -1156,17 +1180,8 @@ def test_list_names(self): self.assertIn("helper", done["names"]) self.assertNotIn("_hidden", done["names"]) self.assertNotIn("rlm", done["names"]) - self.assertEqual(done["names"], sorted(done["names"])) - - def test_list_names_skips_non_string_keys(self): - self.repl.execute("lnk1", "globals()[1] = 1\nbeta = 2") - self.repl.send({"type": "list_names", "id": "lnk2"}) - done = one(self.repl.until_done("lnk2"), "done") - self.assertEqual(done["status"], "ok") - self.assertIn("beta", done["names"]) self.assertNotIn(1, done["names"]) - events = self.repl.execute("lnk3", "'alive'") - self.assertEqual(one(events, "result")["text"], "'alive'") + self.assertEqual(done["names"], sorted(done["names"])) def test_host_request_round_trip(self): code = "\n".join( @@ -1250,11 +1265,6 @@ def test_typed_host_request_rejects_unexpected_envelope_status(self): ) self.assertEqual(one(events, "done")["status"], "error") - def test_host_reply_for_unknown_id_dropped(self): - self.repl.send({"type": "host_reply", "id": "no-such-request", "data": {"status": "ok"}}) - events = self.repl.execute("ok", "'alive'") - self.assertEqual(one(events, "result")["text"], "'alive'") - def test_host_request_cancelled_cell_drops_pending_future(self): code = "\n".join( [ @@ -1293,9 +1303,6 @@ def test_display_from_detached_task_keeps_cell_id(self): self.assertEqual(display["id"], "det") self.assertEqual(display["data"], {"text/plain": "late"}) - def test_shutdown_clean_exit(self): - self.assertEqual(self.repl.shutdown(), 0) - def test_shutdown_after_mcp_import_exits_cleanly(self): events = self.repl.execute("mcp-import", "import rlm.mcp") self.assertEqual(one(events, "done")["status"], "ok") @@ -2197,11 +2204,6 @@ def test_near_cap_payload_skips_tail_instead_of_failing_snapshot(self): with open(self.path, "rb") as fh: self.assertEqual(list(dill.load(fh)), ["a"]) - def test_zero_size_cap_writes_no_empty_payload_overhead(self): - result = self._snap({}, max_bytes=0, max_variable_bytes=0) - self.assertEqual(result, {"error": "write failed: snapshot exceeds aggregate snapshot size cap"}) - self.assertEqual(os.listdir(self.dir), []) - def test_manifest_write_failure_preserves_prior_pair(self): old_payload, old_manifest = self._old_pair() ns = {"keep": 1, "big": b"x" * 100_000}