From 1738f1e9fce93ec9f9c75260f1142d65422c7d82 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 31 Aug 2026 00:36:38 +0000 Subject: [PATCH 1/3] Author managed config's mcp/skills/spend-tiers via `ucode configure` Move the managed-config section authoring from `ucode setup
` onto the role-aware `ucode configure`, so it survives the `ucode setup` removal (AIGTWY-4342) and matches the ug IA (admin authoring stays in the CLI only temporarily, until the Databricks UI/REST/TF surfaces land). - `ucode configure spend-tiers`: new admin-only command wrapping setup_budget_policy_command (no developer form, so it always authors). - `ucode configure mcp` / `ucode configure skills`: now role-aware, mirroring bare `ucode configure`. A workspace admin (with ENABLE_MANAGED_AGENT_CONFIG set) authors the managed config's MCP servers / skills; a developer configures their own tools. Gated by the flag, so with it off these stay developer-only exactly as before. Admins wanting personal MCP servers use `ucode mcp add`/`remove`. - Repoint the authoring flow's "Next steps" list at the `ucode configure` section commands. `ucode setup
` still works; its removal is the separate AIGTWY-4342 deprecation. Co-authored-by: Isaac --- src/ucode/cli.py | 75 ++++++++++++++++++++- src/ucode/managed_wizard.py | 6 +- tests/test_cli.py | 127 +++++++++++++++++++++++++++++++++++ tests/test_managed_wizard.py | 12 ++-- 4 files changed, 210 insertions(+), 10 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 4b0c6967..32c21d7d 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,47 @@ def configure( raise typer.Exit(130) from None +def _authoring_managed_config() -> bool: + """True when a section command should author the workspace's managed config, not local settings. + + `ucode configure mcp`/`skills` are role-aware, mirroring bare `ucode configure`: a workspace + admin authors the config the whole workspace pulls, a developer configures their own tools. + Gated by ENABLE_MANAGED_AGENT_CONFIG, so with the flag off these commands stay developer-only, + exactly as before. Best-effort — a missing workspace or an auth failure returns False and falls + through to the local flow rather than blocking the developer path on an admin check. + """ + 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) + 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 +2896,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 _authoring_managed_config(): + _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 +2956,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 _authoring_managed_config(): + _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 +3018,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. Unlike `configure mcp`/`skills` this + has no developer-facing form, so it always authors (and errors for non-admins). + """ + _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..348392e5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -830,6 +830,133 @@ 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_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", ): From c6271dced3fc3e8c6cd131fb726ff476ed408b44 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 31 Aug 2026 01:15:38 +0000 Subject: [PATCH 2/3] Rename _authoring_managed_config -> _is_managed_config_admin; fix docstrings The old name described how the result is used, not what the function checks; it's a predicate for "is the caller a workspace admin (with the feature on)". Also tighten its docstring and configure_spend_tiers' to state behavior rather than caller context. Co-authored-by: Isaac --- src/ucode/cli.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 32c21d7d..43e70572 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2831,14 +2831,10 @@ def configure( raise typer.Exit(130) from None -def _authoring_managed_config() -> bool: - """True when a section command should author the workspace's managed config, not local settings. - - `ucode configure mcp`/`skills` are role-aware, mirroring bare `ucode configure`: a workspace - admin authors the config the whole workspace pulls, a developer configures their own tools. - Gated by ENABLE_MANAGED_AGENT_CONFIG, so with the flag off these commands stay developer-only, - exactly as before. Best-effort — a missing workspace or an auth failure returns False and falls - through to the local flow rather than blocking the developer path on an admin check. +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 @@ -2902,7 +2898,7 @@ def configure_mcp( 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 _authoring_managed_config(): + if _is_managed_config_admin(): _run_managed_authoring(setup_mcp_command) return # `--services` absent -> None (whole schema); present (even empty) -> the @@ -2962,7 +2958,7 @@ def configure_skills( the download-only flags (``--mcp``/``--path``/``--skill``) don't apply. A developer configures their own tools as described above. """ - if _authoring_managed_config(): + if _is_managed_config_admin(): _run_managed_authoring( lambda: setup_skills_command( None if location is None else _parse_skill_locations(location) @@ -3023,8 +3019,8 @@ 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. Unlike `configure mcp`/`skills` this - has no developer-facing form, so it always authors (and errors for non-admins). + 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) From b2341a67d6dc4bbc7ec32a39f92fd971134a47ca Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 31 Aug 2026 01:49:58 +0000 Subject: [PATCH 3/3] Quiet the role-dispatch admin check so auth prints once MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode configure mcp`/`skills` authenticate once in `_is_managed_config_admin` to pick the admin-vs-developer branch, then the chosen branch authenticates again — printing "Databricks auth already available" twice. Pass quiet=True to the pre-check's ensure_databricks_auth (as setup_command already does) so only the branch reports it. A real login is never silenced by quiet. Co-authored-by: Isaac --- src/ucode/cli.py | 2 +- tests/test_cli.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 43e70572..a8fc448d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2844,7 +2844,7 @@ def _is_managed_config_admin() -> bool: return False profile = state.get("profile") try: - ensure_databricks_auth(workspace, profile) + ensure_databricks_auth(workspace, profile, quiet=True) token = get_databricks_token(workspace, profile) except RuntimeError: return False diff --git a/tests/test_cli.py b/tests/test_cli.py index 348392e5..aa1675fe 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -898,6 +898,17 @@ def test_admin_authors_managed_mcp(self, monkeypatch): 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 (