"""Background task lifecycle for the interactive REPL runtime.""" from __future__ import annotations import asyncio import logging from collections.abc import Callable, Coroutine from typing import Any from rich.console import Console from core.domain.alerts import inbox as _alert_inbox from surfaces.interactive_shell.runtime.background.runner import drain_background_notices from surfaces.interactive_shell.runtime.core.state import ReplState, SpinnerState from surfaces.interactive_shell.session import Session from surfaces.interactive_shell.ui.alerts import drain_and_render_incoming log = logging.getLogger(__name__) class BackgroundTaskManager: """Start background workers and drain their user-visible output.""" def __init__( self, session: Session, state: ReplState, spinner: SpinnerState, inbox: _alert_inbox.AlertInbox | None, prompt_invalidator: Callable[[], None], ) -> None: self.session = session self.state = state self.spinner = spinner self.inbox = inbox self.prompt_invalidator = prompt_invalidator self.tasks: list[tuple[str, asyncio.Task[None]]] = [] self._loop: asyncio.AbstractEventLoop | None = None self._sampler_started = False def start_all( self, processor_coro: Callable[[], Coroutine[Any, Any, None]], ) -> list[tuple[str, asyncio.Task[None]]]: # The fleet sampler (and its psutil dependency) is intentionally NOT # started here: local-agent monitoring only runs once the user actually # opens /fleet, via ``ensure_fleet_sampler_started``. self._loop = asyncio.get_running_loop() self.tasks = [ ("processor", asyncio.create_task(processor_coro())), ("alert watcher", asyncio.create_task(self._alert_watcher())), ("spinner ticker", asyncio.create_task(self._spinner_ticker())), ] return self.tasks def ensure_fleet_sampler_started(self) -> None: """Start the fleet sampler on demand (first live ``/fleet`` use). Safe to call from any thread: a shell turn (and thus the ``/fleet`` handler) runs in a worker thread via ``asyncio.to_thread``, so sampler task creation is marshalled back onto the REPL event loop. Idempotent. """ loop = self._loop if loop is None: return loop.call_soon_threadsafe(self._start_sampler_task) def _start_sampler_task(self) -> None: if self._sampler_started: return # Imported lazily so base REPL startup does not pull the sampler + # psutil into the import path. from tools.system.fleet_monitoring.sampler import start_sampler self._sampler_started = True self.tasks.append(("sampler", start_sampler())) def drain_turn_start_output(self, console: Console) -> None: if self.inbox is not None: try: drain_and_render_incoming(self.session, console, self.inbox) except Exception as exc: log.warning("Error draining alerts at turn start: %s", exc) try: drain_background_notices(self.session, console) except Exception as exc: log.warning("Error draining background notices at turn start: %s", exc) async def _alert_watcher(self) -> None: if self.inbox is None: return alert_console = Console( highlight=False, force_terminal=True, color_system="truecolor", legacy_windows=False, ) drain_and_render_incoming(self.session, alert_console, self.inbox) while not self.state.exit_requested: try: await asyncio.to_thread(self.inbox.pending_event.wait, timeout=1) except asyncio.CancelledError: return try: drain_and_render_incoming(self.session, alert_console, self.inbox) except Exception as exc: log.warning("Error draining incoming alerts: %s", exc) async def _spinner_ticker(self) -> None: # prompt_async's refresh_interval alone is not guaranteed to drive # visible prompt redraws while patch_stdout(raw=True) is active and # the LLM stream is writing rapidly. This task explicitly invalidates # the prompt at 100 ms intervals so the braille glyph cycles smoothly. tick_s = 0.1 was_streaming = False while not self.state.exit_requested: try: await asyncio.sleep(tick_s) except asyncio.CancelledError: return streaming = self.spinner.streaming # Invalidate while streaming, plus one extra tick on the # streaming->idle edge so the prompt repaints without the stale # spinner/phase label instead of waiting for unrelated I/O. if streaming or was_streaming: self.prompt_invalidator() was_streaming = streaming