fix(agent): rehydrate session cost counters on agent construction (#67762)

agent/agent_init.py

@@ -229,6 +229,55 @@ def _custom_provider_runtime_ids(value: Any) -> set[str]:
     return {normalized, f"custom:{normalized}"}
 
 
+def _rehydrate_session_cost(agent) -> None:
+    """Restore ``agent.session_estimated_cost_usd`` and
+    ``agent.session_cost_status`` from SQLite on agent construction
+    (issue #67762).
+
+    Before this helper, both attributes were reset to ``0.0`` / ``"unknown"``
+    by the ``init_agent`` reset block with no read from any persisted source.
+    After a gateway restart mid-session, the live counter would silently drop
+    to ``$0.00`` even though ``session_model_usage`` had the real accumulated
+    cost.
+
+    Fail-open: any unexpected error leaves the agent at the post-reset values
+    (``$0.0`` / ``"unknown"``). The next API call will accumulate from there,
+    so a transient rehydration failure is not catastrophic — just lossy for
+    one turn.
+
+    Note: the rehydrated ``session_cost_status`` is overwritten by the
+    unconditional ``=`` assignment at ``agent/conversation_loop.py:2321``
+    (and the equivalent at ``agent/codex_runtime.py:150``) on the very next
+    API call. The *cost value* survives across a gateway restart; the *status*
+    label reverts to "latest call wins" until issue #67764 (priority ladder)
+    also lands and converts those assignments to sticky ``max()`` semantics.
+    """
+    import sqlite3
+    _session_db = getattr(agent, "_session_db", None)
+    _session_id = getattr(agent, "session_id", None)
+    if _session_db is None or not _session_id:
+        return
+
+    try:
+        row = _session_db.get_session_cost_summary(_session_id)
+    except (sqlite3.Error, AttributeError, TypeError, ValueError) as exc:
+        # Scoped: only the failures we expect from the SQLite read. Anything
+        # more exotic (e.g., a bug in our SQL) should surface, not be swallowed.
+        _ra().logger.debug(
+            "Cost rehydration read failed for session %s: %s",
+            _session_id, exc,
+        )
+        return
+
+    if row is not None:
+        agent.session_estimated_cost_usd = float(
+            row.get("estimated_cost_usd") or 0.0
+        )
+        status = row.get("cost_status")
+        if status:
+            agent.session_cost_status = str(status)
+
+
 def _build_codex_gpt5_autoraise_notice(
     autoraise: Dict[str, Any], context_length: Optional[int] = None
 ) -> str:
@@ -2531,7 +2580,24 @@ def _parse_prune_int(raw, default):
     agent.session_estimated_cost_usd = 0.0
     agent.session_cost_status = "unknown"
     agent.session_cost_source = "none"
-
+
+    # Rehydrate the cost counters from the persisted source so a gateway
+    # restart mid-session doesn't reset the live counter to $0.00 (issue #67762).
+    # SQLite is the source of truth (the JSON SessionEntry can desync if a
+    # write fails between the per-call update_token_counts transaction and
+    # the JSON to_dict save). The JSON fallback below covers the case where
+    # SQLite is unavailable. Fail-open: if both raise, the agent continues
+    # at $0.0 and accumulates from the next API call.
+    #
+    # Note: the rehydrated session_cost_status will be overwritten by the
+    # unconditional `=` assignment at agent/conversation_loop.py:2321 (and
+    # the equivalent at agent/codex_runtime.py:150) on the very next API
+    # call. The *cost value* survives across a gateway restart; the *status*
+    # label reverts to "latest call wins" until issue #67764 (priority
+    # ladder) also lands and converts those assignments to sticky `max()`
+    # semantics.
+    _rehydrate_session_cost(agent)
+
     # ── Ollama num_ctx injection ──
     # Ollama defaults to 2048 context regardless of the model's capabilities.
     # When running against an Ollama server, detect the model's max context

hermes_state.py

@@ -5227,6 +5227,65 @@ def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
             row = cursor.fetchone()
         return dict(row) if row else None
 
+    def get_session_cost_summary(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """Return accumulated cost + status across the session's per-model rows.
+
+        Used by ``agent/agent_init.py`` on resume to rehydrate the in-memory
+        ``agent.session_estimated_cost_usd`` and ``agent.session_cost_status``
+        so a gateway restart doesn't reset the live counter to $0 (issue #67762).
+
+        Returns a dict with ``estimated_cost_usd`` (float, sum of all main-loop
+        ``session_model_usage`` rows for this session — ``task=''`` only, no
+        auxiliary-task rows like title-generation or delegate summaries) and
+        ``cost_status`` (sticky priority across the same rows). Returns ``None`` if the session has
+        no per-model rows yet — callers use this to distinguish "no spend"
+        from "explicit zero spend".
+
+        Read-only; uses the connection's read path (no transaction needed).
+        Single SQL query: ``SUM(CASE WHEN ...)`` is the codebase's preferred
+        conditional-aggregation pattern (see ``agent/insights.py:373-376``).
+        """
+        if self._conn is None:
+            return None
+        with self._lock:
+            cursor = self._conn.execute(
+                """
+                SELECT
+                    COALESCE(SUM(estimated_cost_usd), 0.0) AS estimated_cost_usd,
+                    CASE
+                        WHEN SUM(CASE WHEN cost_status = 'actual' THEN 1 ELSE 0 END) > 0
+                            THEN 'actual'
+                        WHEN SUM(CASE WHEN cost_status = 'included' THEN 1 ELSE 0 END) > 0
+                            THEN 'included'
+                        WHEN SUM(CASE WHEN cost_status = 'unknown' THEN 1 ELSE 0 END) > 0
+                            THEN 'unknown'
+                        ELSE COALESCE(
+                            (SELECT cost_status FROM session_model_usage
+                             WHERE session_id = :sid AND task = ''
+                             ORDER BY last_seen DESC LIMIT 1),
+                            'estimated'
+                        )
+                    END AS cost_status,
+                    COUNT(*) AS row_count
+                FROM session_model_usage
+                WHERE session_id = :sid AND task = ''
+                """,
+                {"sid": session_id},
+            )
+            row = cursor.fetchone()
+        if row is None:
+            return None
+        # row_count lets us distinguish "no rows" from "rows but zero sum"
+        result = dict(row)
+        if int(result.get("row_count") or 0) == 0:
+            return None
+        try:
+            result["estimated_cost_usd"] = float(result["estimated_cost_usd"] or 0.0)
+        except (TypeError, ValueError):
+            result["estimated_cost_usd"] = 0.0
+        result.pop("row_count", None)
+        return result
+
     def resolve_session_id(self, session_id_or_prefix: str) -> Optional[str]:
         """Resolve an exact or uniquely prefixed session ID to the full ID.