From 770811c49dcc2b26ac821f930f3cda9c75723b14 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Mon, 24 Aug 2026 16:20:50 +0000
Subject: [PATCH] Prompt before querying detailed usage

---
 src/ucode/usage.py  | 22 +++++++++++++++++-----
 tests/test_usage.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/src/ucode/usage.py b/src/ucode/usage.py
index 014b7991..dc392796 100644
--- a/src/ucode/usage.py
+++ b/src/ucode/usage.py
@@ -35,6 +35,7 @@
     print_heading,
     print_note,
     print_warning,
+    prompt_yes_no_default,
     render_box_table,
     spinner,
     value,
@@ -727,6 +728,22 @@ def usage(warehouse_id: str | None = None) -> int:
     with spinner("Retrieving Databricks access token..."):
         token = get_databricks_token(workspace, profile)
 
+    # Budget spend comes from AI Gateway directly, so show the useful at-a-glance result before
+    # asking whether to start (and potentially wait for) a SQL warehouse for the detailed report.
+    with spinner("Checking budget spend..."):
+        budget_spend, _ = resolve_current_budget_spend(workspace, token)
+    budget_lines = render_budget_lines(budget_spend)
+    if budget_lines:
+        console.print("\n".join([heading("Usage Budget"), "", *budget_lines]))
+    else:
+        print_note("Budget spend and threshold are unavailable.")
+
+    if not prompt_yes_no_default(
+        "Show token usage and estimated cost details? This queries a SQL warehouse.",
+        default=False,
+    ):
+        return 0
+
     with spinner("Discovering SQL warehouse..."):
         candidates = discover_sql_warehouses(workspace, token, warehouse_id=warehouse_id)
 
@@ -736,10 +753,6 @@ def usage(warehouse_id: str | None = None) -> 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.
-    with spinner("Checking budget spend..."):
-        budget_spend, _ = resolve_current_budget_spend(workspace, token)
-
     # Per-model dollar cost is estimated from tokens × catalog prices; omit cost rather than fail
     # when the price catalog is unreachable.
     with spinner(PRICES_MESSAGE):
@@ -756,7 +769,6 @@ def usage(warehouse_id: str | None = None) -> int:
             records,
             requester_name,
             configured_tool_displays,
-            budget_spend=budget_spend,
             price_lookup=price_lookup,
         )
     )
diff --git a/tests/test_usage.py b/tests/test_usage.py
index f0e97199..6ade58d4 100644
--- a/tests/test_usage.py
+++ b/tests/test_usage.py
@@ -661,6 +661,7 @@ def fake_render_box_table(headers, table_rows, max_widths=None):
         monkeypatch.setattr(
             usage_mod, "resolve_current_budget_spend", lambda *args, **kwargs: (None, "disabled")
         )
+        monkeypatch.setattr(usage_mod, "prompt_yes_no_default", lambda *args, **kwargs: True)
         monkeypatch.setattr(
             usage_mod, "fetch_external_model_prices", lambda *args, **kwargs: ([], "disabled")
         )
@@ -675,6 +676,7 @@ def fake_render_box_table(headers, table_rows, max_widths=None):
         assert "Claude Code · Last 7 Days" in headings
         assert all("Gemini" not in heading for heading in headings)
         assert notes == [
+            "Budget spend and threshold are unavailable.",
             "Using SQL warehouse `wh` (RUNNING).",
             f"No usage for Claude Code in the last {USAGE_BREAKDOWN_DAYS} days.",
         ]
@@ -686,6 +688,47 @@ def fake_render_box_table(headers, table_rows, max_widths=None):
         assert "gemini" not in "\n".join(printed).lower()
         assert "900" not in "\n".join(printed)
 
+    def test_shows_budget_before_prompt_and_skips_sql_when_declined(self, monkeypatch):
+        events: list[str] = []
+
+        class DummyConsole:
+            def print(self, output):
+                events.append(str(output))
+
+        monkeypatch.setattr(
+            usage_mod,
+            "load_state",
+            lambda: {"workspace": "https://workspace", "available_tools": ["codex"]},
+        )
+        monkeypatch.setattr(usage_mod, "ensure_databricks_auth", lambda *args, **kwargs: None)
+        monkeypatch.setattr(usage_mod, "get_databricks_token", lambda *args, **kwargs: "token")
+        monkeypatch.setattr(
+            usage_mod,
+            "resolve_current_budget_spend",
+            lambda *args, **kwargs: ((Decimal("12.34"), Decimal("100")), None),
+        )
+        monkeypatch.setattr(usage_mod, "console", DummyConsole())
+
+        def decline(prompt, *, default):
+            assert "$12.34 of $100.00" in "\n".join(events)
+            assert "SQL warehouse" in prompt
+            assert default is False
+            return False
+
+        monkeypatch.setattr(usage_mod, "prompt_yes_no_default", decline)
+        monkeypatch.setattr(
+            usage_mod,
+            "discover_sql_warehouses",
+            lambda *args, **kwargs: pytest.fail("SQL discovery should not run"),
+        )
+        monkeypatch.setattr(
+            usage_mod,
+            "fetch_external_model_prices",
+            lambda *args, **kwargs: pytest.fail("price lookup should not run"),
+        )
+
+        assert usage() == 0
+
 
 class TestRunQueryOnFirstWorkingWarehouse:
     _COLUMNS = ["requester_name"]
@@ -763,6 +806,7 @@ def fake_discover(workspace, token, *, warehouse_id=None):
         monkeypatch.setattr(
             usage_mod, "resolve_current_budget_spend", lambda *a, **k: (None, "disabled")
         )
+        monkeypatch.setattr(usage_mod, "prompt_yes_no_default", lambda *a, **k: True)
         monkeypatch.setattr(
             usage_mod, "fetch_external_model_prices", lambda *a, **k: ([], "disabled")
         )