diff --git a/hermes_state.py b/hermes_state.py index f693f391f78e4..d6751e21ca69a 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -205,6 +205,14 @@ def _log_wal_fallback_once(db_label: str, exc: Exception) -> None: cache_read_tokens INTEGER DEFAULT 0, cache_write_tokens INTEGER DEFAULT 0, reasoning_tokens INTEGER DEFAULT 0, + -- Size of the LATEST API call's prompt (full input including cached + -- prefixes) and completion. Unlike the cumulative counters above, + -- these are SET on every update — they represent the current + -- context-window snapshot, not lifetime totals. Lets downstream + -- dashboards display "current context size" without having to + -- compute it from cumulative deltas + a cache model. + last_prompt_tokens INTEGER DEFAULT 0, + last_completion_tokens INTEGER DEFAULT 0, billing_provider TEXT, billing_base_url TEXT, billing_mode TEXT, @@ -768,16 +776,29 @@ def update_token_counts( billing_base_url: Optional[str] = None, billing_mode: Optional[str] = None, api_call_count: int = 0, + last_prompt_tokens: Optional[int] = None, + last_completion_tokens: Optional[int] = None, absolute: bool = False, ) -> None: """Update token counters and backfill model if not already set. - When *absolute* is False (default), values are **incremented** — use - this for per-API-call deltas (CLI path). - - When *absolute* is True, values are **set directly** — use this when - the caller already holds cumulative totals (gateway path, where the - cached agent accumulates across messages). + When *absolute* is False (default), the cumulative counters + (input_tokens, output_tokens, cache_*_tokens, reasoning_tokens, + cost, api_call_count) are **incremented** — use this for + per-API-call deltas (CLI path). + + When *absolute* is True, those values are **set directly** — + use this when the caller already holds cumulative totals + (gateway path, where the cached agent accumulates across + messages). + + ``last_prompt_tokens`` / ``last_completion_tokens`` are ALWAYS + set when not None, regardless of *absolute*. They snapshot the + most-recent API call's prompt/completion sizes (i.e. the + current context-window size), which is what downstream + dashboards need to display "context burn" without rebuilding + it from cumulative deltas + a cache model. Pass None to leave + the existing snapshot untouched (e.g. for a cost-only update). """ # Ensure the session row exists so the UPDATE doesn't silently affect # 0 rows. Under concurrent load (cron + kanban + delegate_task) the @@ -803,7 +824,9 @@ def update_token_counts( billing_base_url = COALESCE(billing_base_url, ?), billing_mode = COALESCE(billing_mode, ?), model = COALESCE(model, ?), - api_call_count = ? + api_call_count = ?, + last_prompt_tokens = COALESCE(?, last_prompt_tokens), + last_completion_tokens = COALESCE(?, last_completion_tokens) WHERE id = ?""" else: sql = """UPDATE sessions SET @@ -824,7 +847,9 @@ def update_token_counts( billing_base_url = COALESCE(billing_base_url, ?), billing_mode = COALESCE(billing_mode, ?), model = COALESCE(model, ?), - api_call_count = COALESCE(api_call_count, 0) + ? + api_call_count = COALESCE(api_call_count, 0) + ?, + last_prompt_tokens = COALESCE(?, last_prompt_tokens), + last_completion_tokens = COALESCE(?, last_completion_tokens) WHERE id = ?""" params = ( input_tokens, @@ -843,6 +868,8 @@ def update_token_counts( billing_mode, model, api_call_count, + last_prompt_tokens, + last_completion_tokens, session_id, ) def _do(conn): diff --git a/run_agent.py b/run_agent.py index b60f6c43ce693..ddc330293c5a2 100644 --- a/run_agent.py +++ b/run_agent.py @@ -13251,6 +13251,16 @@ def _stop_spinner(): if cost_result.status == "included" else None, model=self.model, api_call_count=1, + # Snapshot the LATEST call's prompt/completion + # sizes — the cumulative counters above mix + # all calls together, but downstream + # dashboards need the most-recent prompt + # size to show current context-window usage. + # ``prompt_tokens`` is the full input + # (including cached prefixes), so it IS the + # context-window size at this moment. + last_prompt_tokens=prompt_tokens, + last_completion_tokens=completion_tokens, ) except Exception as e: # Log token persistence failures so they're diff --git a/tests/test_hermes_state.py b/tests/test_hermes_state.py index 3bae763b9412a..dbc1e1bdb5b97 100644 --- a/tests/test_hermes_state.py +++ b/tests/test_hermes_state.py @@ -129,6 +129,74 @@ def test_update_token_counts_preserves_existing_model(self, db): session = db.get_session("s1") assert session["model"] == "anthropic/claude-opus-4.6" + def test_last_prompt_tokens_is_set_not_incremented(self, db): + """``last_prompt_tokens`` is the most-recent prompt size — it + MUST be set on every update, never accumulated. Otherwise a + 100-turn session would report a prompt size of N * per_turn.""" + db.create_session(session_id="s1", source="cli") + db.update_token_counts( + "s1", input_tokens=100, output_tokens=50, + last_prompt_tokens=8000, last_completion_tokens=200, + ) + db.update_token_counts( + "s1", input_tokens=100, output_tokens=50, + last_prompt_tokens=12000, last_completion_tokens=300, + ) + session = db.get_session("s1") + # Cumulative counters DO accumulate. + assert session["input_tokens"] == 200 + assert session["output_tokens"] == 100 + # last_* fields snapshot the LATEST call only. + assert session["last_prompt_tokens"] == 12000 + assert session["last_completion_tokens"] == 300 + + def test_last_prompt_tokens_omitted_keeps_previous_value(self, db): + """Passing ``last_prompt_tokens=None`` (the default) must + leave the existing snapshot untouched — needed for cost-only + updates that don't carry a fresh prompt size.""" + db.create_session(session_id="s1", source="cli") + db.update_token_counts( + "s1", input_tokens=100, output_tokens=50, + last_prompt_tokens=8000, last_completion_tokens=200, + ) + # Subsequent update without last_* — must not zero them out. + db.update_token_counts( + "s1", input_tokens=100, output_tokens=50, + ) + session = db.get_session("s1") + assert session["last_prompt_tokens"] == 8000 + assert session["last_completion_tokens"] == 200 + + def test_last_prompt_tokens_in_absolute_mode(self, db): + """The gateway path uses ``absolute=True``. ``last_prompt_tokens`` + must still be SET (not accumulated) in that branch too.""" + db.create_session(session_id="s1", source="gateway") + db.update_token_counts( + "s1", input_tokens=500, output_tokens=200, + last_prompt_tokens=15000, last_completion_tokens=400, + absolute=True, + ) + db.update_token_counts( + "s1", input_tokens=700, output_tokens=350, + last_prompt_tokens=22000, last_completion_tokens=500, + absolute=True, + ) + session = db.get_session("s1") + assert session["input_tokens"] == 700 # absolute + assert session["last_prompt_tokens"] == 22000 # set, not accumulated + assert session["last_completion_tokens"] == 500 + + def test_existing_db_reconciles_last_prompt_tokens_columns(self, db): + """The declarative reconciliation must auto-ADD the new + columns on databases created before this change. Verifies + the columns exist and default to 0.""" + db.create_session(session_id="s1", source="cli") + session = db.get_session("s1") + assert "last_prompt_tokens" in session.keys() + assert "last_completion_tokens" in session.keys() + assert session["last_prompt_tokens"] == 0 + assert session["last_completion_tokens"] == 0 + def test_parent_session(self, db): db.create_session(session_id="parent", source="cli") db.create_session(session_id="child", source="cli", parent_session_id="parent")