Role-aware managed-config paths in ucode configure (AIGTWY-4338)

src/ucode/cli.py

@@ -253,29 +253,45 @@ def _print_managed_summary_abridged(managed: dict, state: dict, tool: str | None
     )
 
 
+def _confirm_managed_config_in_force(managed: dict, workspace: str) -> None:
+    """Confirm the workspace's managed config is already in force, and show what it holds.
+
+    There is nothing to write locally: the launch path applies the config (agents, MCP servers, and
+    skills) on every ``ucode`` run. So both the non-admin ``configure`` path and setup's "Adopt"
+    choice land here — show the admin's config and point the user at ``ucode``.
+    """
+    print_success("A managed config is published for your workspace — you're all set.")
+    _print_managed_summary(managed, {"workspace": workspace}, tool=None)
+    print_note("Configuration is complete. Just run `ucode` to launch with it applied.")
+
+
 def _resolve_workspace_then_maybe_reject(
     workspace_entries: list[tuple[str, str | None]] | None,
 ) -> list[tuple[str, str | None]] | None:
-    """Resolve the workspace ``ucode configure`` targets, then short-circuit if it is managed.
+    """Resolve the workspace ``ucode configure`` targets, then branch on role + managed config.
 
     When managed coding-agent configs are enabled, ``ucode configure`` must still let a developer
     switch workspaces — so resolve the target workspace up front (prompting when the interactive
-    path gave no ``--workspaces``/``--profiles``) and make it current *before* deciding whether to
-    short-circuit. Only then, if that workspace already publishes a managed config, configuring
-    locally would be overridden at launch anyway: show the admin's config and point the developer
-    at `ucode`.
-
-    When there is no managed config the developer's own ``configure`` always proceeds — an admin
-    just sees an FYI that they could publish one with ``ucode setup`` (never a prompt, never a
-    diversion). Returns the resolved entries to configure so the caller reuses them instead of
-    prompting again. Without the feature enabled it returns ``workspace_entries`` unchanged and
-    prompts nothing.
+    path gave no ``--workspaces``/``--profiles``) and make it current *before* deciding what to do.
+    Then, gated entirely by ``ENABLE_MANAGED_AGENT_CONFIG``, the four role/config paths are:
+
+    * **No managed config** → a workspace admin is dropped straight into the ``ucode setup``
+      authoring flow (``configure`` is replacing ``setup``) and the command exits with its code; a
+      non-admin's own ``configure`` proceeds, with the resolved entries returned so the caller
+      reuses them instead of re-prompting.
+    * **Managed config, non-admin** (or admin status unverifiable) → they're already set: the
+      launch path applies the config on every ``ucode`` run, so just show it and point them there.
+    * **Managed config, admin** → drop into the setup flow, whose existing-config menu lets them
+      adopt it (the same "you're all set" confirmation), re-author it, or delete it; the command exits.
+
+    Without the feature enabled it returns ``workspace_entries`` unchanged and prompts nothing.
     """
     if not managed_agent_config_enabled():
         return workspace_entries
     entries = workspace_entries or [_prompt_for_configuration(None)]
     workspace, profile = entries[0]
     set_current_workspace(workspace)
+    ensure_databricks_auth(workspace, profile)
     # Fetch, don't just read the local cache: on a fresh machine (or right after a reinstall) the
     # cache is empty until the first launch, so a cache read would miss a config the workspace does
     # publish and wrongly fall through to the local configure flow. `refresh_managed_config` reaches
@@ -284,24 +300,33 @@ def _resolve_workspace_then_maybe_reject(
         managed, coding_agent_config_feature_disabled = refresh_managed_config(
             {"workspace": workspace, "profile": profile}
         )
-    if not managed and not coding_agent_config_feature_disabled:
-        _maybe_offer_admin_setup(workspace, profile)
     if not managed:
+        if not coding_agent_config_feature_disabled:
+            _maybe_run_admin_setup(workspace, profile)
         return entries
-    print_success("A managed config has been detected for your workspace — you're all set.")
-    _print_managed_summary(managed, load_state(), tool=None)
-    print_note("Configuration is complete. Just run `ucode` to launch with it applied.")
+    is_admin: bool | None = None
+    try:
+        token = get_databricks_token(workspace, profile)
+    except RuntimeError:
+        token = None
+    if token is not None:
+        with spinner("Checking your workspace permissions..."):
+            is_admin = is_workspace_admin(workspace, token)
+    if is_admin:
+        _run_setup_and_exit(workspace, profile, token)
+    _confirm_managed_config_in_force(managed, workspace)
     raise typer.Exit(0)
 
 
-def _maybe_offer_admin_setup(workspace: str, profile: str | None) -> None:
-    """When a workspace admin runs ``configure`` on a workspace with no managed config, drop an FYI
-    that they could publish one with ``ucode setup`` — without interrupting the configure flow.
+def _maybe_run_admin_setup(workspace: str, profile: str | None) -> None:
+    """When a workspace admin runs ``configure`` on a workspace with no managed config, drop straight
+    into the ``ucode setup`` authoring flow — ``configure`` is replacing ``setup``, so the admin
+    never has to invoke it themselves. On completion, exit with setup's own status code.
 
-    Admins are the ones who'd want a managed config, so the note is only shown to them; a plain
-    developer sees nothing. This never prompts and never diverts the command: the developer's own
-    ``configure`` always runs to completion, with the note printed alongside it. The check is
-    best-effort: any failure to determine admin status (auth or SCIM unreachable) silently skips it.
+    A plain developer (and any caller whose admin status can't be verified) instead falls through to
+    the normal local-configure flow — this function just returns for them. The admin check is
+    best-effort: any failure to determine admin status (auth or SCIM unreachable) silently skips
+    setup and returns, so a developer is never blocked behind an authoring flow they can't complete.
     """
     try:
         token = get_databricks_token(workspace, profile)
@@ -312,10 +337,36 @@ def _maybe_offer_admin_setup(workspace: str, profile: str | None) -> None:
     if not is_admin:
         return
     print_note(
-        "✨ New: run `ucode setup` to publish a managed config to a workspace — set agents, models, mcps "
-        "and skills once, and every developer inherits them when running `ucode`. This scales "
-        "delivery of coding agents to all developers without each one setting up ucode themselves."
+        "You're a workspace admin, and no managed coding agent config exists for this workspace "
+        "yet — let's set one up. Choose the agents, models, MCPs, and skills once and every "
+        "developer inherits them when they run `ucode`."
     )
+    _run_setup_and_exit(workspace, profile, token)
+
+
+def _run_setup_and_exit(workspace: str, profile: str | None, token: str | None = None) -> None:
+    """Launch the ``ucode setup`` authoring flow in place, then exit with its status code.
+
+    Reuses the workspace/profile ``configure`` already resolved and authenticated against so setup
+    doesn't prompt for them again, and hands setup the same ``token`` the admin check already used
+    so setup's admin gate can't disagree with the routing decision (e.g. right after a credential
+    switch, where a second token fetch could resolve a different identity). ``setup_command`` handles
+    an already-existing config (offering to adopt or edit it). Its actionable failures and aborts are
+    mapped to clean exit codes rather than bubbling up as unhandled errors.
+    """
+    try:
+        # Brand the flow as `ucode configure`: it was reached through configure, not a bare
+        # `ucode setup`, so its section headers should say so.
+        code = setup_command(
+            workspace=workspace, profile=profile, command_label="ucode configure", token=token
+        )
+    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
+    raise typer.Exit(code or 0)
 
 
 def _print_discovery_diagnostics(state: dict) -> None:
@@ -340,14 +391,18 @@ def _print_discovery_diagnostics(state: dict) -> None:
     print_note("Re-run with `UCODE_DEBUG=1` to log raw discovery responses to ~/.ucode/debug.log.")
 
 
-def _prompt_for_configuration(tool: str | None = None) -> tuple[str, str | None]:
+def _prompt_for_configuration(
+    tool: str | None = None, *, title: str = "ucode configure"
+) -> tuple[str, str | None]:
+    # Defaults to "ucode configure" since every caller in this module is part of the configure
+    # command; the setup flow (managed_wizard) passes title="ucode setup".
     if tool is None:
         desc = "Configure your Databricks workspace"
     else:
         desc = f"Configure {TOOL_SPECS[tool]['display']} to use your Databricks endpoint."
     with spinner("Loading Databricks workspaces and profiles..."):
         profiles = get_databricks_profiles()
-    return prompt_for_workspace(desc, profiles)
+    return prompt_for_workspace(desc, profiles, title=title)
 
 
 def _parse_agents_option(agents: str) -> list[str]:

src/ucode/managed_wizard.py

@@ -1305,12 +1305,18 @@ def _handle_existing_config(workspace: str, token: str) -> tuple[bool, dict | No
     choice = prompt_for_selection(
         "What would you like to do?",
         [
+            ("adopt", "Adopt the published config as-is (it applies when you run `ucode`)"),
             ("create", "Author a new config (replaces the existing one when you publish)"),
             ("delete", "Delete the existing config (removes it from the workspace, leaves none)"),
         ],
     )
     if choice is None:
         raise KeyboardInterrupt
+    if choice == "adopt":
+        from ucode.cli import _confirm_managed_config_in_force
+
+        _confirm_managed_config_in_force(existing, workspace)
+        return False, existing
     if choice == "create":
         # 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.
@@ -1411,9 +1417,9 @@ def _command_line(command: str, description: str, *, marker: str = " ", width: i
 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 _step_banner(index: int, title: str, command_label: str = "ucode setup") -> None:
+    """Announce one phase of the flow as `step N of M`, branded to the invoking command."""
+    print_section(f"{command_label} · step {index} of {len(SETUP_STEP_TITLES)} · {title}")
 
 
 def _selected_recap(workspace: str, enabled_agents: dict, default_agent: str | None) -> None:
@@ -1545,6 +1551,8 @@ def setup_command(
     *,
     workspace: str | None = None,
     profile: str | None = None,
+    command_label: str = "ucode setup",
+    token: str | None = None,
 ) -> int:
     """Author the agents and models half of the workspace's managed coding config interactively.
 
@@ -1557,6 +1565,17 @@ def setup_command(
     workspace hand it in so the admin isn't prompted to pick one again — e.g. `ucode configure`
     launching setup after its admin offer. When ``workspace`` is None the flow prompts as usual.
 
+    ``command_label`` brands the section headers to the invoking command: `ucode configure` passes
+    "ucode configure" so a user who never typed `ucode setup` isn't jarred by it (the standalone
+    `ucode setup` command keeps the default). References to specific sub-commands (`ucode setup
+    mcps`, `ucode apply`, …) stay verbatim — those are real command names, not branding.
+
+    ``token`` lets a caller that already authenticated and admin-checked the workspace (e.g.
+    `ucode configure`) hand its token in, so setup's admin gate uses the *same* token as the routing
+    decision — a second fetch here could resolve a different identity right after a credential
+    switch and reject a caller configure just treated as an admin. When None, setup authenticates
+    and fetches its own token as usual.
+
     Returns a process exit code. Raises RuntimeError for actionable failures (not an admin, no
     agents available) and KeyboardInterrupt when the admin aborts a picker; the CLI maps both.
     """
@@ -1567,17 +1586,19 @@ def setup_command(
     # would be circular.
     from ucode.cli import _prompt_for_configuration, configure_shared_state
 
-    print_section("ucode setup")
+    print_section(command_label)
     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:
-        workspace, profile = _prompt_for_configuration()
+        workspace, profile = _prompt_for_configuration(title="ucode setup")
     # `configure_shared_state` below authenticates too and prints its own success line, so this one
     # stays quiet rather than reporting the same thing twice. It still has to run first: the admin
-    # gate and the existing-config check both need a token before discovery.
-    ensure_databricks_auth(workspace, profile, quiet=True)
-    token = get_databricks_token(workspace, profile)
+    # gate and the existing-config check both need a token before discovery. A token handed in by
+    # the caller is reused as-is (see the docstring); otherwise fetch one here.
+    if token is None:
+        ensure_databricks_auth(workspace, profile, quiet=True)
+        token = get_databricks_token(workspace, profile)
 
     _require_admin(workspace, token)
     keep_going, published = _handle_existing_config(workspace, token)
@@ -1607,7 +1628,7 @@ def setup_command(
     previously_enabled = [
         tool for tool in (previous.get("enabled_agents") or {}) if tool in available
     ]
-    _step_banner(1, SETUP_STEP_TITLES[0])
+    _step_banner(1, SETUP_STEP_TITLES[0], command_label)
     picked = prompt_for_tools(
         [(tool, TOOL_SPECS[tool]["display"]) for tool in available],
         preselected=previously_enabled or None,
@@ -1616,7 +1637,7 @@ def setup_command(
         print_note("No coding agents selected — nothing to configure.")
         return 0
 
-    _step_banner(2, SETUP_STEP_TITLES[1])
+    _step_banner(2, SETUP_STEP_TITLES[1], command_label)
     enabled_agents: dict[str, dict] = {}
     for index, tool in enumerate(picked, start=1):
         print_heading(f"{TOOL_SPECS[tool]['display']}  ({index} of {len(picked)})")
@@ -1642,7 +1663,7 @@ def setup_command(
     # 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])
+    _step_banner(3, SETUP_STEP_TITLES[2], command_label)
     default_agent = picked[0]
     if len(picked) > 1:
         _selected_recap(workspace, enabled_agents, default_agent=None)
@@ -2019,7 +2040,7 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int:
     workspace = state.get("workspace")
     profile = state.get("profile")
     if not workspace:
-        workspace, profile = _prompt_for_configuration()
+        workspace, profile = _prompt_for_configuration(title="ucode setup")
 
     manifest, api_payload = parse_publish_payload(load_publish_payload(file_path), workspace)
 

src/ucode/ui.py

@@ -344,6 +344,8 @@ def normalize_workspace_url(workspace: str) -> str:
 def prompt_for_workspace(
     description: str,
     profiles: list[tuple[str, str]] | None = None,
+    *,
+    title: str = "ucode setup",
 ) -> tuple[str, str | None]:
     """Ask the user for a workspace URL, offering profiles as quick-select.
 
@@ -352,10 +354,10 @@ def prompt_for_workspace(
     profiles pointing at the same workspace) are shown separately; the picker
     returns the exact (host, profile_name) the user selected. Returns
     ``(url, profile_name)``; profile_name is ``None`` when the user typed a
-    URL manually.
+    URL manually. ``title`` names the invoking command in the panel header.
     """
     console.print()
-    console.print(Panel(description, title="ucode setup", style="bold blue", expand=False))
+    console.print(Panel(description, title=title, style="bold blue", expand=False))
 
     if profiles:
         name_header = "Profile Name"