diff --git a/api/background_process.py b/api/background_process.py index 8cf76dfd..52c3b92d 100644 --- a/api/background_process.py +++ b/api/background_process.py @@ -62,6 +62,9 @@ logger = logging.getLogger(__name__) _DRAIN_THREAD: Optional[threading.Thread] = None _DRAIN_STOP = threading.Event() +_PROCESS_RECOVERY_DONE = False +_PROCESS_CHECKPOINT_RECOVERED = False +_PROCESS_RECOVERY_LOCK = threading.Lock() _REAPER_THREAD: Optional[threading.Thread] = None _REAPER_STOP = threading.Event() @@ -1617,6 +1620,58 @@ def _drain_loop() -> None: logger.warning("bg_task_complete event handling failed", exc_info=True) +def recover_processes_for_webui(process_registry=None, get_session_fn=None) -> int: + """Recover core background processes and restore WebUI routing metadata. + + The core gateway performs this during gateway startup, but this WebUI host + previously started only the queue drain. That left checkpointed processes + invisible after a WebUI restart. + """ + global _PROCESS_CHECKPOINT_RECOVERED, _PROCESS_RECOVERY_DONE + if process_registry is None: + try: + from tools.process_registry import process_registry + except ImportError: + # Hermes Agent is optional in isolated WebUI/test environments. + # The drain loop already treats a missing registry as unavailable; + # startup recovery must preserve that fail-soft contract. + logger.debug("process recovery unavailable: Hermes Agent is not installed") + return 0 + if get_session_fn is None: + from api.models import get_session as get_session_fn + + with _PROCESS_RECOVERY_LOCK: + if _PROCESS_RECOVERY_DONE: + return 0 + + recovered = 0 + if not _PROCESS_CHECKPOINT_RECOVERED: + recovered = process_registry.recover_from_checkpoint() + _PROCESS_CHECKPOINT_RECOVERED = True + + for row in process_registry.list_sessions(): + process_id = str(row.get("session_id") or "") + if not process_id: + continue + try: + proc_session = process_registry.get(process_id) + session_key = str(getattr(proc_session, "session_key", "") or "") + if not session_key or get_session_fn(session_key, metadata_only=True) is None: + continue + except Exception: + logger.warning( + "Could not resolve recovered WebUI process %r", + process_id, + exc_info=True, + ) + continue + register_process_session(session_key, session_key) + _PROCESS_RECOVERY_DONE = True + if recovered: + logger.info("Recovered %d background process(es) for WebUI", recovered) + return recovered + + def register_process_session(session_key: str, session_id: str) -> None: """Bind a process-registry session_key to a WebUI session_id. @@ -1663,6 +1718,12 @@ def start_drain_thread() -> bool: with _THREAD_LIFECYCLE_LOCK: if _DRAIN_THREAD is not None and _DRAIN_THREAD.is_alive(): return False + try: + recover_processes_for_webui() + except Exception: + # Recovery is best-effort. A corrupt checkpoint or transient I/O + # error must not disable notifications for newly spawned tasks. + logger.warning("background process recovery failed", exc_info=True) _DRAIN_STOP.clear() _DRAIN_THREAD = threading.Thread( target=_drain_loop, diff --git a/tests/test_background_process_restart_recovery.py b/tests/test_background_process_restart_recovery.py new file mode 100644 index 00000000..0ab76e36 --- /dev/null +++ b/tests/test_background_process_restart_recovery.py @@ -0,0 +1,161 @@ +"""Regression coverage for notify_on_complete across a WebUI restart.""" + +import threading +import time +from types import SimpleNamespace + +import pytest + +from api import background_process as bp + + +class _FakeThread: + def __init__(self, *args, **kwargs): + self.started = False + + def is_alive(self): + return self.started + + def start(self): + self.started = True + + +class _FakeProcessSession: + def __init__(self, session_key): + self.session_key = session_key + + +def test_start_drain_thread_invokes_recovery(monkeypatch): + calls = [] + monkeypatch.setattr(bp, "_DRAIN_THREAD", None) + monkeypatch.setattr(bp, "recover_processes_for_webui", lambda: calls.append("recover") or 0) + monkeypatch.setattr(bp.threading, "Thread", _FakeThread) + + assert bp.start_drain_thread() is True + assert calls == ["recover"] + + +def test_start_drain_thread_survives_recovery_failure(monkeypatch): + def fail_recovery(): + raise OSError("corrupt checkpoint") + + monkeypatch.setattr(bp, "_DRAIN_THREAD", None) + monkeypatch.setattr(bp, "recover_processes_for_webui", fail_recovery) + monkeypatch.setattr(bp.threading, "Thread", _FakeThread) + + assert bp.start_drain_thread() is True + assert bp._DRAIN_THREAD is not None + assert bp._DRAIN_THREAD.is_alive() + + +def test_recovery_runs_once_and_rebuilds_session_mapping(monkeypatch): + calls = {"recover": 0, "registered": []} + + class FakeRegistry: + def recover_from_checkpoint(self): + calls["recover"] += 1 + return 1 + + def list_sessions(self): + return [{ + "session_id": "proc_recovered", + "detached": True, + }] + + def get(self, process_id): + assert process_id == "proc_recovered" + return _FakeProcessSession("webui-session") + + fake_registry = FakeRegistry() + monkeypatch.setattr(bp, "_PROCESS_CHECKPOINT_RECOVERED", False) + monkeypatch.setattr(bp, "_PROCESS_RECOVERY_DONE", False) + monkeypatch.setattr( + bp, + "register_process_session", + lambda key, sid: calls["registered"].append((key, sid)), + ) + + def get_session(sid, metadata_only=False): + return SimpleNamespace(id=sid) + + assert bp.recover_processes_for_webui(fake_registry, get_session) == 1 + assert bp.recover_processes_for_webui(fake_registry, get_session) == 0 + assert calls == { + "recover": 1, + "registered": [("webui-session", "webui-session")], + } + + +def test_partial_recovery_retry_does_not_repeat_checkpoint_adoption(monkeypatch): + calls = {"recover": 0, "list": 0} + + class FlakyRegistry: + def recover_from_checkpoint(self): + calls["recover"] += 1 + return 1 + + def list_sessions(self): + calls["list"] += 1 + if calls["list"] == 1: + raise OSError("transient list failure") + return [] + + registry = FlakyRegistry() + monkeypatch.setattr(bp, "_PROCESS_CHECKPOINT_RECOVERED", False) + monkeypatch.setattr(bp, "_PROCESS_RECOVERY_DONE", False) + + with pytest.raises(OSError, match="transient list failure"): + bp.recover_processes_for_webui(registry, lambda *_args, **_kwargs: None) + + assert bp.recover_processes_for_webui(registry, lambda *_args, **_kwargs: None) == 0 + assert calls == {"recover": 1, "list": 2} + + +def test_concurrent_direct_recovery_runs_once(monkeypatch): + calls = {"recover": 0} + + class FakeRegistry: + def recover_from_checkpoint(self): + time.sleep(0.02) + calls["recover"] += 1 + return 1 + + def list_sessions(self): + return [] + + monkeypatch.setattr(bp, "_PROCESS_CHECKPOINT_RECOVERED", False) + monkeypatch.setattr(bp, "_PROCESS_RECOVERY_DONE", False) + registry = FakeRegistry() + barrier = threading.Barrier(8) + results = [] + + def recover(): + barrier.wait() + results.append( + bp.recover_processes_for_webui(registry, lambda *_args, **_kwargs: None) + ) + + workers = [threading.Thread(target=recover) for _ in range(8)] + for worker in workers: + worker.start() + for worker in workers: + worker.join(timeout=5) + + assert calls["recover"] == 1 + assert sorted(results) == [0] * 7 + [1] + + +def test_recovery_is_fail_soft_without_agent(monkeypatch): + real_import = __import__ + + def fake_import(name, *args, **kwargs): + if name == "tools.process_registry": + raise ImportError("Hermes Agent not installed") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(bp, "_PROCESS_CHECKPOINT_RECOVERED", False) + monkeypatch.setattr(bp, "_PROCESS_RECOVERY_DONE", False) + monkeypatch.setattr("builtins.__import__", fake_import) + + assert bp.recover_processes_for_webui() == 0 + assert bp._PROCESS_RECOVERY_DONE is False \ No newline at end of file