diff --git a/agent_core/core/session/session.py b/agent_core/core/session/session.py index 00ec9c4b..444e8c83 100644 --- a/agent_core/core/session/session.py +++ b/agent_core/core/session/session.py @@ -58,6 +58,8 @@ class Session: (reset when a new run starts). input_tokens/output_tokens/cache_tokens: LLM usage breakdown for the current run. + total_input_tokens/total_output_tokens/total_cache_tokens: the same + breakdown accumulated across every run in this session. """ id: str @@ -83,6 +85,11 @@ class Session: input_tokens: int = 0 output_tokens: int = 0 cache_tokens: int = 0 + # Cumulative session totals — deliberately NOT cleared by + # reset_run_counters(); they span every run in this session. + total_input_tokens: int = 0 + total_output_tokens: int = 0 + total_cache_tokens: int = 0 def touch(self) -> None: """Update last_active_at to now.""" @@ -138,6 +145,9 @@ def to_dict(self) -> Dict[str, Any]: "input_tokens": self.input_tokens, "output_tokens": self.output_tokens, "cache_tokens": self.cache_tokens, + "total_input_tokens": self.total_input_tokens, + "total_output_tokens": self.total_output_tokens, + "total_cache_tokens": self.total_cache_tokens, } @classmethod @@ -165,4 +175,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "Session": input_tokens=data.get("input_tokens", 0), output_tokens=data.get("output_tokens", 0), cache_tokens=data.get("cache_tokens", 0), + total_input_tokens=data.get("total_input_tokens", 0), + total_output_tokens=data.get("total_output_tokens", 0), + total_cache_tokens=data.get("total_cache_tokens", 0), ) diff --git a/app/ui_layer/browser/frontend/src/pages/Dashboard/DashboardPage.tsx b/app/ui_layer/browser/frontend/src/pages/Dashboard/DashboardPage.tsx index 3506b521..18dc66dd 100644 --- a/app/ui_layer/browser/frontend/src/pages/Dashboard/DashboardPage.tsx +++ b/app/ui_layer/browser/frontend/src/pages/Dashboard/DashboardPage.tsx @@ -153,15 +153,23 @@ export function DashboardPage() { // Token metrics - use cached filtered metrics for all periods (including 'total') const tokenFilteredData = filteredMetricsCache[tokenPeriod] - const inputTokens = tokenFilteredData?.token.input ?? (metrics?.token.input ?? 0) + const rawInputTokens = tokenFilteredData?.token.input ?? (metrics?.token.input ?? 0) const outputTokens = tokenFilteredData?.token.output ?? (metrics?.token.output ?? 0) - const totalTokens = tokenFilteredData?.token.total ?? (metrics?.token.total ?? 0) const cachedTokens = tokenFilteredData?.token.cached ?? (metrics?.token.cached ?? 0) + // `token.input` from the API is the full prompt size — cache reads included. + // The Input tile shows only the genuinely new tokens; Cached shows the rest. + const inputTokens = Math.max(0, rawInputTokens - cachedTokens) + // Total counts new tokens only. `token.total` from the API is + // rawInput + output, so it double-counts cache reads — derive instead. + const totalTokens = inputTokens + outputTokens + // Calculate token ratios const inputRatio = totalTokens > 0 ? Math.round((inputTokens / totalTokens) * 100) : 0 const outputRatio = totalTokens > 0 ? Math.round((outputTokens / totalTokens) * 100) : 0 - const cachedRatio = inputTokens > 0 ? Math.min(100, Math.round((cachedTokens / inputTokens) * 100)) : 0 + // Cache hit rate — share of the prompt served from cache. Denominator is the + // full prompt, not totalTokens: output tokens are never cacheable. + const cachedRatio = rawInputTokens > 0 ? Math.min(100, Math.round((cachedTokens / rawInputTokens) * 100)) : 0 const cpuPercent = metrics?.system.cpuPercent ?? 0 const memoryPercent = metrics?.system.memoryPercent ?? 0 diff --git a/app/ui_layer/commands/builtin/__init__.py b/app/ui_layer/commands/builtin/__init__.py index 75a79eed..787d2c44 100644 --- a/app/ui_layer/commands/builtin/__init__.py +++ b/app/ui_layer/commands/builtin/__init__.py @@ -11,6 +11,7 @@ from app.ui_layer.commands.builtin.cred import CredCommand from app.ui_layer.commands.builtin.integrations import IntegrationCommand from app.ui_layer.commands.builtin.update import UpdateCommand +from app.ui_layer.commands.builtin.tokens import TokensCommand from app.ui_layer.commands.builtin.agent_command import AgentCommandWrapper from app.ui_layer.commands.builtin.skill_invoke import SkillInvokeCommand @@ -26,6 +27,7 @@ "CredCommand", "IntegrationCommand", "UpdateCommand", + "TokensCommand", "AgentCommandWrapper", "SkillInvokeCommand", ] diff --git a/app/ui_layer/commands/builtin/tokens.py b/app/ui_layer/commands/builtin/tokens.py new file mode 100644 index 00000000..e61a5520 --- /dev/null +++ b/app/ui_layer/commands/builtin/tokens.py @@ -0,0 +1,59 @@ +"""Tokens command implementation — shows this session's token usage.""" + +from __future__ import annotations + +from typing import List + +from agent_core.core.session import MAIN_SESSION_ID + +from app.ui_layer.commands.base import Command, CommandResult + + +class TokensCommand(Command): + """Show cumulative token usage for the session it was typed in.""" + + @property + def name(self) -> str: + return "/tokens" + + @property + def description(self) -> str: + return "Show this session's token usage" + + async def execute( + self, + args: List[str], + adapter_id: str = "", + session_id: str | None = None, + ) -> CommandResult: + """Report this session's token totals as a chat message.""" + target = session_id or MAIN_SESSION_ID + # Sessions are created lazily on first message, so a brand-new chat + # has no session object yet — that reads as zero usage, not an error. + session = self._controller.agent.session_manager.get(target) + + # Providers report input as the full prompt, cache reads included, so + # `cached` is a subset of `total_input_tokens` — never a sibling. + cached = getattr(session, "total_cache_tokens", 0) or 0 + new_input = max(0, (getattr(session, "total_input_tokens", 0) or 0) - cached) + output = getattr(session, "total_output_tokens", 0) or 0 + total = new_input + output + + message = ( + "Session token usage\n" + f" Input: {new_input:,}\n" + f" Cached: {cached:,}\n" + f" Output: {output:,}\n" + f" Total: {total:,}" + ) + return CommandResult( + success=True, + message=message, + data={ + "session_id": target, + "input": new_input, + "cached": cached, + "output": output, + "total": total, + }, + ) diff --git a/app/ui_layer/controller/ui_controller.py b/app/ui_layer/controller/ui_controller.py index 87152add..7cdddf16 100644 --- a/app/ui_layer/controller/ui_controller.py +++ b/app/ui_layer/controller/ui_controller.py @@ -433,6 +433,7 @@ def _register_builtin_commands(self) -> None: SkillCommand, CredCommand, UpdateCommand, + TokensCommand, ) self._command_registry.register(HelpCommand(self)) @@ -445,6 +446,7 @@ def _register_builtin_commands(self) -> None: self._command_registry.register(SkillCommand(self)) self._command_registry.register(CredCommand(self)) self._command_registry.register(UpdateCommand(self)) + self._command_registry.register(TokensCommand(self)) # Register integration commands self._register_integration_commands() diff --git a/app/usage/task_attribution.py b/app/usage/task_attribution.py index 93dc0e43..7d23f9f3 100644 --- a/app/usage/task_attribution.py +++ b/app/usage/task_attribution.py @@ -44,6 +44,18 @@ def attribute_usage_to_current_task(event: UsageEventData) -> None: event.cached_tokens or 0 ) + # Cumulative session totals — survive reset_run_counters(), which + # clears the per-run counters above at the start of every run. + session.total_input_tokens = (session.total_input_tokens or 0) + int( + event.input_tokens or 0 + ) + session.total_output_tokens = (session.total_output_tokens or 0) + int( + event.output_tokens or 0 + ) + session.total_cache_tokens = (session.total_cache_tokens or 0) + int( + event.cached_tokens or 0 + ) + logger.info( f"[TOKEN_ATTR] session={session.id} +in={event.input_tokens} " f"+out={event.output_tokens} +cached={event.cached_tokens} " diff --git a/tests/test_tokens_command.py b/tests/test_tokens_command.py new file mode 100644 index 00000000..e8b941f5 --- /dev/null +++ b/tests/test_tokens_command.py @@ -0,0 +1,98 @@ +# -*- coding: utf-8 -*- +"""Checks for /tokens: cumulative counters and the display math.""" + +import asyncio +from types import SimpleNamespace + +from agent_core.core.session import Session +from app.ui_layer.commands.builtin.tokens import TokensCommand + + +def _run(session): + """Execute /tokens against a stub controller holding `session`.""" + controller = SimpleNamespace( + agent=SimpleNamespace( + session_manager=SimpleNamespace(get=lambda _id: session) + ) + ) + return asyncio.run(TokensCommand(controller).execute([], session_id="s1")) + + +def test_totals_survive_run_reset(): + """reset_run_counters() clears per-run counters but not the totals.""" + s = Session(id="s1") + s.input_tokens, s.output_tokens, s.cache_tokens = 100, 20, 80 + s.total_input_tokens, s.total_output_tokens, s.total_cache_tokens = 100, 20, 80 + + s.reset_run_counters() + + assert s.input_tokens == 0 + assert (s.total_input_tokens, s.total_output_tokens, s.total_cache_tokens) == ( + 100, + 20, + 80, + ) + + +def test_totals_round_trip_through_dict(): + """Persistence goes through session_json, so to_dict/from_dict must carry them.""" + s = Session(id="s1") + s.total_input_tokens, s.total_output_tokens, s.total_cache_tokens = 100, 20, 80 + + restored = Session.from_dict(s.to_dict()) + + assert restored.total_input_tokens == 100 + assert restored.total_output_tokens == 20 + assert restored.total_cache_tokens == 80 + + +def test_old_sessions_load_as_zero(): + """Sessions persisted before this feature have no totals key.""" + restored = Session.from_dict({"id": "s1"}) + + assert restored.total_input_tokens == 0 + + +def test_input_excludes_cached_and_total_excludes_both(): + """cached is a subset of input; total counts new tokens only.""" + s = Session(id="s1") + s.total_input_tokens, s.total_output_tokens, s.total_cache_tokens = 100, 20, 80 + + d = _run(s).data + + assert d["input"] == 20 # 100 - 80 + assert d["cached"] == 80 + assert d["output"] == 20 + assert d["total"] == 40 # 20 + 20, cached excluded + + +def test_input_clamps_when_cached_exceeds_input(): + """A provider over-reporting cache reads must not yield a negative count.""" + s = Session(id="s1") + s.total_input_tokens, s.total_cache_tokens = 50, 80 + + d = _run(s).data + + assert d["input"] == 0 + assert d["total"] == 0 + + +def test_unsaved_session_reads_as_zero(): + """A brand-new chat has no session yet (id is "new") — report zeros.""" + result = _run(None) + + assert result.success is True + assert result.data["input"] == 0 + assert result.data["cached"] == 0 + assert result.data["output"] == 0 + assert result.data["total"] == 0 + + +if __name__ == "__main__": + test_totals_survive_run_reset() + test_totals_round_trip_through_dict() + test_old_sessions_load_as_zero() + test_input_excludes_cached_and_total_excludes_both() + test_input_clamps_when_cached_exceeds_input() + test_unsaved_session_reads_as_zero() + print("ok")