diff --git a/tests/tui_gateway/test_restore_session_usage.py b/tests/tui_gateway/test_restore_session_usage.py new file mode 100644 index 000000000000..d1a1072c53ef --- /dev/null +++ b/tests/tui_gateway/test_restore_session_usage.py @@ -0,0 +1,160 @@ +"""Tests for _restore_session_usage — token counter restore on session.resume. + +When the TUI gateway resumes a session, a fresh agent is built whose +session_*_tokens counters all start at zero. _restore_session_usage copies +the cumulative counts from the stored session row so that session.info / +session.usage reflects actual usage instead of 0/1.0M-0%. +""" + +import pytest + + +class _FakeAgent: + """Minimal stand-in with the same counter attributes as AIAgent.""" + + def __init__(self): + self.session_input_tokens = 0 + self.session_output_tokens = 0 + self.session_cache_read_tokens = 0 + self.session_cache_write_tokens = 0 + self.session_reasoning_tokens = 0 + self.session_total_tokens = 0 + self.session_api_calls = 0 + self.session_estimated_cost_usd = 0.0 + self.session_cost_status = "unknown" + + +def _import_restore(): + from tui_gateway.server import _restore_session_usage + return _restore_session_usage + + +# ── Unit tests for _restore_session_usage ────────────────────────── + + +def test_restore_basic_token_counts(): + """All five token fields are populated from the stored session row.""" + restore = _import_restore() + agent = _FakeAgent() + stored = { + "input_tokens": 12000, + "output_tokens": 3500, + "cache_read_tokens": 800, + "cache_write_tokens": 200, + "reasoning_tokens": 1500, + "api_call_count": 7, + } + restore(agent, stored) + + assert agent.session_input_tokens == 12000 + assert agent.session_output_tokens == 3500 + assert agent.session_cache_read_tokens == 800 + assert agent.session_cache_write_tokens == 200 + assert agent.session_reasoning_tokens == 1500 + assert agent.session_total_tokens == 12000 + 3500 + 800 + 200 + 1500 + assert agent.session_api_calls == 7 + + +def test_restore_total_is_sum(): + """session_total_tokens equals the sum of all five token components.""" + restore = _import_restore() + agent = _FakeAgent() + stored = { + "input_tokens": 100, + "output_tokens": 200, + "cache_read_tokens": 300, + "cache_write_tokens": 400, + "reasoning_tokens": 500, + } + restore(agent, stored) + assert agent.session_total_tokens == 100 + 200 + 300 + 400 + 500 + + +def test_restore_handles_none_values(): + """Missing or None fields default to zero.""" + restore = _import_restore() + agent = _FakeAgent() + stored = { + "input_tokens": None, + "output_tokens": 50, + } + restore(agent, stored) + + assert agent.session_input_tokens == 0 + assert agent.session_output_tokens == 50 + assert agent.session_cache_read_tokens == 0 + assert agent.session_cache_write_tokens == 0 + assert agent.session_reasoning_tokens == 0 + assert agent.session_total_tokens == 50 + assert agent.session_api_calls == 0 + + +def test_restore_preserves_cost_estimate(): + """estimated_cost_usd and cost_status are carried over when present.""" + restore = _import_restore() + agent = _FakeAgent() + stored = { + "input_tokens": 1000, + "output_tokens": 500, + "estimated_cost_usd": 0.042, + "cost_status": "estimated", + } + restore(agent, stored) + + assert agent.session_estimated_cost_usd == pytest.approx(0.042) + assert agent.session_cost_status == "estimated" + + +def test_restore_skips_cost_when_absent(): + """Cost fields stay at defaults when not in the stored row.""" + restore = _import_restore() + agent = _FakeAgent() + stored = {"input_tokens": 100, "output_tokens": 50} + restore(agent, stored) + + assert agent.session_estimated_cost_usd == 0.0 + assert agent.session_cost_status == "unknown" + + +def test_restore_empty_dict(): + """An empty stored dict leaves the agent at zero counters.""" + restore = _import_restore() + agent = _FakeAgent() + restore(agent, {}) + + assert agent.session_input_tokens == 0 + assert agent.session_output_tokens == 0 + assert agent.session_total_tokens == 0 + assert agent.session_api_calls == 0 + + +# ── _get_usage picks up restored counters ────────────────────────── + + +def test_get_usage_reflects_restored_counters(): + """After restore, _get_usage returns the stored token counts.""" + from tui_gateway.server import _get_usage + + agent = _FakeAgent() + restore = _import_restore() + stored = { + "input_tokens": 5000, + "output_tokens": 2000, + "cache_read_tokens": 100, + "cache_write_tokens": 50, + "reasoning_tokens": 300, + "api_call_count": 3, + "estimated_cost_usd": 0.015, + "cost_status": "estimated", + } + restore(agent, stored) + + usage = _get_usage(agent) + + assert usage["input"] == 5000 + assert usage["output"] == 2000 + assert usage["cache_read"] == 100 + assert usage["cache_write"] == 50 + assert usage["reasoning"] == 300 + assert usage["total"] == 5000 + 2000 + 100 + 50 + 300 + assert usage["calls"] == 3 diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 5df9e51de299..e9cdebebd453 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1935,6 +1935,38 @@ def _get_usage(agent) -> dict: return usage +def _restore_session_usage(agent, stored: dict) -> None: + """Restore cumulative token counters from a stored session row. + + When ``session.resume`` builds a fresh agent its counters start at zero. + The real cumulative usage lives in ``state.db`` (written by + ``update_token_counts`` on every API call). Without this restore the + ``session.info`` / ``session.usage`` payloads report all-zero usage + which shows as ``0/1.0M-0%`` in the Desktop status bar for gateway + sessions viewed after the fact. + """ + agent.session_input_tokens = stored.get("input_tokens") or 0 + agent.session_output_tokens = stored.get("output_tokens") or 0 + agent.session_cache_read_tokens = stored.get("cache_read_tokens") or 0 + agent.session_cache_write_tokens = stored.get("cache_write_tokens") or 0 + agent.session_reasoning_tokens = stored.get("reasoning_tokens") or 0 + agent.session_total_tokens = ( + agent.session_input_tokens + + agent.session_output_tokens + + agent.session_cache_read_tokens + + agent.session_cache_write_tokens + + agent.session_reasoning_tokens + ) + agent.session_api_calls = stored.get("api_call_count") or 0 + # Preserve cost estimate if available + cost_usd = stored.get("estimated_cost_usd") + if cost_usd: + agent.session_estimated_cost_usd = float(cost_usd) + cost_status = stored.get("cost_status") + if cost_status: + agent.session_cost_status = cost_status + + def _probe_credentials(agent) -> str: """Light credential check at session creation — returns warning or ''.""" try: @@ -3585,6 +3617,11 @@ def _(rid, params: dict) -> dict: # state.db; home override is active here so config/skills/model # resolve to the profile too. agent = _make_agent(sid, target, session_id=target, session_db=db) + # Restore cumulative token counters from the stored session so + # that session.info / session.usage shows correct usage instead + # of all zeros when viewing gateway sessions in Desktop. + if found: + _restore_session_usage(agent, found) finally: _clear_session_context(tokens) except Exception as e: