diff --git a/api/models.py b/api/models.py index c666fc1f..b9cd532c 100644 --- a/api/models.py +++ b/api/models.py @@ -1144,6 +1144,13 @@ def _append_journaled_partial_output( # so users do not see "reload to retry" prompts forever. _JOURNAL_RETRY_MAX_ATTEMPTS = 12 _JOURNAL_RETRY_GIVEUP_SECONDS = 24 * 3600 +_JOURNAL_RETRY_LOCKS: dict[str, threading.Lock] = {} +_JOURNAL_RETRY_LOCKS_GUARD = threading.Lock() + + +def _journal_retry_lock_for_sid(sid: str) -> threading.Lock: + with _JOURNAL_RETRY_LOCKS_GUARD: + return _JOURNAL_RETRY_LOCKS.setdefault(str(sid), threading.Lock()) def _build_recovery_marker_with_retry_hook( @@ -1231,6 +1238,21 @@ def _reorder_journal_tail_above_marker(session, marker_idx: int) -> None: session.messages = new_messages +def _try_retry_journal_recovery_in_place(session) -> bool: + sid = str(getattr(session, 'session_id', '') or '') + lock = _journal_retry_lock_for_sid(sid) + if not lock.acquire(blocking=False): + logger.debug("lazy journal-retry already running for session %s", sid) + return False + try: + return _retry_journal_recovery_in_place(session) + finally: + lock.release() + with _JOURNAL_RETRY_LOCKS_GUARD: + if _JOURNAL_RETRY_LOCKS.get(sid) is lock: + _JOURNAL_RETRY_LOCKS.pop(sid, None) + + def _retry_journal_recovery_in_place(session) -> bool: """Re-attempt run-journal recovery for the most recent pending marker. @@ -1630,7 +1652,7 @@ def get_session(sid, metadata_only=False): if cached is not None: if not metadata_only and _session_has_pending_journal_retry(cached): try: - _retry_journal_recovery_in_place(cached) + _try_retry_journal_recovery_in_place(cached) except Exception: logger.debug( "lazy journal-retry failed on cache hit for session %s", @@ -1658,7 +1680,7 @@ def get_session(sid, metadata_only=False): # chance to self-heal on this read. if not repaired and _session_has_pending_journal_retry(s): try: - _retry_journal_recovery_in_place(s) + _try_retry_journal_recovery_in_place(s) except Exception: logger.debug( "lazy journal-retry failed on cold load for session %s", diff --git a/tests/test_session_lost_response_regression.py b/tests/test_session_lost_response_regression.py index eb5c4db0..db129d45 100644 --- a/tests/test_session_lost_response_regression.py +++ b/tests/test_session_lost_response_regression.py @@ -15,6 +15,9 @@ The scenario this test pins down: assistant text/tools into the transcript in the correct chronological position. """ +from concurrent.futures import ThreadPoolExecutor +import threading + import pytest import api.models as models @@ -59,8 +62,10 @@ def _isolate_stream_state(): @pytest.fixture(autouse=True) def _isolate_agent_locks(): config.SESSION_AGENT_LOCKS.clear() + models._JOURNAL_RETRY_LOCKS.clear() yield config.SESSION_AGENT_LOCKS.clear() + models._JOURNAL_RETRY_LOCKS.clear() @pytest.fixture() @@ -97,6 +102,32 @@ def _make_dead_stream_session( return s +def _make_pending_retry_session(session_id: str, *, stream_id: str): + s = Session(session_id=session_id, title="Pending retry", messages=[ + {"role": "user", "content": "q", "timestamp": 1}, + { + "role": "assistant", + "content": models._INTERRUPTED_PENDING_RETRY_WORDING, + "timestamp": 2, + "_error": True, + "type": "interrupted", + "_pending_journal_recovery": True, + "_journal_retry_stream_id": stream_id, + "_journal_retry_attempts": 0, + "_journal_retry_first_seen_ts": int(models.time.time()), + }, + ]) + s.save() + return s + + +def _assert_retry_meta_removed(marker): + assert "_pending_journal_recovery" not in marker + assert "_journal_retry_stream_id" not in marker + assert "_journal_retry_attempts" not in marker + assert "_journal_retry_first_seen_ts" not in marker + + # ── The regression test ──────────────────────────────────────────────────── @@ -196,3 +227,36 @@ def test_lost_response_recovered_on_second_read(hermes_home): assert "_journal_retry_stream_id" not in promoted assert "_journal_retry_attempts" not in promoted assert "_journal_retry_first_seen_ts" not in promoted + + +def test_concurrent_get_session_serializes_lazy_journal_retry(hermes_home, monkeypatch): + sid = "retry_lock_sid" + stream_id = "retry_lock_stream" + s = _make_pending_retry_session(sid, stream_id=stream_id) + models.SESSIONS[sid] = s + + entered = threading.Event() + release = threading.Event() + counter_lock = threading.Lock() + calls = 0 + + def slow_retry(session): + nonlocal calls + with counter_lock: + calls += 1 + entered.set() + assert release.wait(timeout=2), "test timed out waiting to release retry body" + return False + + monkeypatch.setattr(models, "_retry_journal_recovery_in_place", slow_retry) + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(models.get_session, sid) + assert entered.wait(timeout=2), "first caller did not enter retry helper" + second = executor.submit(models.get_session, sid) + assert second.result(timeout=2) is s + release.set() + assert first.result(timeout=2) is s + + assert calls == 1 +