Show budget spend and threshold in ucode usage

README.md

@@ -170,7 +170,7 @@ you to run `ucode <agent>` (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 |

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,

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 "-"

src/ucode/usage.py

@@ -8,22 +8,27 @@
 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 (
     apply_pat_environment,
     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]

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

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