From fd324562d3ad43df7631428df4585cf9de8154f2 Mon Sep 17 00:00:00 2001 From: Austin Pickett Date: Mon, 29 Jun 2026 09:18:10 -0400 Subject: [PATCH] feat(desktop): add context usage breakdown popover Let users click the status bar context indicator to see how tokens are split across system prompt, tools, rules, skills, MCP, and conversation. Co-authored-by: Cursor --- agent/context_breakdown.py | 156 ++++++++++++++++++ .../src/app/shell/context-usage-panel.tsx | 147 +++++++++++++++++ .../app/shell/hooks/use-statusbar-items.tsx | 21 ++- apps/desktop/src/i18n/en.ts | 18 ++ apps/desktop/src/i18n/ja.ts | 18 ++ apps/desktop/src/i18n/types.ts | 18 ++ apps/desktop/src/i18n/zh-hant.ts | 18 ++ apps/desktop/src/i18n/zh.ts | 18 ++ apps/desktop/src/styles.css | 8 + apps/desktop/src/types/hermes.ts | 16 ++ tests/agent/test_context_breakdown.py | 60 +++++++ tui_gateway/server.py | 30 ++++ 12 files changed, 524 insertions(+), 4 deletions(-) create mode 100644 agent/context_breakdown.py create mode 100644 apps/desktop/src/app/shell/context-usage-panel.tsx create mode 100644 tests/agent/test_context_breakdown.py diff --git a/agent/context_breakdown.py b/agent/context_breakdown.py new file mode 100644 index 00000000000..0e2eb772f2f --- /dev/null +++ b/agent/context_breakdown.py @@ -0,0 +1,156 @@ +"""Live session context-window breakdown for UI surfaces. + +Estimates how the next provider request is composed: system prompt tiers, +tool schemas, and conversation history. Uses the same rough char/4 heuristic +as ``agent.model_metadata.estimate_request_tokens_rough`` so numbers align +with compression thresholds — not exact tokenizer counts. +""" + +from __future__ import annotations + +import json +import re +from typing import Any, Dict, List, Optional, Sequence, Tuple + +_SKILLS_BLOCK_RE = re.compile(r".*?", re.DOTALL) + +_SUBAGENT_TOOL_NAMES = frozenset({"delegate_task"}) + +_CATEGORY_COLORS = { + "system_prompt": "var(--context-usage-system)", + "tool_definitions": "var(--context-usage-tools)", + "rules": "var(--context-usage-rules)", + "skills": "var(--context-usage-skills)", + "mcp": "var(--context-usage-mcp)", + "subagent_definitions": "var(--context-usage-subagents)", + "memory": "var(--context-usage-memory)", + "conversation": "var(--context-usage-conversation)", +} + + +def _chars_to_tokens(text: str) -> int: + if not text: + return 0 + return (len(text) + 3) // 4 + + +def _json_tokens(value: Any) -> int: + if not value: + return 0 + return _chars_to_tokens(json.dumps(value, ensure_ascii=False)) + + +def _tool_name(tool: dict) -> str: + fn = tool.get("function") if isinstance(tool, dict) else None + if isinstance(fn, dict): + return str(fn.get("name") or "") + return str(tool.get("name") or "") + + +def _split_tools(tools: Sequence[dict]) -> Tuple[List[dict], List[dict], List[dict]]: + builtin: List[dict] = [] + mcp: List[dict] = [] + subagent: List[dict] = [] + for tool in tools: + name = _tool_name(tool) + if name.startswith("mcp_"): + mcp.append(tool) + elif name in _SUBAGENT_TOOL_NAMES: + subagent.append(tool) + else: + builtin.append(tool) + return builtin, mcp, subagent + + +def _memory_blocks(agent: Any) -> Tuple[str, str]: + memory_block = "" + user_block = "" + store = getattr(agent, "_memory_store", None) + if store is None: + return memory_block, user_block + try: + if getattr(agent, "_memory_enabled", True): + memory_block = store.format_for_system_prompt("memory") or "" + if getattr(agent, "_user_profile_enabled", True): + user_block = store.format_for_system_prompt("user") or "" + except Exception: + pass + return memory_block, user_block + + +def _strip_blocks(text: str, *blocks: str) -> str: + out = text + for block in blocks: + if block: + out = out.replace(block, "") + return out.strip() + + +def compute_session_context_breakdown( + agent: Any, + messages: Optional[List[dict]] = None, +) -> Dict[str, Any]: + """Return a Cursor-style context usage breakdown for one live agent.""" + from agent.model_metadata import estimate_messages_tokens_rough + from agent.system_prompt import build_system_prompt_parts + + parts = build_system_prompt_parts(agent) + stable = parts.get("stable", "") or "" + context = parts.get("context", "") or "" + volatile = parts.get("volatile", "") or "" + + skills_match = _SKILLS_BLOCK_RE.search(stable) + skills_index = skills_match.group(0) if skills_match else "" + + memory_block, user_block = _memory_blocks(agent) + memory_text = "\n\n".join(part for part in (memory_block, user_block) if part).strip() + + system_core = _strip_blocks(stable, skills_index) + system_tail = _strip_blocks(volatile, memory_block, user_block) + system_prompt_text = "\n\n".join(part for part in (system_core, system_tail) if part).strip() + + tools = list(getattr(agent, "tools", None) or []) + builtin_tools, mcp_tools, subagent_tools = _split_tools(tools) + + conversation_tokens = estimate_messages_tokens_rough(messages or []) + + categories = [ + ("system_prompt", "System prompt", _chars_to_tokens(system_prompt_text)), + ("tool_definitions", "Tool definitions", _json_tokens(builtin_tools)), + ("rules", "Rules", _chars_to_tokens(context)), + ("skills", "Skills", _chars_to_tokens(skills_index)), + ("mcp", "MCP", _json_tokens(mcp_tools)), + ("subagent_definitions", "Subagent definitions", _json_tokens(subagent_tools)), + ("memory", "Memory", _chars_to_tokens(memory_text)), + ("conversation", "Conversation", conversation_tokens), + ] + + estimated_total = sum(tokens for _, _, tokens in categories) + + comp = getattr(agent, "context_compressor", None) + context_max = int(getattr(comp, "context_length", 0) or 0) if comp else 0 + measured_used = int(getattr(comp, "last_prompt_tokens", 0) or 0) if comp else 0 + context_used = measured_used if measured_used > 0 else estimated_total + context_percent = ( + max(0, min(100, round(context_used / context_max * 100))) + if context_max + else 0 + ) + + return { + "categories": [ + { + "color": _CATEGORY_COLORS.get(category_id, "var(--ui-text-tertiary)"), + "id": category_id, + "label": label, + "tokens": tokens, + } + for category_id, label, tokens in categories + if tokens > 0 + ], + "context_max": context_max, + "context_percent": context_percent, + "context_used": context_used, + "estimated_total": estimated_total, + "model": getattr(agent, "model", "") or "", + } diff --git a/apps/desktop/src/app/shell/context-usage-panel.tsx b/apps/desktop/src/app/shell/context-usage-panel.tsx new file mode 100644 index 00000000000..5343515ef04 --- /dev/null +++ b/apps/desktop/src/app/shell/context-usage-panel.tsx @@ -0,0 +1,147 @@ +import { useEffect, useMemo, useState } from 'react' + +import { useI18n } from '@/i18n' +import { formatK } from '@/lib/statusbar' +import { cn } from '@/lib/utils' +import type { ContextBreakdown, ContextUsageCategory, UsageStats } from '@/types/hermes' + +interface ContextUsagePanelProps { + currentUsage: UsageStats + requestGateway: (method: string, params?: Record) => Promise + sessionId: string | null +} + +export function ContextUsagePanel({ currentUsage, requestGateway, sessionId }: ContextUsagePanelProps) { + const { t } = useI18n() + const copy = t.shell.statusbar.contextUsagePanel + const [breakdown, setBreakdown] = useState(null) + const [loading, setLoading] = useState(false) + + useEffect(() => { + if (!sessionId) { + setBreakdown(null) + setLoading(false) + return + } + + let cancelled = false + setLoading(true) + + void requestGateway('session.context_breakdown', { session_id: sessionId }) + .then(data => { + if (!cancelled) { + setBreakdown(data) + } + }) + .catch(() => { + if (!cancelled) { + setBreakdown(null) + } + }) + .finally(() => { + if (!cancelled) { + setLoading(false) + } + }) + + return () => { + cancelled = true + } + }, [requestGateway, sessionId]) + + const contextMax = breakdown?.context_max ?? currentUsage.context_max ?? 0 + const contextUsed = breakdown?.context_used ?? currentUsage.context_used ?? 0 + const contextPercent = Math.max( + 0, + Math.min(100, Math.round(breakdown?.context_percent ?? currentUsage.context_percent ?? 0)) + ) + + const categories = useMemo( + () => + (breakdown?.categories ?? []).map(category => ({ + ...category, + label: copy.categories[category.id as keyof typeof copy.categories] ?? category.label + })), + [breakdown?.categories, copy.categories] + ) + + const segmentTotal = categories.reduce((sum, category) => sum + category.tokens, 0) || contextUsed || 1 + + return ( +
+
+

{copy.title}

+ + + {copy.tokenSummary(`~${formatK(contextUsed)}`, formatK(contextMax))} + +
+ +

{copy.percentFull(contextPercent)}

+ + + +
    + {categories.map(category => ( +
  • + + + + {category.label} + + + {formatCategoryTokens(category.tokens)} +
    • + ))} +
+ + {loading &&

{copy.loading}

} + + {!loading && !categories.length &&

{copy.empty}

} +
+ ) +} + +function ContextUsageBar({ + categories, + segmentTotal +}: { + categories: readonly ContextUsageCategory[] + segmentTotal: number +}) { + return ( +
+ {categories.map(category => ( + + ))} +
+ ) +} + +function formatCategoryTokens(value: number): string { + if (!Number.isFinite(value) || value <= 0) { + return '0' + } + + if (value >= 1_000) { + return `${formatK(value)}` + } + + return value.toLocaleString() +} diff --git a/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx b/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx index b6328be6543..c58483ad095 100644 --- a/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx +++ b/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx @@ -4,6 +4,7 @@ import { useCallback, useMemo } from 'react' import type { CommandCenterSection } from '@/app/command-center' import { $terminalTakeover, setTerminalTakeover } from '@/app/right-sidebar/store' import { GatewayMenuPanel } from '@/app/shell/gateway-menu-panel' +import { ContextUsagePanel } from '@/app/shell/context-usage-panel' import { Codicon } from '@/components/ui/codicon' import { GlyphSpinner } from '@/components/ui/glyph-spinner' import { useI18n } from '@/i18n' @@ -365,8 +366,17 @@ export function useStatusbarItems({ hidden: !contextUsage, id: 'context-usage', label: contextUsage, - title: copy.contextUsage, - variant: 'text' + menuAlign: 'end', + menuClassName: 'w-auto border-(--ui-stroke-secondary) p-0', + menuContent: ( + + ), + title: copy.openContextUsage, + variant: 'menu' }, { detail: , @@ -402,18 +412,21 @@ export function useStatusbarItems({ ...(backendVersionItem ? [backendVersionItem] : []) ], [ + activeSessionId, + backendVersionItem, busy, chatOpen, + clientVersionItem, contextBar, contextUsage, copy, + currentUsage, + requestGateway, sessionStartedAt, showYoloToggle, terminalTakeover, toggleYolo, turnStartedAt, - clientVersionItem, - backendVersionItem, yoloActive ] ) diff --git a/apps/desktop/src/i18n/en.ts b/apps/desktop/src/i18n/en.ts index f7239b83aae..2afe579aeec 100644 --- a/apps/desktop/src/i18n/en.ts +++ b/apps/desktop/src/i18n/en.ts @@ -1844,6 +1844,24 @@ export const en: Translations = { turnRunning: 'Running', currentTurnElapsed: 'Current turn elapsed', contextUsage: 'Context usage', + contextUsagePanel: { + categories: { + conversation: 'Conversation', + mcp: 'MCP', + memory: 'Memory', + rules: 'Rules', + skills: 'Skills', + subagent_definitions: 'Subagent definitions', + system_prompt: 'System prompt', + tool_definitions: 'Tool definitions' + }, + empty: 'No context data yet', + loading: 'Loading breakdown…', + percentFull: percent => `${percent}% Full`, + title: 'Context Usage', + tokenSummary: (used, max) => `${used} / ${max} Tokens` + }, + openContextUsage: 'Open context usage breakdown', session: 'Session', runtimeSessionElapsed: 'Runtime session elapsed', yoloOn: 'YOLO on — auto-approving dangerous commands. Click to turn off. Shift+click toggles it globally.', diff --git a/apps/desktop/src/i18n/ja.ts b/apps/desktop/src/i18n/ja.ts index e407c454837..2959c14bfb2 100644 --- a/apps/desktop/src/i18n/ja.ts +++ b/apps/desktop/src/i18n/ja.ts @@ -1964,6 +1964,24 @@ export const ja = defineLocale({ turnRunning: '実行中', currentTurnElapsed: '現在のターン経過時間', contextUsage: 'コンテキスト使用状況', + contextUsagePanel: { + categories: { + conversation: '会話', + mcp: 'MCP', + memory: 'メモリ', + rules: 'ルール', + skills: 'スキル', + subagent_definitions: 'サブエージェント定義', + system_prompt: 'システムプロンプト', + tool_definitions: 'ツール定義' + }, + empty: 'コンテキストデータはまだありません', + loading: '内訳を読み込み中…', + percentFull: percent => `${percent}% 使用中`, + title: 'コンテキスト使用状況', + tokenSummary: (used, max) => `${used} / ${max} Tokens` + }, + openContextUsage: 'コンテキスト使用状況の内訳を開く', session: 'セッション', runtimeSessionElapsed: 'ランタイムセッション経過時間', yoloOn: 'YOLO オン — 危険なコマンドを自動承認中。クリックでオフに。Shift+クリックで全体に切り替え。', diff --git a/apps/desktop/src/i18n/types.ts b/apps/desktop/src/i18n/types.ts index 742bdfa5e7c..c759d471039 100644 --- a/apps/desktop/src/i18n/types.ts +++ b/apps/desktop/src/i18n/types.ts @@ -1497,6 +1497,24 @@ export interface Translations { turnRunning: string currentTurnElapsed: string contextUsage: string + contextUsagePanel: { + categories: { + conversation: string + mcp: string + memory: string + rules: string + skills: string + subagent_definitions: string + system_prompt: string + tool_definitions: string + } + empty: string + loading: string + percentFull: (percent: number) => string + title: string + tokenSummary: (used: string, max: string) => string + } + openContextUsage: string session: string runtimeSessionElapsed: string yoloOn: string diff --git a/apps/desktop/src/i18n/zh-hant.ts b/apps/desktop/src/i18n/zh-hant.ts index e7d607a0772..4b265e3a803 100644 --- a/apps/desktop/src/i18n/zh-hant.ts +++ b/apps/desktop/src/i18n/zh-hant.ts @@ -1903,6 +1903,24 @@ export const zhHant = defineLocale({ turnRunning: '執行中', currentTurnElapsed: '目前回合已用時間', contextUsage: '上下文使用量', + contextUsagePanel: { + categories: { + conversation: '對話', + mcp: 'MCP', + memory: '記憶', + rules: '規則', + skills: '技能', + subagent_definitions: '子代理定義', + system_prompt: '系統提示詞', + tool_definitions: '工具定義' + }, + empty: '尚無上下文資料', + loading: '正在載入明細…', + percentFull: percent => `已用 ${percent}%`, + title: '上下文使用量', + tokenSummary: (used, max) => `${used} / ${max} Tokens` + }, + openContextUsage: '開啟上下文使用量明細', session: '工作階段', runtimeSessionElapsed: '執行時工作階段已用時間', yoloOn: 'YOLO 已開啟 — 自動核准危險指令。點擊關閉。Shift+點擊可全域切換。', diff --git a/apps/desktop/src/i18n/zh.ts b/apps/desktop/src/i18n/zh.ts index aa02579771d..02c2834b31c 100644 --- a/apps/desktop/src/i18n/zh.ts +++ b/apps/desktop/src/i18n/zh.ts @@ -2015,6 +2015,24 @@ export const zh: Translations = { turnRunning: '运行中', currentTurnElapsed: '当前回合已用时间', contextUsage: '上下文用量', + contextUsagePanel: { + categories: { + conversation: '对话', + mcp: 'MCP', + memory: '记忆', + rules: '规则', + skills: '技能', + subagent_definitions: '子代理定义', + system_prompt: '系统提示词', + tool_definitions: '工具定义' + }, + empty: '暂无上下文数据', + loading: '正在加载明细…', + percentFull: percent => `已用 ${percent}%`, + title: '上下文用量', + tokenSummary: (used, max) => `${used} / ${max} Tokens` + }, + openContextUsage: '打开上下文用量明细', session: '会话', runtimeSessionElapsed: '运行时会话已用时间', yoloOn: 'YOLO 已开启 - 自动批准危险命令。点击关闭。Shift+点击可全局切换。', diff --git a/apps/desktop/src/styles.css b/apps/desktop/src/styles.css index 5e9a156b841..636d71c18a3 100644 --- a/apps/desktop/src/styles.css +++ b/apps/desktop/src/styles.css @@ -164,6 +164,14 @@ --ui-cyan: #4c7f8c; --ui-blue: #0053fd; --ui-purple: #9e94d5; + --context-usage-system: color-mix(in srgb, var(--ui-base) 55%, transparent); + --context-usage-tools: var(--ui-purple); + --context-usage-rules: var(--ui-green); + --context-usage-skills: var(--ui-yellow); + --context-usage-mcp: color-mix(in srgb, var(--ui-red) 72%, var(--ui-purple)); + --context-usage-subagents: color-mix(in srgb, var(--ui-blue) 70%, var(--ui-cyan)); + --context-usage-memory: color-mix(in srgb, var(--ui-orange) 80%, var(--ui-yellow)); + --context-usage-conversation: var(--ui-cyan); --ui-bg-chrome: color-mix( in srgb, var(--theme-background-seed) var(--theme-mix-chrome), diff --git a/apps/desktop/src/types/hermes.ts b/apps/desktop/src/types/hermes.ts index 08e29ce4b40..b5ae206c833 100644 --- a/apps/desktop/src/types/hermes.ts +++ b/apps/desktop/src/types/hermes.ts @@ -428,6 +428,22 @@ export interface UsageStats { total: number } +export interface ContextUsageCategory { + color: string + id: string + label: string + tokens: number +} + +export interface ContextBreakdown { + categories: ContextUsageCategory[] + context_max: number + context_percent: number + context_used: number + estimated_total: number + model?: string +} + export interface AnalyticsDailyEntry { actual_cost: number api_calls: number diff --git a/tests/agent/test_context_breakdown.py b/tests/agent/test_context_breakdown.py new file mode 100644 index 00000000000..d8a8c2fc2bb --- /dev/null +++ b/tests/agent/test_context_breakdown.py @@ -0,0 +1,60 @@ +"""Tests for live session context breakdown.""" + +from unittest.mock import MagicMock, patch + +from agent.context_breakdown import compute_session_context_breakdown + + +def _make_agent( + *, + stable: str = "identity and guidance", + context: str = "", + volatile: str = "timestamp line", + tools: list | None = None, + context_length: int = 200_000, + last_prompt_tokens: int = 0, +): + agent = MagicMock() + agent.model = "openai/gpt-5.4" + agent.tools = tools or [ + {"type": "function", "function": {"name": "terminal", "description": "run"}}, + {"type": "function", "function": {"name": "mcp_demo_tool", "description": "mcp"}}, + {"type": "function", "function": {"name": "delegate_task", "description": "spawn"}}, + ] + agent._memory_store = None + agent._memory_enabled = True + agent._user_profile_enabled = True + agent.context_compressor = MagicMock( + context_length=context_length, + last_prompt_tokens=last_prompt_tokens, + ) + return agent, {"stable": stable, "context": context, "volatile": volatile} + + +def test_breakdown_includes_major_categories(): + stable = ( + "base guidance\n" + "\n demo:\n - hello: hi\n" + ) + context = "# Project Context\nFollow AGENTS.md" + volatile = "Current time: now" + history = [{"role": "user", "content": "hello there"}] + agent, parts = _make_agent(stable=stable, context=context, volatile=volatile) + + with patch("agent.system_prompt.build_system_prompt_parts", return_value=parts): + data = compute_session_context_breakdown(agent, history) + + ids = {item["id"] for item in data["categories"]} + assert {"system_prompt", "tool_definitions", "rules", "skills", "mcp", "subagent_definitions", "conversation"} <= ids + assert data["context_max"] == 200_000 + assert data["estimated_total"] > 0 + + +def test_breakdown_uses_measured_context_when_available(): + agent, parts = _make_agent(last_prompt_tokens=42_000) + + with patch("agent.system_prompt.build_system_prompt_parts", return_value=parts): + data = compute_session_context_breakdown(agent, []) + + assert data["context_used"] == 42_000 + assert data["context_percent"] == 21 diff --git a/tui_gateway/server.py b/tui_gateway/server.py index ec873bfba5e..5ac8591f0a7 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -6165,6 +6165,36 @@ def _(rid, params: dict) -> dict: return _ok(rid, usage) +@method("session.context_breakdown") +def _(rid, params: dict) -> dict: + session, err = _sess_nowait(params, rid) + if err: + return err + agent = session.get("agent") + if agent is None: + usage = _get_usage(None) + return _ok( + rid, + { + "categories": [], + "context_max": usage.get("context_max", 0) or 0, + "context_percent": usage.get("context_percent", 0) or 0, + "context_used": usage.get("context_used", 0) or 0, + "estimated_total": 0, + "model": "", + }, + ) + with session["history_lock"]: + history = list(session.get("history", [])) + try: + from agent.context_breakdown import compute_session_context_breakdown + + payload = compute_session_context_breakdown(agent, history) + except Exception as exc: + return _err(rid, 5000, f"Could not compute context breakdown: {exc}") + return _ok(rid, payload) + + def _pet_frame_counts(spritesheet) -> dict: """Real (padding-trimmed) frame count per state, for the desktop canvas.