Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 108 additions & 9 deletions menubar/token_meter_tray.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,17 @@
("today_spend", "Today spend"),
)
DEFAULT_TITLE_METRICS = ("cost", "speed", "today_spend")
STATUS_DISPLAY_MODES = (
("automatic", "Automatic"),
("text", "Cost + speed text"),
("icon", "Icon only"),
)
QUOTA_THRESHOLDS = (80, 90, 95)
DEFAULT_STATE = {
"pinned_session": None,
"selected_tab": "run",
"title_metrics": list(DEFAULT_TITLE_METRICS),
"status_display_mode": "automatic",
"quota_alerts_enabled": True,
"quota_alert_threshold": 80,
"quota_notification_states": {},
Expand All @@ -46,6 +52,18 @@
}


def normalize_status_display_mode(value):
if isinstance(value, str) and value in {mode for mode, _label in STATUS_DISPLAY_MODES}:
return value
return "automatic"


def status_display_label(title, mode):
# AppIndicator has no reliable screen-width signal; Automatic preserves
# the existing Linux text label, while Icon only explicitly hides it.
return "" if normalize_status_display_mode(mode) == "icon" else str(title or "")


def metric_available(availability, metric):
if not isinstance(availability, dict) or metric not in availability:
return True
Expand Down Expand Up @@ -517,6 +535,41 @@ def tray_status_title(state, title_metrics, providers, budget, offline=False):
return base


def tray_status_tooltip(state, title_metrics, providers, budget, offline=False):
if offline:
return "Token Meter · Server offline"
labels = snapshot_labels(state)
today_spend = state.get("today_spend") or {}
today_spend_value = (
f"{money(today_spend.get('total_cost'))}"
f"{' est' if today_spend.get('estimated') else ''}"
if today_spend.get("available") else "--"
)
descriptions = {
"cost": ("Current session cost", labels["cost_label"]),
"speed": ("Output speed", labels["output_speed_label"]),
"context": ("Current session context", labels["context_label"]),
"model": ("Current session model", labels["model"]),
"today_spend": ("Today's spend across sessions", today_spend_value),
}
parts = []
for metric, _label in TITLE_METRICS:
if metric not in title_metrics:
continue
if metric == "limits":
limits = limits_status_title(providers)
if limits:
parts.append(f"Provider limits: {limits}")
continue
description, value = descriptions[metric]
parts.append(f"{description}: {value}")
if not parts:
parts.append("Token Meter")
if budget_any_exceeded(budget):
parts.insert(0, "Monthly budget exceeded")
return "\n".join(parts)


def evaluate_quota_notifications(providers, settings, *, established):
states = dict(settings.get("quota_notification_states") or {})
enabled = bool(settings.get("quota_alerts_enabled", True))
Expand Down Expand Up @@ -711,6 +764,9 @@ def __init__(self):
self.pinned_session = self.settings.get("pinned_session")
self.selected_tab = self.settings.get("selected_tab") or "run"
self.title_metrics = set(self.settings.get("title_metrics") or DEFAULT_TITLE_METRICS)
self.status_display_mode = normalize_status_display_mode(
self.settings.get("status_display_mode")
)
self.quota_alerts_enabled = bool(self.settings.get("quota_alerts_enabled", True))
self.quota_alert_threshold = int(self.settings.get("quota_alert_threshold") or 80)
self.quota_observation_established = bool(
Expand Down Expand Up @@ -857,6 +913,18 @@ def __init__(self):
self.title_metrics_menu.append(item)
self.title_metrics_item.set_submenu(self.title_metrics_menu)
self.settings_menu.append(self.title_metrics_item)
self.display_mode_item, self.display_mode_menu = self._submenu_item("Menu bar display")
self.display_mode_items = {}
display_mode_group = None
for mode, label in STATUS_DISPLAY_MODES:
item = Gtk.RadioMenuItem.new_with_label_from_widget(display_mode_group, label)
display_mode_group = item
item.token_meter_status_display_mode = mode
item.connect("toggled", self._on_display_mode_toggled)
self.display_mode_items[mode] = item
self.display_mode_menu.append(item)
self.display_mode_item.set_submenu(self.display_mode_menu)
self.settings_menu.append(self.display_mode_item)
self.quota_alerts_item = Gtk.CheckMenuItem(label="Quota notifications")
self.quota_alerts_item.connect("toggled", self._on_quota_alerts_toggled)
self.settings_menu.append(self.quota_alerts_item)
Expand Down Expand Up @@ -969,9 +1037,16 @@ def load_state(self):
*merged["title_metrics"], "today_spend",
]))
merged["today_spend_title_preference"] = True
merged["status_display_mode"] = normalize_status_display_mode(
merged.get("status_display_mode")
)
display_mode_migrated = value.get("status_display_mode") != merged["status_display_mode"]
if migrated or display_mode_migrated:
persisted = dict(value)
persisted["title_metrics"] = merged["title_metrics"]
persisted["today_spend_title_preference"] = True
if migrated:
persisted["title_metrics"] = merged["title_metrics"]
persisted["today_spend_title_preference"] = True
persisted["status_display_mode"] = merged["status_display_mode"]
try:
atomic_write_json(STATE_PATH, persisted)
except OSError as exc:
Expand All @@ -986,6 +1061,7 @@ def save_state(self):
"pinned_session": self.pinned_session,
"selected_tab": self.selected_tab,
"title_metrics": [key for key, _ in TITLE_METRICS if key in self.title_metrics],
"status_display_mode": self.status_display_mode,
"today_spend_title_preference": "today_spend" in self.title_metrics,
"quota_alerts_enabled": self.quota_alerts_enabled,
"quota_alert_threshold": self.quota_alert_threshold,
Expand Down Expand Up @@ -1054,14 +1130,24 @@ def _evaluate_notifications(self):

def update_tray_label(self):
if self.error:
title = tray_status_title({}, self.title_metrics, [], None, offline=True)
self.indicator.set_label("TM off", "TM off")
state = {}
providers = []
budget = None
title = tray_status_title(state, self.title_metrics, providers, budget, offline=True)
tooltip = tray_status_tooltip(state, self.title_metrics, providers, budget, offline=True)
label = "TM off"
else:
title = tray_status_title(
self.snapshot, self.title_metrics, self.provider_quotas, self.monthly_budget,
)
self.indicator.set_label(title, "⚠︎ $999.99 · 999 tok/s · 100% ctx")
self.indicator.set_title(title)
state = self.snapshot
providers = self.provider_quotas
budget = self.monthly_budget
title = tray_status_title(state, self.title_metrics, providers, budget)
tooltip = tray_status_tooltip(state, self.title_metrics, providers, budget)
label = title
self.indicator.set_label(
status_display_label(label, self.status_display_mode),
"⚠︎ $999.99 · 999 tok/s · 100% ctx",
)
self.indicator.set_title(tooltip)

def open_url(self, path="", include_pinned_session=True):
url = dashboard_url(path, self.pinned_session, include_pinned_session)
Expand Down Expand Up @@ -1116,6 +1202,17 @@ def _on_title_metric_toggled(self, item):
self.update_tray_label()
self.refresh_menu_content()

def _on_display_mode_toggled(self, item):
if not item.get_active():
return
mode = getattr(item, "token_meter_status_display_mode", None)
if mode not in {available for available, _label in STATUS_DISPLAY_MODES}:
return
self.status_display_mode = mode
self.save_state()
self.update_tray_label()
self.refresh_menu_content()

def _on_quota_alerts_toggled(self, item):
self.quota_alerts_enabled = item.get_active()
self.save_state()
Expand Down Expand Up @@ -1231,6 +1328,8 @@ def _apply_menu_content(self):
item.set_active(threshold == self.quota_alert_threshold)
for metric_id, item in self.title_metric_items.items():
item.set_active(metric_id in self.title_metrics)
for mode, item in self.display_mode_items.items():
item.set_active(mode == self.status_display_mode)

if self.error:
self._set_visible(self.connection_item, True)
Expand Down
125 changes: 125 additions & 0 deletions tests/test_meter.py
Original file line number Diff line number Diff line change
Expand Up @@ -14032,6 +14032,89 @@ def test_dashboard_url_opens_sessions_unless_pinned(self):
"http://127.0.0.1:8722/sessions/abc#summary",
)

def test_status_display_mode_selects_or_hides_the_tray_label(self):
label = "$1 · 12 tok/s"
self.assertEqual(self.tray.normalize_status_display_mode(None), "automatic")
self.assertEqual(self.tray.normalize_status_display_mode("invalid"), "automatic")
self.assertEqual(self.tray.normalize_status_display_mode({"mode": "icon"}), "automatic")
self.assertEqual(self.tray.status_display_label(label, "automatic"), label)
self.assertEqual(self.tray.status_display_label(label, "text"), label)
self.assertEqual(self.tray.status_display_label(label, "icon"), "")

def test_update_tray_label_keeps_title_when_icon_mode_hides_label(self):
tray = object.__new__(self.tray.TokenMeterTray)
tray.error = ""
tray.title_metrics = set()
tray.snapshot = {}
tray.provider_quotas = []
tray.monthly_budget = None
tray.indicator = mock.Mock()
tray.status_display_mode = "icon"

tray.update_tray_label()

tray.indicator.set_label.assert_called_once_with(
"", "⚠︎ $999.99 · 999 tok/s · 100% ctx",
)
tray.indicator.set_title.assert_called_once_with("Token Meter")

def test_update_tray_label_uses_descriptive_tooltip_but_keeps_compact_status(self):
tray = object.__new__(self.tray.TokenMeterTray)
tray.error = ""
tray.title_metrics = {"cost", "speed", "today_spend"}
tray.snapshot = {
"availability": {"cost": True, "throughput": True, "tokens": True},
"source": {"label": "Claude", "token_estimate": False},
"total_cost": 1.25,
"throughput": {"available": True, "output_tps": 12.3},
"today_spend": {"available": True, "total_cost": 2.75, "estimated": True},
}
tray.provider_quotas = []
tray.monthly_budget = None
tray.indicator = mock.Mock()
tray.status_display_mode = "automatic"

tray.update_tray_label()

tray.indicator.set_label.assert_called_once_with(
"$1 · 12 tok/s · $3 est", "⚠︎ $999.99 · 999 tok/s · 100% ctx",
)
tray.indicator.set_title.assert_called_once_with(
"Current session cost: $1.25\nOutput speed: 12.3 tok/s\n"
"Today's spend across sessions: $2.75 est",
)

def test_update_tray_label_explains_offline_status_in_tooltip(self):
tray = object.__new__(self.tray.TokenMeterTray)
tray.error = "server unavailable"
tray.title_metrics = {"cost", "speed"}
tray.indicator = mock.Mock()
tray.status_display_mode = "automatic"

tray.update_tray_label()

tray.indicator.set_label.assert_called_once_with(
"TM off", "⚠︎ $999.99 · 999 tok/s · 100% ctx",
)
tray.indicator.set_title.assert_called_once_with("Token Meter · Server offline")

def test_display_mode_selection_persists_and_refreshes_the_tray(self):
tray = object.__new__(self.tray.TokenMeterTray)
tray.status_display_mode = "automatic"
tray.save_state = mock.Mock()
tray.update_tray_label = mock.Mock()
tray.refresh_menu_content = mock.Mock()
item = mock.Mock()
item.get_active.return_value = True
item.token_meter_status_display_mode = "icon"

tray._on_display_mode_toggled(item)

self.assertEqual(tray.status_display_mode, "icon")
tray.save_state.assert_called_once_with()
tray.update_tray_label.assert_called_once_with()
tray.refresh_menu_content.assert_called_once_with()

def test_tray_status_title_uses_selected_metrics_and_budget_warning(self):
state = {
"availability": {"cost": True, "throughput": True, "cache": True, "tokens": True},
Expand Down Expand Up @@ -14109,6 +14192,34 @@ def test_load_state_persists_today_spend_migration(self):
])
self.assertTrue(persisted["today_spend_title_preference"])

def test_load_state_migrates_and_preserves_status_display_mode(self):
with tempfile.TemporaryDirectory() as tmp:
path = Path(tmp) / "tray.json"
path.write_text(json.dumps({
"today_spend_title_preference": True,
"title_metrics": ["cost", "speed"],
}), encoding="utf-8")
with mock.patch.object(self.tray, "STATE_PATH", str(path)):
loaded = self.tray.TokenMeterTray.load_state(object.__new__(
self.tray.TokenMeterTray,
))
self.assertEqual(loaded["status_display_mode"], "automatic")
self.assertEqual(
json.loads(path.read_text(encoding="utf-8"))["status_display_mode"],
"automatic",
)

path.write_text(json.dumps({
"today_spend_title_preference": True,
"title_metrics": ["cost"],
"status_display_mode": "icon",
}), encoding="utf-8")
with mock.patch.object(self.tray, "STATE_PATH", str(path)):
loaded = self.tray.TokenMeterTray.load_state(object.__new__(
self.tray.TokenMeterTray,
))
self.assertEqual(loaded["status_display_mode"], "icon")

def test_parse_monthly_budget_preserves_unconfigured_runtime_scopes(self):
budget = self.tray.parse_monthly_budget({
"configured": True,
Expand Down Expand Up @@ -14212,6 +14323,20 @@ def test_linux_tray_exposes_tabs_title_metrics_and_notifications(self):
):
self.assertIn(marker, self.source)

def test_linux_tray_exposes_persisted_status_display_modes(self):
for marker in (
'"status_display_mode": "automatic"',
'self.status_display_mode = normalize_status_display_mode',
'"status_display_mode": self.status_display_mode',
'self.display_mode_item, self.display_mode_menu = self._submenu_item("Menu bar display")',
'("automatic", "Automatic")',
'("text", "Cost + speed text")',
'("icon", "Icon only")',
'def _on_display_mode_toggled(self, item):',
'status_display_label(label, self.status_display_mode)',
):
self.assertIn(marker, self.source)


class ProviderQuotaTests(unittest.TestCase):
def tearDown(self):
Expand Down
Loading