Skip to content
Closed
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
23 changes: 23 additions & 0 deletions agent/account_usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,29 @@ def render_account_usage_lines(snapshot: Optional[AccountUsageSnapshot], *, mark
return lines


def render_account_usage_compact(snapshot: Optional[AccountUsageSnapshot]) -> str:
"""Return a terse one-line quota summary suitable for a status bar."""
if not snapshot or snapshot.unavailable_reason:
return ""

preferred_labels = {"session", "current session", "primary", "five hour"}
window = None
for candidate in snapshot.windows:
if (candidate.label or "").strip().lower() in preferred_labels:
window = candidate
break
if window is None and snapshot.windows:
window = snapshot.windows[0]
if window is None or window.used_percent is None:
return ""

remaining = max(0, round(100 - float(window.used_percent)))
provider = (snapshot.provider or "").strip().lower()
if provider == "openai-codex":
return f"quota {remaining}%"
return f"acct {remaining}%"


def _resolve_codex_usage_url(base_url: str) -> str:
normalized = (base_url or "").strip().rstrip("/")
if not normalized:
Expand Down
67 changes: 67 additions & 0 deletions cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -2418,6 +2418,10 @@ def __init__(
self._image_counter = 0
self.preloaded_skills: list[str] = []
self._startup_skills_line_shown = False
self._status_bar_account_quota_label = ""
self._status_bar_account_usage_signature = None
self._status_bar_account_usage_fetched_at = 0.0
self._status_bar_account_usage_fetching = False

# Voice mode state (also reinitialized inside run() for interactive TUI).
self._voice_lock = threading.Lock()
Expand Down Expand Up @@ -2635,6 +2639,7 @@ def _get_status_bar_snapshot(self) -> Dict[str, Any]:
"session_total_tokens": 0,
"session_api_calls": 0,
"compressions": 0,
"account_quota_label": "",
}

if not agent:
Expand All @@ -2659,8 +2664,63 @@ def _get_status_bar_snapshot(self) -> Dict[str, Any]:
if context_length:
snapshot["context_percent"] = max(0, min(100, round((context_tokens / context_length) * 100)))

snapshot["account_quota_label"] = self._get_status_bar_account_quota_label(agent)
return snapshot

def _get_status_bar_account_quota_label(self, agent: Any) -> str:
"""Return cached compact account-quota text for the status bar.

Refreshes asynchronously so the hot render path never blocks on a
provider API call.
"""
provider = str(getattr(agent, "provider", "") or "").strip().lower()
if provider != "openai-codex":
return ""

base_url = str(getattr(agent, "base_url", "") or "")
signature = (provider, base_url)
now = time.time()
ttl_seconds = 60.0

cached_label = getattr(self, "_status_bar_account_quota_label", "") or ""
cached_sig = getattr(self, "_status_bar_account_usage_signature", None)
fetched_at = float(getattr(self, "_status_bar_account_usage_fetched_at", 0.0) or 0.0)
fetching = bool(getattr(self, "_status_bar_account_usage_fetching", False))

stale = (cached_sig != signature) or ((now - fetched_at) > ttl_seconds)
if stale and not fetching:
self._status_bar_account_usage_fetching = True

def _refresh_quota() -> None:
label = ""
try:
from agent.account_usage import fetch_account_usage, render_account_usage_compact

snapshot = fetch_account_usage(
provider,
base_url=getattr(agent, "base_url", None),
api_key=getattr(agent, "api_key", None),
)
label = render_account_usage_compact(snapshot)
except Exception:
label = ""
finally:
self._status_bar_account_quota_label = label
self._status_bar_account_usage_signature = signature
self._status_bar_account_usage_fetched_at = time.time()
self._status_bar_account_usage_fetching = False
try:
if getattr(self, "_app", None) is not None:
self._app.invalidate()
except Exception:
pass

threading.Thread(target=_refresh_quota, daemon=True, name="status-bar-account-usage").start()

if cached_sig == signature:
return cached_label
return ""

@staticmethod
def _status_bar_display_width(text: str) -> int:
"""Return terminal cell width for status-bar text.
Expand Down Expand Up @@ -2852,6 +2912,9 @@ def _build_status_bar_text(self, width: Optional[int] = None) -> str:

parts = [f"⚕ {snapshot['model_short']}", context_label, percent_label]
parts.append(duration_label)
quota_label = snapshot.get("account_quota_label")
if quota_label:
parts.append(quota_label)
prompt_elapsed = snapshot.get("prompt_elapsed")
if prompt_elapsed:
parts.append(prompt_elapsed)
Expand Down Expand Up @@ -2914,6 +2977,10 @@ def _get_status_bar_fragments(self):
("class:status-bar-dim", " │ "),
("class:status-bar-dim", duration_label),
]
quota_label = snapshot.get("account_quota_label")
if quota_label:
frags.append(("class:status-bar-dim", " │ "))
frags.append(("class:status-bar-dim", quota_label))
# Position 7: per-prompt elapsed timer (live or frozen)
prompt_elapsed = snapshot.get("prompt_elapsed")
if prompt_elapsed:
Expand Down
20 changes: 20 additions & 0 deletions tests/cli/test_cli_status_bar.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,26 @@ def test_build_status_bar_text_handles_missing_agent(self):
assert "⚕" in text
assert "claude-sonnet-4-20250514" in text

def test_build_status_bar_text_includes_cached_quota_label_for_codex(self):
cli_obj = _attach_agent(
_make_cli(model="gpt-5.4-codex"),
prompt_tokens=10000,
completion_tokens=2400,
total_tokens=12400,
api_calls=7,
context_tokens=12400,
context_length=200_000,
)
cli_obj.agent.provider = "openai-codex"
cli_obj.agent.base_url = "https://chatgpt.com/backend-api/codex"
cli_obj._status_bar_account_quota_label = "quota 9%"
cli_obj._status_bar_account_usage_signature = ("openai-codex", "https://chatgpt.com/backend-api/codex")
cli_obj._status_bar_account_usage_fetched_at = 10**12

text = cli_obj._build_status_bar_text(width=120)

assert "quota 9%" in text

def test_minimal_tui_chrome_threshold(self):
cli_obj = _make_cli()

Expand Down
15 changes: 15 additions & 0 deletions tests/test_account_usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
AccountUsageSnapshot,
AccountUsageWindow,
fetch_account_usage,
render_account_usage_compact,
render_account_usage_lines,
)

Expand Down Expand Up @@ -118,6 +119,20 @@ def test_render_account_usage_lines_includes_reset_and_provider():
assert "Credits balance: $9.99" in lines[3]


def test_render_account_usage_compact_prefers_session_window():
snapshot = AccountUsageSnapshot(
provider="openai-codex",
source="usage_api",
fetched_at=datetime.now(timezone.utc),
windows=(
AccountUsageWindow(label="Weekly", used_percent=18),
AccountUsageWindow(label="Session", used_percent=91),
),
)

assert render_account_usage_compact(snapshot) == "quota 9%"


def test_fetch_account_usage_openrouter_uses_limit_remaining_and_ignores_deprecated_rate_limit(monkeypatch):
monkeypatch.setattr(
"agent.account_usage.resolve_runtime_provider",
Expand Down