From 58472d803a32edd19773bb2ed7426490981a636e Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Tue, 1 Sep 2026 21:07:37 +0530
Subject: [PATCH] refactor(state): flatten registry acquire flow; restore
 mirror cleanup assertion
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Post-merge review follow-ups on #100201:

- acquire(): drop the never-iterating 'while True' and the redundant
  'existing is not generation' half of the race check — after the retire
  path, generation is always None on the fresh-open leg, so 'existing is
  not None' is the complete condition. Same behavior, flat flow.
- test_mirror: the conversion to the shared registry dropped the
  cleanup assertion entirely; restore it by patching
  hermes_state.release_or_close and asserting the handle is released
  exactly once after _append_to_sqlite.
---
 hermes_state_registry.py     | 74 ++++++++++++++++++------------------
 tests/gateway/test_mirror.py | 19 +++++----
 2 files changed, 47 insertions(+), 46 deletions(-)

diff --git a/hermes_state_registry.py b/hermes_state_registry.py
index 2fb1d5a7f2..3c5b6ec24e 100644
--- a/hermes_state_registry.py
+++ b/hermes_state_registry.py
@@ -134,45 +134,43 @@ def acquire(db_path: Optional[Path] = None) -> "SessionDB":
 
     path = Path(db_path) if db_path is not None else Path(_default_db_path())
 
-    while True:
-        with _lock:
-            generation = _generations.get(path)
-            if generation is not None:
-                current = _stat_db_file_identity(path)
-                if (
-                    current is not None
-                    and generation.identity is not None
-                    and current != generation.identity
-                ):
-                    # File replaced: retire the live generation (its
-                    # holders keep it until they release) and fall
-                    # through to opening a fresh one below.
-                    _retire_generation_locked(path, generation)
-                    generation = None
-                else:
-                    generation.refcount += 1
-                    return generation.db
-
-        # Open a fresh generation OUTSIDE the lock: construction can
-        # take seconds (write-lock patience) and must not block every
-        # other state.db acquisition in the process.
-        db = _open_session_db(path)
-        db._shared_registry_owned = True
-        identity = _stat_db_file_identity(path)
-        with _lock:
-            existing = _generations.get(path)
-            if existing is not None and existing is not generation:
-                # Someone else opened a generation while we were
-                # constructing.  Ours loses — close it (outside the
-                # lock) and use theirs.
-                existing.refcount += 1
-                winner = existing.db
+    with _lock:
+        generation = _generations.get(path)
+        if generation is not None:
+            current = _stat_db_file_identity(path)
+            if (
+                current is not None
+                and generation.identity is not None
+                and current != generation.identity
+            ):
+                # File replaced: retire the live generation (its
+                # holders keep it until they release) and fall
+                # through to opening a fresh one below.
+                _retire_generation_locked(path, generation)
             else:
-                _generations[path] = _Generation(db, identity)
-                winner = db
-        if winner is not db:
-            _teardown(db)
-        return winner
+                generation.refcount += 1
+                return generation.db
+
+    # Open a fresh generation OUTSIDE the lock: construction can
+    # take seconds (write-lock patience) and must not block every
+    # other state.db acquisition in the process.
+    db = _open_session_db(path)
+    db._shared_registry_owned = True
+    identity = _stat_db_file_identity(path)
+    with _lock:
+        existing = _generations.get(path)
+        if existing is not None:
+            # Someone else opened a generation while we were
+            # constructing (or retired ours and installed a new one).
+            # Ours loses — close it (outside the lock) and use theirs.
+            existing.refcount += 1
+            winner = existing.db
+        else:
+            _generations[path] = _Generation(db, identity)
+            winner = db
+    if winner is not db:
+        _teardown(db)
+    return winner
 
 
 def _retire_generation_locked(path: Path, generation: _Generation) -> None:
diff --git a/tests/gateway/test_mirror.py b/tests/gateway/test_mirror.py
index a2890c8b94..671d9e9883 100644
--- a/tests/gateway/test_mirror.py
+++ b/tests/gateway/test_mirror.py
@@ -119,18 +119,21 @@ class TestMirrorToSession:
 
 
 class TestAppendToSqlite:
-    def test_connection_is_closed_after_use(self, tmp_path):
-        """Verify _append_to_sqlite releases the shared SessionDB handle."""
+    def test_connection_is_released_after_use(self, tmp_path):
+        """Verify _append_to_sqlite returns the shared SessionDB reference."""
         from gateway.mirror import _append_to_sqlite
         mock_db = MagicMock()
+        released = []
 
-        with patch("hermes_state.get_shared_session_db", return_value=mock_db):
+        with patch("hermes_state.get_shared_session_db", return_value=mock_db), \
+             patch(
+                 "hermes_state.release_or_close",
+                 side_effect=lambda db: released.append(db),
+             ):
             _append_to_sqlite("sess_1", {"role": "assistant", "content": "hello"})
 
         mock_db.append_message.assert_called_once()
-        # Shared instances are released (not closed) — the registry owns close().
-        # release_shared_session_db is a module-level function, so verify the
-        # mock was passed to it by checking that append_message was called
-        # (the real release_shared_session_db on a MagicMock is a no-op since
-        # the mock isn't in the registry).
+        assert released == [mock_db], (
+            "the shared handle must be released exactly once after use"
+        )