feat: Role aware ucode configure for mcp, skills, spend-tiers subcommands

src/ucode/cli.py

@@ -5,6 +5,7 @@
 
 import os
 import shutil
+from collections.abc import Callable
 from typing import Annotated
 
 import typer
@@ -2830,6 +2831,43 @@ def configure(
         raise typer.Exit(130) from None
 
 
+def _is_managed_config_admin() -> bool:
+    """True when the managed-config feature is enabled and the current caller is a workspace admin.
+
+    Best-effort: returns False when the workspace is unknown or authentication fails.
+    """
+    if not managed_agent_config_enabled():
+        return False
+    state = load_state()
+    workspace = state.get("workspace")
+    if not workspace:
+        return False
+    profile = state.get("profile")
+    try:
+        ensure_databricks_auth(workspace, profile, quiet=True)
+        token = get_databricks_token(workspace, profile)
+    except RuntimeError:
+        return False
+    return bool(is_workspace_admin(workspace, token))
+
+
+def _run_managed_authoring(run: Callable[[], int]) -> None:
+    """Run a managed-config authoring wizard with the CLI's standard error/exit mapping."""
+    # `typer.Exit` subclasses RuntimeError, so it must be raised outside the try — inside, the
+    # `except RuntimeError` below would swallow it and report a clean exit as an error.
+    try:
+        install_databricks_cli()
+        code = run()
+    except RuntimeError as exc:
+        print_err(str(exc))
+        raise typer.Exit(1) from None
+    except KeyboardInterrupt:
+        print_err("Interrupted.")
+        raise typer.Exit(130) from None
+    if code:
+        raise typer.Exit(code)
+
+
 @configure_app.command("mcp")
 def configure_mcp(
     location: Annotated[
@@ -2854,7 +2892,15 @@ def configure_mcp(
         ),
     ] = None,
 ) -> None:
-    """Add Databricks MCP servers to installed coding tools."""
+    """Add Databricks MCP servers to installed coding tools.
+
+    Role-aware: a workspace admin (with ENABLE_MANAGED_AGENT_CONFIG set) authors the managed
+    config's MCP servers for the whole workspace; a developer configures their own tools. Admins
+    who want their own personal MCP servers use `ucode mcp add`/`remove`.
+    """
+    if _is_managed_config_admin():
+        _run_managed_authoring(setup_mcp_command)
+        return
     # `--services` absent -> None (whole schema); present (even empty) -> the
     # explicit subset, so `--services ""` deselects everything.
     selected = None if services is None else {s.strip() for s in services.split(",") if s.strip()}
@@ -2906,7 +2952,19 @@ def configure_skills(
     registers the MCP connection with utility tools only. ``--skill`` narrows a
     download to a named subset of a single schema's skills (requires exactly one
     ``--location``).
+
+    Role-aware: a workspace admin (with ENABLE_MANAGED_AGENT_CONFIG set) authors the managed
+    config's skills for the whole workspace — ``--location`` (or the prompt) names the schemas, and
+    the download-only flags (``--mcp``/``--path``/``--skill``) don't apply. A developer configures
+    their own tools as described above.
     """
+    if _is_managed_config_admin():
+        _run_managed_authoring(
+            lambda: setup_skills_command(
+                None if location is None else _parse_skill_locations(location)
+            )
+        )
+        return
     try:
         locations = _parse_skill_locations(location)
         # `--skill` absent -> None (whole schema); present (even empty) -> the
@@ -2956,6 +3014,17 @@ def configure_tracing(
         raise typer.Exit(130) from None
 
 
+@configure_app.command("spend-tiers")
+def configure_spend_tiers() -> None:
+    """Route developers to cheaper agents as the workspace spends its budget (admins only).
+
+    Authors the managed config's tiered spend policy — the workspace-wide config `ucode publish`
+    sends to every developer, not this machine's own settings. Always authors and errors for
+    non-admins (there is no developer-facing form).
+    """
+    _run_managed_authoring(setup_budget_policy_command)
+
+
 @setup_app.callback(invoke_without_command=True)
 def setup(
     ctx: typer.Context,

src/ucode/managed_wizard.py

@@ -1400,10 +1400,10 @@ def setup_from_file(path: str) -> int:
 # The sections that have their own `ucode setup <thing>` command, in the order the checklist lists
 # them: the command, the label the summary uses, and how to tell whether the manifest has one.
 SETUP_SECTIONS: list[tuple[str, str, Callable[[dict], bool]]] = [
-    ("ucode setup mcps", "MCP servers", lambda m: bool(m.get("mcp_servers"))),
-    ("ucode setup skills", "Skills", lambda m: bool((m.get("skills") or {}).get("names"))),
+    ("ucode configure mcp", "MCP servers", lambda m: bool(m.get("mcp_servers"))),
+    ("ucode configure skills", "Skills", lambda m: bool((m.get("skills") or {}).get("names"))),
     (
-        "ucode setup spend-tiers",
+        "ucode configure spend-tiers",
         "Tiered Spend Policy",
         lambda m: isinstance(m.get("budget_policy"), dict),
     ),

tests/test_cli.py

@@ -830,6 +830,144 @@ def test_path_without_location_exit_1(self):
         mock_download.assert_not_called()
 
 
+class TestConfigureSpendTiersCommand:
+    """`ucode configure spend-tiers` authors the managed config's tiered spend policy (admin)."""
+
+    def test_registered_and_calls_the_wizard(self):
+        with (
+            patch("ucode.cli.install_databricks_cli"),
+            patch("ucode.cli.setup_budget_policy_command", return_value=0) as fn,
+        ):
+            result = runner.invoke(app, ["configure", "spend-tiers"])
+        assert result.exit_code == 0, result.output
+        assert fn.called
+        assert "ERROR" not in _strip_ansi(result.output)
+
+    def test_runtime_error_exits_1(self):
+        with (
+            patch("ucode.cli.install_databricks_cli"),
+            patch(
+                "ucode.cli.setup_budget_policy_command",
+                side_effect=RuntimeError("not an admin"),
+            ),
+        ):
+            result = runner.invoke(app, ["configure", "spend-tiers"])
+        assert result.exit_code == 1
+
+    def test_interrupt_exits_130(self):
+        with (
+            patch("ucode.cli.install_databricks_cli"),
+            patch("ucode.cli.setup_budget_policy_command", side_effect=KeyboardInterrupt),
+        ):
+            result = runner.invoke(app, ["configure", "spend-tiers"])
+        assert result.exit_code == 130
+
+    def test_nonzero_wizard_code_propagates(self):
+        # `setup_budget_policy_command` returns a process exit code; a non-zero one must surface as
+        # the command's exit code, not be swallowed into a success.
+        with (
+            patch("ucode.cli.install_databricks_cli"),
+            patch("ucode.cli.setup_budget_policy_command", return_value=3),
+        ):
+            result = runner.invoke(app, ["configure", "spend-tiers"])
+        assert result.exit_code == 3
+
+
+class TestConfigureMcpSkillsRoleAware:
+    """`ucode configure mcp`/`skills` are role-aware: an admin (with the flag set) authors the
+    managed config; a developer configures their own tools. Gated by ENABLE_MANAGED_AGENT_CONFIG."""
+
+    def _role(self, monkeypatch, *, is_admin):
+        monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1")
+        monkeypatch.setattr(
+            "ucode.cli.load_state", lambda: {"workspace": "https://w", "profile": None}
+        )
+        monkeypatch.setattr("ucode.cli.ensure_databricks_auth", lambda *a, **k: None)
+        monkeypatch.setattr("ucode.cli.get_databricks_token", lambda *a, **k: "tok")
+        monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda *a, **k: is_admin)
+
+    def test_admin_authors_managed_mcp(self, monkeypatch):
+        self._role(monkeypatch, is_admin=True)
+        with (
+            patch("ucode.cli.install_databricks_cli"),
+            patch("ucode.cli.setup_mcp_command", return_value=0) as author,
+            patch("ucode.cli.configure_mcp_command") as local,
+        ):
+            result = runner.invoke(app, ["configure", "mcp"])
+        assert result.exit_code == 0, result.output
+        assert author.called
+        assert not local.called
+
+    def test_admin_check_authenticates_quietly(self, monkeypatch):
+        self._role(monkeypatch, is_admin=True)
+        with (
+            patch("ucode.cli.ensure_databricks_auth") as auth,
+            patch("ucode.cli.install_databricks_cli"),
+            patch("ucode.cli.setup_mcp_command", return_value=0),
+        ):
+            result = runner.invoke(app, ["configure", "mcp"])
+        assert result.exit_code == 0, result.output
+        assert auth.call_args.kwargs.get("quiet") is True
+
+    def test_developer_configures_local_mcp(self, monkeypatch):
+        self._role(monkeypatch, is_admin=False)
+        with (
+            patch("ucode.cli.setup_mcp_command") as author,
+            patch("ucode.cli.configure_mcp_command") as local,
+        ):
+            result = runner.invoke(app, ["configure", "mcp"])
+        assert result.exit_code == 0, result.output
+        assert local.called
+        assert not author.called
+
+    def test_flag_off_stays_local_mcp(self, monkeypatch):
+        # No ENABLE_MANAGED_AGENT_CONFIG -> never an admin check, always the developer flow.
+        monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False)
+        with (
+            patch("ucode.cli.setup_mcp_command") as author,
+            patch("ucode.cli.configure_mcp_command") as local,
+        ):
+            result = runner.invoke(app, ["configure", "mcp"])
+        assert result.exit_code == 0, result.output
+        assert local.called
+        assert not author.called
+
+    def test_admin_authors_managed_skills(self, monkeypatch):
+        self._role(monkeypatch, is_admin=True)
+        with (
+            patch("ucode.cli.install_databricks_cli"),
+            patch("ucode.cli.setup_skills_command", return_value=0) as author,
+            patch("ucode.cli.configure_skills_mcp_command") as local_mcp,
+            patch("ucode.cli.configure_skills_download_command") as local_dl,
+        ):
+            result = runner.invoke(app, ["configure", "skills"])
+        assert result.exit_code == 0, result.output
+        assert author.called
+        assert not local_mcp.called
+        assert not local_dl.called
+
+    def test_admin_skills_forwards_location(self, monkeypatch):
+        self._role(monkeypatch, is_admin=True)
+        with (
+            patch("ucode.cli.install_databricks_cli"),
+            patch("ucode.cli.setup_skills_command", return_value=0) as author,
+        ):
+            result = runner.invoke(app, ["configure", "skills", "--location", "main.a,main.b"])
+        assert result.exit_code == 0, result.output
+        assert author.call_args.args[0] == ["main.a", "main.b"]
+
+    def test_developer_configures_local_skills(self, monkeypatch):
+        self._role(monkeypatch, is_admin=False)
+        with (
+            patch("ucode.cli.setup_skills_command") as author,
+            patch("ucode.cli.configure_skills_mcp_command") as local_mcp,
+        ):
+            result = runner.invoke(app, ["configure", "skills"])
+        assert result.exit_code == 0, result.output
+        assert local_mcp.called
+        assert not author.called
+
+
 class TestApplyManagedSkills:
     """The launch path both registers the skills MCP connection and downloads bundles to disk."""
 

tests/test_managed_wizard.py

@@ -2205,9 +2205,9 @@ def test_marks_configured_and_unconfigured_sections(self, capsys):
         manifest = {**AGENTS_ONLY, "skills": {"names": ["main.default"]}}
         wizard._print_next_steps(manifest)
         out = capsys.readouterr().out
-        assert "ucode setup mcps" in out
-        assert "ucode setup skills" in out
-        assert "ucode setup spend-tiers" in out
+        assert "ucode configure mcp" in out
+        assert "ucode configure skills" in out
+        assert "ucode configure spend-tiers" in out
         assert "ucode publish" in out
 
     def test_dry_run_says_nothing_was_saved(self, capsys, monkeypatch):
@@ -2380,9 +2380,9 @@ def test_lists_every_setup_command(self, capsys):
         out = capsys.readouterr().out
         for command in (
             "ucode setup",
-            "ucode setup mcps",
-            "ucode setup skills",
-            "ucode setup spend-tiers",
+            "ucode configure mcp",
+            "ucode configure skills",
+            "ucode configure spend-tiers",
             "ucode setup show",
             "ucode publish",
         ):