From 68d3b5b5ef5017677ffb9f320eb207dff7a871bd Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:26:07 +0000 Subject: [PATCH] Show budget spend and threshold in `ucode usage` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode usage` reported tokens only, with no dollar figure. Surface the caller's live coding-agent budget spend against its alert threshold, from the AI Gateway `coding-agent-configs:resolveCurrentBudgetSpend` endpoint. Budget spend: $12.34 of $100.00 (12%) [███░░░░░░░░░░░░░░░░░░░░░░░░░░░] - databricks.py: `resolve_current_budget_spend` POSTs an empty body (caller and workspace come from the bearer) and parses both amounts as `Decimal`. Absence is the common case — the endpoint is behind a per-org SAFE flag (default off), needs a coding-agent config, and returns both fields unset when no budget matches — so all of it comes back as a reason string, never an exception. A spend without a threshold counts as no spend, mirroring the server's `BudgetSpend.fromProto`. - ui.py: `format_usd` (half-up to cents) and `format_meter` (clamped to [0, 1]; a nonzero fraction fills at least one cell so a small real spend doesn't read as an empty bar). - usage.py: `render_budget_lines` + a `budget_spend` arg on `render_usage_summary`. Unavailable spend omits the lines rather than failing the report; a zero threshold shows the bare amount, since there is no whole to be a fraction of. Co-authored-by: Isaac --- README.md | 2 +- src/ucode/databricks.py | 43 +++++++++++++++++++++ src/ucode/ui.py | 15 ++++++++ src/ucode/usage.py | 35 ++++++++++++++++- tests/test_databricks.py | 81 ++++++++++++++++++++++++++++++++++++++++ tests/test_ui.py | 42 +++++++++++++++++++++ tests/test_usage.py | 52 ++++++++++++++++++++++++++ 7 files changed, 268 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 03527575..c98b2fab 100644 --- a/README.md +++ b/README.md @@ -170,7 +170,7 @@ you to run `ucode ` (existing agent sessions need a restart before the MC | Command | Description | |---------|-------------| | `ucode status` | Show current workspace, base URLs, managed config files, and selected models | -| `ucode usage` | Show AI Gateway usage summary | +| `ucode usage` | Show AI Gateway usage summary, plus your budget spend against its alert threshold when the workspace reports one | | `ucode revert` | Clear saved state and restore backed-up config files | | `ucode configure --dry-run` | Preview config files without writing them | | `ucode configure --agents claude,codex` | Configure specific agents without the interactive picker | diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 1d32f319..723fd5b2 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -23,6 +23,7 @@ from concurrent.futures import ( TimeoutError as FutureTimeoutError, ) +from decimal import Decimal, InvalidOperation from pathlib import Path from typing import Literal, cast, overload from urllib import error as urllib_error @@ -2234,6 +2235,48 @@ def _looks_like_auth_failure(reason: str) -> bool: return False +CODING_AGENT_BUDGET_SPEND_PATH = "/api/ai-gateway/v2/coding-agent-configs:resolveCurrentBudgetSpend" + + +def resolve_current_budget_spend( + workspace: str, + token: str, + *, + timeout: int = 10, +) -> tuple[tuple[Decimal, Decimal] | None, str | None]: + """Fetch the caller's coding-agent budget spend and alert threshold. + + Returns `((spend, threshold), None)` or `(None, reason)`. Absence is + routine — the endpoint needs a per-org SAFE flag (default off) and a + coding-agent config — so it never raises. + """ + url = f"https://{workspace_hostname(workspace)}{CODING_AGENT_BUDGET_SPEND_PATH}" + payload, reason = _http_post_json(url, token, {}, timeout=timeout) + if payload is None: + return None, reason or "unknown error" + if not isinstance(payload, dict): + return None, "response was not a JSON object" + + # Per the server's BudgetSpend.fromProto, a spend with no threshold to + # measure against counts as no spend. + spend = _parse_decimal(payload.get("current_spend")) + threshold = _parse_decimal(payload.get("effective_threshold")) + if spend is None or threshold is None: + return None, "workspace reported no coding-agent budget spend" + return (spend, threshold), None + + +def _parse_decimal(value: object) -> Decimal | None: + if isinstance(value, str) and value.strip(): + try: + return Decimal(value.strip()) + except InvalidOperation: + return None + if isinstance(value, int): + return Decimal(value) + return None + + def discover_sql_warehouse_http_path( workspace: str, token: str, diff --git a/src/ucode/ui.py b/src/ucode/ui.py index 81ffc2e0..10849250 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -10,6 +10,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager from datetime import timedelta +from decimal import ROUND_HALF_UP, Decimal import questionary from rich.console import Console @@ -204,6 +205,20 @@ def format_token_count(token_count: int) -> str: return str(token_count) +def format_usd(amount: Decimal) -> str: + return f"${amount.quantize(Decimal('0.01'), rounding=ROUND_HALF_UP):,}" + + +def format_meter(fraction: float, width: int = 30) -> str: + """Text meter for `fraction` of a whole, clamped to [0, 1].""" + clamped = min(max(fraction, 0.0), 1.0) + filled = int(clamped * width) + # A small-but-real fraction shouldn't read as empty. + if clamped > 0: + filled = max(filled, 1) + return "[" + "█" * filled + "░" * (width - filled) + "]" + + def format_duration(duration_value: timedelta | None) -> str: if not duration_value or duration_value.total_seconds() <= 0: return "-" diff --git a/src/ucode/usage.py b/src/ucode/usage.py index 0e30aa5b..3f464463 100644 --- a/src/ucode/usage.py +++ b/src/ucode/usage.py @@ -8,6 +8,7 @@ import json from collections.abc import Mapping from datetime import date, datetime, timedelta +from decimal import Decimal from typing import cast from ucode.databricks import ( @@ -15,15 +16,19 @@ discover_sql_warehouse_http_path, ensure_databricks_auth, get_databricks_token, + resolve_current_budget_spend, run_usage_query, ) from ucode.state import load_state from ucode.ui import ( console, format_duration, + format_meter, format_token_count, + format_usd, heading, label, + muted, print_heading, print_note, render_box_table, @@ -370,10 +375,27 @@ def find_requester_name( return "current user" +def render_budget_lines(budget_spend: tuple[Decimal, Decimal] | None) -> list[str]: + """Spend-against-threshold lines, or nothing when unavailable.""" + if budget_spend is None: + return [] + spend, threshold = budget_spend + # No whole to be a fraction of; dividing would raise. + if threshold <= 0: + return [f"{label('Budget spend:')} {value(format_usd(spend))}"] + fraction = float(spend / threshold) + summary = f"{format_usd(spend)} of {format_usd(threshold)} ({fraction:.0%})" + return [ + f"{label('Budget spend:')} {value(summary)}", + muted(format_meter(fraction)), + ] + + def render_usage_summary( records: list[dict[str, object]], requester_name: str, tool_displays: dict[str, str], + budget_spend: tuple[Decimal, Decimal] | None = None, ) -> str: today = date.today() week_start = today - timedelta(days=USAGE_BREAKDOWN_DAYS - 1) @@ -434,6 +456,7 @@ def render_usage_summary( for model_name, token_total in top_models ) lines.append(f"{label('Top models this week:')} {value(models_text)}") + lines.extend(render_budget_lines(budget_spend)) return "\n".join(lines) @@ -465,12 +488,22 @@ def usage() -> int: records = parse_usage_rows(columns, rows) requester_name = find_requester_name(workspace, resolved_http_path, token, records) + # Opt-in per workspace: omit the lines rather than fail the report. + budget_spend, _ = resolve_current_budget_spend(workspace, token) + tool_displays = {tool: spec["display"] for tool, spec in TOOL_SPECS.items()} configured_tools = configured_usage_tools(state, tool_displays) configured_tool_displays = {tool: tool_displays[tool] for tool in configured_tools} records = filter_records_for_tools(records, configured_tools) - console.print(render_usage_summary(records, requester_name, configured_tool_displays)) + console.print( + render_usage_summary( + records, + requester_name, + configured_tool_displays, + budget_spend=budget_spend, + ) + ) table_headers = ["Date", "Day", "Tokens", "Sessions", "Duration", "Models"] table_widths = [8, 5, 10, 8, 8, 24] diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7e1a73a1..27f4b50f 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -5,12 +5,14 @@ import json import os import subprocess +from decimal import Decimal import pytest import ucode.databricks as db_mod from ucode.databricks import ( AI_GATEWAY_V2_DOCS_URL, + CODING_AGENT_BUDGET_SPEND_PATH, _format_subprocess_result, _parse_databricks_cli_version, _run_databricks_cli_installer, @@ -30,6 +32,7 @@ list_databricks_apps, list_databricks_connections, list_genie_spaces, + resolve_current_budget_spend, workspace_hostname, ) @@ -1940,3 +1943,81 @@ def test_failure_surfaces_cli_stderr(self, monkeypatch): install_ai_tools(["copilot"]) assert len(warnings) == 1 assert "copilot: cli-not-on-path: could not resolve copilot" in warnings[0] + + +class TestResolveCurrentBudgetSpend: + def test_parses_spend_and_threshold(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, payload, timeout=10: ( + {"current_spend": "12.34", "effective_threshold": "100"}, + None, + ), + ) + spend, reason = resolve_current_budget_spend("https://ws", "token") + assert spend == (Decimal("12.34"), Decimal("100")) + assert reason is None + + def test_posts_empty_body_to_coding_agent_path(self, monkeypatch): + captured = {} + + def fake_post(url, token, payload, timeout=10): + captured["url"] = url + captured["payload"] = payload + return {"current_spend": "1", "effective_threshold": "2"}, None + + monkeypatch.setattr(db_mod, "_http_post_json", fake_post) + resolve_current_budget_spend("https://ws.example.com", "token") + assert captured["url"] == (f"https://ws.example.com{CODING_AGENT_BUDGET_SPEND_PATH}") + assert captured["payload"] == {} + + def test_feature_disabled_returns_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, payload, timeout=10: ( + None, + "HTTP 400 Bad Request: FEATURE_DISABLED", + ), + ) + spend, reason = resolve_current_budget_spend("https://ws", "token") + assert spend is None + assert "FEATURE_DISABLED" in reason + + def test_unset_fields_treated_as_no_spend(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_post_json", lambda url, token, payload, timeout=10: ({}, None) + ) + spend, reason = resolve_current_budget_spend("https://ws", "token") + assert spend is None + assert "no coding-agent budget spend" in reason + + def test_spend_without_threshold_is_no_spend(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, payload, timeout=10: ({"current_spend": "12.34"}, None), + ) + spend, _ = resolve_current_budget_spend("https://ws", "token") + assert spend is None + + def test_malformed_decimal_is_no_spend(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, payload, timeout=10: ( + {"current_spend": "not-a-number", "effective_threshold": "100"}, + None, + ), + ) + spend, _ = resolve_current_budget_spend("https://ws", "token") + assert spend is None + + def test_non_object_payload_is_no_spend(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_post_json", lambda url, token, payload, timeout=10: ([], None) + ) + spend, reason = resolve_current_budget_spend("https://ws", "token") + assert spend is None + assert "not a JSON object" in reason diff --git a/tests/test_ui.py b/tests/test_ui.py index 3c0fbc89..58300b75 100644 --- a/tests/test_ui.py +++ b/tests/test_ui.py @@ -3,13 +3,16 @@ from __future__ import annotations from datetime import timedelta +from decimal import Decimal from unittest.mock import patch import pytest from ucode.ui import ( format_duration, + format_meter, format_token_count, + format_usd, normalize_workspace_url, prompt_for_workspace, prompt_yes_no_default, @@ -225,3 +228,42 @@ def test_no_profiles_goes_straight_to_manual_prompt(self): url, profile = prompt_for_workspace("desc", profiles=None) assert url == "https://example.databricks.com" assert profile is None + + +class TestFormatUsd: + def test_rounds_to_cents(self): + assert format_usd(Decimal("12.345")) == "$12.35" + assert format_usd(Decimal("12.344")) == "$12.34" + + def test_pads_to_two_decimals(self): + assert format_usd(Decimal("5")) == "$5.00" + + def test_thousands_separator(self): + assert format_usd(Decimal("1234567.5")) == "$1,234,567.50" + + def test_zero(self): + assert format_usd(Decimal("0")) == "$0.00" + + +class TestFormatMeter: + def test_empty(self): + assert format_meter(0.0, width=10) == "[" + "\u2591" * 10 + "]" + + def test_full(self): + assert format_meter(1.0, width=10) == "[" + "\u2588" * 10 + "]" + + def test_half(self): + assert format_meter(0.5, width=10) == "[" + "\u2588" * 5 + "\u2591" * 5 + "]" + + def test_tiny_nonzero_fills_one_cell(self): + assert format_meter(0.001, width=10) == "[\u2588" + "\u2591" * 9 + "]" + + def test_clamps_above_one(self): + assert format_meter(2.5, width=10) == "[" + "\u2588" * 10 + "]" + + def test_clamps_below_zero(self): + assert format_meter(-1.0, width=10) == "[" + "\u2591" * 10 + "]" + + def test_width_is_constant(self): + for fraction in (0.0, 0.13, 0.5, 0.99, 1.0): + assert len(format_meter(fraction)) == 32 diff --git a/tests/test_usage.py b/tests/test_usage.py index d3c36bc0..e3f20377 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -3,8 +3,10 @@ from __future__ import annotations from datetime import date, datetime, timedelta +from decimal import Decimal import ucode.usage as usage_mod +from ucode.ui import label, value from ucode.usage import ( USAGE_BREAKDOWN_DAYS, USAGE_SUMMARY_DAYS, @@ -20,6 +22,7 @@ filter_records_for_tools, has_tool_usage_last_week, parse_usage_rows, + render_budget_lines, render_usage_summary, simplify_model_name, summarize_model_tokens, @@ -301,6 +304,37 @@ def test_structure(self): assert row["models"] == "-" +class TestRenderBudgetLines: + def test_no_lines_when_unavailable(self): + assert render_budget_lines(None) == [] + + def test_shows_spend_threshold_and_percent(self): + lines = render_budget_lines((Decimal("12.34"), Decimal("100"))) + assert "$12.34" in lines[0] + assert "$100.00" in lines[0] + assert "12%" in lines[0] + + def test_renders_meter(self): + lines = render_budget_lines((Decimal("50"), Decimal("100"))) + assert len(lines) == 2 + assert "█" in lines[1] + assert "░" in lines[1] + + def test_zero_threshold_omits_percent_and_meter(self): + lines = render_budget_lines((Decimal("5"), Decimal("0"))) + assert lines == [f"{label('Budget spend:')} {value('$5.00')}"] + + def test_spend_over_threshold_clamps_meter(self): + lines = render_budget_lines((Decimal("250"), Decimal("100"))) + assert "250%" in lines[0] + assert "░" not in lines[1] + + def test_thousands_separator(self): + lines = render_budget_lines((Decimal("1234.5"), Decimal("10000"))) + assert "$1,234.50" in lines[0] + assert "$10,000.00" in lines[0] + + class TestRenderUsageSummary: def _make_record(self, days_ago: int, tool: str, tokens: int, model: str = "") -> dict: d = date.today() - timedelta(days=days_ago) @@ -341,6 +375,21 @@ def test_top_models_listed(self): result = render_usage_summary(records, "user", {"claude": "Claude Code"}) assert "sonnet-4" in result + def test_includes_budget_spend_when_available(self): + records = [self._make_record(0, "claude", 1000)] + result = render_usage_summary( + records, + "user", + {"claude": "Claude Code"}, + budget_spend=(Decimal("12.34"), Decimal("100")), + ) + assert "$12.34 of $100.00" in result + + def test_omits_budget_spend_by_default(self): + records = [self._make_record(0, "claude", 1000)] + result = render_usage_summary(records, "user", {"claude": "Claude Code"}) + assert "Budget spend" not in result + def test_top_models_uses_per_model_token_totals(self): records = [ { @@ -470,6 +519,9 @@ def fake_render_box_table(headers, table_rows, max_widths=None): lambda *args, **kwargs: "/sql/1.0/warehouses/abc", ) monkeypatch.setattr(usage_mod, "run_usage_query", lambda *args, **kwargs: (columns, rows)) + monkeypatch.setattr( + usage_mod, "resolve_current_budget_spend", lambda *args, **kwargs: (None, "disabled") + ) monkeypatch.setattr(usage_mod, "console", DummyConsole()) monkeypatch.setattr(usage_mod, "print_heading", headings.append) monkeypatch.setattr(usage_mod, "print_note", notes.append)