Fix budget spend bar: show zero-spend, hide budgets without per-user …

src/ucode/databricks.py

@@ -409,12 +409,29 @@ 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"
+
+
+def _has_per_user_alert(entry: dict) -> bool:
+    """Whether a raw budget entry carries a per-user alert threshold."""
+    for alert in entry.get("alert_configurations") or []:
+        if isinstance(alert, dict) and alert.get("scope_type") == _PER_USER_ALERT_SCOPE:
+            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": ...}``.
+    Returns ``(budgets, reason)`` where each budget is ``{"id", "display_name", "has_per_user_alert"}``.
     ``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.
+    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``).
     """
     hostname = workspace_hostname(workspace)
     url = f"https://{hostname}{_WORKSPACE_BUDGETS_API_PATH}"
@@ -438,6 +455,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),
             }
         )
     if not budgets:
@@ -2712,12 +2730,22 @@ def resolve_current_budget_spend(
     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"))
+    # The threshold is what anchors the spend: a caller with a per-user threshold but no spend yet
+    # this period gets `effective_threshold` set and `current_spend` omitted (see the server's
+    # per-user spend resolution). Treat an *absent* spend as $0 rather than "no budget", so a
+    # developer who hasn't spent anything still sees their budget instead of a blank. With no
+    # threshold there is nothing to measure against, so that genuinely counts as no spend.
     threshold = _parse_decimal(payload.get("effective_threshold"))
-    if spend is None or threshold is None:
+    if threshold is None:
         return None, "workspace reported no coding-agent budget spend"
+    raw_spend = payload.get("current_spend")
+    if raw_spend is None:
+        spend: Decimal | None = Decimal(0)
+    else:
+        # Present but unparseable is corrupt data, not zero spend — don't silently mask it.
+        spend = _parse_decimal(raw_spend)
+        if spend is None:
+            return None, "workspace reported no coding-agent budget spend"
     return (spend, threshold), None
 
 

src/ucode/managed_wizard.py

@@ -528,11 +528,26 @@ 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")]
+    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`."
+        )
+        return None
+    print_note(
+        "Showing only budgets with a per-user threshold configured, which spend routing needs."
+    )
+
     budget_id = prompt_for_selection(
         "Which budget should this policy track?",
         [
             (budget["id"], f"{budget['display_name'] or budget['id']} ({budget['id']})")
-            for budget in budgets
+            for budget in usable
         ],
         searchable=True,
     )

tests/conftest.py

@@ -31,6 +31,9 @@ def _isolate_ucode_state(tmp_path, monkeypatch):
     state_dir.mkdir()
     monkeypatch.setattr(state_mod, "STATE_PATH", state_dir / "state.json")
     monkeypatch.setattr(config_io_mod, "APP_DIR", state_dir)
+    # Isolate the managed-config opt-in from the developer's own shell: leaving it set changes what
+    # `ucode`/`ucode configure` do mid-test. Tests that exercise the managed path set it explicitly.
+    monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False)
     # The model-services listing is memoized for the life of the process, so without this a cached
     # result would leak into the next test and make a stubbed listing look like it was never called.
     databricks_mod.clear_model_services_cache()

tests/test_databricks.py

@@ -35,6 +35,7 @@
     list_databricks_apps,
     list_databricks_connections,
     list_genie_spaces,
+    list_workspace_budgets,
     resolve_current_budget_spend,
     workspace_hostname,
 )
@@ -2610,6 +2611,18 @@ def test_spend_without_threshold_is_no_spend(self, monkeypatch):
         spend, _ = resolve_current_budget_spend("https://ws", "token")
         assert spend is None
 
+    def test_threshold_without_spend_is_zero_spend(self, monkeypatch):
+        # A per-user threshold with no spend yet this period is $0 spent, not "no budget" — the
+        # developer should still see their budget rather than a blank.
+        monkeypatch.setattr(
+            db_mod,
+            "_http_post_json",
+            lambda url, token, payload, timeout=10: ({"effective_threshold": "100"}, None),
+        )
+        spend, reason = resolve_current_budget_spend("https://ws", "token")
+        assert spend == (Decimal(0), Decimal("100"))
+        assert reason is None
+
     def test_malformed_decimal_is_no_spend(self, monkeypatch):
         monkeypatch.setattr(
             db_mod,
@@ -2631,6 +2644,55 @@ def test_non_object_payload_is_no_spend(self, monkeypatch):
         assert "not a JSON object" in reason
 
 
+class TestListWorkspaceBudgets:
+    PER_USER = "ALERT_CONFIGURATION_SCOPE_TYPE_PER_USER"
+    SHARED = "ALERT_CONFIGURATION_SCOPE_TYPE_SHARED"
+
+    def _stub(self, monkeypatch, payload):
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None)
+        )
+
+    def test_flags_per_user_alert_presence(self, monkeypatch):
+        self._stub(
+            monkeypatch,
+            {
+                "workspace_ai_gateway_budgets": [
+                    {
+                        "budget_configuration_id": "with",
+                        "display_name": "has per-user",
+                        "alert_configurations": [
+                            {"scope_type": self.SHARED},
+                            {"scope_type": self.PER_USER},
+                        ],
+                    },
+                    {
+                        "budget_configuration_id": "without",
+                        "display_name": "shared only",
+                        "alert_configurations": [{"scope_type": self.SHARED}],
+                    },
+                ]
+            },
+        )
+        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
+
+    def test_missing_alert_configs_is_not_per_user(self, monkeypatch):
+        self._stub(
+            monkeypatch,
+            {
+                "workspace_ai_gateway_budgets": [
+                    {"budget_configuration_id": "b", "display_name": "x"}
+                ]
+            },
+        )
+        budgets, _ = list_workspace_budgets("https://ws", "token")
+        assert budgets[0]["has_per_user_alert"] is False
+
+
 class TestDiscoverSqlWarehouses:
     def _payload(self, *entries: dict) -> dict:
         return {"warehouses": list(entries)}

tests/test_managed_wizard.py

@@ -964,8 +964,44 @@ 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}]
+        with (
+            patch.object(wizard, "prompt_yes_no_default", return_value=True),
+            patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)),
+            patch.object(wizard, "print_warning") as warn,
+        ):
+            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.
+        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},
+        ]
+        with (
+            patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]),
+            patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)),
+            patch.object(
+                wizard,
+                "prompt_for_selection",
+                side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"],
+            ) as select,
+            patch.object(wizard, "prompt_for_text", return_value="tiered"),
+            patch.object(wizard, "prompt_for_percentage", return_value=0.8),
+        ):
+            policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE)
+        # First selection call is the budget picker; only the per-user budget is offered.
+        offered = [value for value, _ in select.call_args_list[0][0][1]]
+        assert offered == [BUDGET_ID]
+        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"}]
+        budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}]
         with (
             patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]),
             patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)),
@@ -1000,7 +1036,7 @@ def test_offers_only_the_models_the_agent_was_configured_with(self):
                 }
             }
         }
-        budgets = [{"id": BUDGET_ID, "display_name": "eng"}]
+        budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}]
         with (
             patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]),
             patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)),
@@ -1029,7 +1065,7 @@ def test_claude_family_slots_are_flattened_for_the_picker(self):
                 }
             }
         }
-        budgets = [{"id": BUDGET_ID, "display_name": "eng"}]
+        budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}]
         with (
             patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]),
             patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)),
@@ -1048,7 +1084,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"}]
+        budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}]
         with (
             patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]),
             patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)),
@@ -1065,7 +1101,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"}]
+        budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_alert": True}]
         with (
             patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]),
             patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)),
@@ -1371,7 +1407,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"}]
+        budgets = [{"id": "budget-1", "display_name": "eng", "has_per_user_alert": True}]
         searchable_prompts: list[str] = []
 
         def fake_sel(prompt, options, **kwargs):