fix(cli): keep CLI cache followers from hanging or dropping rows (#4842)

api/models.py

@@ -48,6 +48,7 @@ _CLI_SESSIONS_CACHE_LOCK = threading.Lock()
 _CLI_SESSIONS_CACHE_INFLIGHT: "dict[tuple, threading.Event]" = {}
 _CLI_SESSIONS_CACHE_INVALIDATION_VERSION = 0
 _CLI_SESSIONS_CACHE = {}
+_CLI_SESSIONS_CACHE_WAIT_SECONDS = 0.25
 # Event waits that keep stale rows visible while a rebuild is in flight.
 _CLI_SESSIONS_CACHE_STALE_WAIT_SECONDS = 0.10
 
@@ -4264,6 +4265,54 @@ def _cache_cli_sessions_if_current(
     return True
 
 
+def _copy_fresh_cli_sessions_cache_entry(cache_key: tuple):
+    with _CLI_SESSIONS_CACHE_LOCK:
+        cached_entry = _CLI_SESSIONS_CACHE.get(cache_key)
+        if cached_entry is None:
+            return None
+        if len(cached_entry) == 3:
+            cached_expires_at, cached_stamp, cached_sessions = cached_entry
+        else:
+            cached_expires_at, cached_sessions = cached_entry
+            cached_stamp = _CLI_SESSIONS_CACHE_INVALIDATION_VERSION
+        if cached_stamp != _CLI_SESSIONS_CACHE_INVALIDATION_VERSION:
+            _CLI_SESSIONS_CACHE.pop(cache_key, None)
+            return None
+        if cached_expires_at <= time.monotonic():
+            return None
+        return _copy_cli_sessions(cached_sessions)
+
+
+def _load_and_cache_cli_sessions(
+    *,
+    cache_key: tuple,
+    ttl: float,
+    invalidation_stamp: int,
+    load_sessions,
+    stale_sessions,
+    stale_stamp,
+    all_profiles: bool,
+    db_path,
+) -> list:
+    try:
+        sessions = load_sessions()
+    except Exception as _cli_err:
+        logger.warning(
+            "get_cli_sessions() failed — check state.db schema or path (%s): %s",
+            "all profiles" if all_profiles else db_path, _cli_err,
+        )
+        if stale_sessions is not None and stale_stamp == _cli_sessions_cache_invalidation_stamp():
+            return stale_sessions
+        return []
+    _cache_cli_sessions_if_current(
+        cache_key,
+        ttl,
+        invalidation_stamp,
+        sessions,
+    )
+    return _copy_cli_sessions(sessions)
+
+
 def _reload_cli_sessions_after_inflight(
     *,
     cache_key: tuple,
@@ -4274,47 +4323,50 @@ def _reload_cli_sessions_after_inflight(
     all_profiles: bool,
     db_path: str,
 ) -> list:
-    event, is_owner = _cli_sessions_cache_claim_rebuild(cache_key)
-    if not is_owner:
+    while True:
+        event, is_owner = _cli_sessions_cache_claim_rebuild(cache_key)
+        if is_owner:
+            break
+        wait_finished = False
         try:
-            event.wait()
+            wait_finished = bool(
+                event.wait(
+                    _CLI_SESSIONS_CACHE_STALE_WAIT_SECONDS
+                    if stale_sessions is not None
+                    else _CLI_SESSIONS_CACHE_WAIT_SECONDS
+                )
+            )
         except Exception:
             pass
-        with _CLI_SESSIONS_CACHE_LOCK:
-            cached_entry = _CLI_SESSIONS_CACHE.get(cache_key)
-            if cached_entry is not None:
-                if len(cached_entry) == 3:
-                    cached_expires_at, cached_stamp, cached_sessions = cached_entry
-                else:
-                    cached_expires_at, cached_sessions = cached_entry
-                    cached_stamp = _CLI_SESSIONS_CACHE_INVALIDATION_VERSION
-                if (
-                    cached_stamp == _CLI_SESSIONS_CACHE_INVALIDATION_VERSION
-                    and cached_expires_at > time.monotonic()
-                ):
-                    return _copy_cli_sessions(cached_sessions)
+        cached_sessions = _copy_fresh_cli_sessions_cache_entry(cache_key)
+        if cached_sessions is not None:
+            return cached_sessions
         if stale_sessions is not None and stale_stamp == _cli_sessions_cache_invalidation_stamp():
             return stale_sessions
-        return []
-    try:
-        fallback_invalidation_stamp = _cli_sessions_cache_invalidation_stamp()
-        try:
-            sessions = load_sessions()
-        except Exception as _cli_err:
-            logger.warning(
-                "get_cli_sessions() failed — check state.db schema or path (%s): %s",
-                "all profiles" if all_profiles else db_path, _cli_err,
+        if not wait_finished:
+            fallback_invalidation_stamp = _cli_sessions_cache_invalidation_stamp()
+            return _load_and_cache_cli_sessions(
+                cache_key=cache_key,
+                ttl=ttl,
+                invalidation_stamp=fallback_invalidation_stamp,
+                load_sessions=load_sessions,
+                stale_sessions=stale_sessions,
+                stale_stamp=stale_stamp,
+                all_profiles=all_profiles,
+                db_path=db_path,
             )
-            if stale_sessions is not None and stale_stamp == _cli_sessions_cache_invalidation_stamp():
-                return stale_sessions
-            return []
-        _cache_cli_sessions_if_current(
-            cache_key,
-            ttl,
-            fallback_invalidation_stamp,
-            sessions,
+    try:
+        invalidation_stamp = _cli_sessions_cache_invalidation_stamp()
+        return _load_and_cache_cli_sessions(
+            cache_key=cache_key,
+            ttl=ttl,
+            invalidation_stamp=invalidation_stamp,
+            load_sessions=load_sessions,
+            stale_sessions=stale_sessions,
+            stale_stamp=stale_stamp,
+            all_profiles=all_profiles,
+            db_path=db_path,
         )
-        return _copy_cli_sessions(sessions)
     finally:
         _cli_sessions_cache_done(cache_key, event)
 
@@ -4927,46 +4979,18 @@ def get_cli_sessions(source_filter=None, *, all_profiles: bool = False) -> list:
         if is_owner:
             try:
                 invalidation_stamp = _cli_sessions_cache_invalidation_stamp()
-                try:
-                    sessions = _load_sessions()
-                except Exception as _cli_err:
-                    logger.warning(
-                        "get_cli_sessions() failed — check state.db schema or path (%s): %s",
-                        "all profiles" if all_profiles else db_path, _cli_err,
-                    )
-                    return []
-                _cache_cli_sessions_if_current(
-                    cache_key,
-                    ttl,
-                    invalidation_stamp,
-                    sessions,
+                return _load_and_cache_cli_sessions(
+                    cache_key=cache_key,
+                    ttl=ttl,
+                    invalidation_stamp=invalidation_stamp,
+                    load_sessions=_load_sessions,
+                    stale_sessions=stale_sessions,
+                    stale_stamp=stale_stamp,
+                    all_profiles=all_profiles,
+                    db_path=db_path,
                 )
-                return _copy_cli_sessions(sessions)
             finally:
                 _cli_sessions_cache_done(cache_key, event)
-        try:
-            event.wait(
-                _CLI_SESSIONS_CACHE_STALE_WAIT_SECONDS
-                if stale_sessions is not None
-                else None
-            )
-        except Exception:
-            pass
-        with _CLI_SESSIONS_CACHE_LOCK:
-            cached_entry = _CLI_SESSIONS_CACHE.get(cache_key)
-            if cached_entry is not None:
-                if len(cached_entry) == 3:
-                    cached_expires_at, cached_stamp, cached_sessions = cached_entry
-                else:
-                    cached_expires_at, cached_sessions = cached_entry
-                    cached_stamp = _CLI_SESSIONS_CACHE_INVALIDATION_VERSION
-                if (
-                    cached_stamp == _CLI_SESSIONS_CACHE_INVALIDATION_VERSION
-                    and cached_expires_at > time.monotonic()
-                ):
-                    return _copy_cli_sessions(cached_sessions)
-        if stale_sessions is not None and stale_stamp == _cli_sessions_cache_invalidation_stamp():
-            return stale_sessions
         return _reload_cli_sessions_after_inflight(
             cache_key=cache_key,
             ttl=ttl,
@@ -4976,25 +5000,6 @@ def get_cli_sessions(source_filter=None, *, all_profiles: bool = False) -> list:
             all_profiles=all_profiles,
             db_path=db_path,
         )
-        fallback_invalidation_stamp = _cli_sessions_cache_invalidation_stamp()
-        try:
-            sessions = _load_sessions()
-        except Exception as _cli_err:
-            logger.warning(
-                "get_cli_sessions() failed — check state.db schema or path (%s): %s",
-                "all profiles" if all_profiles else db_path, _cli_err,
-            )
-            if stale_sessions is not None and stale_stamp == _cli_sessions_cache_invalidation_stamp():
-                return stale_sessions
-            return []
-        if _cli_sessions_cache_invalidation_stamp() == fallback_invalidation_stamp:
-            with _CLI_SESSIONS_CACHE_LOCK:
-                _CLI_SESSIONS_CACHE[cache_key] = (
-                    time.monotonic() + ttl,
-                    _CLI_SESSIONS_CACHE_INVALIDATION_VERSION,
-                    _copy_cli_sessions(sessions),
-                )
-        return _copy_cli_sessions(sessions)
 
     try:
         return _load_sessions()
@@ -5005,7 +5010,6 @@ def get_cli_sessions(source_filter=None, *, all_profiles: bool = False) -> list:
         )
         return []
 
-
 def _json_loads_if_string(value):
     if not isinstance(value, str):
         return value

tests/test_issue4842_cli_cache_streaming.py

@@ -153,6 +153,61 @@ def test_get_cli_sessions_cold_followers_join_single_rebuild(monkeypatch, tmp_pa
     assert results.get("follower") == [{"session_id": "fresh", "title": "fresh-row"}]
 
 
+def test_get_cli_sessions_cold_follower_times_out_to_independent_rebuild(monkeypatch, tmp_path):
+    hermes_home = tmp_path / "hermes"
+    hermes_home.mkdir()
+    monkeypatch.setattr(profiles, "get_active_hermes_home", lambda: str(hermes_home))
+    monkeypatch.setattr(profiles, "get_active_profile_name", lambda: "default")
+    models.clear_cli_sessions_cache()
+    monkeypatch.setattr(models, "_CLI_SESSIONS_CACHE_TTL_SECONDS", 60.0, raising=False)
+    monkeypatch.setattr(models, "_CLI_SESSIONS_CACHE_WAIT_SECONDS", 0.05, raising=False)
+
+    owner_started = threading.Event()
+    owner_block = threading.Event()
+    fallback_started = threading.Event()
+    load_count = {"value": 0}
+    load_count_lock = threading.Lock()
+    results = {}
+
+    def _blocking_loader(*_args, **_kwargs):
+        with load_count_lock:
+            load_count["value"] += 1
+            call_number = load_count["value"]
+        if call_number == 1:
+            owner_started.set()
+            owner_block.wait()
+        else:
+            fallback_started.set()
+        return [{"session_id": "fresh", "title": "fresh-row"}]
+
+    monkeypatch.setattr(models, "_load_cli_sessions_uncached", _blocking_loader)
+
+    owner = threading.Thread(
+        target=lambda: results.setdefault("owner", models.get_cli_sessions()),
+        daemon=True,
+    )
+    follower = threading.Thread(
+        target=lambda: results.setdefault("follower", models.get_cli_sessions()),
+        daemon=True,
+    )
+
+    owner.start()
+    assert owner_started.wait(1.0), "owner did not start"
+    follower.start()
+    assert fallback_started.wait(1.0), "follower did not start fallback rebuild"
+    follower.join(1.0)
+
+    assert not follower.is_alive()
+    assert owner.is_alive()
+    assert load_count["value"] == 2
+    assert results.get("follower") == [{"session_id": "fresh", "title": "fresh-row"}]
+
+    owner_block.set()
+    owner.join(2.0)
+    assert not owner.is_alive()
+    assert results.get("owner") == [{"session_id": "fresh", "title": "fresh-row"}]
+
+
 def test_get_cli_sessions_clear_during_rebuild_does_not_restore_stale_rows(monkeypatch, tmp_path):
     hermes_home = tmp_path / "hermes"
     hermes_home.mkdir()
@@ -268,6 +323,71 @@ def test_get_cli_sessions_clear_during_rebuild_preserves_joiners(monkeypatch, tm
     assert results.get("follower_b") == [{"session_id": "fresh", "title": "fresh-row"}]
 
 
+def test_get_cli_sessions_clear_during_rebuild_reclaims_after_invalidated_wait(monkeypatch, tmp_path):
+    hermes_home = tmp_path / "hermes"
+    hermes_home.mkdir()
+    monkeypatch.setattr(profiles, "get_active_hermes_home", lambda: str(hermes_home))
+    monkeypatch.setattr(profiles, "get_active_profile_name", lambda: "default")
+    models.clear_cli_sessions_cache()
+    monkeypatch.setattr(models, "_CLI_SESSIONS_CACHE_TTL_SECONDS", 60.0, raising=False)
+    monkeypatch.setattr(models, "_CLI_SESSIONS_CACHE_STALE_WAIT_SECONDS", 0.2, raising=False)
+
+    (_, _, _, cache_key) = models._resolve_cli_sessions_context(None)
+    cache_stamp = models._cli_sessions_cache_invalidation_stamp()
+    with models._CLI_SESSIONS_CACHE_LOCK:
+        models._CLI_SESSIONS_CACHE[cache_key] = (
+            time.monotonic() - 1.0,
+            cache_stamp,
+            [{"session_id": "stale", "title": "stale-row"}],
+        )
+
+    owner_started = threading.Event()
+    owner_block = threading.Event()
+    second_started = threading.Event()
+    load_count = {"value": 0}
+    load_count_lock = threading.Lock()
+    results = {}
+
+    def _blocking_loader(*_args, **_kwargs):
+        with load_count_lock:
+            load_count["value"] += 1
+            call_number = load_count["value"]
+        if call_number == 1:
+            owner_started.set()
+            owner_block.wait()
+            return [{"session_id": "owner", "title": "owner-row"}]
+        second_started.set()
+        return [{"session_id": "fresh", "title": "fresh-row"}]
+
+    monkeypatch.setattr(models, "_load_cli_sessions_uncached", _blocking_loader)
+
+    owner = threading.Thread(
+        target=lambda: results.setdefault("owner", models.get_cli_sessions()),
+        daemon=True,
+    )
+    follower = threading.Thread(
+        target=lambda: results.setdefault("follower", models.get_cli_sessions()),
+        daemon=True,
+    )
+
+    owner.start()
+    assert owner_started.wait(1.0), "owner did not enter rebuild"
+    follower.start()
+    time.sleep(0.05)
+    models.clear_cli_sessions_cache()
+    owner_block.set()
+
+    assert second_started.wait(1.0), "follower did not reclaim the rebuild"
+    owner.join(2.0)
+    follower.join(2.0)
+
+    assert not owner.is_alive()
+    assert not follower.is_alive()
+    assert load_count["value"] == 2
+    assert results.get("owner") == [{"session_id": "owner", "title": "owner-row"}]
+    assert results.get("follower") == [{"session_id": "fresh", "title": "fresh-row"}]
+
+
 def test_cache_cli_sessions_if_current_skips_stale_store(monkeypatch, tmp_path):
     hermes_home = tmp_path / "hermes"
     hermes_home.mkdir()