feat(/usage): append account limits section in CLI and gateway

Wires the agent/account_usage module from the preceding commit into
/usage so users see provider-side quota/credit info alongside the
existing session token report.

CLI:
- `_show_usage` appends account lines under the token table. Fetch
  runs in a 1-worker ThreadPoolExecutor with a 10s timeout so a slow
  provider API can never hang the prompt.

Gateway:
- `_handle_usage_command` resolves provider from the live agent when
  available, else from the persisted billing_provider/billing_base_url
  on the SessionDB row, so /usage still returns account info between
  turns when no agent is resident. Fetch runs via asyncio.to_thread.
- Account section is appended to all three return branches: running
  agent, no-agent-with-history, and the new no-agent-no-history path
  (falls back to account-only output instead of "no data").

Tests:
- 2 new tests in tests/gateway/test_usage_command.py cover the live-
  agent account section and the persisted-billing fallback path.

Salvaged from PR #2486 by @kshitijk4poor. The original branch had
drifted ~2615 commits behind main and rewrote _show_usage wholesale,
which would have dropped the rate-limit and cached-agent blocks added
in PRs #6541 and #7038. This commit re-adds only the new behavior on
top of current main.
This commit is contained in:
Teknium 2026-04-21 01:54:10 -07:00 committed by Teknium
parent 8a11b0a204
commit bcc5d7b67d
3 changed files with 149 additions and 6 deletions

View file

@ -175,3 +175,79 @@ class TestUsageCachedAgent:
result = await runner._handle_usage_command(event)
assert "Cost: included" in result
class TestUsageAccountSection:
"""Account-limits section appended to /usage output (PR #2486)."""
@pytest.mark.asyncio
async def test_usage_command_includes_account_section(self, monkeypatch):
agent = _make_mock_agent(provider="openai-codex")
agent.base_url = "https://chatgpt.com/backend-api/codex"
agent.api_key = "unused"
runner = _make_runner(SK, cached_agent=agent)
event = MagicMock()
monkeypatch.setattr(
"gateway.run.fetch_account_usage",
lambda provider, base_url=None, api_key=None: object(),
)
monkeypatch.setattr(
"gateway.run.render_account_usage_lines",
lambda snapshot, markdown=False: [
"📈 **Account limits**",
"Provider: openai-codex (Pro)",
"Session: 85% remaining (15% used)",
],
)
with patch("agent.rate_limit_tracker.format_rate_limit_compact", return_value="RPM: 50/60"), \
patch("agent.usage_pricing.estimate_usage_cost") as mock_cost:
mock_cost.return_value = MagicMock(amount_usd=None, status="included")
result = await runner._handle_usage_command(event)
assert "📊 **Session Token Usage**" in result
assert "📈 **Account limits**" in result
assert "Provider: openai-codex (Pro)" in result
@pytest.mark.asyncio
async def test_usage_command_uses_persisted_provider_when_agent_not_running(self, monkeypatch):
runner = _make_runner(SK)
runner._session_db = MagicMock()
runner._session_db.get_session.return_value = {
"billing_provider": "openai-codex",
"billing_base_url": "https://chatgpt.com/backend-api/codex",
}
session_entry = MagicMock()
session_entry.session_id = "sess-1"
runner.session_store.get_or_create_session.return_value = session_entry
runner.session_store.load_transcript.return_value = [
{"role": "user", "content": "earlier"},
]
calls = {}
async def _fake_to_thread(fn, *args, **kwargs):
calls["args"] = args
calls["kwargs"] = kwargs
return fn(*args, **kwargs)
monkeypatch.setattr("gateway.run.asyncio.to_thread", _fake_to_thread)
monkeypatch.setattr(
"gateway.run.fetch_account_usage",
lambda provider, base_url=None, api_key=None: object(),
)
monkeypatch.setattr(
"gateway.run.render_account_usage_lines",
lambda snapshot, markdown=False: [
"📈 **Account limits**",
"Provider: openai-codex (Pro)",
],
)
event = MagicMock()
result = await runner._handle_usage_command(event)
assert calls["args"] == ("openai-codex",)
assert calls["kwargs"]["base_url"] == "https://chatgpt.com/backend-api/codex"
assert "📊 **Session Info**" in result
assert "📈 **Account limits**" in result