diff --git a/src/ucode/usage.py b/src/ucode/usage.py index 153ecfba..bf85ac1d 100644 --- a/src/ucode/usage.py +++ b/src/ucode/usage.py @@ -35,6 +35,7 @@ print_heading, print_note, print_warning, + prompt_for_selection, prompt_yes_no_default, render_box_table, spinner, @@ -503,7 +504,13 @@ class ToolUsageTotals(NamedTuple): cost: Decimal | None -TOOL_MODEL_TABLE_HEADERS = ["Model", "Requests", "Input (incl. cache)", "Output", "Cost (USD)"] +TOOL_MODEL_TABLE_HEADERS = [ + "Model", + "Requests", + "Input (incl. cache)", + "Output", + "Est. Cost (USD)", +] def aggregate_tool_model_usage(records: list[dict[str, object]], tool: str) -> list[ModelUsage]: @@ -692,6 +699,49 @@ def run_query_on_first_working_warehouse( raise last_error or RuntimeError("No SQL warehouse could run the usage query.") +def select_sql_warehouse(candidates: list[SqlWarehouse]) -> SqlWarehouse | None: + """Ask which discovered warehouse should run the detailed usage query.""" + selected_path = prompt_for_selection( + "Select a SQL warehouse for the usage query:", + [ + (warehouse.http_path, f"{warehouse.label} ({warehouse.state})") + for warehouse in candidates + ], + ) + return next( + (warehouse for warehouse in candidates if warehouse.http_path == selected_path), + None, + ) + + +def select_and_run_usage_query( + workspace: str, + token: str, + candidates: list[SqlWarehouse], + query: str, +) -> tuple[str, list[str], list[tuple]] | None: + """Let the user choose warehouses until one runs the query or they cancel. + + A failed warehouse is removed before the picker is shown again so the user cannot + accidentally retry the same unusable option forever. Returns ``None`` when the picker + is cancelled and raises the final warehouse error when no candidates remain. + """ + remaining = list(candidates) + last_error: RuntimeError | None = None + while remaining: + selected = select_sql_warehouse(remaining) + if selected is None: + return None + try: + return run_query_on_first_working_warehouse(workspace, token, [selected], query) + except RuntimeError as exc: + last_error = exc + remaining = [ + warehouse for warehouse in remaining if warehouse.http_path != selected.http_path + ] + raise last_error or RuntimeError("No SQL warehouse could run the usage query.") + + def _query_with_progress( workspace: str, token: str, @@ -746,10 +796,18 @@ def usage(warehouse_id: str | None = None) -> int: with spinner("Discovering SQL warehouse..."): candidates = discover_sql_warehouses(workspace, token, warehouse_id=warehouse_id) - - resolved_http_path, columns, rows = run_query_on_first_working_warehouse( - workspace, token, candidates, build_usage_report_query() - ) + if warehouse_id is None: + query_result = select_and_run_usage_query( + workspace, token, candidates, build_usage_report_query() + ) + if query_result is None: + print_note("Usage details cancelled.") + return 0 + resolved_http_path, columns, rows = query_result + else: + resolved_http_path, columns, rows = run_query_on_first_working_warehouse( + workspace, token, candidates, build_usage_report_query() + ) records = parse_usage_rows(columns, rows) requester_name = find_requester_name(workspace, resolved_http_path, token, records) @@ -793,6 +851,6 @@ def usage(warehouse_id: str | None = None) -> int: console.print(f"{label('Requests:')} {value(f'{totals.requests:,}')}") console.print(f"{label('Total tokens:')} {value(f'{totals.tokens:,}')}") if totals.cost is not None: - console.print(f"{label('Cost (USD):')} {value(format_cost_usd(totals.cost))}") + console.print(f"{label('Est. Cost (USD):')} {value(format_cost_usd(totals.cost))}") console.print(render_box_table(TOOL_MODEL_TABLE_HEADERS, rows, max_widths=table_widths)) return 0 diff --git a/tests/test_usage.py b/tests/test_usage.py index f0f0e02f..a1e92ef1 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -12,6 +12,7 @@ from ucode.databricks import SqlWarehouse from ucode.ui import label, value from ucode.usage import ( + TOOL_MODEL_TABLE_HEADERS, USAGE_BREAKDOWN_DAYS, USAGE_SUMMARY_DAYS, ModelPrice, @@ -33,6 +34,8 @@ render_budget_lines, render_usage_summary, run_query_on_first_working_warehouse, + select_and_run_usage_query, + select_sql_warehouse, simplify_model_name, summarize_models, usage, @@ -464,6 +467,9 @@ def test_totals_cost_none_when_nothing_priced(self): _, totals = build_tool_model_rows(records, "claude", self._lookup()) assert totals.cost is None + def test_cost_header_is_explicitly_estimated(self): + assert TOOL_MODEL_TABLE_HEADERS[-1] == "Est. Cost (USD)" + class TestRenderBudgetLines: def test_no_lines_when_unavailable(self): @@ -662,6 +668,9 @@ def fake_render_box_table(headers, table_rows, max_widths=None): 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, "prompt_for_selection", lambda prompt, options: options[0][0] + ) monkeypatch.setattr( usage_mod, "fetch_external_model_prices", lambda *args, **kwargs: ([], "disabled") ) @@ -688,6 +697,53 @@ 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_queries_only_the_selected_warehouse(self, monkeypatch): + queried_paths: list[str] = [] + candidates = [ + SqlWarehouse("/sql/1.0/warehouses/first", "First", "RUNNING"), + SqlWarehouse("/sql/1.0/warehouses/second", "Second", "STOPPED"), + ] + + monkeypatch.setattr( + usage_mod, + "load_state", + lambda: {"workspace": "https://workspace", "available_tools": []}, + ) + 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("1"), Decimal("10")), None), + ) + monkeypatch.setattr(usage_mod, "prompt_yes_no_default", lambda *args, **kwargs: True) + monkeypatch.setattr( + usage_mod, "discover_sql_warehouses", lambda *args, **kwargs: candidates + ) + + def choose_second(prompt, options): + assert options == [ + (candidates[0].http_path, "First (RUNNING)"), + (candidates[1].http_path, "Second (STOPPED)"), + ] + return candidates[1].http_path + + monkeypatch.setattr(usage_mod, "prompt_for_selection", choose_second) + + def fake_query(workspace, http_path, token, query, on_connected=None): + queried_paths.append(http_path) + return ["requester_name"], [("user@example.com",)] + + monkeypatch.setattr(usage_mod, "run_usage_query", fake_query) + monkeypatch.setattr( + usage_mod, "fetch_external_model_prices", lambda *args, **kwargs: ([], None) + ) + monkeypatch.setattr(usage_mod, "print_note", lambda *args: None) + monkeypatch.setattr(usage_mod, "console", type("C", (), {"print": lambda *args: None})()) + + assert usage() == 0 + assert queried_paths == [candidates[1].http_path] + def test_shows_budget_before_prompt_and_skips_sql_when_declined(self, monkeypatch): events: list[str] = [] @@ -730,6 +786,22 @@ def decline(prompt, *, default): assert usage() == 0 +class TestSelectSqlWarehouse: + def test_returns_selected_candidate(self, monkeypatch): + candidates = [ + SqlWarehouse("/sql/1.0/warehouses/a", "Alpha", "RUNNING"), + SqlWarehouse("/sql/1.0/warehouses/b", "Beta", "STOPPED"), + ] + monkeypatch.setattr( + usage_mod, "prompt_for_selection", lambda prompt, options: candidates[1].http_path + ) + assert select_sql_warehouse(candidates) == candidates[1] + + def test_returns_none_when_cancelled(self, monkeypatch): + monkeypatch.setattr(usage_mod, "prompt_for_selection", lambda prompt, options: None) + assert select_sql_warehouse([SqlWarehouse("/path", "Warehouse", "RUNNING")]) is None + + class TestRunQueryOnFirstWorkingWarehouse: _COLUMNS = ["requester_name"] _ROWS = [("user@example.com",)] @@ -788,6 +860,54 @@ def test_raises_when_no_candidates(self, monkeypatch): run_query_on_first_working_warehouse("https://ws", "token", [], "SELECT 1") +class TestSelectAndRunUsageQuery: + _COLUMNS = ["requester_name"] + _ROWS = [("user@example.com",)] + + def _warehouses(self) -> list[SqlWarehouse]: + return [ + SqlWarehouse("/sql/1.0/warehouses/dead", "Dead", "RUNNING"), + SqlWarehouse("/sql/1.0/warehouses/alive", "Alive", "RUNNING"), + ] + + def test_prompts_again_without_failed_warehouse(self, monkeypatch): + candidates = self._warehouses() + picker_options: list[list[str]] = [] + + def choose(remaining): + picker_options.append([warehouse.label for warehouse in remaining]) + return remaining[0] + + def query(workspace, http_path, token, query, on_connected=None): + if http_path.endswith("dead"): + raise RuntimeError("warehouse unavailable") + return self._COLUMNS, self._ROWS + + monkeypatch.setattr(usage_mod, "select_sql_warehouse", choose) + monkeypatch.setattr(usage_mod, "run_usage_query", query) + monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) + monkeypatch.setattr(usage_mod, "print_warning", lambda *a: None) + + result = select_and_run_usage_query("https://ws", "token", candidates, "SELECT 1") + + assert result == (candidates[1].http_path, self._COLUMNS, self._ROWS) + assert picker_options == [["Dead", "Alive"], ["Alive"]] + + def test_can_cancel_after_a_warehouse_fails(self, monkeypatch): + candidates = self._warehouses() + choices = iter([candidates[0], None]) + monkeypatch.setattr(usage_mod, "select_sql_warehouse", lambda remaining: next(choices)) + monkeypatch.setattr( + usage_mod, + "run_usage_query", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("warehouse unavailable")), + ) + monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) + monkeypatch.setattr(usage_mod, "print_warning", lambda *a: None) + + assert select_and_run_usage_query("https://ws", "token", candidates, "SELECT 1") is None + + class TestUsageWarehouseIdPassthrough: def test_forwards_warehouse_id_to_discovery(self, monkeypatch): captured = {}