[RSI, performance] perf(kernel): defer the event-loop import stack past the ready event

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.

prime-agent-runtime/src/rlm/bash.py

@@ -2,34 +2,29 @@
 
 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
 import time
 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

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,14 +1184,21 @@ def main() -> None:
     user_module.__dict__["__builtins__"] = __builtins__
     sys.modules["__main__"] = user_module
 
+    signal.signal(signal.SIGINT, _sigint_handler)
+    _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__))
     # A KeyboardInterrupt escaping a cell or background task stops
     # run_until_complete; the interrupt is already recorded, so resume serving.

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,27 @@ async def test_await_returns_result(self):
         awaited = await handle
         self.assertEqual(handle.poll(), awaited)
 
+    def test_handle_construction_binds_asyncio_without_event_loop(self):
+        # rlm.bash defers its asyncio import and binds it on first handle
+        # construction: a fresh interpreter (no event loop, no bash() call yet)
+        # must construct a handle and reap it without NameError.
+        code = (
+            "import rlm, sys, time\n"
+            "assert 'asyncio' not in sys.modules\n"
+            "handle = rlm.BashHandle('exit 7')\n"
+            "for _ in range(250):\n"
+            "    result = handle.poll()\n"
+            "    if result is not None:\n"
+            "        break\n"
+            "    time.sleep(0.02)\n"
+            "assert result is not None, 'handle never completed'\n"
+            "assert result.exit_code == 7\n"
+            "sys.exit(0)"
+        )
+        # The runtime test venv has the rlm package installed, so the child
+        # interpreter resolves it the same way this process does.
+        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 +622,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 +673,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

prime-agent-runtime/test/test_repl.py

@@ -137,6 +137,26 @@ 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_event_loop_stack(self):
+        # `import rlm` (the pre-ready boot path of a kernel) must stay lean:
+        # the asyncio stack loads after the ready event, not during package
+        # import. A regression here reintroduces the boot-time import cost.
+        code = (
+            "import rlm, sys; "
+            "assert 'asyncio' not in sys.modules, 'rlm import must defer asyncio'; "
+            "assert 'secrets' not in sys.modules, 'rlm import must defer secrets'; "
+            "sys.exit(0)"
+        )
+        env = {**os.environ, "PYTHONPATH": SRC + os.pathsep + os.environ.get("PYTHONPATH", "")}
+        subprocess.run([sys.executable, "-c", code], env=env, check=True, timeout=30)
+
+    def test_serving_kernel_loads_asyncio_after_ready(self):
+        # The deferral must not break serving: by the first executed cell the
+        # event loop stack is resident and drives cell execution as usual.
+        events = self.repl.execute("serving", "import sys\n'asyncio' in sys.modules")
+        self.assertEqual(one(events, "result")["text"], "True")
+        self.assertEqual(one(events, "done")["status"], "ok")
+
     def test_result_echo(self):
         events = self.repl.execute("a", "1+1")
         self.assertEqual(one(events, "result")["text"], "2")