diff --git a/menubar/token_meter_tray.py b/menubar/token_meter_tray.py index 82ddd9d..4f037a1 100755 --- a/menubar/token_meter_tray.py +++ b/menubar/token_meter_tray.py @@ -11,6 +11,7 @@ BASE_URL = os.environ.get("TOKEN_METER_URL", "http://127.0.0.1:8722").rstrip("/") STATE_URL = BASE_URL + "/menubar" +TOKENOMICS_URL = "https://www.splunk.com/en_us/products/tokenomics.html" CONFIG_HOME = os.path.expanduser(os.environ.get("XDG_CONFIG_HOME", "~/.config")) STATE_PATH = os.path.join(CONFIG_HOME, "token-meter", "tray.json") MAX_RECENT_SESSIONS = 10 @@ -29,8 +30,9 @@ ("context", "Context"), ("model", "Model"), ("limits", "Limits"), + ("today_spend", "Today spend"), ) -DEFAULT_TITLE_METRICS = ("cost", "speed") +DEFAULT_TITLE_METRICS = ("cost", "speed", "today_spend") QUOTA_THRESHOLDS = (80, 90, 95) DEFAULT_STATE = { "pinned_session": None, @@ -50,6 +52,22 @@ def metric_available(availability, metric): return bool(availability.get(metric)) +def atomic_write_json(path, payload): + os.makedirs(os.path.dirname(path), mode=0o700, exist_ok=True) + temporary = path + ".tmp" + try: + with open(temporary, "w", encoding="utf-8") as handle: + json.dump(payload, handle) + os.chmod(temporary, 0o600) + os.replace(temporary, path) + except OSError: + try: + os.unlink(temporary) + except OSError: + pass + raise + + def money(value): value = float(value or 0) if abs(value) >= 1: @@ -172,16 +190,19 @@ def parse_monthly_budget(payload): if not isinstance(row, dict): continue allocation = float(row.get("allocation") or 0) - percent = row.get("percent") provider = row.get("provider") - if allocation <= 0 or percent is None or not provider: + if not provider: continue + percent = ( + float(row.get("percent") or 0) * 100 + if allocation > 0 else 0 + ) scopes.append({ "id": str(provider), "label": str(row.get("label") or str(provider).title()), "spend": float(row.get("spend") or 0), "budget": allocation, - "percent": float(percent) * 100, + "percent": percent, }) if configured: scopes.insert(0, { @@ -311,10 +332,18 @@ def snapshot_labels(state): context = state.get("context") or {} cache = state.get("cache") or {} throughput = state.get("throughput") or {} + live_throughput = state.get("live_throughput") or {} cost_available = metric_available(availability, "cost") tokens_available = metric_available(availability, "tokens") cache_available = metric_available(availability, "cache") throughput_available = metric_available(availability, "throughput") + selected_throughput = ( + live_throughput + if throughput_available and live_throughput.get("available") + else throughput + ) + today_spend = state.get("today_spend") or {} + today_spend_available = bool(today_spend.get("available")) estimated_cost = bool(state.get("cost_approx") or source.get("approximate_cost")) estimated_tokens = bool(source.get("token_estimate")) context_pct = context.get("latest_pct") @@ -329,9 +358,9 @@ def snapshot_labels(state): if pct <= 1: pct *= 100 context_label = f"{int(round(pct))}% ctx" - if throughput_available and throughput.get("available") and throughput.get("output_tps"): + if throughput_available and selected_throughput.get("available") and selected_throughput.get("output_tps"): output_speed_label = ( - f"{format_token_rate(throughput['output_tps'])} tok/s" + f"{format_token_rate(selected_throughput['output_tps'])} tok/s" f"{' est' if estimated_tokens else ''}" ) else: @@ -347,19 +376,42 @@ def snapshot_labels(state): f"{compact_number(state.get('total_tokens'))}{' est' if estimated_tokens else ''}" if tokens_available else "--" ) + today_spend_label = ( + f"${float(today_spend.get('total_cost') or 0):.0f}" + f"{' est' if today_spend.get('estimated') else ''}" + if today_spend_available else "--" + ) + menu_bar_context_label = "--%" + if context_pct is not None: + pct = float(context_pct) + if pct <= 1: + pct *= 100 + menu_bar_context_label = f"{int(round(pct))}%" return { "model": str(state.get("model") or source.get("model") or "unknown"), "cost_label": cost_label, + "menu_bar_cost_label": ( + f"${float(state.get('total_cost') or 0):.0f}" + if cost_available else "--" + ), "context_label": context_label, + "menu_bar_context_label": menu_bar_context_label, "output_speed_label": output_speed_label, + "menu_bar_output_speed_label": ( + f"{float(selected_throughput['output_tps']):.0f} tok/s" + if throughput_available and selected_throughput.get("available") and selected_throughput.get("output_tps") + else "-- tok/s" + ), "cache_label": cache_label, "tokens_label": tokens_label, + "today_spend_label": today_spend_label, "estimated_cost": estimated_cost, "estimated_tokens": estimated_tokens, "cost_available": cost_available, "tokens_available": tokens_available, "cache_available": cache_available, "throughput_available": throughput_available, + "today_spend_available": today_spend_available, "last_turn_cost": money(state.get("last_turn_cost")) if cost_available else "--", "turns": int(state.get("turns") or 0), "context_tokens": compact_number(context.get("latest")), @@ -401,17 +453,19 @@ def tray_status_title(state, title_metrics, providers, budget, offline=False): if key not in title_metrics: continue if key == "cost": - parts.append(labels["cost_label"]) + parts.append(labels["menu_bar_cost_label"]) elif key == "speed": - parts.append(labels["output_speed_label"]) + parts.append(labels["menu_bar_output_speed_label"]) elif key == "context": - parts.append(labels["context_label"]) + parts.append(labels["menu_bar_context_label"]) elif key == "model": parts.append(labels["model"]) elif key == "limits": limits = limits_status_title(providers) if limits: parts.append(limits) + elif key == "today_spend": + parts.append(labels["today_spend_label"]) base = " · ".join(parts) if parts else "TM" if budget_any_exceeded(budget): return f"⚠︎ {base}" @@ -704,7 +758,29 @@ def __init__(self): lambda *_: self.open_url("#capabilities", include_pinned_session=False), ), } + self.budgets_item, self.budgets_menu = self._submenu_item("Budgets") + self.budget_empty_item = self._metric_item("") + self.budget_overall_item = self._action_item( + "", lambda *_: self.open_url("#settings-budgets", include_pinned_session=False), + ) + self.budget_separator = Gtk.SeparatorMenuItem() + self.budget_scope_items = [ + self._action_item( + "", lambda *_: self.open_url("#settings-budgets", include_pinned_session=False), + ) + for _ in range(MAX_PROVIDERS) + ] + self.budgets_menu.append(self.budget_empty_item) + self.budgets_menu.append(self.budget_overall_item) + self.budgets_menu.append(self.budget_separator) + for item in self.budget_scope_items: + self.budgets_menu.append(item) self.settings_item, self.settings_menu = self._submenu_item("Settings") + self.open_settings_item = self._action_item( + "Open Settings", + lambda *_: self.open_url("#settings", include_pinned_session=False), + ) + self.settings_menu.append(self.open_settings_item) self.model_prices_item = self._action_item( "Model Prices", lambda *_: self.open_url("#model-pricing", include_pinned_session=False), @@ -736,6 +812,10 @@ def __init__(self): self.threshold_menu.append(item) self.threshold_item.set_submenu(self.threshold_menu) self.settings_menu.append(self.threshold_item) + self.enterprise_item = self._action_item( + "Get Enterprise Tokenomics", + lambda *_: self.open_external_url(TOKENOMICS_URL), + ) self.quit_separator = Gtk.SeparatorMenuItem() self.quit_item = self._action_item("Quit Token Meter Tray", lambda *_: Gtk.main_quit()) self._assemble_menu() @@ -784,10 +864,12 @@ def _assemble_menu(self): self.menu.append(window["pace"]) self.menu.append(self.provider_detail_items["coverage"]) self.menu.append(self.provider_detail_items["footer"]) + self.menu.append(self.budgets_item) self.menu.append(self.actions_separator) for item in self.action_items.values(): self.menu.append(item) self.menu.append(self.settings_item) + self.menu.append(self.enterprise_item) self.menu.append(self.quit_separator) self.menu.append(self.quit_item) self.menu.show_all() @@ -820,6 +902,19 @@ def load_state(self): merged.update(value) if not merged.get("title_metrics"): merged["title_metrics"] = list(DEFAULT_TITLE_METRICS) + migrated = "today_spend_title_preference" not in value + if migrated: + merged["title_metrics"] = list(dict.fromkeys([ + *merged["title_metrics"], "today_spend", + ])) + merged["today_spend_title_preference"] = True + persisted = dict(value) + persisted["title_metrics"] = merged["title_metrics"] + persisted["today_spend_title_preference"] = True + try: + atomic_write_json(STATE_PATH, persisted) + except OSError as exc: + print(f"Token Meter could not migrate tray state: {exc}", file=sys.stderr) threshold = int(merged.get("quota_alert_threshold") or 80) if threshold not in QUOTA_THRESHOLDS: merged["quota_alert_threshold"] = 80 @@ -830,6 +925,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], + "today_spend_title_preference": "today_spend" in self.title_metrics, "quota_alerts_enabled": self.quota_alerts_enabled, "quota_alert_threshold": self.quota_alert_threshold, "quota_notification_states": self.settings.get("quota_notification_states") or {}, @@ -840,12 +936,7 @@ def save_state(self): "quota_observation_established": self.quota_observation_established, } try: - os.makedirs(os.path.dirname(STATE_PATH), mode=0o700, exist_ok=True) - temporary = STATE_PATH + ".tmp" - with open(temporary, "w", encoding="utf-8") as handle: - json.dump(payload, handle) - os.chmod(temporary, 0o600) - os.replace(temporary, STATE_PATH) + atomic_write_json(STATE_PATH, payload) self.settings.update(payload) except OSError as exc: print(f"Token Meter could not save tray state: {exc}", file=sys.stderr) @@ -911,6 +1002,9 @@ def update_tray_label(self): def open_url(self, path="", include_pinned_session=True): url = dashboard_url(path, self.pinned_session, include_pinned_session) + self.open_external_url(url) + + def open_external_url(self, url): try: subprocess.Popen( ["xdg-open", url], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, @@ -954,6 +1048,7 @@ def _on_title_metric_toggled(self, item): self.title_metrics.add(metric) else: self.title_metrics.discard(metric) + self.settings["today_spend_title_preference"] = "today_spend" in self.title_metrics self.save_state() self.update_tray_label() self.refresh_menu_content() @@ -1010,6 +1105,47 @@ def _hide_provider_detail_section(self): for key in window: self._set_visible(window[key], False) + def _apply_budget_menu(self): + budget = self.monthly_budget + has_budget = isinstance(budget, dict) + self._set_visible(self.budget_empty_item, not has_budget) + self._set_visible(self.budget_overall_item, has_budget) + self._set_visible(self.budget_separator, False) + for item in self.budget_scope_items: + self._set_visible(item, False) + if not has_budget: + self._set_item_label(self.budget_empty_item, "Budget data unavailable") + return + if budget.get("configured"): + overall = ( + f"Overall · {money(budget.get('spend'))} of {money(budget.get('budget'))} · " + f"{int(round(float(budget.get('percent') or 0)))}% used" + ) + else: + overall = "Overall · Not set" + self._set_item_label(self.budget_overall_item, overall) + scopes = [scope for scope in budget.get("scopes") or [] if scope.get("id") != "overall"] + self._set_visible(self.budget_separator, bool(scopes)) + for index, scope in enumerate(scopes): + if index >= len(self.budget_scope_items): + item = self._action_item( + "", lambda *_: self.open_url("#settings-budgets", include_pinned_session=False), + ) + self.budget_scope_items.append(item) + self.budgets_menu.append(item) + item = self.budget_scope_items[index] + allocation = float(scope.get("budget") or 0) + label = str(scope.get("label") or scope.get("id") or "Runtime") + if allocation > 0: + value = ( + f"{label} · {money(scope.get('spend'))} of {money(allocation)} · " + f"{int(round(float(scope.get('percent') or 0)))}% used" + ) + else: + value = f"{label} · Not set" + self._set_item_label(item, value) + self._set_visible(item, True) + def _apply_menu_content(self): for tab_id, item in self.tab_items.items(): item.set_active(tab_id == self.selected_tab) @@ -1025,6 +1161,7 @@ def _apply_menu_content(self): self._hide_run_section() self._hide_overview_section() self._hide_provider_detail_section() + self._set_visible(self.budgets_item, False) self._set_visible(self.tab_separator, False) self._set_visible(self.view_item, False) return False @@ -1035,6 +1172,8 @@ def _apply_menu_content(self): self._hide_run_section() self._hide_overview_section() self._hide_provider_detail_section() + self._set_visible(self.budgets_item, True) + self._apply_budget_menu() if self.selected_tab == "run": self._apply_run_tab() diff --git a/tests/test_meter.py b/tests/test_meter.py index e90fcdf..56102ba 100644 --- a/tests/test_meter.py +++ b/tests/test_meter.py @@ -10416,10 +10416,76 @@ def test_tray_status_title_uses_selected_metrics_and_budget_warning(self): state, {"cost", "speed", "limits"}, providers, budget, ) self.assertTrue(title.startswith("⚠︎ ")) - self.assertIn("$1.25", title) - self.assertIn("12.3 tok/s", title) + self.assertIn("$1", title) + self.assertIn("12 tok/s", title) self.assertIn("Claude 88% · weekly", title) + def test_tray_status_title_includes_today_spend_with_native_compact_labels(self): + state = { + "availability": {"cost": True, "throughput": True, "tokens": True}, + "source": {"label": "Claude", "token_estimate": False}, + "total_cost": 1.25, + "cost_approx": True, + "context": {"latest_pct": 0.42}, + "throughput": {"available": True, "output_tps": 12.3}, + "live_throughput": {"available": True, "output_tps": 24.6}, + "today_spend": {"available": True, "total_cost": 2.75, "estimated": True}, + } + title = self.tray.tray_status_title( + state, {"cost", "speed", "context", "today_spend"}, [], None, + ) + self.assertEqual(title, "$1 · 25 tok/s · 42% · $3 est") + self.assertNotIn("$1 est", title) + + unavailable_context = dict(state) + unavailable_context["context"] = {} + unavailable_title = self.tray.tray_status_title( + unavailable_context, {"context"}, [], None, + ) + self.assertEqual(unavailable_title, "--%") + + def test_load_state_persists_today_spend_migration(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "tray.json" + path.write_text(json.dumps({ + "title_metrics": ["cost", "speed"], + "selected_tab": "run", + }), 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.assertIn("today_spend", loaded["title_metrics"]) + persisted = json.loads(path.read_text(encoding="utf-8")) + self.assertEqual(persisted["title_metrics"], [ + "cost", "speed", "today_spend", + ]) + self.assertTrue(persisted["today_spend_title_preference"]) + + def test_parse_monthly_budget_preserves_unconfigured_runtime_scopes(self): + budget = self.tray.parse_monthly_budget({ + "configured": True, + "month": "2026-07", + "spend": 6.5, + "budget": 1000, + "percent": 0.0065, + "runtimes": [ + { + "provider": "claude", "label": "Claude", "spend": 0, + "allocation": 0, "percent": None, + }, + { + "provider": "codex", "label": "Codex", "spend": 6.5, + "allocation": 1000, "percent": 0.0065, + }, + ], + }) + self.assertEqual([scope["id"] for scope in budget["scopes"]], [ + "overall", "claude", "codex", + ]) + self.assertEqual(budget["scopes"][1]["percent"], 0) + self.assertAlmostEqual(budget["scopes"][2]["percent"], 0.65) + def test_budget_notifications_fire_on_threshold_crossing(self): budget = { "configured": True, @@ -10463,8 +10529,14 @@ def test_linux_tray_matches_native_dashboard_routing_and_actions(self): 'self.open_url("#summary")', "Open Budget Settings", "#settings-budgets", + "Open Settings", + "#settings", "Model Prices", "#model-pricing", + "Budgets", + "Get Enterprise Tokenomics", + "TOKENOMICS_URL", + "open_external_url", ): self.assertIn(marker, self.source) @@ -10476,6 +10548,11 @@ def test_linux_tray_exposes_tabs_title_metrics_and_notifications(self): for marker in ( '("overview", "All")', '("claude", "Claude")', + '("today_spend", "Today spend")', + 'DEFAULT_TITLE_METRICS = ("cost", "speed", "today_spend")', + "_apply_budget_menu", + "for index, scope in enumerate(scopes):", + "self.budget_scope_items.append(item)", "Menu bar title", "Quota notifications", "Warn at",