mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-29 18:46:59 +00:00
174ad45939
Every API call in the tool loop persisted its token/cost delta by calling SessionDB.update_token_counts() synchronously on the turn thread — a BEGIN IMMEDIATE sessions UPDATE plus a session_model_usage upsert, measured in production at p50 3.3ms / p95 70.4ms per call and up to 299ms against a cold multi-GB state.db. The tool loop stalls for that long between calls, N times per multi-tool turn. SessionDB gains queue_token_counts(): same signature and semantics as update_token_counts(), but the critical path is a deque append plus a condvar notify. A lazily started daemon thread applies deltas in enqueue order through the existing update_token_counts -> _execute_write path, so the established self._lock / BEGIN IMMEDIATE / jitter-retry discipline is unchanged. When a backlog forms, adjacent same-route incremental deltas coalesce into one UPDATE: token and api-call fields sum, cost fields sum None-preservingly (an all-None run stays None so COALESCE keeps the stored value), and absolute=True deltas never merge and act as ordering barriers. Route equality is required for a merge because those fields feed COALESCE backfill, the last-non-None-wins status fields, and the per-model usage attribution key — a merged apply is row-equivalent to sequential applies. Correctness and durability: - flush_token_counts() gives read-your-writes to token/cost readers (get_session, list_sessions_rich, _get_session_rich_row, list_gateway_sessions, InsightsEngine.generate) — a plain attribute check when nothing is queued. The writer sets its busy flag before popping the queue so the lock-free fast path can never miss an in-flight batch. - update_session_model / update_session_billing_route / update_session_meta write the sessions row synchronously, bypassing the queue, so they flush it first: a still-queued first-of-session delta carries the pre-switch route, and applying it after the switch UPDATE would trip the first_accounted_route branch (api_call_count == 0 plus a route mismatch) and resurrect the old model/provider. - AIAgent._persist_session flushes at turn finalize and every error-exit persist point; close() stops and drains the writer before the WAL checkpoint; an atexit hook (registered on first enqueue, unregistered on close so closed instances are not pinned until interpreter exit) drains at shutdown. Worst-case crash loss is the in-flight call's delta — the same window as the old inline write. - A flush trusts a live stop-flagged writer (its loop drains before exiting) and only drains on the caller's thread when the writer is dead or never started, claiming the same busy flag so concurrent flushes wait instead of racing an in-flight batch. - After close() has stopped the writer, queue_token_counts applies the delta inline instead of parking it on a queue nothing will drain; a closed-connection failure then raises at the call site, which already guards for it, exactly like the old synchronous path. - Writer apply failures are logged and never raise into a turn; the writer thread survives and keeps applying. Call sites switched to the queue: the per-call site in agent/conversation_loop.py and both codex app-server sites in agent/codex_runtime.py. In-memory per-turn counters (agent.session_estimated_cost_usd etc.) stay synchronous, so live turn displays never see the queue. Tests: tests/agent/test_async_token_accounting.py (19 tests: enqueue ordering, absolute-as-barrier, backlog coalescing with exact sums, coalesced-vs-sequential row equivalence, merge unit rules, None-cost preservation, read-your-writes, flush vs stop-flagged/concurrent drains, inline apply after writer stop, close/atexit durability, _persist_session drain, writer failure isolation); tests/run_agent/test_token_persistence_non_cli.py updated to the queue_token_counts contract.
97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
import json
|
|
import sys
|
|
|
|
from run_agent import AIAgent
|
|
|
|
|
|
def _mock_response(*, usage: dict, content: str = "done"):
|
|
msg = SimpleNamespace(content=content, tool_calls=None)
|
|
choice = SimpleNamespace(message=msg, finish_reason="stop")
|
|
return SimpleNamespace(
|
|
choices=[choice],
|
|
model="test/model",
|
|
usage=SimpleNamespace(**usage),
|
|
)
|
|
|
|
|
|
def _make_agent(session_db, *, platform: str):
|
|
with (
|
|
patch("run_agent.get_tool_definitions", return_value=[]),
|
|
patch("run_agent.check_toolset_requirements", return_value={}),
|
|
patch("run_agent.OpenAI"),
|
|
):
|
|
agent = AIAgent(
|
|
api_key="test-key",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
quiet_mode=True,
|
|
skip_context_files=True,
|
|
skip_memory=True,
|
|
session_db=session_db,
|
|
session_id=f"{platform}-session",
|
|
platform=platform,
|
|
)
|
|
agent.client = MagicMock()
|
|
agent.client.chat.completions.create.return_value = _mock_response(
|
|
usage={
|
|
"prompt_tokens": 11,
|
|
"completion_tokens": 7,
|
|
"total_tokens": 18,
|
|
}
|
|
)
|
|
return agent
|
|
|
|
|
|
def test_run_conversation_persists_tokens_for_telegram_sessions():
|
|
session_db = MagicMock()
|
|
agent = _make_agent(session_db, platform="telegram")
|
|
|
|
result = agent.run_conversation("hello")
|
|
|
|
assert result["final_response"] == "done"
|
|
# Per-call deltas are enqueued for the SessionDB background writer
|
|
# (queue_token_counts) rather than written inline on the turn thread.
|
|
session_db.queue_token_counts.assert_called_once()
|
|
assert session_db.queue_token_counts.call_args.args[0] == "telegram-session"
|
|
|
|
|
|
def test_run_conversation_persists_tokens_for_cron_sessions():
|
|
session_db = MagicMock()
|
|
agent = _make_agent(session_db, platform="cron")
|
|
|
|
result = agent.run_conversation("hello")
|
|
|
|
assert result["final_response"] == "done"
|
|
session_db.queue_token_counts.assert_called_once()
|
|
assert session_db.queue_token_counts.call_args.args[0] == "cron-session"
|
|
|
|
|
|
def test_session_search_lazily_opens_db_when_entrypoint_did_not_pass_one(monkeypatch):
|
|
sentinel_db = object()
|
|
captured = {}
|
|
|
|
class FakeSessionDB:
|
|
def __new__(cls):
|
|
return sentinel_db
|
|
|
|
hermes_state = ModuleType("hermes_state")
|
|
hermes_state.SessionDB = FakeSessionDB
|
|
monkeypatch.setitem(sys.modules, "hermes_state", hermes_state)
|
|
|
|
session_search_mod = ModuleType("tools.session_search_tool")
|
|
|
|
def fake_session_search(**kwargs):
|
|
captured.update(kwargs)
|
|
return json.dumps({"success": True, "results": []})
|
|
|
|
session_search_mod.session_search = fake_session_search
|
|
monkeypatch.setitem(sys.modules, "tools.session_search_tool", session_search_mod)
|
|
|
|
agent = _make_agent(None, platform="acp")
|
|
result = json.loads(agent._invoke_tool("session_search", {"query": "Hermes"}, "task-id"))
|
|
|
|
assert result["success"] is True
|
|
assert captured["db"] is sentinel_db
|
|
assert captured["query"] == "Hermes"
|
|
assert agent._session_db is sentinel_db
|