diff --git a/README.md b/README.md index d0a457c..c657210 100644 --- a/README.md +++ b/README.md @@ -179,24 +179,35 @@ you to run `ucode ` (existing agent sessions need a restart before the MC ### Managed config for a workspace (admins) +Author the coding config your developers pick up automatically, instead of asking each of them to +run `ucode configure` by hand. Restricted to workspace admins. `ucode setup help` prints the whole +sequence; the short version is one command for the agents and models, then a command per optional +section, then publish: + ```bash -ucode setup +ucode setup # agents and models (start here) +ucode setup mcp # managed MCP servers +ucode setup skills # managed skills +ucode setup budget-policy # spend-based routing +ucode apply # publish it to the workspace ``` -Author the coding config your developers pick up automatically, instead of asking each of them to -run `ucode configure` by hand. Restricted to workspace admins. +`ucode setup` walks through the agents to enable and which one bare `ucode` launches, then per agent: +Databricks-hosted models or an external Model Provider Service, the models to expose, and (for Claude +Code and Codex) whether the config writes the agent's own OS-level settings file or a ucode-only one. +Claude Code is asked one model per family (opus/sonnet/haiku/fable), since it selects models by family +alias; any family can be skipped. -The flow walks through the agents to enable and which one bare `ucode` launches, then per agent: -Databricks-hosted models or an external Model Provider Service, the models to expose, and whether -the config applies machine-wide or per user. Claude Code is asked one model per family -(opus/sonnet/haiku/fable), since Claude Code selects models by family alias; any family can be -skipped. It then offers tracing, managed MCP servers, skills, and a spend-based budget policy that -switches the default agent and model as the workspace burns through a budget. +The optional sections each edit their own part of the same config, so you can add an MCP server or +change a budget tier later without walking the whole flow. `ucode setup skills --location +main.default,other.schema` skips the prompt. `ucode setup budget-policy` sets a spend-based policy +that switches the default agent and model as the workspace burns through a budget. Answering these +also runs the matching `ucode configure` step, which does configure this machine. -The result is written to `~/.ucode/managed-state.json` — the one local managed-config file — which -`ucode apply` publishes to the workspace. Your own agent configs are left alone, with one exception: -answering yes to tracing, MCP servers, or skills runs the matching `ucode configure` step, which -does configure this machine. +Everything is written to `~/.ucode/managed-state.json` — the one local managed-config file — which +`ucode apply` publishes. Re-running `ucode setup` keeps the MCP servers, skills, tracing table, and +budget policy already authored, rather than clearing them; to drop one, edit the file and reload it +with `ucode setup --from-file`. ```bash # Review the manifest and the exact payload `ucode apply` would publish. @@ -209,7 +220,7 @@ ucode setup --from-file ./managed-config.json Once the manifest looks right, publish it: ```bash -# Validate, show what would change, and ask before publishing. +# Validate, show a diff against what's live, and ask before publishing. ucode apply # Publish without the confirmation prompt (for CI). @@ -217,9 +228,11 @@ ucode apply --yes ``` `apply` updates the workspace's existing config in place rather than replacing it, so a failed -publish leaves the current config intact. It is still a whole-manifest write: every field ucode -authors is sent, so anything skipped in a re-run is cleared rather than carried over. Developers -pick the new config up on their next ucode run. +publish leaves the current config intact. It shows a diff of exactly what changes against the +published config before asking to confirm, and does nothing when the two already match. It is a +whole-manifest write — every field ucode authors is sent — but because `ucode setup` carries the +other sections forward, a re-run no longer silently drops them. Developers pick the new config up on +their next ucode run. --- @@ -247,10 +260,14 @@ pick the new config up on their next ucode run. | `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | | `ucode configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | | `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | -| `ucode setup` | Author the workspace's managed coding config (workspace admins only) | +| `ucode setup` | Author the managed config's agents and models (workspace admins only) | +| `ucode setup mcp` | Add or change the managed config's MCP servers | +| `ucode setup skills [--location a.b,c.d]` | Add or change the managed config's skills | +| `ucode setup budget-policy` | Set the managed config's spend-based routing policy | +| `ucode setup help` | Walk through the whole setup sequence, marking what's already configured | | `ucode setup show` | Print the authored config and the payload `ucode apply` would publish | | `ucode setup --from-file ` | Load a hand-written managed config instead of running the prompts | -| `ucode apply` | Publish the authored managed config to the workspace (workspace admins only) | +| `ucode apply` | Publish the authored managed config to the workspace, after a diff and confirmation (admins only) | | `ucode apply --yes` | Publish without the confirmation prompt | ## Managed Local Files diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0f16465..2eeb621 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -80,7 +80,15 @@ recommended_agent, resolve_state, ) -from ucode.managed_wizard import apply_command, setup_command, show_command +from ucode.managed_wizard import ( + apply_command, + setup_budget_policy_command, + setup_command, + setup_help_command, + setup_mcp_command, + setup_skills_command, + show_command, +) from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -263,8 +271,9 @@ def _maybe_offer_admin_setup(workspace: str, profile: str | None) -> None: if not is_admin: return print_note( - "✨ New: as a workspace admin you can publish a managed config with `ucode setup` — set " - "the agents, models, MCPs, and skills once, and every developer picks them up automatically." + "✨ New: as a workspace admin you can publish a managed config with `ucode setup` — set the " + "agents and models once (then MCP servers and skills with `ucode setup mcp` / `skills`), and " + "every developer picks them up automatically." ) if prompt_yes_no("Set one up now with `ucode setup`?"): # Launch the setup flow in place rather than telling them to re-run a command. Reuse the @@ -714,7 +723,7 @@ def _use_databricks() -> dict: return _use_databricks() choice = prompt_for_selection( - f"How should {display} be configured?", + f"How should {display} get its models?", [ ("databricks", "Databricks Hosted"), ("mps", "External Models"), @@ -1044,7 +1053,9 @@ def revert() -> int: app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") setup_app = typer.Typer(add_completion=False, no_args_is_help=False) app.add_typer( - setup_app, name="setup", help="Author the workspace's managed coding config (admins only)." + setup_app, + name="setup", + help="Author the workspace's managed coding config (admins only). See `ucode setup help`.", ) @@ -2614,7 +2625,10 @@ def setup( ), ] = None, ) -> None: - """Author the managed coding config for your workspace (workspace admins only).""" + """Choose the agents and models for your workspace's managed config (admins only). + + MCP servers, skills, and the budget policy have their own commands — see `ucode setup help`. + """ if ctx.invoked_subcommand is not None: return # `typer.Exit` subclasses RuntimeError, so it must be raised outside the try — inside, the @@ -2632,6 +2646,80 @@ def setup( raise typer.Exit(code) +@setup_app.command("mcp") +def setup_mcp_cmd() -> None: + """Choose the MCP servers the managed config gives developers (admins only).""" + # Same `typer.Exit`/RuntimeError ordering trap as the `setup` callback above. + try: + install_databricks_cli() + code = setup_mcp_command() + 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) + + +@setup_app.command("skills") +def setup_skills_cmd( + location: Annotated[ + str | None, + typer.Option( + "--location", + help="Skill schemas to publish as `.` (comma-separated for several). " + "Skips the prompt.", + ), + ] = None, +) -> None: + """Choose the skills the managed config gives developers (admins only).""" + try: + install_databricks_cli() + # None means "prompt"; an explicit `--location` is parsed to the list to publish. + locations = None if location is None else _parse_skill_locations(location) + code = setup_skills_command(locations) + 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) + + +@setup_app.command("budget-policy") +def setup_budget_policy_cmd() -> None: + """Route developers to cheaper agents as the workspace spends its budget (admins only).""" + try: + install_databricks_cli() + code = setup_budget_policy_command() + 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) + + +@setup_app.command("help") +def setup_help_cmd() -> None: + """Walk through the managed-config setup: every command, in order, and what's already done.""" + # No auth and no CLI install: this reads the local draft only, so it works before `ucode + # configure` and on a machine without the Databricks CLI. + try: + code = setup_help_command() + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + if code: + raise typer.Exit(code) + + @setup_app.command("show") def setup_show_cmd() -> None: """Print the authored managed config and the payload `ucode apply` would publish.""" diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 88935dd..89ef09d 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1681,7 +1681,7 @@ def discover_model_services( - ``claude_models`` maps ``fable``/``opus``/``sonnet``/``haiku`` to the newest matching ``system.ai.claude-*`` id (mirrors ``discover_claude_models``). - - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids. + - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids, newest first. - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids. @@ -1708,7 +1708,7 @@ def discover_model_services( # newest-wins once the router accepts opus-5 (PR databricks-eng/universe#2365446). _prefer_opus_4_8(claude_models, ids) - codex_models = [m for m in ids if "gpt-" in m] + codex_models = sorted([m for m in ids if "gpt-" in m], key=model_version_sort_key) gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)] @@ -2839,7 +2839,11 @@ def discover_gemini_models(workspace: str, token: str) -> tuple[list[str], str | def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str | None]: - return discover_endpoints_with_api_type(workspace, token, "openai/v1/responses") + # Order newest model version first (like `discover_gemini_models`), so the picker's top choice + # and default is e.g. gpt-5-4 rather than the alphabetically-first gpt-5. + return discover_endpoints_with_api_type( + workspace, token, "openai/v1/responses", sort_key=model_version_sort_key + ) def fetch_gemini_models(workspace: str, token: str) -> list[str]: diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 2940b8b..82c9d94 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -1,24 +1,30 @@ """Interactive `ucode setup`: author the workspace's managed coding-agent config. -Workspace admins run this to build the ``CodingAgentConfig`` their developers will pull. It walks -the admin through agents, per-agent models, tracing, MCP servers, skills, and a spend-routing budget -policy, then writes the manifest to ``~/.ucode/managed-state.json`` (the one local managed-config -file, owned by :mod:`ucode.managed_config`). Publishing it to the workspace is ``ucode apply`` (a -separate command, so an admin can review the file first). +Workspace admins run this to build the ``CodingAgentConfig`` their developers will pull, then publish +it with ``ucode apply`` (a separate command, so the manifest can be reviewed first). The config lives +at ``~/.ucode/managed-state.json`` (the one local managed-config file, owned by +:mod:`ucode.managed_config`). + +Authoring is split across commands so an admin can change one part without walking the whole flow: +``ucode setup`` picks the agents and models, and ``ucode setup mcp`` / ``skills`` / ``budget-policy`` +each edit their own section of the same manifest. ``ucode setup`` carries the other sections forward +untouched (:func:`_carry_forward_sections`), and ``ucode setup help`` prints the whole sequence. Serialization, validation, and the per-agent model catalogs live in :mod:`ucode.managed_setup`; this -module is the interaction layer on top of them. Sub-flows an admin already knows — tracing, MCP, -skills — are delegated to the existing ``ucode configure `` commands and their results read -back out of ``state.json``, so there is exactly one picker per concern in the codebase. +module is the interaction layer on top of them. Sub-flows an admin already knows — MCP, skills — are +delegated to the existing ``ucode configure `` commands and their results read back out of +``state.json``, so there is exactly one picker per concern in the codebase. """ from __future__ import annotations import json +from collections.abc import Callable from decimal import Decimal from pathlib import Path from typing import cast +from ucode import config_io from ucode.agents import GLOBAL_SETTINGS_AGENTS, TOOL_SPECS, check_gateway_endpoint from ucode.databricks import ( ANTHROPIC_FAMILIES, @@ -42,6 +48,7 @@ get_managed_config, load_managed_state, managed_state_workspace, + normalize_managed_config, save_managed_state, ) from ucode.managed_setup import ( @@ -65,6 +72,7 @@ print_section, print_success, print_warning, + print_warning_panel, prompt_for_multi_selection, prompt_for_percentage, prompt_for_selection, @@ -75,27 +83,26 @@ ) # The OS-level managed settings file `use_as_global_settings` writes for each agent — named in the -# prompt so an admin sees exactly what answering "yes" touches. +# prompt so an admin sees exactly what answering "yes" touches. Yes writes this file (needs sudo +# once) so a bare `claude`/`codex` reaches the gateway on its own; no keeps it ucode-only. GLOBAL_SETTINGS_FILES = { - "claude": "Claude Code's managed-settings.json", - "codex": "Codex's managed_config.toml", + "claude": "managed-settings.json", + "codex": "managed_config.toml", } -# What `use_as_global_settings` actually does, in plain terms. `{binary}` is filled in per agent. -GLOBAL_SETTINGS_BLURB = ( - "Answer yes to write the gateway config into that file (needs sudo once), so a bare `{binary}` " - "reaches the Databricks gateway on its own — you don't have to launch it through ucode. Answer " - "no to write a ucode-only settings file instead, so `{binary}` uses the gateway only when " - "started with `ucode {binary}`." -) - BUDGET_POLICY_BLURB = ( - "A budget policy moves developers onto cheaper agents and models as the workspace spends " - "against a budget — for example Claude Code on Opus by default, then Sonnet at 80%, then " - "OpenCode on Kimi at 100%. It only changes the default; developers can still pick any Model " - "Service to which they have access. Hard caps stay with the budget's own blocking threshold." + "As the workspace spends more of a budget, a policy automatically switches everyone's default " + "agent and model to a cheaper one — for example Claude Code / Opus normally, Claude Code / " + "Sonnet once spend passes 80%, OpenCode / Kimi past 100%.\n\n" + "It only moves the default. Developers can still pick any model they have access to, and the " + "budget's own hard block is what actually caps spend." ) +# Agents not offered in `ucode setup`'s picker, even when the workspace serves their models. +# `ucode gemini` still works as a launch target; it's just not part of the managed config authored +# here. Serialize/validate keep supporting it, so a `--from-file` manifest can still name it. +SETUP_EXCLUDED_AGENTS = frozenset({"gemini"}) + def _tracing_table_from_state(state: dict) -> str | None: """The UC table `ucode configure tracing` wired up, or None when tracing is off. @@ -401,6 +408,39 @@ def _prompt_models_for_agent(tool: str, state: dict, provider_service: dict | No _SKIP_FAMILY = "__skip__" +def _confirm_agent(tool: str, agent_config: dict) -> None: + """One consistent closing line per agent in step 2, whatever its model shape. + + Every agent — a single-model codex, a multi-model opencode, a family-slotted claude — ends its + block with the same `✔ configured — · ` line, so the step reads as a + uniform checklist rather than each agent's picker trailing off differently. + """ + display = TOOL_SPECS.get(tool, {}).get("display", tool) + model_config = agent_config.get("model_config") or {} + detail = model_config.get("default_model") or "no model" + provider = model_config.get("model_provider_service") + if provider: + detail = f"{detail} via {provider}" + if tool in GLOBAL_SETTINGS_AGENTS: + scope = "global settings" if agent_config.get("use_as_global_settings") else "ucode-only" + detail = f"{detail} · {scope}" + print_success(f"{display} configured — {detail}") + + +def _render_family_slots(slots: dict[str, str]) -> None: + """Recap the Claude family → model slots just chosen, before the overall-default question. + + Same "form filling in" motif as :func:`_selected_recap`: the per-family answers scrolled by one at + a time, so gathering them into one box makes "which of these is the overall default?" a choice + over something the admin can see rather than recall. + """ + lines = [ + kv_line(slot.removeprefix("default_").removesuffix("_model"), model) + for slot, model in slots.items() + ] + print_panel("Claude Code models", lines) + + def _prompt_claude_models(state: dict) -> dict: """Build Claude's ``model_config`` one family slot at a time. @@ -468,8 +508,9 @@ def _prompt_claude_models(state: dict) -> dict: # A one-option prompt is a wasted keystroke, but skipping it silently reads as a dropped # step — say what was inferred so the admin knows the default is set, and to what. model_config["default_model"] = chosen[0] - print_success(f"Overall default for {display}: {chosen[0]} (the only model configured)") + print_note(f"Only one model configured, so it's {display}'s overall default.") else: + _render_family_slots(slots) model_config["default_model"] = _require_selection( f"Which of those is {display}'s overall default?", [(m, m) for m in chosen] ) @@ -553,7 +594,8 @@ def _prompt_claude_provider_family_models(targets: list[str], service_name: str) for fam in ANTHROPIC_FAMILIES if CLAUDE_SLOT_FOR_FAMILY[fam] in slots ) - print_success(f"{display}: {summary} (default: {default_family})") + # A note, not a success line: the loop's `_confirm_agent` prints the single ✔ for the agent. + print_note(f"Quick setup — {summary} (default: {default_family}).") return model_config print_note( @@ -581,8 +623,10 @@ def _prompt_claude_provider_family_models(targets: list[str], service_name: str) chosen = list(dict.fromkeys(slots.values())) if len(chosen) == 1: model_config["default_model"] = chosen[0] - print_success(f"Overall default for {display}: {chosen[0]} (the only model configured)") + print_note(f"Only one model configured, so it's {display}'s overall default.") else: + if slots: + _render_family_slots(slots) # Offered over every target, not just the slots: `default_model` needn't be a family model, # and a mixed-catalog service may expose one an admin wants as the overall default. options = chosen or list(targets) @@ -808,31 +852,67 @@ def configured_models_for_agent(agent_config: dict) -> list[str]: return list(dict.fromkeys(models)) +def _render_tier_ladder(tiers: list[dict], threshold: object, *, base_default: str = "") -> None: + """Show the tiers built so far as spend ranges, so the fallback ladder reads at a glance. + + A tier activates once spend passes its percentage and the highest passed tier wins, so each + tier really owns the range from its own percentage up to the next tier's. Rendering those ranges + ("50–90%", "90%+") rather than bare thresholds ("at 50%", "at 90%") is what makes the ladder + legible — the admin sees which agent a developer actually gets at any level of spend. Reprinted + as the ladder grows, so the sequence forms in front of them instead of in their head. + """ + ordered = sorted(tiers, key=lambda t: t["spending_percentage"]) + lines: list[str] = [] + if base_default: + # Below the first tier the manifest's own default applies; naming it anchors the sequence. + first = ordered[0]["spending_percentage"] * 100 + lines.append(kv_line(f"under {first:g}%", f"{base_default} (default)")) + for i, tier in enumerate(ordered): + low = tier["spending_percentage"] * 100 + agent = TOOL_SPECS.get(tier["default_agent"], {}).get("display", tier["default_agent"]) + if i + 1 < len(ordered): + span = f"{low:g}–{ordered[i + 1]['spending_percentage'] * 100:g}%" + else: + span = f"{low:g}%+" + lines.append(kv_line(span, f"{agent} / {tier['default_model']}")) + print_panel("Budget tiers so far", lines) + + def _prompt_budget_policy( - workspace: str, token: str, enabled_agents: dict[str, dict], state: dict + workspace: str, + token: str, + enabled_agents: dict[str, dict], + state: dict, + *, + base_default: str = "", ) -> dict | None: - """Author a spend-routing ``budget_policy``, or None when the admin declines or can't. + """Author a spend-routing ``budget_policy``, or None when the admin backs out or can't. Budgets themselves are created in the Databricks console (they're account-level objects), so the admin picks an existing one here. Tiers are prompted in percent and stored as fractions, which is what the API validates. - A tier's model choices come from what the admin configured for that agent earlier in this run — - not the workspace catalog. Offering the catalog would let a tier point an agent at a model it - wasn't given, which neither this validation nor the server's would reject: the tier would + ``enabled_agents`` is what the manifest gives each agent, so a tier's model choices come from + that rather than the workspace catalog. Offering the catalog would let a tier point an agent at a + model it wasn't given, which neither this validation nor the server's would reject: the tier would activate and hand the developer a model their agent doesn't have. + + Asks no "set up a budget policy?" gate — running `ucode setup budget-policy` is the answer to that + question, the same way `ucode configure ` needs no confirmation. """ print_section("Budget policy") - print_note(BUDGET_POLICY_BLURB) - if not prompt_yes_no_default("Set up a budget policy for this workspace?", default=False): - return None + # Check for attachable budgets before anything else: budgets are created in the Databricks + # console, so if there are none (or none that can enforce routing) there is nothing to do here. + # Bail with a boxed warning and skip the explanatory blurb — no point explaining a feature the + # workspace can't use yet. with spinner("Listing workspace budgets..."): budgets, reason = list_workspace_budgets(workspace, token) if reason is not None or not budgets: - print_warning( + print_warning_panel( "No AI Gateway budgets are visible for this workspace, so there is nothing to attach a " - "policy to. Create a budget in the Databricks console first, then re-run `ucode setup`." + "policy to. Create a budget in the Databricks console first, then re-run " + "`ucode setup budget-policy`." ) return None @@ -842,12 +922,17 @@ def _prompt_budget_policy( # 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( + print_warning_panel( "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`." + "block action to a budget in the Databricks console, then re-run " + "`ucode setup budget-policy`." ) return None + + # Budgets exist — now explain what a policy does, before asking the admin to pick one. Boxed so + # the concept is read as a unit rather than skimmed as one more bullet. + print_panel("What is a budget policy?", [BUDGET_POLICY_BLURB]) print_note( "Showing only budgets with a per-user hard block configured, which spend routing enforces." ) @@ -889,15 +974,21 @@ def _prompt_budget_policy( seen_percentages: set[float] = set() seen_combos: set[tuple[str, str]] = set() print_note( - "Add one tier per step-down. Each tier activates once spend reaches its percentage, and " - "the highest activated tier wins." + "Add a tier for each step down: once spend passes the percentage you set, everyone's " + "default switches to the cheaper agent and model you pick." ) while True: index = len(tiers) + 1 - fraction = prompt_for_percentage(f"Tier {index}: activates at what percent of budget?") - if fraction in seen_percentages: - print_err("That percentage is already used by another tier; pick a different one.") - continue + + # Percentage first, in its own retry loop so a duplicate here re-asks only the percentage. + while True: + fraction = prompt_for_percentage( + f"Tier {index}: switch once spend passes what % of the budget?" + ) + if fraction in seen_percentages: + print_err("That percentage is already used by another tier; pick a different one.") + continue + break if threshold is not None: # Echo the dollars this percentage stands for, so the admin can sanity-check the tier # against the real per-user cap instead of reasoning about percentages in a vacuum. @@ -905,34 +996,47 @@ def _prompt_budget_policy( f" {fraction * 100:g}% of {format_usd(threshold)} is " f"{format_usd(threshold * Decimal(str(fraction)))}." ) - agent = prompt_for_selection( - f"Tier {index}: which agent becomes the default?", - [(tool, TOOL_SPECS[tool]["display"]) for tool in enabled_agents], - ) - if not agent: - break - # Only what this agent was actually configured with; the workspace catalog would offer - # models the agent doesn't have. - options = configured_models_for_agent(enabled_agents.get(agent) or {}) - if not options: - options = model_options_for_agent(agent, state) - if options: - model = prompt_for_selection( - f"Tier {index}: which model?", [(m, m) for m in options], searchable=True + + # Agent + model in their own retry loop: a duplicate agent/model re-asks just these two, so + # the admin doesn't have to retype the percentage they already entered for this tier. + agent = model = None + while True: + agent = prompt_for_selection( + f"Tier {index}: switch the default to which agent?", + [(tool, TOOL_SPECS[tool]["display"]) for tool in enabled_agents], ) - else: - model = prompt_for_text(f"Tier {index}: which model?") - if not model: + if not agent: + break + # Only what this agent was actually configured with; the workspace catalog would offer + # models the agent doesn't have. + options = configured_models_for_agent(enabled_agents.get(agent) or {}) + if not options: + options = model_options_for_agent(agent, state) + if options: + model = prompt_for_selection( + f"Tier {index}: using which model?", + [(m, m) for m in options], + searchable=True, + ) + else: + model = prompt_for_text(f"Tier {index}: using which model?") + if not model: + break + if (agent, model) in seen_combos: + # The highest crossed tier wins, so a second tier on the same agent+model never + # changes what the lower one already selected — a step-down that doesn't step down. + # Reject it rather than build a policy with a silently inert tier; only the agent and + # model are re-asked, the percentage above is kept. + print_err( + f"{TOOL_SPECS[agent]['display']} / {model} is already used by another tier; a " + "repeated agent/model makes this tier do nothing. Pick a different one." + ) + continue break - if (agent, model) in seen_combos: - # The highest crossed tier wins, so a second tier on the same agent+model never changes - # what the lower one already selected — it is a step-down that doesn't step down. Reject - # it here rather than let the admin build a policy with a silently inert tier. - print_err( - f"{TOOL_SPECS[agent]['display']} / {model} is already used by another tier; a " - "repeated agent/model makes this tier a no-op. Pick a different one." - ) - continue + # Cancelling the agent or model picker abandons this tier and stops adding more. + if not agent or not model: + break + seen_percentages.add(fraction) seen_combos.add((agent, model)) tiers.append( @@ -942,6 +1046,7 @@ def _prompt_budget_policy( "default_model": model, } ) + _render_tier_ladder(tiers, threshold, base_default=base_default) if not prompt_yes_no_default("Add another tier?", default=False): break @@ -999,7 +1104,10 @@ def _render_summary(workspace: str, manifest: dict) -> None: ) skills = (manifest.get("skills") or {}).get("names") or [] lines.append(kv_line("Skills", ", ".join(skills) if skills else "none")) - lines.append(kv_line("Tracing", manifest.get("tracing_table") or "disabled")) + # Managed tracing isn't offered by the flow yet, so a "disabled" line is just noise. Only surface + # it when a `--from-file` config actually set a table. + if manifest.get("tracing_table"): + lines.append(kv_line("Tracing", str(manifest["tracing_table"]))) policy = manifest.get("budget_policy") if isinstance(policy, dict): @@ -1019,6 +1127,119 @@ def _render_summary(workspace: str, manifest: dict) -> None: print_panel("Configuration summary", lines) +def _config_facts(manifest: dict) -> list[tuple[str, str, str]]: + """Flatten a normalized config into ordered ``(key, label, value)`` facts, for diffing. + + Each fact is one thing an admin would think of as a single setting — the default agent, an agent's + model, its settings scope, an MCP server, a skill, the tracing table, a budget tier. The ``key`` is + a stable identity so the same setting lines up across two configs even when values differ; the + ``label`` is what the admin reads. Deliberately mirrors what :func:`_render_summary` chooses to + show, so the diff and the summary never disagree about what's in a config. + """ + facts: list[tuple[str, str, str]] = [] + + default_agent = manifest.get("default_agent") + if isinstance(default_agent, str): + display = TOOL_SPECS.get(default_agent, {}).get("display", default_agent) + facts.append(("default_agent", "Default agent", display)) + + for tool, agent_config in (manifest.get("enabled_agents") or {}).items(): + display = TOOL_SPECS.get(tool, {}).get("display", tool) + model_config = agent_config.get("model_config") or {} + detail = model_config.get("default_model") or "no model" + provider = model_config.get("model_provider_service") + if provider: + detail = f"{detail} via {provider}" + facts.append((f"agent:{tool}", display, detail)) + models = model_config.get("models") + if isinstance(models, dict): + for slot, model in models.items(): + family = slot.removeprefix("default_").removesuffix("_model") + facts.append((f"agent:{tool}:model:{family}", f"{display} ({family})", str(model))) + elif isinstance(models, list) and len(models) > 1: + facts.append((f"agent:{tool}:models", f"{display} models", ", ".join(map(str, models)))) + if tool in GLOBAL_SETTINGS_AGENTS: + scope = ( + "global settings" if agent_config.get("use_as_global_settings") else "ucode-only" + ) + facts.append((f"agent:{tool}:scope", f"{display} settings", scope)) + + for server in manifest.get("mcp_servers") or []: + name = str(server.get("name")) + facts.append((f"mcp:{name}", f"MCP server {name}", str(server.get("type") or ""))) + + for skill in (manifest.get("skills") or {}).get("names") or []: + facts.append((f"skill:{skill}", f"Skill {skill}", "published")) + + tracing = manifest.get("tracing_table") + if tracing: + facts.append(("tracing_table", "Tracing table", str(tracing))) + + policy = manifest.get("budget_policy") + if isinstance(policy, dict): + facts.append( + ( + "budget:id", + "Budget", + policy.get("budget_display_name") or policy.get("budget_id", ""), + ) + ) + if policy.get("display_name"): + facts.append(("budget:name", "Policy name", str(policy["display_name"]))) + for tier in policy.get("tiers") or []: + agent = tier.get("default_agent") + agent_display = TOOL_SPECS.get(agent, {}).get("display", agent) + percent = float(tier.get("spending_percentage", 0)) * 100 + facts.append( + ( + f"budget:tier:{percent:g}", + f"Budget tier at {percent:g}%", + f"{agent_display} / {tier.get('default_model')}", + ) + ) + + return facts + + +def _render_config_diff(existing: dict | None, incoming: dict, workspace: str) -> bool: + """Show what publishing ``incoming`` changes versus the ``existing`` published config. + + Returns True when there is a difference. Lists only what changes — labelled ADD, DELETE, or + CHANGE (``old → new``) — since the full config was just printed by :func:`_render_summary` above; + repeating the unchanged rows here would bury the actual delta. Both configs are in ucode's + normalized shape (the caller round-trips the local manifest through serialize/normalize first), + so the comparison is field-for-field with what the workspace holds. + """ + old = {key: (label, value) for key, label, value in _config_facts(existing or {})} + new = {key: (label, value) for key, label, value in _config_facts(incoming)} + + # Fixed-width verbs so the labels line up in a column and the eye can scan one kind of change. + add = "[green]ADD [/green]" + delete = "[red]DELETE[/red]" + change = "[yellow]CHANGE[/yellow]" + + # Incoming order first (added/changed read top-down like the summary), then removed keys. + ordered = list(new) + [key for key in old if key not in new] + rows: list[str] = [] + for key in ordered: + if key in new and key not in old: + label, value = new[key] + rows.append(f" {add} {label}: {value}") + elif key in old and key not in new: + label, value = old[key] + rows.append(f" {delete} {label}: {value}") + elif old[key][1] != new[key][1]: + label, old_value = old[key] + rows.append(f" {change} {label}: {old_value} → {new[key][1]}") + + if not rows: + return False + print_heading(f"Changes to publish on {workspace}") + for row in rows: + console.print(row) + return True + + def _require_admin(workspace: str, token: str) -> None: """Stop unless the caller is a workspace admin. @@ -1041,24 +1262,25 @@ def _require_admin(workspace: str, token: str) -> None: print_success("Admin permissions verified") -def _handle_existing_config(workspace: str, token: str) -> bool: +def _handle_existing_config(workspace: str, token: str) -> tuple[bool, dict | None]: """Decide what to do when the workspace already has a published config. - Returns True to keep authoring a new config (the wizard continues; publishing later replaces the - existing one) and False to stop (no config exists, the check failed, or the admin chose to delete - the existing one instead of authoring a replacement). + Returns ``(keep_going, existing)``: ``keep_going`` is True to continue authoring (publishing later + replaces the existing config) and False to stop (the admin chose to delete it instead). ``existing`` + is the published config when one was read, so the caller can carry its MCP servers / skills / + tracing / budget policy forward — the local draft may be missing on a fresh machine or after + ``ucode revert``, and without this those sections would be silently dropped on the next publish. Deliberately doesn't itemize what the existing config holds. The admin doesn't need an inventory - to act on this — the instruction is the same either way ("include everything you want to keep") - — and `ucode setup show` prints the real thing for anyone who wants to compare. + to act on this, and `ucode setup show` prints the real thing for anyone who wants to compare. """ with spinner("Checking for an existing managed config..."): existing, reason = get_managed_config(workspace, token) if reason is not None: print_note(f"Could not check for an existing config: {reason}") - return True + return True, None if existing is None: - return True + return True, None print_warning( "This workspace already has a managed configuration — one config covers every agent, MCP " @@ -1074,11 +1296,12 @@ def _handle_existing_config(workspace: str, token: str) -> bool: if choice is None: raise KeyboardInterrupt if choice == "create": - print_note("Make sure this run includes everything you want to keep.") - return True + # The agent/model half is re-authored here; the other sections carry forward from `existing` + # (see `_carry_forward_sections`), so no need to warn the admin to re-enter them. + return True, existing _delete_existing_config(workspace, token, existing) - return False + return False, existing def _delete_existing_config(workspace: str, token: str, existing: dict) -> None: @@ -1145,14 +1368,140 @@ def setup_from_file(path: str) -> int: save_managed_state(workspace, manifest) _render_summary(workspace, manifest) print_success(f"Saved to {manifest_path.name} -> ~/.ucode/managed-state.json") - _print_next_steps() + _print_next_steps(manifest) return 0 -def _print_next_steps() -> None: +# 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 mcp", "MCP servers", lambda m: bool(m.get("mcp_servers"))), + ("ucode setup skills", "Skills", lambda m: bool((m.get("skills") or {}).get("names"))), + ( + "ucode setup budget-policy", + "Budget policy", + lambda m: isinstance(m.get("budget_policy"), dict), + ), +] + + +def _command_line(command: str, description: str, *, marker: str = " ", width: int = 0) -> str: + """A `` `` row, for command lists that read as a column.""" + return f" {marker} [bold]{command.ljust(width)}[/bold] {description}" + + +# `ucode setup` walks these phases in order; the banners announce each one so the admin can see how +# far along the flow they are, the way a multi-page form numbers its pages. +SETUP_STEP_TITLES = ["Coding agents", "Models & settings", "Default agent"] + + +def _step_banner(index: int, title: str) -> None: + """Announce one phase of `ucode setup` as `step N of M`.""" + print_section(f"ucode setup · step {index} of {len(SETUP_STEP_TITLES)} · {title}") + + +def _selected_recap(workspace: str, enabled_agents: dict, default_agent: str | None) -> None: + """A compact panel of what's chosen so far, reprinted as the flow advances. + + Turns the run of prompts into something that reads like a form filling in: each phase reprints + the growing set of decisions before asking the next question. Agents still mid-configuration show + a `…` placeholder for their model. + """ + lines = [kv_line("Workspace", workspace)] + for tool, config in enabled_agents.items(): + model = (config.get("model_config") or {}).get("default_model") or "…" + lines.append(kv_line(TOOL_SPECS.get(tool, {}).get("display", tool), model)) + if default_agent: + lines.append( + kv_line("Default", TOOL_SPECS.get(default_agent, {}).get("display", default_agent)) + ) + print_panel("Selected so far", lines) + + +def _section_status_lines(manifest: dict, width: int = 0) -> list[str]: + """One row per optional section: its command, what it covers, and whether it's configured.""" + width = width or max(len(command) for command, _, _ in SETUP_SECTIONS) + lines: list[str] = [] + for command, label, configured in SETUP_SECTIONS: + if configured(manifest): + marker, state = "[green]✔[/green]", "[green]configured[/green]" + else: + marker, state = "[dim]○[/dim]", "[dim]not configured[/dim]" + lines.append(_command_line(command, f"{label} — {state}", marker=marker, width=width)) + return lines + + +def _print_next_steps(manifest: dict) -> None: + """List the setup commands still worth running, then the publish step. + + Printed rather than prompted: each section is its own command now, so the admin drives the rest of + the setup themselves instead of being walked through a chain they mostly want to skip. Showing + what is already configured keeps a re-run from looking like it lost the other sections — it + didn't; `setup` carries them forward. + """ console.print() print_heading("Next steps") - print_note("Publish it to the workspace: ucode apply") + if config_io.is_dry_run(): + # Under --dry-run nothing was written, so the section commands (which read the saved draft) + # and `apply` have nothing to act on. Say so rather than send the admin to commands that + # would report "run `ucode setup` first". + print_note("Dry run — nothing was saved. Re-run without --dry-run to author the config.") + return + # These sections aren't required to publish — call them out as optional so an admin doesn't read + # a config with none configured as unfinished. + print_note("[dim]Optional — configure any of these, or skip straight to publishing:[/dim]") + for line in _section_status_lines(manifest): + console.print(line) + print_panel( + "All done?", + ["Publish with [bold]ucode apply[/bold] so all developers use this configuration."], + ) + + +# The sections `ucode setup` carries forward instead of prompting for, and how to rebuild each one. +CARRIED_SECTIONS: list[tuple[str, str, str]] = [ + ("mcp_servers", "MCP servers", "ucode setup mcp"), + ("skills", "Skills", "ucode setup skills"), + ("tracing_table", "Tracing table", "ucode setup --from-file"), + ("budget_policy", "Budget policy", "ucode setup budget-policy"), +] + + +def _carry_forward_sections(previous: dict, manifest: dict) -> None: + """Copy the sections `ucode setup` no longer prompts for out of a previously authored config. + + `setup` writes the whole manifest, so without this a re-run would silently clear the MCP servers, + skills, tracing table, and budget policy an admin authored with the other commands — they'd have + to redo every one of them just to change a model. + + Each section is probe-validated before it's carried, and dropped with a warning if it no longer + fits. Otherwise a carried section could make the manifest invalid and block the save outright, + with no way out: the commands that could repair a section read the very manifest that can't be + written. The live case is a budget-policy tier naming an agent the admin just de-selected, but + hand-edited drafts and configs authored by an older ucode can trip the others the same way. + """ + # Validating against no inventory keeps this to structural checks, which is all that's at stake + # here: the models were just picked from the workspace's own catalog a few prompts ago. + baseline = validate_manifest(manifest, None) + for key, label, rebuild in CARRIED_SECTIONS: + if key not in previous: + continue + candidate = previous[key] + new_errors = [ + error + for error in validate_manifest({**manifest, key: candidate}, None) + if error not in baseline + ] + if not new_errors: + manifest[key] = candidate + continue + print_warning( + f"{label} from the existing config no longer fits what you just picked, so it was left " + "out:" + ) + for error in new_errors: + print_note(error) + print_note(f"Rebuild it with `{rebuild}`.") def setup_command( @@ -1161,7 +1510,12 @@ def setup_command( workspace: str | None = None, profile: str | None = None, ) -> int: - """Author the workspace's managed coding-agent config interactively. + """Author the agents and models half of the workspace's managed coding config interactively. + + Agents and per-agent models only. MCP servers, skills, and the budget policy each have their own + command (`ucode setup mcp` / `skills` / `budget-policy`), so an admin changing one of them doesn't + have to walk the whole flow again — and this command carries whatever they already authored + forward untouched rather than clearing it (:func:`_carry_forward_sections`). ``workspace``/``profile`` let a caller that has already resolved (and authenticated against) a workspace hand it in so the admin isn't prompted to pick one again — e.g. `ucode configure` @@ -1178,7 +1532,7 @@ def setup_command( from ucode.cli import _prompt_for_configuration, configure_shared_state print_section("ucode setup") - print_note("Author the managed coding config for this workspace.") + print_note("Choose the coding agents and models for this workspace's managed config.") print_note("Developers pull it automatically when they run ucode.") if workspace is None: @@ -1190,7 +1544,8 @@ def setup_command( token = get_databricks_token(workspace, profile) _require_admin(workspace, token) - if not _handle_existing_config(workspace, token): + keep_going, published = _handle_existing_config(workspace, token) + if not keep_going: return 0 # Discover the workspace's models and gateway URLs. This also logs in and persists local state. @@ -1198,17 +1553,25 @@ def setup_command( workspace = state.get("workspace") or workspace profile = state.get("profile") or profile - available = [tool for tool in TOOL_SPECS if check_gateway_endpoint(state, tool)] + available = [ + tool + for tool in TOOL_SPECS + if tool not in SETUP_EXCLUDED_AGENTS and check_gateway_endpoint(state, tool) + ] if not available: raise RuntimeError( f"No coding agents are available on {workspace}. Check that the workspace's AI Gateway " "serves models for at least one agent." ) - previous = load_managed_state(workspace) or {} + # The local draft is the carry-forward source, falling back to what's published on the workspace: + # a fresh machine (or one after `ucode revert`) has no draft, and without the fallback the next + # publish would silently wipe the workspace's MCP servers, skills, tracing, and budget policy. + previous = load_managed_state(workspace) or published or {} previously_enabled = [ - tool for tool in (previous.get("enabled_agents") or {}) if tool in TOOL_SPECS + tool for tool in (previous.get("enabled_agents") or {}) if tool in available ] + _step_banner(1, SETUP_STEP_TITLES[0]) picked = prompt_for_tools( [(tool, TOOL_SPECS[tool]["display"]) for tool in available], preselected=previously_enabled or None, @@ -1217,20 +1580,10 @@ def setup_command( print_note("No coding agents selected — nothing to configure.") return 0 - default_agent = picked[0] - if len(picked) > 1: - chosen = prompt_for_selection( - "Which agent should launch when a developer runs `ucode`?", - [(tool, TOOL_SPECS[tool]["display"]) for tool in picked], - ) - if not chosen: - raise KeyboardInterrupt - default_agent = chosen - print_success(f"Default agent set to {TOOL_SPECS[default_agent]['display']}") - + _step_banner(2, SETUP_STEP_TITLES[1]) enabled_agents: dict[str, dict] = {} - for tool in picked: - print_heading(TOOL_SPECS[tool]["display"]) + for index, tool in enumerate(picked, start=1): + print_heading(f"{TOOL_SPECS[tool]['display']} ({index} of {len(picked)})") provider_service = _select_provider_service(tool, workspace, token) # Always set: `_prompt_models_for_agent` re-prompts rather than returning empty, so every # enabled agent carries a default_model and any of them can be the default_agent. @@ -1243,50 +1596,36 @@ def setup_command( if tool in GLOBAL_SETTINGS_AGENTS: binary = TOOL_SPECS[tool]["binary"] agent_config["use_as_global_settings"] = prompt_yes_no_default( - f"Write {TOOL_SPECS[tool]['display']}'s config to {GLOBAL_SETTINGS_FILES[tool]}? " - f"({GLOBAL_SETTINGS_BLURB.format(binary=binary)})", + f"Route `{binary}` through the gateway too, not just `ucode {binary}`? " + f"(writes {GLOBAL_SETTINGS_FILES[tool]}, needs sudo once)", default=False, ) enabled_agents[tool] = agent_config + _confirm_agent(tool, agent_config) + + # Pick the default after configuring each agent, not before: by now the admin has seen every + # agent's models go by, so "which is the default?" is a choice among things they've just set up + # rather than a bare list up front. The recap reprints those picks so the choice is informed. + _step_banner(3, SETUP_STEP_TITLES[2]) + default_agent = picked[0] + if len(picked) > 1: + _selected_recap(workspace, enabled_agents, default_agent=None) + chosen = prompt_for_selection( + "Which coding agent should be the default?", + [(tool, TOOL_SPECS[tool]["display"]) for tool in picked], + ) + if not chosen: + raise KeyboardInterrupt + default_agent = chosen + print_success(f"Default agent set to {TOOL_SPECS[default_agent]['display']}") manifest: dict = {"default_agent": default_agent, "enabled_agents": enabled_agents} # Tracing is intentionally not prompted here: the managed-tracing path isn't working yet, so # asking would author a `tracing_table` the workspace can't honor. The manifest field and its # serialize/validate support stay in place, so a hand-written `--from-file` config can still set - # it once the backend is ready. Re-add the section below when it is. - - print_section("MCP servers") - if prompt_yes_no_default("Set up managed MCP servers for this workspace?", default=False): - from ucode.mcp import configure_mcp_command - - # Managed configs can't carry a Databricks app (its host isn't reconstructable from the - # workspace), so hide apps from the picker rather than let an admin pick one that is then - # dropped from the published config. - configure_mcp_command(exclude_sources={"apps"}) - mcp_servers = _mcp_servers_from_state(load_state()) - if mcp_servers: - manifest["mcp_servers"] = mcp_servers - print_success(f"{len(mcp_servers)} MCP server(s) added to the managed config") - - print_section("Skills") - if prompt_yes_no_default("Set up managed skills for this workspace?", default=False): - locations = prompt_for_text( - "Skill schemas to publish, comma-separated `catalog.schema` (blank to skip)", - default="", - ) - parsed: list[str] = [item.strip() for item in (locations or "").split(",") if item.strip()] - if parsed: - from ucode.mcp import configure_skills_mcp_command - - configure_skills_mcp_command(parsed) - skill_names = _skill_names_from_state(load_state()) or parsed - manifest["skills"] = {"names": skill_names} - print_success(f"{len(skill_names)} skill schema(s) added to the managed config") - - budget_policy = _prompt_budget_policy(workspace, token, enabled_agents, state) - if budget_policy: - manifest["budget_policy"] = budget_policy + # it once the backend is ready. Re-add a `ucode setup tracing` command when it is. + _carry_forward_sections(previous, manifest) errors = validate_manifest(manifest, state) if errors: @@ -1301,7 +1640,240 @@ def setup_command( _render_summary(workspace, manifest) console.print() print_success("Saved to ~/.ucode/managed-state.json") - _print_next_steps() + _print_next_steps(manifest) + return 0 + + +def _resolve_admin_workspace() -> tuple[str, str | None, str]: + """Resolve the workspace a section command edits, authenticate, and gate on admin. + + Returns ``(workspace, profile, token)``. Unlike `ucode setup`, this doesn't prompt for a workspace + and takes it strictly from local state rather than falling back to the draft file's workspace: the + MCP and skills pickers re-read ``current_workspace`` themselves (via ``setup_mcp_clients``), so a + mismatch would have them operate against one workspace while the manifest is saved for another. + Requiring ``ucode configure`` to have set the current workspace keeps the two in lockstep. It also + skips :func:`_handle_existing_config` — the create-or-delete choice belongs to authoring a config, + not to changing one section of it. + """ + state = load_state() + workspace = state.get("workspace") + if not workspace: + raise RuntimeError( + "No workspace is configured. Run `ucode configure` first, then `ucode setup` to author " + "this workspace's managed config." + ) + profile = state.get("profile") + ensure_databricks_auth(workspace, profile) + token = get_databricks_token(workspace, profile) + _require_admin(workspace, token) + return workspace, profile, token + + +def _manifest_for_edit(workspace: str) -> dict: + """The authored manifest a section command edits. Raises when `ucode setup` hasn't run. + + An empty ``enabled_agents`` counts as "hasn't run": a launch records ``{}`` for a workspace with no + managed config (see ``refresh_managed_config``), so the file existing is not proof an admin + authored anything. Requiring agents first also keeps the budget-policy tiers honest — they can only + name agents the manifest enables. + """ + manifest = load_managed_state(workspace) + if not (manifest or {}).get("enabled_agents"): + raise RuntimeError( + f"No managed config has been authored for {workspace} yet. Run `ucode setup` first to " + "pick the agents and models, then re-run this command." + ) + return cast(dict, manifest) + + +def _save_section_update(workspace: str, manifest: dict) -> int: + """Validate the edited manifest structurally, save it, and show what's left to do. + + Validated with no model inventory (``state=None``), so only structure is checked here — not model + availability. That's deliberate: a section command doesn't touch agents or models, so re-checking + them would only reject a legitimately pinned older Claude model (`load_state` keeps just the newest + per family) or, worse, wrongly flag a codex/gemini model whenever the re-fetched inventory happens + to be Claude-only. `ucode apply` runs the full model check against the live catalog before + publishing, which is where it belongs. + """ + errors = validate_manifest(manifest, None) + if errors: + print_err("The updated config is not valid:") + for error in errors: + print_note(error) + return 1 + + save_managed_state(workspace, manifest) + _render_summary(workspace, manifest) + console.print() + print_success("Saved to ~/.ucode/managed-state.json") + _print_next_steps(manifest) + return 0 + + +def setup_mcp_command() -> int: + """Author the managed config's MCP servers (`ucode setup mcp`).""" + workspace, _, _ = _resolve_admin_workspace() + manifest = _manifest_for_edit(workspace) + + print_section("Managed MCP servers") + print_note("Developers get these MCP servers registered automatically when they run ucode.") + from ucode.mcp import configure_mcp_command + + # Snapshot the managed-shaped servers before the picker so a cancelled or no-op run leaves the + # section exactly as it was: the picker returns 0 on Esc, and re-reading local state would + # otherwise let whatever is registered locally overwrite the manifest — including deleting the + # section when nothing is registered. + before = _mcp_servers_from_state(load_state()) + # Managed configs can't carry a Databricks app (its host isn't reconstructable from the + # workspace), so hide apps from the picker rather than let an admin pick one that is then + # dropped from the published config. + configure_mcp_command(exclude_sources={"apps"}) + after = _mcp_servers_from_state(load_state()) + if after == before: + print_note("No changes to the MCP servers — the managed config is unchanged.") + return 0 + + if after: + manifest["mcp_servers"] = after + print_success(f"{len(after)} MCP server(s) in the managed config") + else: + # Deregistering every server locally is how an admin clears the section — there is no + # separate "remove them all" flag. + manifest.pop("mcp_servers", None) + print_note("No MCP servers are registered, so the managed config now carries none.") + return _save_section_update(workspace, manifest) + + +def setup_skills_command(locations: list[str] | None = None) -> int: + """Author the managed config's skills (`ucode setup skills`). + + ``locations`` comes from ``--location`` (already parsed to `.` refs); when None + the admin is prompted and the answer is parsed the same way. + """ + workspace, _, _ = _resolve_admin_workspace() + manifest = _manifest_for_edit(workspace) + + print_section("Managed skills") + print_note("Developers get these skills downloaded automatically when they run ucode.") + if locations is None: + answer = prompt_for_text( + "Skill schemas to publish, comma-separated `catalog.schema` (blank to leave unchanged)", + default="", + ) + # Route the interactive answer through the same parser as `--location` so `main` (missing the + # schema) is rejected here rather than published as a bogus skill name. + from ucode.cli import _parse_skill_locations + + locations = _parse_skill_locations(answer) + # A blank answer / empty `--location` means "leave the skills alone". Returning before delegating + # matters: `configure_skills_mcp_command([])` is not a no-op — it registers the schema-less skills + # MCP connection into the admin's own agents. + if not locations: + print_note("No skill schemas given — the managed config's skills are unchanged.") + return 0 + + from ucode.mcp import configure_skills_mcp_command + + configure_skills_mcp_command(locations) + skill_names = _skill_names_from_state(load_state()) or locations + manifest["skills"] = {"names": skill_names} + print_success(f"{len(skill_names)} skill schema(s) in the managed config") + return _save_section_update(workspace, manifest) + + +def setup_budget_policy_command() -> int: + """Author the managed config's spend-routing budget policy (`ucode setup budget-policy`).""" + workspace, _, token = _resolve_admin_workspace() + manifest = _manifest_for_edit(workspace) + + # The manifest's own default, shown as the "under the first rung" row of the fallback ladder so + # the admin sees what developers get before any tier kicks in. + default_agent = manifest.get("default_agent") + default_config = (manifest.get("enabled_agents") or {}).get(default_agent) or {} + default_model = (default_config.get("model_config") or {}).get("default_model") + base_default = ( + f"{TOOL_SPECS.get(default_agent, {}).get('display', default_agent)} / {default_model}" + if default_agent and default_model + else "" + ) + + # `_prompt_budget_policy` returns None on an environmental dead end too (no budgets, no budget with + # a per-user block, or the admin backing out of a picker), not only on an explicit decline. Leave + # any existing policy untouched in every one of those cases — never pop it — so a transient budget + # listing failure can't silently delete a policy the admin already published. + policy = _prompt_budget_policy( + workspace, token, manifest["enabled_agents"], load_state(), base_default=base_default + ) + if not policy: + print_note("The managed config's budget policy is unchanged.") + return 0 + manifest["budget_policy"] = policy + return _save_section_update(workspace, manifest) + + +def setup_help_command() -> int: + """Walk through the whole managed-config setup, marking what this machine has authored. + + Hand-written rather than left to `--help`: the point is the *order* of the commands and the fact + that nothing reaches developers until `ucode apply`, neither of which a flag listing conveys. Reads + the manifest but never authenticates, so it works before `ucode configure`. + """ + print_section("ucode setup") + print_note( + "A managed config is the coding setup your developers pull automatically — they run ucode " + "and get the agents, models, MCP servers, and skills you chose here. Admins only." + ) + print_note( + "Each command below edits your local draft; nothing reaches the workspace until " + "`ucode apply`." + ) + + workspace = load_state().get("workspace") or managed_state_workspace() + manifest = load_managed_state(workspace) or {} + agents_done = bool(manifest.get("enabled_agents")) + # One column width across all three groups, so the commands line up as a single list. + width = max(len(command) for command, _, _ in SETUP_SECTIONS) + width = max(width, len("ucode setup --from-file ")) + + print_heading("1. Start here") + console.print( + _command_line( + "ucode setup", + "Agents and models — " + + ("[green]configured[/green]" if agents_done else "[yellow]not configured[/yellow]"), + marker="[green]✔[/green]" if agents_done else "[yellow]○[/yellow]", + width=width, + ) + ) + if not agents_done: + print_note("The commands below edit that config, so they need this one to have run.") + + print_heading("2. Then any of these, in any order") + for line in _section_status_lines(manifest, width): + console.print(line) + + print_heading("3. Review and publish") + console.print( + _command_line("ucode setup show", "The draft, and the payload `apply` sends", width=width) + ) + console.print(_command_line("ucode apply", "Publish it to the workspace", width=width)) + + print_heading("Also") + console.print( + _command_line( + "ucode setup --from-file ", + "Load a hand-written manifest instead of prompting", + width=width, + ) + ) + print_note( + f"The draft lives in ~/.ucode/managed-state.json (workspace: {workspace or 'none'})." + ) + print_note( + "Re-running `ucode setup` keeps the sections in step 2; to drop one, edit the draft and " + "reload it with `ucode setup --from-file`." + ) return 0 @@ -1446,11 +2018,17 @@ def apply_command(*, yes: bool = False) -> int: if existing is None: print_note(f"This will create a new managed config on {workspace}.") else: - agents = ", ".join((existing.get("enabled_agents") or {}).keys()) or "no agents" - print_warning( - f"This will replace the config already published on {workspace} (currently: {agents}). " - "Every developer picks the new one up on their next ucode run." - ) + # Diff against what's live, normalized the same way, so the admin sees exactly what changes + # rather than a bare "this replaces the current config". Comparing the round-tripped payload + # (not the raw manifest) shows the real post-publish state — any field serialization drops + # won't appear as a phantom change. + changed = _render_config_diff(existing, normalize_managed_config(payload), workspace) + if not changed: + print_success(f"{workspace}'s published config already matches this one.") + print_note("Nothing to publish.") + return 0 + console.print() + print_warning("This takes effect for every developer on their next `ucode` run.") if not yes and not prompt_yes_no_default("Publish this config?", default=False): print_note("Nothing was published.") return 1 @@ -1468,8 +2046,17 @@ def apply_command(*, yes: bool = False) -> int: name = (published or {}).get("name") or existing_name or "coding-agent-configs/?" print_success(f"Published {name} to {workspace}") - print_note("Developers pick this up on their next ucode run.") + print_note("Developers get it automatically the next time they run `ucode`.") return 0 -__all__ = ["apply_command", "setup_command", "setup_from_file", "show_command"] +__all__ = [ + "apply_command", + "setup_budget_policy_command", + "setup_command", + "setup_from_file", + "setup_help_command", + "setup_mcp_command", + "setup_skills_command", + "show_command", +] diff --git a/src/ucode/ui.py b/src/ucode/ui.py index f338e8a..e53761a 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -92,6 +92,19 @@ def print_panel(title: str, lines: list[str]) -> None: console.print(Panel("\n".join(lines), title=title, style="blue", expand=False)) +def print_warning_panel(message: str, *, title: str = "Warning") -> None: + """Render a warning as a boxed panel, for a blocker that should stand out from inline notes. + + A `!` marker keeps it visually of a kind with :func:`print_warning`, but the box gives a + dead-end message (e.g. "no budgets exist, nothing to do") the weight to be read before the flow + exits, rather than scrolling past as one more line. + """ + console.print() + console.print( + Panel(f"[bold yellow]![/bold yellow] {message}", title=title, style="yellow", expand=False) + ) + + def print_note(text: str) -> None: console.print(f"[dim]•[/dim] {text}") diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 64b4dcc..a7a060d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1134,9 +1134,9 @@ def test_returns_newest_flash_first(self, monkeypatch): assert reason is None assert models[0] == "databricks-gemini-3-5-flash" - def test_codex_discovery_keeps_alphabetical_order(self, monkeypatch): - # Codex passes no sort_key, so ordering must stay the plain alphabetical - # default — guarding against the gemini change leaking across tools. + def test_codex_discovery_orders_newest_version_first(self, monkeypatch): + # Codex orders newest model version first (like gemini), so the picker's top choice and + # default is the newest gpt, not the alphabetically-first one. payload = { "endpoints": [ { @@ -1152,7 +1152,7 @@ def test_codex_discovery_keeps_alphabetical_order(self, monkeypatch): ] }, } - for name in ["databricks-gpt-5-2-codex", "databricks-gpt-4-1"] + for name in ["databricks-gpt-4-1", "databricks-gpt-5-2-codex"] ] } monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) @@ -1160,7 +1160,7 @@ def test_codex_discovery_keeps_alphabetical_order(self, monkeypatch): models, reason = db_mod.discover_codex_models(WS, "token") assert reason is None - assert models == ["databricks-gpt-4-1", "databricks-gpt-5-2-codex"] + assert models == ["databricks-gpt-5-2-codex", "databricks-gpt-4-1"] class TestResolvePatToken: diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 4872b79..9cf96c7 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -20,7 +20,7 @@ import ucode.managed_config as managed_config_mod import ucode.managed_wizard as wizard from ucode.cli import app -from ucode.managed_setup import validate_manifest +from ucode.managed_setup import serialize_managed_config, validate_manifest runner = CliRunner() @@ -231,7 +231,8 @@ def test_continue_when_no_config_exists(self): patch.object(wizard, "prompt_for_selection") as select, patch.object(wizard, "print_warning") as warn, ): - assert wizard._handle_existing_config(WORKSPACE, "token") is True + # No published config, so nothing to carry forward. + assert wizard._handle_existing_config(WORKSPACE, "token") == (True, None) assert not select.called assert not warn.called @@ -242,7 +243,7 @@ def test_read_failure_continues_with_a_note(self): patch.object(wizard, "prompt_for_selection") as select, patch.object(wizard, "print_note") as note, ): - assert wizard._handle_existing_config(WORKSPACE, "token") is True + assert wizard._handle_existing_config(WORKSPACE, "token") == (True, None) assert not select.called assert note.called @@ -255,7 +256,11 @@ def test_choosing_create_continues_authoring(self): ), patch.object(wizard, "prompt_for_selection", return_value="create"), ): - assert wizard._handle_existing_config(WORKSPACE, "token") is True + # Continues authoring, and hands back the published config to carry its sections forward. + assert wizard._handle_existing_config(WORKSPACE, "token") == ( + True, + {"name": "x", "enabled_agents": {}}, + ) def test_warning_does_not_itemize_the_existing_config(self): # The warning is the same whatever the config holds: an inventory doesn't change what the @@ -282,7 +287,8 @@ def test_choosing_delete_stops_and_deletes(self): patch.object(wizard, "prompt_yes_no_default", return_value=True), patch.object(wizard, "delete_coding_agent_config", return_value=None) as delete, ): - assert wizard._handle_existing_config(WORKSPACE, "token") is False + keep_going, _ = wizard._handle_existing_config(WORKSPACE, "token") + assert keep_going is False delete.assert_called_once_with(WORKSPACE, "token", "cfg/1") def test_delete_declined_leaves_config_intact(self): @@ -297,7 +303,8 @@ def test_delete_declined_leaves_config_intact(self): patch.object(wizard, "delete_coding_agent_config") as delete, ): # Still stops the wizard: the admin chose the delete path, not the author path. - assert wizard._handle_existing_config(WORKSPACE, "token") is False + keep_going, _ = wizard._handle_existing_config(WORKSPACE, "token") + assert keep_going is False assert not delete.called def test_delete_failure_raises(self): @@ -453,7 +460,8 @@ def test_claude_falls_back_to_text_when_nothing_discovered(self): def test_single_slot_announces_the_inferred_default(self): # The one-option prompt is skipped, but silence reads as a dropped step — the admin has to - # learn that the default is set, and to what. + # learn that the default was inferred. Announced as a note; the per-agent ✔ (`_confirm_agent`) + # in the setup loop carries the final confirmation. candidates = {"opus": ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"]} def fake_sel(prompt, options, **kwargs): @@ -464,14 +472,13 @@ def fake_sel(prompt, options, **kwargs): with ( patch.object(wizard, "_claude_candidates", return_value=candidates), patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success") as success, + patch.object(wizard, "print_note") as note, ): config = wizard._prompt_models_for_agent("claude", STATE, None) assert config["default_model"] == "system.ai.claude-opus-4-8" - assert success.called - assert "system.ai.claude-opus-4-8" in success.call_args[0][0] + notes = " ".join(str(call.args[0]) for call in note.call_args_list) + assert "overall default" in notes def test_claude_all_families_skipped_still_picks_from_the_candidates(self): # Skipping every slot is a legitimate minimal config — `models` is optional and each unset @@ -1346,30 +1353,40 @@ def test_malformed_targets_yield_nothing(self): class TestBudgetPolicy: - def test_declining_yields_none(self): - with patch.object(wizard, "prompt_yes_no_default", return_value=False): - assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None + def test_no_up_front_gate(self): + # Running `ucode setup budget-policy` is the consent, so the flow asks no "set up a policy?" + # question — it goes straight to listing budgets. (The only yes/no it asks is "add another + # tier?", after a tier is built.) + with ( + patch.object(wizard, "prompt_yes_no_default") as ask, + patch.object(wizard, "list_workspace_budgets", return_value=([], "none found")), + patch.object(wizard, "print_warning_panel"), + ): + wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) + assert not ask.called - def test_no_budgets_warns_and_yields_none(self): + def test_no_budgets_warns_in_a_box_and_skips_the_blurb(self): + # No attachable budget: show the dead-end warning as a box and don't explain a feature the + # workspace can't use yet. with ( - patch.object(wizard, "prompt_yes_no_default", return_value=True), patch.object(wizard, "list_workspace_budgets", return_value=([], "none found")), - patch.object(wizard, "print_warning") as warn, + patch.object(wizard, "print_warning_panel") as warn_box, + patch.object(wizard, "print_note") as note, ): assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None - assert warn.called + assert warn_box.called + assert not note.called # the BUDGET_POLICY_BLURB note is skipped 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)), - patch.object(wizard, "print_warning") as warn, + patch.object(wizard, "print_warning_panel") as warn_box, ): assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None - assert warn.called + assert warn_box.called 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 @@ -1379,7 +1396,7 @@ def test_only_per_user_block_budgets_are_offered(self): {"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, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1398,7 +1415,7 @@ def test_only_per_user_block_budgets_are_offered(self): def test_percentages_are_stored_as_fractions(self): 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, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1434,7 +1451,7 @@ def test_shows_per_user_threshold_and_tier_dollars(self): } ] with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1454,7 +1471,7 @@ def test_missing_threshold_skips_the_dollar_hints(self): # A budget whose threshold couldn't be read still works; the prompt just omits the dollars. 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, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1483,7 +1500,7 @@ def test_offers_only_the_models_the_agent_was_configured_with(self): } 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, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1512,7 +1529,7 @@ def test_claude_family_slots_are_flattened_for_the_picker(self): } 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, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1531,7 +1548,7 @@ def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): # catalog than nothing at all. 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, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1548,7 +1565,7 @@ def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): def test_authored_policy_validates(self): 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, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1568,7 +1585,8 @@ def test_authored_policy_validates(self): def test_a_repeated_agent_model_pair_is_rejected_and_re_prompted(self): # The highest crossed tier wins, so a second tier on the same agent+model is inert. The loop - # must reject the repeat and re-prompt, the way it already does for a repeated percentage. + # rejects the repeat and re-asks only the agent/model — the percentage already entered for + # this tier is kept, not re-prompted. two_models = { "claude": { "model_config": { @@ -1585,7 +1603,7 @@ def test_a_repeated_agent_model_pair_is_rejected_and_re_prompted(self): # 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, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, @@ -1595,7 +1613,7 @@ def test_a_repeated_agent_model_pair_is_rejected_and_re_prompted(self): # Tier 1: claude / opus. "claude", "system.ai.claude-opus-4-8", - # Tier 2 first attempt: claude / opus again — rejected, so the loop re-asks. + # Tier 2 first attempt: claude / opus again — rejected, so just agent/model re-ask. "claude", "system.ai.claude-opus-4-8", # Tier 2 retry: a genuine step-down. @@ -1604,7 +1622,9 @@ def test_a_repeated_agent_model_pair_is_rejected_and_re_prompted(self): ], ), patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", side_effect=[0.5, 0.9, 0.9]), + # Percentage asked once per tier: 0.5 for tier 1, 0.9 for tier 2. The combo retry does + # not re-ask it. + patch.object(wizard, "prompt_for_percentage", side_effect=[0.5, 0.9]), patch.object(wizard, "print_err") as err, ): policy = wizard._prompt_budget_policy(WORKSPACE, "token", two_models, STATE) @@ -1612,7 +1632,7 @@ def test_a_repeated_agent_model_pair_is_rejected_and_re_prompted(self): ("claude", "system.ai.claude-opus-4-8"), ("claude", "system.ai.claude-sonnet-4-6"), ] - assert any("no-op" in call.args[0] for call in err.call_args_list) + assert any("do nothing" in call.args[0] for call in err.call_args_list) class TestConfiguredModelsForAgent: @@ -1973,7 +1993,7 @@ def fake_sel(prompt, options, **kwargs): return "system.ai.claude-opus-4-8" with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), patch.object(wizard, "prompt_for_text", return_value="tiered"), @@ -1985,6 +2005,265 @@ def fake_sel(prompt, options, **kwargs): assert any("model" in p for p in searchable_prompts), searchable_prompts +# A minimal authored manifest (agents + models only), the shape `ucode setup` now writes. +AGENTS_ONLY = { + "default_agent": "claude", + "enabled_agents": {"claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}}, +} + + +class TestCarryForwardSections: + def test_all_optional_sections_survive_a_rerun(self): + previous = { + **AGENTS_ONLY, + "mcp_servers": [{"name": "system.ai.github", "type": "mcp-service"}], + "skills": {"names": ["main.default"]}, + "tracing_table": "main.default.traces", + "budget_policy": { + "budget_id": BUDGET_ID, + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-8", + } + ], + }, + } + manifest = dict(AGENTS_ONLY) + wizard._carry_forward_sections(previous, manifest) + assert manifest["mcp_servers"] == previous["mcp_servers"] + assert manifest["skills"] == previous["skills"] + assert manifest["tracing_table"] == previous["tracing_table"] + assert manifest["budget_policy"] == previous["budget_policy"] + + def test_empty_previous_adds_nothing(self): + manifest = dict(AGENTS_ONLY) + wizard._carry_forward_sections({}, manifest) + assert set(manifest) == set(AGENTS_ONLY) + + def test_budget_policy_naming_a_dropped_agent_is_left_out_with_a_warning(self): + # The admin re-ran `setup` and de-selected codex; the saved policy still routes to it, which + # would fail `validate_manifest` and block the whole save. Drop just the policy, loudly. + previous = { + **AGENTS_ONLY, + "budget_policy": { + "budget_id": BUDGET_ID, + "tiers": [ + { + "spending_percentage": 0.9, + "default_agent": "codex", + "default_model": "system.ai.gpt-5", + } + ], + }, + } + manifest = dict(AGENTS_ONLY) + with patch.object(wizard, "print_warning") as warn: + wizard._carry_forward_sections(previous, manifest) + assert "budget_policy" not in manifest + assert warn.called + # The rest of the manifest is untouched and still valid. + assert validate_manifest(manifest, None) == [] + + +class TestNextSteps: + 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 mcp" in out + assert "ucode setup skills" in out + assert "ucode setup budget-policy" in out + assert "ucode apply" in out + + def test_dry_run_says_nothing_was_saved(self, capsys, monkeypatch): + monkeypatch.setattr(config_io_mod, "_dry_run", True) + wizard._print_next_steps(AGENTS_ONLY) + out = capsys.readouterr().out + assert "Dry run" in out + assert "ucode apply" not in out + + +class TestSectionCommands: + """The `ucode setup mcp` / `skills` / `budget-policy` section commands.""" + + @staticmethod + def _admin(**overrides): + """Patch the auth/admin boundary the section commands resolve through.""" + defaults = { + "load_state": lambda: {"workspace": WORKSPACE, "profile": "p", **STATE}, + "ensure_databricks_auth": lambda *a, **k: None, + "get_databricks_token": lambda *a, **k: "tok", + "is_workspace_admin": lambda *a, **k: True, + } + defaults.update(overrides) + return [patch.object(wizard, name, value) for name, value in defaults.items()] + + def _run(self, fn, *, admin_overrides=None, **patches): + import contextlib + + with contextlib.ExitStack() as stack: + for p in self._admin(**(admin_overrides or {})): + stack.enter_context(p) + for name, value in patches.items(): + stack.enter_context(patch.object(wizard, name, value)) + return fn() + + def test_mcp_requires_an_authored_config(self): + # No manifest on disk → the command can't edit a section that doesn't exist. + with pytest.raises(RuntimeError, match="ucode setup"): + self._run(wizard.setup_mcp_command) + + def test_mcp_requires_enabled_agents(self): + # A launch stores `{}` to mean "no managed config"; that must not count as authored. + managed_config_mod.save_managed_state(WORKSPACE, {}) + with pytest.raises(RuntimeError, match="ucode setup"): + self._run(wizard.setup_mcp_command) + + def test_mcp_writes_only_its_section(self): + managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + servers = [{"name": "system.ai.github", "type": "mcp-service"}] + # The picker (imported lazily inside the command) registers servers into local state; fake + # that by having the before/after reads bracket a change. + reads = iter([[], servers]) + with ( + patch("ucode.mcp.configure_mcp_command", return_value=0) as picker, + patch.object(wizard, "_mcp_servers_from_state", side_effect=lambda *_: next(reads)), + ): + code = self._run(wizard.setup_mcp_command) + assert code == 0 + assert picker.call_args.kwargs == {"exclude_sources": {"apps"}} + saved = managed_config_mod.load_managed_state(WORKSPACE) + assert saved["mcp_servers"] == servers + assert saved["enabled_agents"] == AGENTS_ONLY["enabled_agents"] + + def test_mcp_cancel_is_a_no_op(self): + # Picker cancelled / nothing changed → the section is left exactly as it was. + managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + with ( + patch("ucode.mcp.configure_mcp_command", return_value=0), + patch.object(wizard, "_mcp_servers_from_state", return_value=[]), + patch.object(wizard, "save_managed_state") as save, + ): + code = self._run(wizard.setup_mcp_command) + assert code == 0 + assert not save.called + + def test_mcp_not_admin_raises(self): + managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + with pytest.raises(RuntimeError, match="not an admin"): + self._run( + wizard.setup_mcp_command, + admin_overrides={"is_workspace_admin": lambda *a, **k: False}, + ) + + def test_skills_location_bypasses_the_prompt(self): + managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + with ( + patch("ucode.mcp.configure_skills_mcp_command", return_value=0) as configure, + patch.object(wizard, "_skill_names_from_state", return_value=["main.default"]), + patch.object(wizard, "prompt_for_text") as prompt, + ): + code = self._run(lambda: wizard.setup_skills_command(["main.default"])) + assert code == 0 + assert not prompt.called + configure.assert_called_once_with(["main.default"]) + assert managed_config_mod.load_managed_state(WORKSPACE)["skills"] == { + "names": ["main.default"] + } + + def test_skills_blank_answer_writes_nothing(self): + # A blank answer must not delegate: `configure_skills_mcp_command([])` is not a no-op. + managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + with ( + patch("ucode.mcp.configure_skills_mcp_command") as configure, + patch.object(wizard, "prompt_for_text", return_value=""), + patch.object(wizard, "save_managed_state") as save, + ): + code = self._run(wizard.setup_skills_command) + assert code == 0 + assert not configure.called + assert not save.called + + def test_budget_policy_offers_only_the_manifests_agents(self): + managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + captured = {} + + def fake_prompt(workspace, token, enabled_agents, state, **kwargs): + captured["agents"] = enabled_agents + return None # decline / dead end → leave the policy unchanged + + with patch.object(wizard, "_prompt_budget_policy", side_effect=fake_prompt): + code = self._run(wizard.setup_budget_policy_command) + assert code == 0 + assert captured["agents"] == AGENTS_ONLY["enabled_agents"] + + def test_budget_policy_none_leaves_existing_untouched(self): + # A transient budget-listing failure returns None; it must never delete a saved policy. + seeded = { + **AGENTS_ONLY, + "budget_policy": { + "budget_id": BUDGET_ID, + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-8", + } + ], + }, + } + managed_config_mod.save_managed_state(WORKSPACE, seeded) + with patch.object(wizard, "_prompt_budget_policy", return_value=None): + code = self._run(wizard.setup_budget_policy_command) + assert code == 0 + assert ( + managed_config_mod.load_managed_state(WORKSPACE)["budget_policy"] + == seeded["budget_policy"] + ) + + +class TestSetupHelp: + def test_lists_every_setup_command(self, capsys): + wizard.setup_help_command() + out = capsys.readouterr().out + for command in ( + "ucode setup", + "ucode setup mcp", + "ucode setup skills", + "ucode setup budget-policy", + "ucode setup show", + "ucode apply", + ): + assert command in out + + +class TestApplyDiff: + def test_lists_added_removed_and_changed(self, capsys): + existing = { + "name": "cfg/1", + **AGENTS_ONLY, + "mcp_servers": [{"name": "system.ai.slack", "type": "mcp-service"}], + } + incoming = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-9"}} + }, + "mcp_servers": [{"name": "system.ai.github", "type": "mcp-service"}], + } + changed = wizard._render_config_diff(existing, incoming, WORKSPACE) + out = capsys.readouterr().out + assert changed is True + assert "CHANGE" in out and "claude-opus-4-8" in out and "claude-opus-4-9" in out + assert "ADD" in out and "system.ai.github" in out # added server + assert "DELETE" in out and "system.ai.slack" in out # removed server + + def test_identical_configs_report_no_change(self, capsys): + assert wizard._render_config_diff(AGENTS_ONLY, AGENTS_ONLY, WORKSPACE) is False + + class TestApplyCommand: MANIFEST = { "default_agent": "claude", @@ -2068,6 +2347,30 @@ def fake_create(*a, **k): assert updated["name"] == "coding-agent-configs/abc" assert created["called"] is False + def test_no_publish_when_the_published_config_already_matches(self): + # Publishing a config identical to what's live is a no-op; say so and skip the write rather + # than PATCH the same bytes back. + managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + # What's live is the manifest normalized the same way `apply` will send it. + existing = { + "name": "coding-agent-configs/abc", + **managed_config_mod.normalize_managed_config(serialize_managed_config(self.MANIFEST)), + } + updated = {"called": False} + + def fake_update(*a, **k): + updated["called"] = True + return {}, None + + assert ( + self._run( + get_managed_config=lambda *a, **k: (existing, None), + update_coding_agent_config=fake_update, + ) + == 0 + ) + assert updated["called"] is False + def test_invalid_manifest_is_not_published(self): # `default_agent` names an agent that isn't enabled. managed_config_mod.save_managed_state( @@ -2327,6 +2630,68 @@ def test_show_exits_zero(self): result = runner.invoke(app, ["setup", "show"]) assert result.exit_code == 0 + @pytest.mark.parametrize( + ("command", "target"), + [ + ("mcp", "setup_mcp_command"), + ("skills", "setup_skills_command"), + ("budget-policy", "setup_budget_policy_command"), + ], + ) + def test_section_subcommands_are_registered_and_called(self, command, target): + with ( + patch("ucode.cli.install_databricks_cli"), + patch(f"ucode.cli.{target}", return_value=0) as fn, + ): + result = runner.invoke(app, ["setup", command]) + assert result.exit_code == 0 + assert fn.called + assert "ERROR" not in _out(result) + + def test_setup_skills_declares_location(self): + group = typer.main.get_command(app).commands["setup"] # type: ignore[attr-defined] + skills = group.commands["skills"] # type: ignore[attr-defined] + declared = {opt for param in skills.params for opt in param.opts} + assert "--location" in declared + + def test_setup_skills_location_is_parsed_to_a_list(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_skills_command", return_value=0) as fn, + ): + runner.invoke(app, ["setup", "skills", "--location", "main.a,main.b"]) + assert fn.call_args.args[0] == ["main.a", "main.b"] + + def test_setup_help_needs_no_auth(self): + # `ucode setup help` reads the local draft only — it must not shell out to install the CLI. + with ( + patch("ucode.cli.install_databricks_cli") as install, + patch("ucode.cli.setup_help_command", return_value=0) as fn, + ): + result = runner.invoke(app, ["setup", "help"]) + assert result.exit_code == 0 + assert fn.called + assert not install.called + + def test_section_command_runtime_error_exits_1(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch( + "ucode.cli.setup_mcp_command", side_effect=RuntimeError("run `ucode setup` first") + ), + ): + result = runner.invoke(app, ["setup", "mcp"]) + assert result.exit_code == 1 + assert "ucode setup" in _out(result) + + def test_section_command_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, ["setup", "budget-policy"]) + assert result.exit_code == 130 + def _out(result) -> str: """CliRunner output with stderr folded in, since print_err writes to a stderr console."""