diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 705b5519..7848a02e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -409,29 +409,32 @@ def is_workspace_admin(workspace: str, token: str) -> bool | None: _WORKSPACE_BUDGETS_API_PATH = "/api/ai-gateway/v2/workspace-metrics/budgets" -# Alert-config scope that carries a per-user threshold. A budget's coding-agent spend routing only -# works when it has one: the gateway's `recommendModel` measures the caller's spend against a -# per-user threshold, so a budget with only a shared (workspace-wide) alert reports no spend and -# leaves every tier inert. The listing exposes `scope_type` but not the alert's action, so ucode can -# only check for the scope's presence; the server enforces the (block) action on config create. _PER_USER_ALERT_SCOPE = "ALERT_CONFIGURATION_SCOPE_TYPE_PER_USER" +_BLOCK_ACTION_TYPE = "BLOCK_USAGE" -def _has_per_user_alert(entry: dict) -> bool: - """Whether a raw budget entry carries a per-user alert threshold.""" +def _has_per_user_block(entry: dict) -> bool: + """Whether a raw budget entry has a per-user alert threshold that hard-blocks usage. + + True only when some alert is per-user scoped *and* carries a ``BLOCK_USAGE`` action; a per-user + alert with only an email notification does not enforce spend routing. + """ for alert in entry.get("alert_configurations") or []: - if isinstance(alert, dict) and alert.get("scope_type") == _PER_USER_ALERT_SCOPE: - return True + if not isinstance(alert, dict) or alert.get("scope_type") != _PER_USER_ALERT_SCOPE: + continue + for action in alert.get("action_configurations") or []: + if isinstance(action, dict) and action.get("action_type") == _BLOCK_ACTION_TYPE: + return True return False def list_workspace_budgets(workspace: str, token: str) -> tuple[list[dict], str | None]: """List the AI Gateway budgets that apply to this workspace. - Returns ``(budgets, reason)`` where each budget is ``{"id", "display_name", "has_per_user_alert"}``. + Returns ``(budgets, reason)`` where each budget is ``{"id", "display_name", "has_per_user_block"}``. ``reason`` is None on success, otherwise it explains why the list is empty. ucode never creates - budgets — an admin picks an existing one to attach a spend-routing policy to. ``has_per_user_alert`` - lets the picker hide budgets that can't drive spend routing (see ``_PER_USER_ALERT_SCOPE``). + budgets — an admin picks an existing one to attach a spend-routing policy to. ``has_per_user_block`` + lets the picker hide budgets that can't enforce spend routing (see ``_has_per_user_block``). """ hostname = workspace_hostname(workspace) url = f"https://{hostname}{_WORKSPACE_BUDGETS_API_PATH}" @@ -455,7 +458,7 @@ def list_workspace_budgets(workspace: str, token: str) -> tuple[list[dict], str { "id": budget_id, "display_name": display_name if isinstance(display_name, str) else "", - "has_per_user_alert": _has_per_user_alert(entry), + "has_per_user_block": _has_per_user_block(entry), } ) if not budgets: diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index f5da862b..c58e8206 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -709,19 +709,20 @@ def _prompt_budget_policy( ) return None - # Spend routing only works on a budget with a per-user threshold; without one the gateway reports - # no spend and every tier stays inert. The listing can't reveal the alert's action, so this hides - # the clearly-unusable budgets and the server rejects the rest on create. - usable = [budget for budget in budgets if budget.get("has_per_user_alert")] + # Spend routing only works on a budget with a per-user threshold that hard-blocks: without a + # per-user threshold the gateway reports no spend and every tier stays inert, and without a + # BLOCK_USAGE action the policy is never enforced (an email-only alert does not gate spend). The + # listing now exposes each alert's action, so hide the budgets that can't enforce routing. + usable = [budget for budget in budgets if budget.get("has_per_user_block")] if not usable: print_warning( - "None of this workspace's AI Gateway budgets have a per-user threshold configured, which " - "spend routing requires. Add a per-user alert threshold to a budget in the Databricks " - "console, then re-run `ucode setup`." + "None of this workspace's AI Gateway budgets have a per-user threshold with a usage " + "block configured, which spend routing enforces. Add a per-user alert threshold with a " + "block action to a budget in the Databricks console, then re-run `ucode setup`." ) return None print_note( - "Showing only budgets with a per-user threshold configured, which spend routing needs." + "Showing only budgets with a per-user hard block configured, which spend routing enforces." ) budget_id = prompt_for_selection( @@ -736,6 +737,13 @@ def _prompt_budget_policy( return None policy: dict = {"budget_id": budget_id} + # Remember the budget's own name so the summary can show it beside the policy name. It's a local + # display aid only — `_budget_policy_payload` doesn't serialize it, so it never reaches the API. + budget_display_name = next( + (budget["display_name"] for budget in usable if budget["id"] == budget_id), "" + ) + if budget_display_name: + policy["budget_display_name"] = budget_display_name display_name = prompt_for_text("Policy name", default="coding-agents-tiered-routing") if display_name: policy["display_name"] = display_name @@ -847,8 +855,9 @@ def _render_summary(workspace: str, manifest: dict) -> None: if isinstance(policy, dict): tiers = policy.get("tiers") or [] lines.append( - kv_line("Budget policy", policy.get("display_name") or policy.get("budget_id") or "set") + kv_line("Budget", policy.get("budget_display_name") or policy.get("budget_id") or "set") ) + lines.append(kv_line("Policy name", policy.get("display_name") or "unnamed")) for tier in tiers: agent = tier.get("default_agent") display = TOOL_SPECS.get(agent, {}).get("display", agent) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 0b51f8bc..9d2b0546 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2666,23 +2666,46 @@ def _stub(self, monkeypatch, payload): db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) ) - def test_flags_per_user_alert_presence(self, monkeypatch): + BLOCK = "BLOCK_USAGE" + EMAIL = "EMAIL_NOTIFICATION" + + def test_flags_per_user_block_presence(self, monkeypatch): self._stub( monkeypatch, { "workspace_ai_gateway_budgets": [ { - "budget_configuration_id": "with", - "display_name": "has per-user", + "budget_configuration_id": "blocks", + "display_name": "per-user block", "alert_configurations": [ {"scope_type": self.SHARED}, - {"scope_type": self.PER_USER}, + { + "scope_type": self.PER_USER, + "action_configurations": [{"action_type": self.BLOCK}], + }, + ], + }, + { + "budget_configuration_id": "email_only", + "display_name": "per-user email only", + "alert_configurations": [ + { + "scope_type": self.PER_USER, + "action_configurations": [ + {"action_type": self.EMAIL, "target": "a@b.com"} + ], + } ], }, { - "budget_configuration_id": "without", - "display_name": "shared only", - "alert_configurations": [{"scope_type": self.SHARED}], + "budget_configuration_id": "shared_block", + "display_name": "shared block only", + "alert_configurations": [ + { + "scope_type": self.SHARED, + "action_configurations": [{"action_type": self.BLOCK}], + } + ], }, ] }, @@ -2690,10 +2713,12 @@ def test_flags_per_user_alert_presence(self, monkeypatch): budgets, reason = list_workspace_budgets("https://ws", "token") assert reason is None by_id = {b["id"]: b for b in budgets} - assert by_id["with"]["has_per_user_alert"] is True - assert by_id["without"]["has_per_user_alert"] is False + # Only a per-user threshold that also carries a BLOCK_USAGE action enforces spend routing. + assert by_id["blocks"]["has_per_user_block"] is True + assert by_id["email_only"]["has_per_user_block"] is False + assert by_id["shared_block"]["has_per_user_block"] is False - def test_missing_alert_configs_is_not_per_user(self, monkeypatch): + def test_missing_alert_configs_is_not_per_user_block(self, monkeypatch): self._stub( monkeypatch, { @@ -2703,7 +2728,7 @@ def test_missing_alert_configs_is_not_per_user(self, monkeypatch): }, ) budgets, _ = list_workspace_budgets("https://ws", "token") - assert budgets[0]["has_per_user_alert"] is False + assert budgets[0]["has_per_user_block"] is False class TestDiscoverSqlWarehouses: diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index cc4b8e50..bbe8353f 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -1250,10 +1250,10 @@ def test_no_budgets_warns_and_yields_none(self): assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None assert warn.called - def test_no_per_user_budgets_warns_and_yields_none(self): - # Spend routing needs a per-user threshold; a workspace whose only budgets lack one has - # nothing usable to attach a policy to. - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": False}] + def test_no_per_user_block_budgets_warns_and_yields_none(self): + # Spend routing needs a per-user threshold that hard-blocks; a workspace whose only budgets + # lack one has nothing usable to attach a policy to. + budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": False}] with ( patch.object(wizard, "prompt_yes_no_default", return_value=True), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), @@ -1262,12 +1262,12 @@ def test_no_per_user_budgets_warns_and_yields_none(self): assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None assert warn.called - def test_only_per_user_budgets_are_offered(self): - # The picker hides budgets without a per-user threshold rather than letting the admin pick - # one that would leave every tier inert. + def test_only_per_user_block_budgets_are_offered(self): + # The picker hides budgets without a per-user hard block rather than letting the admin pick + # one that would leave every tier inert or unenforced. budgets = [ - {"id": "no-per-user", "display_name": "shared-only", "has_per_user_alert": False}, - {"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}, + {"id": "no-block", "display_name": "email-only", "has_per_user_block": False}, + {"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}, ] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), @@ -1287,7 +1287,7 @@ def test_only_per_user_budgets_are_offered(self): assert policy is not None and policy["budget_id"] == BUDGET_ID def test_percentages_are_stored_as_fractions(self): - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}] + budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), @@ -1303,6 +1303,8 @@ def test_percentages_are_stored_as_fractions(self): policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) assert policy is not None assert policy["budget_id"] == BUDGET_ID + # The picked budget's name is remembered for the summary. + assert policy["budget_display_name"] == "eng" assert policy["tiers"] == [ { "spending_percentage": 0.8, @@ -1322,7 +1324,7 @@ def test_offers_only_the_models_the_agent_was_configured_with(self): } } } - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}] + budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), @@ -1351,7 +1353,7 @@ def test_claude_family_slots_are_flattened_for_the_picker(self): } } } - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}] + budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), @@ -1370,7 +1372,7 @@ def test_claude_family_slots_are_flattened_for_the_picker(self): def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): # An agent configured through a provider service has no enumerable list; better to offer the # catalog than nothing at all. - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}] + budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), @@ -1387,7 +1389,7 @@ def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): assert offered == ["system.ai.gemini-3-flash"] def test_authored_policy_validates(self): - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}] + budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), @@ -1421,10 +1423,10 @@ def test_a_repeated_agent_model_pair_is_rejected_and_re_prompted(self): } } } - # `has_per_user_alert` is required since the budget-threshold gate landed on main: spend - # routing needs a per-user threshold, so a budget without one is filtered out and the policy - # flow returns before the tier loop this test exercises. - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}] + # `has_per_user_block` is required since the budget-threshold gate landed on main: spend + # routing needs a per-user threshold that hard-blocks, so a budget without one is filtered + # out and the policy flow returns before the tier loop this test exercises. + budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), @@ -1641,6 +1643,47 @@ def test_a_bracketed_policy_name_survives_the_summary(self, capsys): ) assert "[prod] tiered routing" in capsys.readouterr().out + def test_shows_both_budget_and_policy_names(self, capsys): + # An admin checks the policy against two distinct things: which budget it tracks and what the + # policy itself is called. The summary must surface both, not collapse to one. + wizard._render_summary( + WORKSPACE, + { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} + }, + "budget_policy": { + "budget_id": "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57", + "budget_display_name": "eng-budget", + "display_name": "tiered routing", + "tiers": [], + }, + }, + ) + out = capsys.readouterr().out + assert "eng-budget" in out + assert "tiered routing" in out + + def test_falls_back_to_budget_id_without_a_budget_name(self, capsys): + # `--from-file` and server-read manifests carry no `budget_display_name`, so the budget id is + # all there is to show. + wizard._render_summary( + WORKSPACE, + { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} + }, + "budget_policy": { + "budget_id": "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57", + "display_name": "tiered routing", + "tiers": [], + }, + }, + ) + assert "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57" in capsys.readouterr().out + class TestCancelledPromptsAbort: """A dismissed prompt must abort, not re-ask an input that can't answer.""" @@ -1741,7 +1784,7 @@ def fake_sel(prompt, options, **kwargs): assert seen[0].get("searchable") is True def test_budget_and_tier_pickers_are_searchable(self): - budgets = [{"id": "budget-1", "display_name": "eng", "has_per_user_alert": True}] + budgets = [{"id": "budget-1", "display_name": "eng", "has_per_user_block": True}] searchable_prompts: list[str] = [] def fake_sel(prompt, options, **kwargs):