diff --git a/hermes_state.py b/hermes_state.py index f19e36ea7f24..cd05ae2a1083 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -5143,10 +5143,15 @@ class SessionDB: ) if not writer_stopped: self._token_queue.append((session_id, kwargs)) - if thread is None: + if thread is None or not thread.is_alive(): # Daemon so process exit never hangs on accounting; the # atexit hook drains anything still queued at interpreter # shutdown (registered once per instance, on first use). + # ``not is_alive()`` (rather than ``is None`` only) + # respawns the writer if it ever died from an unexpected + # escape — otherwise a dead thread object would block + # respawn forever and deltas would pile up on the deque + # until a reader's flush drained them synchronously. thread = threading.Thread( target=self._token_writer_loop, name="session-db-token-writer", @@ -5240,7 +5245,18 @@ class SessionDB: def _apply_token_batch(self, batch: List[Tuple[str, Dict[str, Any]]]) -> None: """Apply queued deltas in order, coalescing where safe. Never raises.""" - for session_id, kwargs in self._coalesce_token_deltas(batch): + try: + coalesced = self._coalesce_token_deltas(batch) + except Exception as exc: + # Coalescing must never kill the writer thread (a dead writer + # can't be observed by callers). Fall back to applying the raw + # batch delta-by-delta — the merge is an optimization only. + logger.warning( + "async token accounting: coalesce failed, applying raw " + "batch: %s", exc, + ) + coalesced = batch + for session_id, kwargs in coalesced: try: self.update_token_counts(session_id, **kwargs) except Exception as exc: diff --git a/tests/agent/test_async_token_accounting.py b/tests/agent/test_async_token_accounting.py index daa1dbfa2e41..70cd5d8d34c2 100644 --- a/tests/agent/test_async_token_accounting.py +++ b/tests/agent/test_async_token_accounting.py @@ -562,3 +562,49 @@ class TestWriterFailure: db.update_token_counts = original assert _totals(db, "s-f")["input_tokens"] == 7 + + def test_coalesce_failure_falls_back_to_raw_batch(self, db, caplog): + """A coalescing bug must never kill the writer: the batch is applied + raw (delta-by-delta) and the failure is logged.""" + db.create_session("s-co", "test") + + original = db._coalesce_token_deltas + + def broken(batch): + raise TypeError("unclassified kwarg broke the merge") + + db._coalesce_token_deltas = broken + try: + with caplog.at_level("WARNING", logger="hermes_state"): + db.queue_token_counts("s-co", input_tokens=3, api_call_count=1) + db.queue_token_counts("s-co", input_tokens=4, api_call_count=1) + assert db.flush_token_counts() + assert any( + "coalesce failed" in rec.getMessage() for rec in caplog.records + ) + finally: + db._coalesce_token_deltas = original + + totals = _totals(db, "s-co") + assert totals["input_tokens"] == 7 + assert totals["api_call_count"] == 2 + + def test_dead_writer_respawns_on_next_enqueue(self, db): + """If the writer thread ever dies unexpectedly, the next enqueue + must respawn it instead of parking deltas on a queue forever.""" + db.create_session("s-respawn", "test") + db.queue_token_counts("s-respawn", input_tokens=1, api_call_count=1) + assert db.flush_token_counts() + first = db._token_writer_thread + assert first is not None + + # Simulate an unexpected writer death: a finished dummy thread. + dead = threading.Thread(target=lambda: None) + dead.start() + dead.join() + db._token_writer_thread = dead + + db.queue_token_counts("s-respawn", input_tokens=2, api_call_count=1) + assert db._token_writer_thread is not dead + assert db.flush_token_counts() + assert _totals(db, "s-respawn")["input_tokens"] == 3