Split ucode setup into per-section commands + apply diff + wizard UX

README.md

@@ -201,24 +201,35 @@ you to run `ucode <agent>` (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.
@@ -231,17 +242,19 @@ 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).
 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.
 
 ---
 
@@ -271,10 +284,14 @@ pick the new config up on their next ucode run.
 | `ucode configure skills --location main.default [--path <dir>]` | Download a schema's skills to disk (under `<dir>`, 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 <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

src/ucode/cli.py

@@ -81,7 +81,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,
@@ -265,8 +273,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
@@ -716,7 +725,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"),
@@ -1047,7 +1056,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`.",
 )
 
 
@@ -2661,7 +2672,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
@@ -2679,6 +2693,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 `<catalog>.<schema>` (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."""

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]: