Filter budget picker to per-user hard blocks

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:

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)

tests/test_databricks.py

@@ -2666,34 +2666,59 @@ 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}],
+                            }
+                        ],
                     },
                 ]
             },
         )
         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:

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):