From 196ef8cadcc944ea5e4495185bdccca74b2b3b49 Mon Sep 17 00:00:00 2001 From: David Metcalfe Date: Mon, 22 Jun 2026 20:40:47 -0700 Subject: [PATCH] feat: add provider-agnostic balance/credits display to Hermes Desktop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a general system for displaying AI provider balance/credits information in the Desktop statusbar, starting with Kilo AI. Backend: - agent/balance_provider.py: BalanceProvider ABC, ProviderBalance dataclass, BalanceConfig, and thread-safe BalanceProviderRegistry with TTL caching - plugins/model-providers/kilocode/balance.py: KiloBalanceProvider fetching GET https://api.kilo.ai/api/profile/balance - tui_gateway/server.py: @method('balance.view') RPC handler — resolves the active runtime provider, merges config.yaml overrides with per-class defaults, and returns cached or fresh balance data Desktop frontend: - apps/desktop/src/types/hermes.ts: ProviderBalance and BalanceViewResponse types - apps/desktop/src/store/provider-balance.ts: nanostore atom for balance state - apps/desktop/src/lib/hooks/use-provider-balance.ts: React hook with 2min auto-poll, gateway-connect trigger, and click-to-force-refresh - apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx: statusbar item between session-timer and YOLO toggle (color-coded: red=depleted, amber=low) - apps/desktop/src/app/desktop-controller.tsx: wired useProviderBalance Design: provider-agnostic from day one. Adding a new provider is one Python file + one-line BalanceProviderRegistry.register() call — no frontend changes. Config keys under providers..balance: { endpoint, enabled, cache_ttl_seconds }. Fail-open on all error paths. Closes #51175 --- agent/balance_provider.py | 238 ++++++++++++++++++ apps/desktop/src/app/desktop-controller.tsx | 5 + .../app/shell/hooks/use-statusbar-items.tsx | 32 ++- .../src/lib/hooks/use-provider-balance.ts | 68 +++++ apps/desktop/src/store/provider-balance.ts | 43 ++++ apps/desktop/src/types/hermes.ts | 20 ++ plugins/model-providers/kilocode/__init__.py | 5 + plugins/model-providers/kilocode/balance.py | 49 ++++ tui_gateway/server.py | 77 ++++++ 9 files changed, 536 insertions(+), 1 deletion(-) create mode 100644 agent/balance_provider.py create mode 100644 apps/desktop/src/lib/hooks/use-provider-balance.ts create mode 100644 apps/desktop/src/store/provider-balance.ts create mode 100644 plugins/model-providers/kilocode/balance.py diff --git a/agent/balance_provider.py b/agent/balance_provider.py new file mode 100644 index 000000000000..434bf981db0d --- /dev/null +++ b/agent/balance_provider.py @@ -0,0 +1,238 @@ +""" +Provider Balance/Credits ABC +============================= + +Defines the pluggable interface for fetching balance/credit information from +an AI provider. Providers register subclasses via ``BalanceProviderRegistry``; +the Desktop frontend displays the active provider's balance via the +``balance.view`` RPC method. + +Adding a new provider is one file + one registration call — no frontend changes +needed. See ``plugins/model-providers/kilocode/balance.py`` for an example. +""" + +from __future__ import annotations + +import logging +import threading +import time +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, ClassVar, Optional + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Data types +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ProviderBalance: + """Canonical balance data returned by every BalanceProvider. + + ``value`` is a float denominated in the provider's currency (usually USD). + ``currency`` is the ISO 4217 code (default ``"USD"``). ``label`` is a + human-friendly short string like ``"Kilo AI"``. + + ``is_depleted`` is a hint from the provider API (not derived from + ``value > 0`` — some providers have a separate flag). ``fetched_at`` is + set by the registry, not the provider subclass. + + When ``error`` is non-None the fetch failed and ``value`` should not + be displayed. + """ + + provider_name: str + label: str + value: float + currency: str = "USD" + is_depleted: bool = False + fetched_at: float = 0.0 # unix timestamp + error: Optional[str] = None # non-None → fetch failed + + def __str__(self) -> str: + """Compact display string, e.g. ``"$6.61"`` or ``"125,000"``.""" + if self.error: + return "" + if self.currency == "USD": + return f"${self.value:.2f}" + return f"{self.value:.2f}" + + +@dataclass(frozen=True) +class BalanceConfig: + """Per-provider config snippet from ``config.yaml providers..balance:``. + + Default values match the Kilo endpoint; providers override as needed. + """ + + endpoint: str = "" + api_key_env: str = "" + enabled: bool = True + # How long (in seconds) cached balance data is considered fresh. + cache_ttl_seconds: float = 60.0 + + +# --------------------------------------------------------------------------- +# ABC +# --------------------------------------------------------------------------- + + +class BalanceProvider(ABC): + """One registered balance provider. + + Subclass name must match the Hermes provider slug (e.g. + ``KiloBalanceProvider`` lives in the ``kilocode`` plugin; its + ``provider_slug`` is ``"kilocode"``). + """ + + # Must match the slug used in ``providers: {}`` config key and plugin name. + provider_slug: ClassVar[str] = "" + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if not cls.provider_slug: + raise TypeError( + f"{cls.__name__} must define a non-empty 'provider_slug' class variable" + ) + + @abstractmethod + def fetch(self, api_key: str, config: BalanceConfig) -> ProviderBalance: + """Fetch current balance from the provider API. + + Called on a thread off the event loop (blocking I/O is fine). Must + return a ``ProviderBalance`` — if the API is unreachable, set + ``error`` on the return value (do not raise). + """ + + @classmethod + def default_config(cls) -> BalanceConfig: + """Default ``BalanceConfig`` when none is specified in config.yaml.""" + return BalanceConfig() + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + + +class BalanceProviderRegistry: + """Holds all registered ``BalanceProvider`` subclasses. + + Registration is automatic via ``register()`` called alongside the existing + ``register_provider()`` in each model provider plugin's ``__init__.py``. + """ + + _providers: dict[str, type[BalanceProvider]] = {} + _cache: dict[str, ProviderBalance] = {} + _last_fetch: dict[str, float] = {} # provider_slug → unix timestamp + _fetch_in_flight: set[str] = set() # provider_slugs currently being fetched + _lock = threading.Lock() + + @classmethod + def register(cls, provider_cls: type[BalanceProvider]) -> None: + """Register a ``BalanceProvider`` subclass.""" + slug = provider_cls.provider_slug + if not slug: + raise ValueError(f"{provider_cls.__name__} must set provider_slug") + with cls._lock: + cls._providers[slug] = provider_cls + logger.info("BalanceProvider registered: %s", slug) + + @classmethod + def registered_slugs(cls) -> frozenset[str]: + """Return all registered provider slugs.""" + with cls._lock: + return frozenset(cls._providers) + + @classmethod + def get(cls, slug: str) -> type[BalanceProvider] | None: + """Return the registered provider class for *slug*, or None.""" + with cls._lock: + return cls._providers.get(slug) + + @classmethod + def get_cached(cls, slug: str) -> ProviderBalance | None: + """Return cached balance for *slug*, or None.""" + with cls._lock: + return cls._cache.get(slug) + + @classmethod + def cached_or_fetch( + cls, + slug: str, + api_key: str, + config: BalanceConfig, + *, + force: bool = False, + ) -> tuple[ProviderBalance, bool]: + """Return cached balance if fresh, else fetch and cache. + + Returns ``(ProviderBalance, was_cached)`` where ``was_cached`` is + True when the result came from the in-memory cache (fresh enough + per ``cache_ttl_seconds``). + + If another thread is already fetching the same slug, the second + caller receives stale cache rather than issuing a duplicate HTTP + request. + """ + now = time.time() + + # Return cached if fresh enough and not forced. + with cls._lock: + if slug in cls._fetch_in_flight: + # Another thread is already fetching; return stale cache. + cached = cls._cache.get(slug) + if cached: + return cached, True + + if not force: + last = cls._last_fetch.get(slug, 0.0) + cached = cls._cache.get(slug) + if cached and (now - last) < config.cache_ttl_seconds: + return cached, True + + # Mark in-flight so concurrent callers see stale cache. + cls._fetch_in_flight.add(slug) + + # Fetch (outside the lock to avoid holding during I/O). + provider_cls = cls.get(slug) + if provider_cls is None: + with cls._lock: + cls._fetch_in_flight.discard(slug) + return ProviderBalance( + provider_name=slug, + label=slug, + value=0.0, + error=f"No BalanceProvider registered for '{slug}'", + ), False + + try: + balance = provider_cls().fetch(api_key, config) + balance = ProviderBalance( + provider_name=balance.provider_name, + label=balance.label, + value=balance.value, + currency=balance.currency, + is_depleted=balance.is_depleted, + fetched_at=now, + error=balance.error, + ) + except Exception as exc: + logger.debug("BalanceProvider '%s' fetch failed: %s", slug, exc) + balance = ProviderBalance( + provider_name=slug, + label=slug, + value=0.0, + fetched_at=now, + error=str(exc), + ) + + with cls._lock: + cls._cache[slug] = balance + cls._last_fetch[slug] = now + cls._fetch_in_flight.discard(slug) + + return balance, False \ No newline at end of file diff --git a/apps/desktop/src/app/desktop-controller.tsx b/apps/desktop/src/app/desktop-controller.tsx index ced02523d22f..0acfb266f8c2 100644 --- a/apps/desktop/src/app/desktop-controller.tsx +++ b/apps/desktop/src/app/desktop-controller.tsx @@ -125,6 +125,7 @@ import { AppShell } from './shell/app-shell' import { useOverlayRouting } from './shell/hooks/use-overlay-routing' import { useStatusSnapshot } from './shell/hooks/use-status-snapshot' import { useStatusbarItems } from './shell/hooks/use-statusbar-items' +import { useProviderBalance } from '@/lib/hooks/use-provider-balance' import { ModelMenuPanel } from './shell/model-menu-panel' import type { StatusbarItem } from './shell/statusbar-controls' import type { TitlebarTool } from './shell/titlebar-controls' @@ -909,18 +910,22 @@ export function DesktopController() { startFreshSessionDraft }) + const { balance: providerBalance, fetchBalance } = useProviderBalance(requestGateway, gatewayState) + const { leftStatusbarItems, statusbarItems } = useStatusbarItems({ agentsOpen, chatOpen, commandCenterOpen, extraLeftItems: statusbarItemGroups.flat.left, extraRightItems: statusbarItemGroups.flat.right, + fetchBalance, gatewayLogLines, gatewayState, inferenceStatus, openAgents, freshDraftReady, openCommandCenterSection, + providerBalance, requestGateway, statusSnapshot, toggleCommandCenter 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 a95ac3217f55..263d6e0eb569 100644 --- a/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx +++ b/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx @@ -45,7 +45,7 @@ import { $updateStatus, openUpdateOverlayFor } from '@/store/updates' -import type { StatusResponse } from '@/types/hermes' +import type { ProviderBalance, StatusResponse } from '@/types/hermes' import { CRON_ROUTE } from '../../routes' import type { StatusbarItem, StatusbarSelectModifiers } from '../statusbar-controls' @@ -56,12 +56,14 @@ interface StatusbarItemsOptions { commandCenterOpen: boolean extraLeftItems: readonly StatusbarItem[] extraRightItems: readonly StatusbarItem[] + fetchBalance: (force?: boolean) => Promise gatewayLogLines: readonly string[] gatewayState: string inferenceStatus: RuntimeReadinessResult | null openAgents: () => void openCommandCenterSection: (section: CommandCenterSection) => void freshDraftReady: boolean + providerBalance: ProviderBalance | null requestGateway: (method: string, params?: Record) => Promise statusSnapshot: StatusResponse | null toggleCommandCenter: () => void @@ -73,12 +75,14 @@ export function useStatusbarItems({ commandCenterOpen, extraLeftItems, extraRightItems, + fetchBalance, gatewayLogLines, gatewayState, inferenceStatus, openAgents, openCommandCenterSection, freshDraftReady, + providerBalance, requestGateway, statusSnapshot, toggleCommandCenter @@ -399,6 +403,30 @@ export function useStatusbarItems({ title: copy.runtimeSessionElapsed, variant: 'text' }, + // ── Provider balance (hidden when unavailable) ──────────── + ...(providerBalance + ? [ + { + detail: `${providerBalance.currency === 'USD' ? '$' : ''}${providerBalance.value.toFixed(2)}`, + hidden: false, + icon: providerBalance.is_depleted ? ( + + ) : providerBalance.value < 5 ? ( + + ) : undefined, + id: 'provider-balance', + label: providerBalance.label, + onSelect: () => void fetchBalance(true), + title: [ + `${providerBalance.label}: ${providerBalance.currency === 'USD' ? '$' : ''}${providerBalance.value.toFixed(2)}`, + providerBalance.is_depleted ? ' — Depleted' : '', + !providerBalance.is_depleted && providerBalance.value < 5 ? ' — Low balance' : '', + ].filter(Boolean).join(''), + variant: 'action' as const, + } as StatusbarItem + ] + : []), + // ── YOLO toggle ───────────────────────────────────────── { className: cn('px-1', yoloActive && 'bg-(--chrome-action-hover)'), hidden: !showYoloToggle, @@ -437,6 +465,8 @@ export function useStatusbarItems({ turnStartedAt, clientVersionItem, backendVersionItem, + fetchBalance, + providerBalance, yoloActive ] ) diff --git a/apps/desktop/src/lib/hooks/use-provider-balance.ts b/apps/desktop/src/lib/hooks/use-provider-balance.ts new file mode 100644 index 000000000000..bee97f5463a0 --- /dev/null +++ b/apps/desktop/src/lib/hooks/use-provider-balance.ts @@ -0,0 +1,68 @@ +import { useStore } from '@nanostores/react' +import { useCallback, useEffect, useRef } from 'react' + +import { + $providerBalance, + setProviderBalance, + setProviderBalanceError, + setProviderBalanceLoading, +} from '@/store/provider-balance' +import type { BalanceViewResponse, ProviderBalance } from '@/types/hermes' + +const POLL_MS = 120_000 // 2 minutes + +type GatewayRequester = (method: string, params?: Record) => Promise + +export function useProviderBalance( + requestGateway: GatewayRequester, + gatewayState: string | undefined, +): { balance: ProviderBalance | null; fetchBalance: (force?: boolean) => Promise } { + const state = useStore($providerBalance) + const pollRef = useRef | null>(null) + + const fetchBalance = useCallback( + async (force = false) => { + setProviderBalanceLoading(true) + try { + const res = await requestGateway('balance.view', { force }) + if (res.ok && res.balance) { + setProviderBalance(res.balance) + } else { + setProviderBalanceError(res.error ?? 'unknown error') + } + } catch (err) { + setProviderBalanceError(err instanceof Error ? err.message : String(err)) + } + }, + [requestGateway], + ) + + // Fetch on mount and when gateway opens. + useEffect(() => { + if (gatewayState === 'open') { + void fetchBalance() + } + }, [gatewayState, fetchBalance]) + + // Poll while gateway is open. + useEffect(() => { + if (gatewayState !== 'open') { + if (pollRef.current) { + clearInterval(pollRef.current) + pollRef.current = null + } + return + } + + pollRef.current = setInterval(() => void fetchBalance(), POLL_MS) + + return () => { + if (pollRef.current) { + clearInterval(pollRef.current) + pollRef.current = null + } + } + }, [gatewayState, fetchBalance]) + + return { balance: state.balance, fetchBalance } +} \ No newline at end of file diff --git a/apps/desktop/src/store/provider-balance.ts b/apps/desktop/src/store/provider-balance.ts new file mode 100644 index 000000000000..aabc82389a9e --- /dev/null +++ b/apps/desktop/src/store/provider-balance.ts @@ -0,0 +1,43 @@ +import { atom } from 'nanostores' + +import type { ProviderBalance } from '@/types/hermes' + +export interface ProviderBalanceState { + balance: ProviderBalance | null + loading: boolean + error: string | null + fetchedAtMs: number | null +} + +export const $providerBalance = atom({ + balance: null, + loading: false, + error: null, + fetchedAtMs: null, +}) + +export function setProviderBalance(balance: ProviderBalance): void { + $providerBalance.set({ + balance, + loading: false, + error: null, + fetchedAtMs: Date.now(), + }) +} + +export function setProviderBalanceError(error: string): void { + const current = $providerBalance.get() + $providerBalance.set({ + ...current, + loading: false, + error, + fetchedAtMs: Date.now(), + }) +} + +export function setProviderBalanceLoading(loading: boolean): void { + $providerBalance.set({ + ...$providerBalance.get(), + loading, + }) +} \ No newline at end of file diff --git a/apps/desktop/src/types/hermes.ts b/apps/desktop/src/types/hermes.ts index 1dc2d6be50e5..13fe24b4557f 100644 --- a/apps/desktop/src/types/hermes.ts +++ b/apps/desktop/src/types/hermes.ts @@ -765,3 +765,23 @@ export interface ModelAssignmentResponse { stale_aux?: StaleAuxAssignment[] tasks?: string[] } + + +// ── Provider balance (provider-agnostic) ──────────────────────────── + + +export interface ProviderBalance { + provider_name: string + label: string + value: number + currency: string + is_depleted: boolean + fetched_at: number // unix timestamp +} + +export interface BalanceViewResponse { + ok: boolean + balance: ProviderBalance | null + error: string | null + cached: boolean +} diff --git a/plugins/model-providers/kilocode/__init__.py b/plugins/model-providers/kilocode/__init__.py index 23123966aaca..e4ebb472819b 100644 --- a/plugins/model-providers/kilocode/__init__.py +++ b/plugins/model-providers/kilocode/__init__.py @@ -3,6 +3,10 @@ from providers import register_provider from providers.base import ProviderProfile +from agent.balance_provider import BalanceProviderRegistry + +from .balance import KiloBalanceProvider + kilocode = ProviderProfile( name="kilocode", aliases=("kilo-code", "kilo", "kilo-gateway"), @@ -12,3 +16,4 @@ ) register_provider(kilocode) +BalanceProviderRegistry.register(KiloBalanceProvider) diff --git a/plugins/model-providers/kilocode/balance.py b/plugins/model-providers/kilocode/balance.py new file mode 100644 index 000000000000..ff6a5026f935 --- /dev/null +++ b/plugins/model-providers/kilocode/balance.py @@ -0,0 +1,49 @@ +"""Kilo AI balance provider — fetches from the Kilo profile API.""" + +from __future__ import annotations + +import httpx + +from agent.balance_provider import BalanceConfig, BalanceProvider, ProviderBalance + + +class KiloBalanceProvider(BalanceProvider): + provider_slug = "kilocode" + + def fetch(self, api_key: str, config: BalanceConfig) -> ProviderBalance: + endpoint = config.endpoint or "https://api.kilo.ai/api/profile/balance" + headers = { + "Authorization": f"Bearer {api_key}", + "Accept": "application/json", + } + try: + with httpx.Client(timeout=10.0) as client: + resp = client.get(endpoint, headers=headers) + resp.raise_for_status() + data = resp.json() + except Exception as exc: + return ProviderBalance( + provider_name=self.provider_slug, + label="Kilo AI", + value=0.0, + error=f"{type(exc).__name__}: {exc}", + ) + + balance = float(data.get("balance", 0.0)) + depleted = bool(data.get("isDepleted", False)) + + return ProviderBalance( + provider_name=self.provider_slug, + label="Kilo AI", + value=balance, + is_depleted=depleted, + ) + + @classmethod + def default_config(cls) -> BalanceConfig: + return BalanceConfig( + endpoint="https://api.kilo.ai/api/profile/balance", + api_key_env="KILOCODE_API_KEY", + enabled=True, + cache_ttl_seconds=60.0, + ) \ No newline at end of file diff --git a/tui_gateway/server.py b/tui_gateway/server.py index ad3ea68cdd43..889e28908cb4 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -5475,6 +5475,83 @@ def _(rid, params: dict) -> dict: return _ok(rid, {"logged_in": False, "balance_lines": [], "identity_line": None, "topup_url": None, "depleted": False}) +@method("balance.view") +def _(rid, params: dict) -> dict: + """Fetch balance for the current active provider. + + Returns a ``ProviderBalance``-shaped dict, or an error envelope. + The Desktop client uses this via ``requestGateway('balance.view')``. + Fail-open: no active provider or no registered balance provider + returns ``{"ok": False}`` — the UI hides the balance item. + """ + try: + import os + + from agent.balance_provider import BalanceConfig, BalanceProviderRegistry + from hermes_cli.config import read_config + from hermes_cli.runtime_provider import resolve_runtime_provider + + # 1. Determine active provider slug. + provider = resolve_runtime_provider() + if not provider or not provider.get("provider"): + return _ok(rid, {"ok": False, "error": "no_active_provider"}) + + provider_slug = provider["provider"] + + # 2. Read per-provider config from config.yaml. + config_data = read_config() + providers_cfg = config_data.get("providers", {}) or {} + pcfg = providers_cfg.get(provider_slug, {}) or {} + balance_cfg_data = pcfg.get("balance", {}) or {} + + # 3. Get defaults from registered provider class. + bal_cls = BalanceProviderRegistry.get(provider_slug) + if bal_cls: + defaults = bal_cls.default_config() + else: + defaults = BalanceConfig() + + # Merge config.yaml overrides on top of defaults. + cfg = BalanceConfig( + endpoint=balance_cfg_data.get("endpoint", defaults.endpoint), + api_key_env=balance_cfg_data.get("api_key_env", defaults.api_key_env), + enabled=balance_cfg_data.get("enabled", defaults.enabled), + cache_ttl_seconds=float( + balance_cfg_data.get("cache_ttl_seconds", defaults.cache_ttl_seconds) + ), + ) + + if not cfg.enabled: + return _ok(rid, {"ok": False, "error": "balance_disabled"}) + + # 4. Resolve API key. + env_var = (cfg.api_key_env or f"{provider_slug.upper()}_API_KEY").strip() + api_key = os.environ.get(env_var) or pcfg.get("api_key", "") + + # 5. Fetch (or return cached). + force = bool(params.get("force")) + result, was_cached = BalanceProviderRegistry.cached_or_fetch( + provider_slug, api_key, cfg, force=force + ) + + return _ok(rid, { + "ok": result.error is None, + "balance": { + "provider_name": result.provider_name, + "label": result.label, + "value": result.value, + "currency": result.currency, + "is_depleted": result.is_depleted, + "fetched_at": result.fetched_at, + } if result.error is None else None, + "error": result.error, + "cached": was_cached, + }) + except Exception as exc: + logger.debug("balance.view failed", exc_info=True) + return _ok(rid, {"ok": False, "error": str(exc)}) + + # =========================================================================== # Phase 2b terminal billing RPC methods # ===========================================================================