fix(wakeup): recover terminal process completions after restart (#6287)

Recover checkpointed core background processes and rebuild PROCESS_SESSION_INDEX on WebUI startup, so an ordinary terminal(background=True, notify_on_complete=True) proc_* completion that outlives a WebUI restart can still wake its original session. Complements #6283 (which covered async_delegation completions). Thanks @allenliang2022.

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,

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