diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 4b0c6967..a8fc448d 100644 --- a/src/ucode/cli.py +++ b/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, diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 3710c68b..c7fe43ae 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -1400,10 +1400,10 @@ def setup_from_file(path: str) -> int: # The sections that have their own `ucode setup ` 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), ), diff --git a/tests/test_cli.py b/tests/test_cli.py index 605c4327..aa1675fe 100644 --- a/tests/test_cli.py +++ b/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.""" diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 8d3513dd..41a0b990 100644 --- a/tests/test_managed_wizard.py +++ b/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", ):