From 927633272fcb0bbd67bd9b02e963b013fca543c9 Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Tue, 28 Jul 2026 18:05:37 +0500
Subject: [PATCH] fix: claim busy before clearing queue in _stop_token_writer
 drain
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Follow-up to PR #64171. The writer loop and flush_token_counts'
caller-drain both set _token_writer_busy BEFORE popping the queue —
that ordering is what makes flush's lock-free fast path (reads queue
then busy, no cond held) sound. _stop_token_writer's leftover drain did
it backwards (clear queue, then set busy), leaving a few-bytecode window
where a concurrent flush could observe 'empty and idle' and return True
with the popped batch still unapplied. Shutdown-only staleness, no data
loss — but the protocol now matches at all three drain sites.

Test: concurrent flush during a stop-drain mid-apply must time out
(False), never report drained.
---
 hermes_state.py                            |  8 ++++-
 tests/agent/test_async_token_accounting.py | 37 ++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/hermes_state.py b/hermes_state.py
index cd05ae2a1083..c1b9d55e32dd 100644
--- a/hermes_state.py
+++ b/hermes_state.py
@@ -5333,10 +5333,16 @@ class SessionDB:
                     )
                     return
                 self._token_queue_cond.wait(remaining)
+            # busy is claimed BEFORE the queue is cleared — same ordering
+            # as the writer loop and the flush caller-drain. The lock-free
+            # fast path in flush_token_counts() reads queue-then-busy
+            # without the cond, so clearing first would let a concurrent
+            # flush observe "empty and idle" and return True while this
+            # popped batch is still unapplied.
             batch = list(self._token_queue)
-            self._token_queue.clear()
             if batch:
                 self._token_writer_busy = True
+                self._token_queue.clear()
         if batch:
             try:
                 self._apply_token_batch(batch)
diff --git a/tests/agent/test_async_token_accounting.py b/tests/agent/test_async_token_accounting.py
index 70cd5d8d34c2..181fa795a65d 100644
--- a/tests/agent/test_async_token_accounting.py
+++ b/tests/agent/test_async_token_accounting.py
@@ -608,3 +608,40 @@ class TestWriterFailure:
         assert db._token_writer_thread is not dead
         assert db.flush_token_counts()
         assert _totals(db, "s-respawn")["input_tokens"] == 3
+
+    def test_stop_drain_claims_busy_before_clearing_queue(self, db):
+        """_stop_token_writer's leftover drain must follow the same
+        busy-before-clear ordering as the writer loop: a concurrent flush's
+        lock-free fast path (queue-then-busy, no cond held) must never
+        observe 'empty and idle' while the popped batch is unapplied."""
+        db.create_session("s-stopdrain", "test")
+        db.flush_token_counts()
+        db._stop_token_writer()  # writer dead, connection open
+
+        applied = threading.Event()
+        gate = threading.Event()
+        original = db.update_token_counts
+
+        def gated(session_id, **kwargs):
+            applied.set()
+            assert gate.wait(timeout=10)
+            return original(session_id, **kwargs)
+
+        db.update_token_counts = gated
+        try:
+            db._token_queue.append(
+                ("s-stopdrain", dict(input_tokens=6, api_call_count=1))
+            )
+            t = threading.Thread(target=db._stop_token_writer)
+            t.start()
+            assert applied.wait(timeout=10)
+            # Stop-drain is mid-apply with the queue popped: the fast path
+            # must see busy=True and wait (timing out), not return True.
+            assert db.flush_token_counts(timeout=0.3) is False
+            gate.set()
+            t.join(timeout=10)
+            assert db.flush_token_counts()
+        finally:
+            db.update_token_counts = original
+
+        assert _totals(db, "s-stopdrain")["input_tokens"] == 6