From 3f70031a472427cd6e73665f471a817a809a49bd Mon Sep 17 00:00:00 2001 From: David Liu Date: Tue, 1 Sep 2026 23:26:04 +0000 Subject: [PATCH] Replace `ucode setup` with a `ucode configure`-centered workflow Managed config had its own command tree: `ucode setup` to author, `ucode setup show`, `ucode setup help`, `ucode setup spend-tiers`, `ucode setup --from-file`. An admin had to know that `ucode configure` set up their own machine while `ucode setup` authored everyone else's, and a developer running `ucode configure` against a workspace that publishes a config was walked through prompts whose answers the managed config was about to override. `ucode configure` is now the single entry point. It fetches the workspace's published config once, applies it to this machine when there is one (showing the drift inline rather than asking per setting), and only then offers authoring, to admins. `ucode setup` and its subcommands are gone; `spend-tiers` and `--from-file` move onto `configure`, and `ucode export` reads the draft slot. Authoring writes the draft slot only, so nothing reaches the workspace until `ucode publish`. A workspace whose managed-config backend is unavailable still configures the machine and simply does not offer to publish. Co-authored-by: Isaac --- README.md | 119 +++---- src/ucode/cli.py | 479 ++++++++++++++++---------- src/ucode/databricks.py | 17 +- src/ucode/managed_config.py | 58 ++-- src/ucode/managed_export.py | 30 +- src/ucode/managed_resolve.py | 4 +- src/ucode/managed_setup.py | 6 +- src/ucode/managed_wizard.py | 468 ++++++++----------------- src/ucode/ui.py | 53 ++- tests/test_cli.py | 528 +++++++++++++++++++---------- tests/test_databricks.py | 6 +- tests/test_managed_config.py | 144 +++++--- tests/test_managed_export.py | 10 +- tests/test_managed_publish.py | 6 + tests/test_managed_wizard.py | 618 ++++++++++++++-------------------- tests/test_ui.py | 37 +- 16 files changed, 1357 insertions(+), 1226 deletions(-) diff --git a/README.md b/README.md index bbc4a9c9..f56d8ff8 100644 --- a/README.md +++ b/README.md @@ -217,52 +217,58 @@ ucode skills --location main.default,ml.prod --mcp Each run prints the registered server, its URL, the configured agents, and its tools, and reminds you to run `ucode ` (existing agent sessions need a restart before the MCP tools load). -### Managed config for a workspace (admins) +### Managed config for a workspace + +`ucode configure` is the single entry point for both configuring your own machine and (for admins) +authoring the config your developers pick up automatically. It figures out which you need: + +1. It selects and authenticates your workspace and fetches the workspace's managed config once. +2. **If the workspace already has a managed config**, `ucode configure` applies it to this machine + automatically — showing the drift against your current settings inline, without asking you to + approve each change. For a developer, that's it: you're configured. +3. **If you're a workspace admin**, `ucode configure` then asks whether you want to update the + workspace's configuration. If yes (or if no managed config exists yet), it walks you through the + agents to enable and which one bare `ucode` launches, then per agent: Databricks-hosted models or + an external Model Provider Service and the models to expose. Claude Code is asked one model per + family (opus/sonnet/haiku/fable), since it selects models by family alias; any family can be + skipped. + +Interactive Claude Code and Codex configuration installs gateway-critical values in the OS-managed +settings scope so enterprise settings cannot silently override ucode. Non-interactive and CI runs use +local files without invoking `sudo`, and stop with an actionable error if an existing managed value +conflicts. Claude subscription relay is local-only because its loopback proxy exists only for that +session. -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 the spend-tier command, -then publish: +Authoring only ever saves a local **draft** — nothing reaches the workspace until you run +`ucode publish`. `ucode configure` reports when your machine is configured, then (for admins) advises +publishing. + +A managed config carries **agents, models, the default agent, and a tiered spend policy only**. +MCP servers and skills are personal, per-developer configuration (`ucode mcp` / +`ucode skills`) and are never part of it. ```bash -ucode setup # agents and models (start here) -ucode setup spend-tiers # spend-based routing -ucode publish # publish it to the workspace +ucode configure # configure this machine; admins are offered authoring +ucode configure spend-tiers # (admins) edit the tiered spend policy — the one section command +ucode publish # (admins) publish the draft to the workspace ``` -A managed config carries **agents, models, the default agent, and a tiered spend policy only**. -MCP servers and skills are personal, per-developer configuration (`ucode mcp` / `ucode skills`) and -are never part of it. - -`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 and the models to expose. Interactive -Claude Code and Codex configuration installs gateway-critical values in the OS-managed settings -scope so enterprise settings cannot silently override ucode. Non-interactive and CI runs use local -files without invoking `sudo`, and stop with an actionable error if an existing managed value -conflicts. Claude subscription relay is local-only because its loopback proxy exists only for that -session. -Claude Code is asked one model per family (opus/sonnet/haiku/fable), since it selects models by family -alias; any family can be skipped. - -`ucode setup spend-tiers` edits just its own part of the same config, so you can change a spend tier -later without walking the whole flow. It sets a tiered spend policy that switches the default agent -and model as the workspace burns through a budget, and offers to publish right away so you can apply -changes incrementally. +`ucode configure spend-tiers` is the only managed-config section command. It sets a tiered spend +policy that switches everyone's default agent and model as the workspace burns through a budget, and +edits just that section of the draft. It is strictly admin-only — a non-admin gets an actionable +error and a non-zero exit. -Everything is written to `~/.ucode/managed-state.json` — the one local managed-config file — which -`ucode publish` publishes. Re-running `ucode setup` keeps the tracing table and tiered spend policy -already authored, rather than clearing them; to drop one, edit the file and reload -it with `ucode setup --from-file`. +The draft and the last-fetched published snapshot live side by side in `~/.ucode/managed-state.json`, +in separate slots: a launch refreshes the published snapshot but never touches your unpublished +draft. Admins who keep the config in version control can load a hand-written manifest as the draft +instead of running the prompts: ```bash -# Review the manifest and the exact payload `ucode publish` would publish. -ucode setup show - -# Skip the prompts and load a hand-written config instead (validated before saving). -ucode setup --from-file ./managed-config.json +# (admins) Load a hand-written managed config (ucode's manifest shape) as the draft; nothing published. +ucode configure --from-file ./managed-config.json ``` -Once the manifest looks right, publish it: +Once the draft looks right, publish it: ```bash # Validate, show a diff against what's live, and ask before publishing. @@ -271,17 +277,18 @@ ucode publish # Publish without the confirmation prompt (for CI). ucode publish --yes -# Publish a config file exported with `ucode export` instead of the locally authored one. +# Publish a config file exported with `ucode export` instead of the locally authored draft. ucode publish -f ./managed-config.json ucode publish --file ./managed-config.json --yes ``` `publish` updates the workspace's existing config in place rather than replacing it, so a failed 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. +published config before asking to confirm, and does nothing when the two already match. Developers +pick the new config up on their next ucode run. + +If the workspace's managed-config backend isn't available, `ucode configure` still configures your +machine locally and simply doesn't offer to publish. With `-f`/`--file`, `publish` reads a config file produced by `ucode export` and publishes it through the same validation, diff, and confirmation flow. The file's `workspace` must match the configured @@ -292,13 +299,14 @@ file are ignored rather than rejected, since the managed config no longer carrie ### Exporting the config -Any user (not only admins) can print the workspace's managed config as portable JSON with `ucode -export`. The output leads with the source `workspace` URL and a `spec_version` (the export format -version), followed by the canonical external config; credentials and server-assigned fields (the -resource name, timestamps, user ids) are excluded. Without `--file` the JSON is written to stdout; -with `--file`/`-f` the same bytes are written to a file (atomically, and the destination's parent -directory must already exist) while stdout stays empty. The exported file is exactly what `ucode -publish -f ` consumes. +`ucode export` prints the admin's local managed-config **draft** as portable JSON. The output leads +with the source `workspace` URL and a `spec_version` (the export format version), followed by the +canonical external config; credentials and server-assigned fields (the resource name, timestamps, +user ids) are excluded. It reads the draft only — never your personal settings, and never the fetched +published snapshot — so there's nothing to export until you've authored one with `ucode configure`. +Without `--file` the JSON is written to stdout; with `--file`/`-f` the same bytes are written to a +file (atomically, and the destination's parent directory must already exist) while stdout stays +empty. The exported file is exactly what `ucode publish -f ` consumes. ```bash # Print the managed config as JSON. @@ -326,7 +334,7 @@ The output looks like: | Command | Description | |---------|-------------| | `ucode status` | Show current workspace, base URLs, managed config files, and selected models | -| `ucode export` | Print the workspace's managed config as portable JSON (`--file ` / `-f` to write a file) | +| `ucode export` | Print the admin's local managed-config draft as portable JSON (`--file ` / `-f` to write a file) | | `ucode doctor` | Diagnose local issues (uv, npm, Databricks CLI, workspace, credentials, agent CLIs, tracing) and offer to fix any problems found | | `ucode usage` | Show AI Gateway usage summary, plus your budget spend against its alert threshold when the workspace reports one | | `ucode usage --warehouse-id ` | Query a specific SQL warehouse instead of discovering one | @@ -353,13 +361,10 @@ The output looks like: | `ucode 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 skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | | `ucode skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | -| `ucode setup` | Author the managed config's agents and models (workspace admins only) | -| `ucode setup spend-tiers` | Set the managed config's tiered spend 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 publish` would publish | -| `ucode setup --from-file ` | Load a hand-written managed config instead of running the prompts | -| `ucode publish` | Publish the authored managed config to the workspace, after a diff and confirmation (admins only) | -| `ucode publish -f ` | Publish a config file exported with `ucode export` instead of the locally authored one | +| `ucode configure spend-tiers` | Set the managed config's tiered spend routing policy (workspace admins only) | +| `ucode configure --from-file ` | Load a hand-written managed config as the draft, skipping the prompts (workspace admins only) | +| `ucode publish` | Publish the managed-config draft to the workspace, after a diff and confirmation (admins only) | +| `ucode publish -f ` | Publish a config file exported with `ucode export` instead of the locally authored draft | | `ucode publish --yes` | Publish without the confirmation prompt | Databricks AI Tools are installed only by `ucode configure`, never by `ucode ` launches. @@ -381,7 +386,7 @@ control the installation. | `~/.copilot/.env` | GitHub Copilot CLI | | `~/.pi/agent/models.json` | Pi | | `~/.cursor/mcp.json` | Cursor Agent (MCP servers only) | -| `~/.ucode/managed-state.json` | The managed config — authored by `ucode setup` (admins) and refreshed from the workspace on launch | +| `~/.ucode/managed-state.json` | The managed config — the admin's unpublished draft (authored by `ucode configure`) and the last-fetched published snapshot, in separate slots; refreshed from the workspace on launch | | `~/.ucode/managed-state.json.pre-v2.bak` | One-time copy of a pre-slots `managed-state.json`, kept when it is first migrated | | `~/.ucode/managed-backups/` | Baseline backups for OS-managed files changed by ucode | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3da6fb60..ce637ff4 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -68,8 +68,10 @@ ) from ucode.managed_config import ( MANAGED_CONFIG_ENV_VAR, + fetch_published_config, get_model_recommendation, - load_managed_state, + load_draft_config, + load_published_config, managed_agent_config_enabled, refresh_managed_config, ) @@ -80,17 +82,18 @@ managed_launch_model, managed_provider_family_models, managed_provider_service, + managed_state_overrides, managed_supplies_models, managed_unservable_models, recommended_agent, resolve_state, ) from ucode.managed_wizard import ( + author_managed_config, + configure_from_file, + configure_spend_tiers_command, + print_managed_next_steps, publish_command, - setup_budget_policy_command, - setup_command, - setup_help_command, - show_command, ) from ucode.mcp import ( MCP_CLIENTS, @@ -133,6 +136,7 @@ prompt_for_tools, prompt_for_workspace, prompt_yes_no, + prompt_yes_no_default, set_verbosity, spinner, status_badge, @@ -234,115 +238,241 @@ def _print_managed_summary_abridged(managed: dict, state: dict, tool: str | None ) -def _confirm_managed_config_applied(managed: dict, workspace: str) -> None: - print_success("A managed config is published for your workspace — you're all set.") - _print_managed_summary(managed, {"workspace": workspace}, tool=None) - print_note("Run `ucode` to launch with your managed settings.") +def _managed_drift_rows(managed: dict, state: dict, tool: str) -> list[str]: + """The human-readable settings the managed config would change for ``tool`` on this machine. + + Compares the admin's config against the developer's current local state for the settings a launch + actually overlays (:func:`managed_state_overrides` plus the provider service), and returns one + ``old -> new`` line per real difference. Empty when nothing changes. + """ + rows: list[str] = [] + current_provider = get_provider_service(state, tool) + managed_provider = managed_provider_service(managed, tool) + if managed_provider and managed_provider != current_provider: + was = current_provider or "your discovered models" + rows.append(f"provider {was} -> {managed_provider}") + managed_model = managed_default_model(managed, tool) + if managed_model and not managed_provider: + overrides = managed_state_overrides(managed, tool) + if overrides: + rows.append(f"default model -> {managed_model}") + return rows + + +def _render_managed_drift(managed: dict, state: dict, tools: list[str]) -> None: + """Print, without prompting, what applying the managed config changes on this machine.""" + lines: list[str] = [] + for tool in tools: + rows = _managed_drift_rows(managed, state, tool) + display = TOOL_SPECS[tool]["display"] + if rows: + lines.append(f"[bold]{display}[/bold]: {'; '.join(rows)}") + else: + lines.append(f"[bold]{display}[/bold]: [dim]already matches[/dim]") + if lines: + console.print(Panel("\n".join(lines), title="Applying managed config", style="green")) + + +def apply_managed_config_locally(managed: dict, workspace: str, profile: str | None) -> bool: + """Configure this machine's agents from the workspace's managed config, without prompting. + + The launch path overlays the managed config on every ``ucode`` run (``resolve_state`` then + ``configure_tool``); this does the same at configure time so a developer is set up immediately. + Per enabled tool it discovers the workspace's models, overlays the admin's config, and writes the + agent's settings — the managed overlay is transient in ``state.json`` (``save_state`` restores the + developer's own values), exactly as at launch. Best-effort per tool: one agent's failure warns and + the rest still configure. + + Returns whether the managed config governs this machine afterwards: false only when every enabled + agent failed, so a caller does not report completion right after warning about each failure. A + config that enables no agent ucode supports leaves nothing to apply and counts as applied. + """ + tools = [t for t in managed_enabled_tools(managed) if t in TOOL_SPECS] + if not tools: + return True + configure_shared_state(workspace, profile=profile, tools=tools, force_login=False) + _render_managed_drift(managed, load_state(), tools) + applied = 0 + for tool in tools: + display = TOOL_SPECS[tool]["display"] + try: + overlaid = resolve_state(managed, load_state(), tool) + provider = managed_provider_service(managed, tool) + if provider: + provider_models, error, relayed = resolve_provider_models(tool, overlaid, provider) + if error: + print_warning(f"Could not apply the managed config for {display}: {error}") + continue + if tool == "claude": + authored = managed_provider_family_models(managed) + if authored: + provider_models = authored + result = configure_tool( + tool, + overlaid, + None, + provider=provider, + provider_models=provider_models, + relayed=relayed, + ) + else: + overlaid, resolved = resolve_launch_model( + tool, overlaid, managed_default_model(managed, tool) + ) + result = configure_tool(tool, overlaid, resolved) + except RuntimeError as exc: + print_warning(f"Could not apply the managed config for {display}: {exc}") + continue + save_state(result) + applied += 1 + return applied > 0 + + +def _report_managed_apply(applied: bool, subject: str, *, final: bool = True) -> None: + """Report the outcome of :func:`apply_managed_config_locally` without over-claiming. + + ``final`` is false when an admin authoring offer follows, where "Configuration complete" reads + wrong ahead of more configuration. + """ + if not applied: + print_warning(f"Could not apply {subject} to this machine; see the warnings above.") + return + if final: + print_success(f"Configuration complete — this machine now uses {subject}.") + else: + print_success(f"This machine now uses {subject}.") -def _resolve_workspace_then_maybe_reject( +def _prompt_admin_update_workspace() -> bool: + """Ask an admin whether to (re)author the workspace's managed config. Never auto-publishes. + + Defaults to no on an empty answer or closed stdin: the machine is already configured by the time + this is asked, so a scripted run must finish cleanly rather than abort, and declining is the safe + default when the "yes" branch opens an interactive authoring session. + """ + return prompt_yes_no_default( + "Update this workspace's managed configuration? (admins only)", default=False + ) + + +def _run_managed_configure_flow( workspace_entries: list[tuple[str, str | None]] | None, + *, + allow_interactive_authoring: bool = True, ) -> list[tuple[str, str | None]] | None: - """Resolve the workspace ``ucode configure`` targets, then branch on role + managed config. - - Enablement is both client- and server-side: the client-side ``ENABLE_MANAGED_AGENT_CONFIG`` env - var must be set for ``ucode`` to run any of this (the opt-in bug-bash gate below), and the - workspace's gateway must not report the feature disabled (``FEATURE_DISABLED``) — a config only - exists to adopt when the server side is on too. - - 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 what to do. - Then, gated by the client-side ``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 client-side flag set it returns ``workspace_entries`` unchanged and prompts nothing. + """The managed-config-aware core of ``ucode configure``: resolve workspace, apply, offer authoring. + + Gated on the client-side ``ENABLE_MANAGED_AGENT_CONFIG`` opt-in; without it this returns + ``workspace_entries`` unchanged and the caller runs the plain local configure. When enabled, it + resolves and authenticates the target workspace, determines admin access **once**, and fetches the + published config with a clean feature-availability signal, then branches: + + * **Backend feature unavailable** → applies the config last saved for this workspace when there + is one, as the launch path does, and exits; otherwise returns the entries so local + configuration still runs. Neither branch offers authoring or advises publishing. + * **A published config exists** → applies it to this machine and shows the drift inline (no + per-change approval), reports the outcome, and exits. An admin is then asked whether to update + the workspace config; yes runs the guided authoring, which saves a draft, applies that draft + here, and advises ``ucode publish`` — so completion is not claimed ahead of that offer. + * **No published config (feature on)** → a non-admin falls through to local configure (returns + entries). An admin runs the guided authoring, unless explicit local-configure options were + supplied; those options are honored by falling through to local configure rather than silently + ignoring them. + + ``allow_interactive_authoring`` is false when the caller supplied options that the interactive + managed-authoring flow cannot honor (for example ``--agents`` or ``--skip-validate``). It keeps + those options effective by returning to the ordinary local-configure flow. + + Returns the resolved entries when the caller should continue into local configure, or raises + ``typer.Exit`` when the managed flow has fully handled the run. """ if not managed_agent_config_enabled(): return workspace_entries + # A managed config is per-workspace and this flow resolves exactly one, so several explicit + # workspaces go to local configure rather than silently configuring only the first. + if workspace_entries is not None and len(workspace_entries) > 1: + 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 - # the workspace and never raises — it falls back to the persisted copy, then None, on failure. - with spinner("Loading..."): - managed, coding_agent_config_feature_disabled = refresh_managed_config( - {"workspace": workspace, "profile": profile} - ) - if not managed: - if not coding_agent_config_feature_disabled: - _maybe_run_admin_setup(workspace, profile) - return entries - is_admin: bool | None = None try: token = get_databricks_token(workspace, profile) except RuntimeError: - token = None - if token is not None: + return entries + with spinner("Loading your workspace's managed config..."): + published, reason, feature_disabled = fetch_published_config(workspace, token) + if reason is not None: + published = load_published_config(workspace) + summary = " ".join(reason.split()) + summary = summary if len(summary) <= 160 else summary[:157] + "..." + if not published: + if not feature_disabled: + print_warning( + f"Could not read your workspace's managed config ({summary}); configuring your " + "own settings for now. Re-run `ucode configure` once the workspace is reachable." + ) + return entries + if feature_disabled: + print_note("Applying the managed config last saved for this workspace.") + else: + print_warning( + f"Could not read your workspace's managed config ({summary}); applying the last one " + "saved for this workspace." + ) + # Skipped when the feature is off: nothing can be authored or published, so admin status cannot + # change the outcome. + is_admin = False + if not feature_disabled: 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_applied(managed, workspace) - raise typer.Exit(0) - -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. + if published: + applied = apply_managed_config_locally(published, workspace, profile) + offer_authoring = bool(is_admin) and allow_interactive_authoring + _report_managed_apply(applied, "your workspace's managed config", final=not offer_authoring) + if is_admin: + if allow_interactive_authoring: + if _prompt_admin_update_workspace(): + _author_and_advise(workspace, profile, token, published) + else: + print_note( + "To update the workspace configuration, re-run `ucode configure` without " + "local configuration options." + ) + raise typer.Exit(0) - 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) - except RuntimeError: - return - with spinner("Checking your workspace permissions..."): - is_admin = is_workspace_admin(workspace, token) if not is_admin: - return - print_note( - "You're a workspace admin, and no managed coding agent config exists for this workspace " - "yet — let's set one up. Choose the agents and models once and every developer " - "inherits them when they run `ucode`." - ) - _run_setup_and_exit(workspace, profile, token) - + return entries + if not allow_interactive_authoring: + print_note( + "Using the requested local configuration options. To author a workspace configuration, " + "re-run `ucode configure` without them." + ) + return entries + _author_and_advise(workspace, profile, token, None) + raise typer.Exit(0) -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. +def _author_and_advise( + workspace: str, profile: str | None, token: str, published: dict | None +) -> None: + """Run the guided managed-config authoring (draft only), apply it locally, and exit. + + Shared by both admin branches. Authoring saves a draft (it does not advise publishing itself); + the draft is then applied to this machine, so a `ucode configure` run always ends with the admin + on the config they just authored, and that outcome is reported here — *then* the publish advice is + printed, so completion always precedes it. Applying is a local preview, not a promotion: launches + overlay the *published* config, so the next `ucode ` puts this machine back on the + published one until `ucode publish` promotes the draft. Never publishes. """ try: - # Brand the flow as "Configure Unity Gateway": it was reached through `ucode configure`, - # not a bare `ucode setup`, so its section headers use the product name rather than the - # bare command. - code = setup_command( + code = author_managed_config( workspace=workspace, profile=profile, - command_label="Configure Unity Gateway", token=token, + published=published, + command_label="ucode configure", ) except RuntimeError as exc: print_err(str(exc)) @@ -350,6 +480,12 @@ def _run_setup_and_exit(workspace: str, profile: str | None, token: str | None = except KeyboardInterrupt: print_err("Interrupted.") raise typer.Exit(130) from None + if code == 0: + draft = load_draft_config(workspace) + if draft: + applied = apply_managed_config_locally(draft, workspace, profile) + _report_managed_apply(applied, "your authored config") + print_managed_next_steps(draft or {}) raise typer.Exit(code or 0) @@ -382,7 +518,10 @@ def _prompt_for_configuration(tool: str | None = None) -> tuple[str, str | None] 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) + state = load_state() + workspace = state.get("workspace") + preselect = (workspace, state.get("profile")) if workspace else None + return prompt_for_workspace(desc, profiles, preselect=preselect) def _parse_agents_option(agents: str) -> list[str]: @@ -993,7 +1132,7 @@ def status() -> int: # the local cache (no network): status is a quick, offline-safe glance, and the cache is what the # last launch persisted for this workspace. if workspace and managed_agent_config_enabled(): - managed = load_managed_state(workspace) + managed = load_published_config(workspace) if managed: _print_managed_summary(managed, state, None) @@ -1138,12 +1277,6 @@ def revert() -> int: app.add_typer(configure_app, name="configure", help="Configure workspace and tool settings.") mcp_app = typer.Typer(add_completion=False, no_args_is_help=False) app.add_typer(mcp_app, name="mcp", help="Register Databricks MCP servers on your coding tools.") -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). See `ucode setup help`.", -) def _version_callback(value: bool) -> None: @@ -2166,9 +2299,9 @@ def _launch_managed_default( if not current: raise RuntimeError("No workspace configured. Run `ucode configure` first.") apply_pat_environment(state) - # --dry-run avoids the fetch but still applies the last saved config. + coding_agent_config_feature_disabled = False if dry_run: - managed = load_managed_state(current) + managed = load_published_config(current) else: with spinner("Loading..."): managed, coding_agent_config_feature_disabled = refresh_managed_config(state) @@ -2210,10 +2343,11 @@ def _print_no_managed_config_guidance(workspace: str, profile: str | None) -> No with spinner("Checking your workspace permissions..."): is_admin = is_workspace_admin(workspace, token) if is_admin is False: - print_note("Ask a workspace admin to set one up with `ucode setup`.") + print_note("Ask a workspace admin to set one up with `ucode configure`.") else: - # None means the admin check itself failed; point at setup rather than a dead end. - print_note("Run `ucode setup` to configure one for your workspace, then `ucode publish`.") + print_note( + "Run `ucode configure` to configure one for your workspace, then `ucode publish`." + ) @app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) @@ -2538,6 +2672,15 @@ def configure( ), ] = False, skip_managed_config: SkipManagedConfigOption = False, + from_file: Annotated[ + str | None, + typer.Option( + "--from-file", + help="Admins only: load a hand-written managed config (JSON, in ucode's manifest shape) " + "as the local draft instead of running the interactive flow. Validated before it is " + "saved; nothing is published. Publish it with `ucode publish`.", + ), + ] = None, verbose: Annotated[ str, typer.Option( @@ -2556,6 +2699,45 @@ def configure( raise typer.Exit(2) set_dry_run(dry_run) set_verbosity(verbose) + if from_file is not None: + ignored = [ + flag + for flag, given in ( + ("--agent", agent is not None), + ("--agents", agents is not None), + ("--workspaces", workspaces is not None), + ("--profiles", profiles is not None), + ("--use-pat", use_pat), + ("--skip-validate", skip_validate), + ("--skip-unavailable", skip_unavailable), + ("--enable-fable/--disable-fable", enable_fable is not None), + ( + "--enable-databricks-ai-tools/--disable-databricks-ai-tools", + enable_databricks_ai_tools is not None, + ), + ("--tracing", tracing), + ("--skip-upgrade", skip_upgrade), + ) + if given + ] + if ignored: + print_err( + f"--from-file can't be combined with {', '.join(ignored)}. It saves a managed draft " + "for the current workspace and configures no agent. Run `ucode configure " + f"--from-file {from_file}` on its own, then re-run `ucode configure` with those " + "options." + ) + raise typer.Exit(2) + try: + install_databricks_cli() + code = configure_from_file(from_file) + 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) prompt_optional_updates = not skip_upgrade try: install_databricks_cli() @@ -2580,14 +2762,24 @@ def configure( workspace_entries = _parse_workspaces_option(workspaces) if workspaces is not None else None if profiles is not None: workspace_entries = _parse_profiles_option(profiles) - # Whether the user named the workspace(s) via flags, captured before the resolver below - # may fill `workspace_entries` from a prompt — this, not the resolved value, decides - # whether the optional-setup step is offered. flag_driven_workspace = workspace_entries is not None - # Under a managed config, resolve (prompting when interactive) and set the target workspace - # first, so the developer can switch workspaces; only then short-circuit if that workspace - # is already managed. Returns the resolved entries so the flow below doesn't prompt again. - workspace_entries = _resolve_workspace_then_maybe_reject(workspace_entries) + allow_interactive_managed_authoring = not any( + ( + agent is not None, + agents is not None, + use_pat, + skip_validate, + skip_unavailable, + enable_fable is not None, + enable_databricks_ai_tools is not None, + tracing, + skip_upgrade, + ) + ) + workspace_entries = _run_managed_configure_flow( + workspace_entries, + allow_interactive_authoring=allow_interactive_managed_authoring, + ) # Only forward the opt-in flags when set so existing call expectations # (and defaults) stay unchanged for the common interactive path. skip_kwargs: dict = {} @@ -2607,7 +2799,6 @@ def configure( agent = "claude" if enable_databricks_ai_tools is not None: skip_kwargs["databricks_ai_tools_enabled"] = enable_databricks_ai_tools - combined_optional_setup = False if agent is not None: tool = normalize_tool(agent) install_tool_binary( @@ -2840,29 +3031,19 @@ def configure_tracing( raise typer.Exit(130) from None -@setup_app.callback(invoke_without_command=True) -def setup( - ctx: typer.Context, - from_file: Annotated[ - str | None, - typer.Option( - "--from-file", - help="Skip the interactive flow and load a hand-written managed config (JSON, in " - "ucode's manifest shape) instead. Validated before it is saved.", - ), - ] = None, -) -> None: - """Choose the agents and models for your workspace's managed config (admins only). +@configure_app.command("spend-tiers") +def configure_spend_tiers() -> None: + """Route developers to cheaper agents as the workspace spends its budget (admins only). - The tiered spend policy has its own command — see `ucode setup help`. + The one managed-config section command: it edits the tiered spend policy in the local managed + draft, then advise `ucode publish` to apply it. Strictly admin-only — a non-admin gets an + actionable error and a non-zero exit. """ - if ctx.invoked_subcommand is not None: - return # `typer.Exit` subclasses RuntimeError, so it must be raised outside the try — inside, the # `except RuntimeError` below would swallow it and report the exit code as an error message. try: install_databricks_cli() - code = setup_command(from_file=from_file) + code = configure_spend_tiers_command() except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -2873,48 +3054,6 @@ def setup( raise typer.Exit(code) -@setup_app.command("spend-tiers") -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 publish` would publish.""" - try: - code = show_command() - except RuntimeError as exc: - print_err(str(exc)) - raise typer.Exit(1) from None - if code: - raise typer.Exit(code) - - @app.command("publish") def publish_cmd( file_path: Annotated[ @@ -2934,11 +3073,9 @@ def publish_cmd( """Publish this workspace's managed coding config (workspace admins only). Always validates the manifest before publishing (and shows what would change, then confirms), so - there is no separate dry-run: `ucode setup` only ever writes a valid manifest, and a + there is no separate dry-run: `ucode configure` only ever writes a valid draft, and a hand-editing admin sees any error here before anything reaches the workspace. """ - # See the `setup` callback: `typer.Exit` subclasses RuntimeError, so it must be raised after - # the try block or the handler below would report a successful exit as an error. try: install_databricks_cli() code = publish_command(file_path=file_path, yes=yes) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 7fe3911c..5ba8c295 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -431,8 +431,6 @@ def _http_get_bytes(url: str, token: str, *, timeout: int = 10) -> tuple[bytes | return None, f"network error: {exc.reason}" -# Workspace group whose members are workspace admins. `ucode setup` / `ucode publish` are restricted -# to this group because the coding-agent-config CRUD API enforces the same check server-side. WORKSPACE_ADMIN_GROUP = "admins" @@ -447,10 +445,11 @@ def is_workspace_admin(workspace: str, token: str) -> bool | None: """Whether the caller is a workspace admin, via their SCIM `Me` group membership. Returns True/False, or None when the check itself could not be made (SCIM unreachable or a - malformed response). Callers should treat None as "unknown" and proceed optimistically rather - than blocking: the API enforces the same check server-side, so a false negative here would - needlessly stop a legitimate admin, while a false positive just surfaces the server's - PERMISSION_DENIED later. + malformed response). How to treat None is the caller's call, and both conventions are in use: + ``ucode publish`` proceeds optimistically, since the API enforces the same check server-side and a + false negative would needlessly stop a legitimate admin while a false positive just surfaces the + server's PERMISSION_DENIED later; ``ucode configure`` and the admin-only sections decline instead, + since opening an authoring session that cannot end in a publish is worse than asking for a retry. """ payload = _scim_me(workspace, token) if payload is None: @@ -1594,13 +1593,13 @@ def _get_model_services_page( # Successful model-service listings for this process, keyed by workspace. The listing is a paginated # walk of the whole metastore catalog, and several callers want different views of the same result # (`discover_model_services` buckets it per family, `discover_claude_models_unbucketed` keeps the raw -# Claude ids), so a single `ucode setup` run would otherwise page it twice. Cached per process, not +# Claude ids), so a single `ucode configure` run would otherwise page it twice. Cached per process, not # persisted: a long-lived process is not a thing here, and a new model appearing mid-command is not # worth a second walk. Failures are never cached, so a transient error still retries. _MODEL_SERVICES_CACHE: dict[str, list[str]] = {} # Same idea for the Model Provider Service listing (a different endpoint). It is workspace-wide and -# filtered per agent afterwards, so `ucode setup` would otherwise re-list it once per MPS-capable +# filtered per agent afterwards, so `ucode configure` would otherwise re-list it once per MPS-capable # agent. Keyed by ``(workspace, parent)`` — a schema-scoped listing is a different result set than # the metastore-wide one, so they must not share an entry. _MODEL_PROVIDER_SERVICES_CACHE: dict[tuple[str, str], list[dict]] = {} @@ -2145,7 +2144,7 @@ def list_model_provider_services( A successful result is memoized per workspace for the life of the process, like the model-services listing: the listing is workspace-wide (filtered per agent afterwards by - :func:`service_usable_for_tool`), so without the memo `ucode setup` re-lists it once per + :func:`service_usable_for_tool`), so without the memo `ucode configure` re-lists it once per MPS-capable agent. Pass ``use_cache=False`` to force a fresh call. """ # Keyed by workspace *and* parent: a `parent`-scoped listing holds only that schema's services, diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index dd7b279d..3a6309b7 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -6,15 +6,18 @@ - fetching the raw manifest (via :func:`ucode.databricks.fetch_managed_coding_agent_configs`), - normalizing the proto-JSON into a stable internal dict keyed by ucode's own tool names, -- persisting it via :func:`save_managed_state` / :func:`load_managed_state` — the admin-write side - (``managed_setup`` / ``managed_wizard``) authors the manifest here, and the launch path pulls the - published copy back into the same file, and +- persisting it via :func:`save_draft_config` / :func:`save_published_config` (and the matching + loaders) — the admin-write side (``managed_setup`` / ``managed_wizard``) authors the manifest here, + and the launch path pulls the published copy back into the same file, and - re-reading it on each launch, falling back to the persisted copy when the read fails. -There is deliberately one file, not a separate authored ``managed-settings.json``: the workspace is -the source of truth, so an authored draft and the pulled copy are the same shape and coexist in -``managed-state.json``. ``ucode setup`` authors the draft; ``ucode publish`` publishes it; a launch -then pulls the published copy back into the same file. +One file holds two clearly separated slots per workspace. ``managed-state.json`` is a versioned +per-workspace map — ``{"version": 2, "workspaces": {: {"published": ..., "draft": ...}}}`` — so +the admin's locally authored, unpublished ``draft`` never shares a slot with the launch-fetched +``published`` snapshot. ``ucode configure`` (admin) authors the ``draft``; ``ucode export`` / +``ucode publish`` read the ``draft`` only; a launch refreshes the ``published`` slot and never +touches the ``draft``. Keeping a map (not a single slot) means refreshing one workspace can't clobber +another workspace's draft. :func:`refresh_managed_config` is the launch path's entry point. It is called before model discovery, because the manifest decides whether that discovery is needed at all; the launch path then hands the @@ -456,16 +459,6 @@ def managed_state_workspace() -> str | None: return next(iter(workspaces)) if len(workspaces) == 1 else None -def save_managed_state(workspace: str, config: dict) -> None: - """Deprecated alias for :func:`save_published_config`, kept while callers migrate to the slots.""" - save_published_config(workspace, config) - - -def load_managed_state(workspace: str | None) -> dict | None: - """Deprecated alias for :func:`load_published_config`, kept while callers migrate to the slots.""" - return load_published_config(workspace) - - def refresh_managed_config(state: dict) -> tuple[dict | None, bool]: """Fetch the workspace's managed config and persist it, returning ``(manifest, coding_agent_config_feature_disabled)``. @@ -480,7 +473,7 @@ def refresh_managed_config(state: dict) -> tuple[dict | None, bool]: ``coding_agent_config_feature_disabled`` is True when the gateway returned ``FEATURE_DISABLED`` and there was no persisted config to fall back on — the coding-agent-configs feature isn't enabled server-side, - so callers suppress the ``ucode setup`` recommendation. + so callers suppress the ``ucode configure`` publish recommendation. """ workspace = state.get("workspace") if not workspace: @@ -497,14 +490,33 @@ def refresh_managed_config(state: dict) -> tuple[dict | None, bool]: return fallback, _is_feature_disabled(reason) and fallback is None if managed is None: # Record that this workspace has no config, rather than leaving an earlier one on disk: - # the file doubles as the fallback above, so a removed policy would otherwise come back - # into force after the next transient outage. - save_managed_state(workspace, {}) + # the published slot doubles as the fallback above, so a removed policy would otherwise come + # back into force after the next transient outage. The admin's draft, if any, is untouched. + save_published_config(workspace, {}) return None, False - save_managed_state(workspace, managed) + save_published_config(workspace, managed) return managed, False +def fetch_published_config(workspace: str, token: str) -> tuple[dict | None, str | None, bool]: + """Fetch the workspace's published config for ``ucode configure`` and persist the published slot. + + Returns ``(published_or_None, read_error_or_None, feature_disabled)``. Unlike + :func:`refresh_managed_config`, ``feature_disabled`` is reported independently of any persisted + fallback, so ``ucode configure`` can branch on "the managed-config backend is unavailable" — and + then skip publish advice — even when a stale published snapshot is still on disk. A successful + read persists the published slot (recording emptiness as ``{}``); the admin's draft is untouched. + """ + managed, reason = get_managed_config(workspace, token) + if reason is not None: + return None, reason, _is_feature_disabled(reason) + if managed is None: + save_published_config(workspace, {}) + return None, None, False + save_published_config(workspace, managed) + return managed, None, False + + def _is_feature_disabled(reason: str) -> bool: return "feature_disabled" in reason.lower() @@ -520,7 +532,7 @@ def _persisted_fallback(workspace: str, reason: str, *, refused: bool = False) - """ # An empty persisted config means the last successful read found none, so there is no admin # policy to fall back to — treat it the same as having no file at all. - persisted = load_managed_state(workspace) + persisted = load_published_config(workspace) if not persisted: return None summary = _summarize_read_failure(reason) diff --git a/src/ucode/managed_export.py b/src/ucode/managed_export.py index 511d331d..5c5689aa 100644 --- a/src/ucode/managed_export.py +++ b/src/ucode/managed_export.py @@ -1,14 +1,15 @@ """`ucode export`: serialize the workspace's managed coding-agent config to portable JSON. -Reads the local managed config (the one file :mod:`ucode.managed_config` owns, authored by -``ucode setup`` and refreshed by a launch), validates and serializes it through the same path -``ucode publish`` uses, and writes the external proto-JSON ``CodingAgentConfig`` — prefixed with the -source ``workspace`` and a ``spec_version`` envelope, the format ``ucode publish -f `` consumes -— to stdout or a file. - -Deliberately read-only and offline: no auth, no admin check, no discovery, no publish, and no write -except the explicitly requested ``--file`` output. That makes it role-agnostic (any developer can -run it) and keeps the machine-readable stream on stdout uncontaminated by Rich output. +Reads the local managed config **draft** (the ``draft`` slot of the one file +:mod:`ucode.managed_config` owns, authored by ``ucode configure``), validates and serializes it +through the same path ``ucode publish`` uses, and writes the external proto-JSON ``CodingAgentConfig`` +— prefixed with the source ``workspace`` and a ``spec_version`` envelope, the format ``ucode publish +-f `` consumes — to stdout or a file. + +Draft-only: the launch-fetched published snapshot is deliberately not an export source (a cached +publication is not authored work). Read-only and offline otherwise: no auth, no admin check, no +discovery, no publish, and no write except the explicitly requested ``--file`` output, keeping the +machine-readable stream on stdout uncontaminated by Rich output. """ from __future__ import annotations @@ -19,7 +20,7 @@ import tempfile from pathlib import Path -from ucode.managed_config import load_managed_state, managed_state_workspace +from ucode.managed_config import load_draft_config, managed_state_workspace from ucode.managed_setup import serialize_managed_config, validate_manifest from ucode.state import load_state @@ -37,16 +38,17 @@ def build_export_payload() -> dict: actionable message when no config is authored locally or the config fails structural validation. """ workspace = load_state().get("workspace") or managed_state_workspace() - manifest = load_managed_state(workspace) + manifest = load_draft_config(workspace) if not manifest: raise RuntimeError( - "No managed coding-agent config found locally. Run `ucode setup` to author one, or run " - "`ucode` against a workspace that publishes one, then re-run `ucode export`." + "No managed config draft found locally. Only the draft you author with `ucode configure` " + "(admins only) can be exported or published — a fetched published config is not an " + "authoring source. Run `ucode configure` to author one, then re-run this command." ) errors = validate_manifest(manifest, None) if errors: detail = "\n".join(f" - {error}" for error in errors) - raise RuntimeError(f"The managed config is not valid, so it was not exported:\n{detail}") + raise RuntimeError(f"The managed config draft is not valid:\n{detail}") config = serialize_managed_config(manifest) for field in _SERVER_OWNED_FIELDS: config.pop(field, None) diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 947b174b..62425158 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -1,6 +1,6 @@ """Resolve the effective agent settings from the managed config plus local ucode state. -The managed config (``~/.ucode/managed-state.json`` — authored by ``ucode setup`` and refreshed +The managed config (``~/.ucode/managed-state.json`` — authored via ``ucode configure`` and refreshed from the workspace at launch, both through :mod:`ucode.managed_config`) and the developer's own ucode state (``~/.ucode/state.json``) stay separate files — they are never merged on disk. Instead this module resolves them *per key* at config-write time: whatever the manifest specifies wins, and @@ -198,7 +198,7 @@ def managed_provider_family_models(managed: dict) -> dict[str, str] | None: Model Provider Service. The launch path pins each ``ANTHROPIC_DEFAULT__MODEL`` from this so a *managed* launch - uses exactly the versions the admin chose in ``ucode setup`` — rather than + uses exactly the versions the admin chose in ``ucode configure`` — rather than ``resolve_provider_models`` re-deriving "newest per family" from the service's live targets. It returns the manifest's own family slots (``{opus: id, sonnet: id, ...}``), i.e. what the wizard's per-family prompt authored. diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index 94c3ca5e..5181cb99 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -14,9 +14,9 @@ managed config carries agents/models/global policy/spend tiers only — MCP servers, skills, and tracing are personal configuration and are not part of this manifest. -Local persistence is not duplicated here: the authored manifest is saved to and loaded from the one -local file, ``~/.ucode/managed-state.json``, via :func:`ucode.managed_config.save_managed_state` and -:func:`ucode.managed_config.load_managed_state` — the same file the launch path pulls into. +Local persistence is not duplicated here: the authored draft is saved to and loaded from the one +local file, ``~/.ucode/managed-state.json``, via :func:`ucode.managed_config.save_draft_config` and +:func:`ucode.managed_config.load_draft_config` — a slot kept separate from the launch-fetched copy. The interactive wizard that calls these helpers, and the publish step, live in :mod:`ucode.managed_wizard`. diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 6dddbb4d..1f8afe37 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -1,15 +1,15 @@ -"""Interactive `ucode setup`: author the workspace's managed coding-agent config. - -Workspace admins run this to build the ``CodingAgentConfig`` their developers will pull, then publish -it with ``ucode publish`` (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 spend-tiers`` edits the tiered spend -policy of the same manifest. ``ucode setup`` carries the other sections forward untouched -(:func:`_carry_forward_sections`), and ``ucode setup help`` prints the whole sequence. The managed -config carries agents/models/global policy/spend tiers only — MCP servers and skills are personal +"""Interactive managed-config authoring behind `ucode configure` (admin path). + +Workspace admins reach this through ``ucode configure`` to build the ``CodingAgentConfig`` their +developers will pull, then publish it with ``ucode publish`` (a separate, explicit command — authoring +only ever saves a draft, never publishes). The draft lives in the ``draft`` slot of +``~/.ucode/managed-state.json`` (owned by :mod:`ucode.managed_config`), kept separate from the +launch-fetched published snapshot. + +:func:`author_managed_config` picks the agents and models; ``ucode configure spend-tiers`` +(:func:`configure_spend_tiers_command`) edits the tiered spend policy. Authoring carries the spend +policy and tracing table forward untouched (:func:`_carry_forward_sections`). The managed config +carries agents/models/global policy/spend tiers only — MCP servers and skills are personal configuration (`ucode mcp` / `ucode skills`), not part of it. Serialization, validation, and the per-agent model catalogs live in :mod:`ucode.managed_setup`; this @@ -30,7 +30,6 @@ ANTHROPIC_FAMILIES, all_users_can_use_schema, create_coding_agent_config, - delete_coding_agent_config, discover_claude_models_unbucketed, ensure_databricks_auth, get_databricks_token, @@ -46,16 +45,16 @@ ) from ucode.managed_config import ( get_managed_config, - load_managed_state, - managed_state_workspace, - save_managed_state, + load_draft_config, + load_published_config, + save_draft_config, + save_published_config, ) from ucode.managed_setup import ( CLAUDE_SLOT_FOR_FAMILY, claude_family_candidates, claude_family_for_model, model_options_for_agent, - serialize_managed_config, supports_provider_service, validate_manifest, ) @@ -95,7 +94,7 @@ "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. +# Agents not offered in `ucode configure`'s managed picker, even when the workspace serves them. # `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"}) @@ -297,7 +296,7 @@ def _prompt_models_for_agent(tool: str, state: dict, provider_service: dict | No # Nothing pre-checked: the first option is whatever discovery sorted first, not a # recommendation — for pi it is a Claude model, for codex the oldest GPT. Pre-checking it made # "hit Enter" produce an arbitrary config. (A worthwhile follow-up is to pre-check the models - # this workspace was configured with last time, which `load_managed_state` already loads for + # this workspace was configured with last time, which `load_draft_config` already loads for # the agent picker, so a re-run becomes an edit rather than a re-entry.) picked = _select_hosted_models_multi(f"Select models for {display}:", options, state, custom) if len(picked) == 1: @@ -810,7 +809,7 @@ def _prompt_budget_policy( 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 tiered spend policy?" gate — running `ucode setup spend-tiers` is the answer to + Asks no "set up a tiered spend policy?" gate — running `ucode configure spend-tiers` is the answer to that question, the same way `ucode configure ` needs no confirmation. """ print_section("Tiered Spend Policy") @@ -825,7 +824,7 @@ def _prompt_budget_policy( 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 spend-tiers`. Currently, only AI Gateway budgets with hard blocks are " + "`ucode configure spend-tiers`. Currently, only AI Gateway budgets with hard blocks are " "eligible to be associated with Tiered Spend Policies." ) return None @@ -840,7 +839,7 @@ def _prompt_budget_policy( "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 spend-tiers`." + "`ucode configure spend-tiers`." ) return None @@ -1060,6 +1059,7 @@ def _config_facts(manifest: dict) -> list[tuple[str, str, str]]: 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)))) + tracing = manifest.get("tracing_table") if tracing: facts.append(("tracing_table", "Tracing table", str(tracing))) @@ -1129,20 +1129,28 @@ def _render_config_diff(existing: dict | None, incoming: dict, workspace: str) - return True -def _require_admin(workspace: str, token: str) -> None: +def _require_admin(workspace: str, token: str, *, strict: bool = False) -> None: """Stop unless the caller is a workspace admin. - An unverifiable check (SCIM unreachable) warns and continues: the API enforces the same rule, so - the worst case is a clear PERMISSION_DENIED at publish time rather than a false block here. + By default an unverifiable check (SCIM unreachable) warns and continues: the API enforces the same + rule, so the worst case is a clear PERMISSION_DENIED at publish time rather than a false block. A + ``strict`` gate instead refuses when admin status can't be determined — used by commands that must + be admin-only up front (`ucode configure spend-tiers`) rather than relying on a later publish. """ with spinner("Checking workspace admin permissions..."): admin = is_workspace_admin(workspace, token) if admin is False: raise RuntimeError( - f"You are not an admin of {workspace}. `ucode setup` authors the workspace-wide " - "coding config, so it is restricted to workspace admins." + f"You are not an admin of {workspace}. Authoring the workspace-wide coding config is " + "restricted to workspace admins." ) if admin is None: + if strict: + raise RuntimeError( + f"Could not verify that you are an admin of {workspace}, and this command is " + "admin-only. Check your workspace access (the SCIM `Me` API must be reachable) and " + "try again." + ) print_warning( "Could not verify workspace admin permissions. Continuing — `ucode publish` will fail " "if you lack them." @@ -1151,92 +1159,13 @@ def _require_admin(workspace: str, token: str) -> None: print_success("Admin permissions verified") -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 ``(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 tracing table and 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. +def configure_from_file(path: str) -> int: + """Validate an admin-written manifest and save it as the local draft, skipping the wizard. - Deliberately doesn't itemize what the existing config holds. The admin doesn't need an inventory - 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: - if "feature_disabled" in reason.lower(): - # Authoring a draft that the workspace cannot publish only leads the admin through a - # dead-end wizard. Stop before model discovery and point them to per-user setup instead. - raise RuntimeError(CODING_AGENT_CONFIGS_DISABLED_MESSAGE) - print_note(f"Could not check for an existing config: {reason}") - return True, None - if existing is None: - return True, None - - print_warning( - "This workspace already has a managed configuration — one config covers every agent, " - "tracing table, and budget policy for the whole workspace." - ) - choice = prompt_for_selection( - "What would you like to do?", - [ - ( - "adopt", - "Adopt the published config as your current settings. To invoke, 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_applied - - _confirm_managed_config_applied(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. - return True, existing - - _delete_existing_config(workspace, token, existing) - return False, existing - - -def _delete_existing_config(workspace: str, token: str, existing: dict) -> None: - """Delete the workspace's published config after confirming. Raises RuntimeError on failure. - - Deleting leaves the workspace with no managed config, so every developer falls back to their own - settings on their next ucode run — confirm before doing it. - """ - name = existing.get("name") - if not isinstance(name, str): - raise RuntimeError( - "This workspace has a managed config but the API didn't return its resource name, so " - "ucode can't delete it. Delete it in the workspace directly." - ) - print_warning( - "Deleting removes the managed config entirely. Every developer falls back to their own " - "settings on their next ucode run." - ) - if not prompt_yes_no_default("Delete the existing managed config?", default=False): - print_note("Nothing was deleted.") - return - with spinner("Deleting the managed config..."): - delete_reason = delete_coding_agent_config(workspace, token, name) - if delete_reason is not None: - raise RuntimeError(f"Could not delete the managed config on {workspace}: {delete_reason}.") - print_success(f"Deleted the managed config from {workspace}") - - -def setup_from_file(path: str) -> int: - """Validate an admin-written manifest and save it, skipping the interactive flow. - - The non-interactive path for CI and for admins who'd rather keep the JSON in version control. - Reads ucode's own manifest shape (the same thing the wizard writes), not proto-JSON. + The non-interactive path for CI and for admins who keep the JSON in version control (backing + `ucode configure --from-file`). Reads ucode's own manifest shape (the same thing the wizard + writes), not proto-JSON, and writes the draft slot only — nothing is published. Admin-only, since + it authors the workspace's managed config; a non-admin gets an actionable error. """ manifest_path = Path(path).expanduser() try: @@ -1259,6 +1188,9 @@ def setup_from_file(path: str) -> int: "No workspace is configured. Run `ucode configure` first so ucode knows which " "workspace this manifest is for." ) + profile = state.get("profile") + ensure_databricks_auth(workspace, profile) + _require_admin(workspace, get_databricks_token(workspace, profile), strict=True) errors = validate_manifest(manifest, state) if errors: @@ -1267,18 +1199,16 @@ def setup_from_file(path: str) -> int: print_note(error) return 1 - save_managed_state(workspace, manifest) + save_draft_config(workspace, manifest) _render_summary(workspace, manifest) - print_success(f"Saved to {manifest_path.name} -> ~/.ucode/managed-state.json") - _print_next_steps(manifest) + print_success(f"Saved draft from {manifest_path.name} -> ~/.ucode/managed-state.json") + print_managed_next_steps(manifest) return 0 -# 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 spend-tiers", + "ucode configure spend-tiers", "Tiered Spend Policy", lambda m: isinstance(m.get("budget_policy"), dict), ), @@ -1290,12 +1220,10 @@ def _command_line(command: str, description: str, *, marker: str = " ", width: i 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, command_label: str = "ucode setup") -> None: +def _step_banner(index: int, title: str, command_label: str = "ucode configure") -> 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}") @@ -1331,25 +1259,22 @@ def _section_status_lines(manifest: dict, width: int = 0) -> list[str]: return lines -def _print_next_steps(manifest: dict) -> None: - """List the setup commands still worth running, then the publish step. +def print_managed_next_steps(manifest: dict) -> None: + """Report that the draft is saved, list the optional spend-tiers command, then advise publishing. - 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. + Printed rather than prompted: the admin flow only ever saves a draft — publishing is always an + explicit, separate `ucode publish` step (never offered inline). Showing what is already configured + keeps a re-run from looking like it lost a section — it didn't; authoring carries sections forward. """ console.print() print_heading("Next steps") if config_io.is_dry_run(): - # Under --dry-run nothing was written, so the section commands (which read the saved draft) + # Under --dry-run nothing was written, so the section command (which reads the saved draft) # and `publish` have nothing to act on. Say so rather than send the admin to commands that - # would report "run `ucode setup` first". + # would report "run `ucode configure` 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]") + print_note("[dim]Optional — configure this too, or skip straight to publishing:[/dim]") for line in _section_status_lines(manifest): console.print(line) print_panel( @@ -1358,39 +1283,18 @@ def _print_next_steps(manifest: dict) -> None: ) -def _offer_publish() -> None: - """Offer to publish the saved draft right away, so an admin can apply changes incrementally. - - Each `ucode setup` command only writes a local draft. Without this an admin has to remember to run - `ucode publish` separately, and a `ucode setup` re-run in the meantime is easy to mistake for having - lost the change. Answering yes runs `publish_command`, which shows the diff against the published - config as it publishes; declining leaves the draft for a later `ucode publish`. Skipped under - --dry-run, where nothing was saved to publish. - """ - if config_io.is_dry_run(): - return - console.print() - if not prompt_yes_no_default( - "Publish these changes to the workspace now? (runs `ucode publish`)", default=False - ): - print_note("Draft saved. Run `ucode publish` when you're ready to publish.") - return - publish_command(yes=True) - - -# The sections `ucode setup` carries forward instead of prompting for, and how to rebuild each one. CARRIED_SECTIONS: list[tuple[str, str, str]] = [ - ("tracing_table", "Tracing table", "ucode setup --from-file"), - ("budget_policy", "Tiered Spend Policy", "ucode setup spend-tiers"), + ("tracing_table", "Tracing table", "ucode configure --from-file"), + ("budget_policy", "Tiered Spend Policy", "ucode configure spend-tiers"), ] def _carry_forward_sections(previous: dict, manifest: dict) -> None: - """Copy the sections `ucode setup` no longer prompts for out of a previously authored config. + """Copy the sections authoring 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 tracing - table and budget policy an admin authored with the other commands — they'd have to redo both - just to change a model. + Authoring writes the whole manifest, so without this a re-run would silently clear the tracing + table and budget policy an admin set with the other command — they'd have to redo 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, @@ -1422,68 +1326,43 @@ def _carry_forward_sections(previous: dict, manifest: dict) -> None: print_note(f"Rebuild it with `{rebuild}`.") -def setup_command( - from_file: str | None = None, +def author_managed_config( *, - 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. - - Agents and per-agent models only. The tiered spend policy has its own command (`ucode setup - spend-tiers`), so an admin changing it 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` - 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 - "Configure Unity Gateway" 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 spend-tiers`, `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. + workspace: str, + profile: str | None, + token: str, + published: dict | None = None, + command_label: str = "ucode configure", +) -> int | None: + """Guided agent/model authoring for the workspace's managed config draft. + + The shared step-by-step picker behind ``ucode configure``'s admin path: choose the agents, then + per-agent models / provider service / settings scope, then the default agent. Agents and models + only — the tiered spend policy has its own command (`ucode configure spend-tiers`), and it (plus + the tracing table) is carried forward untouched from the existing draft or the ``published`` + snapshot rather than cleared (:func:`_carry_forward_sections`). + + Saves the result to the local **draft** slot only — never publishes — then advises `ucode + publish`. The caller (`ucode configure`) has already resolved and authenticated the workspace, + determined admin access, and fetched the ``published`` snapshot, and hands them all in, so this + never re-checks admin or re-fetches the workspace's config. ``token`` is reused for provider- + service discovery; ``command_label`` brands the step banners. + + Returns a process exit code, or ``None`` when the admin selected no agents: nothing was saved, so + the caller must not treat an earlier draft as the result of this run. Raises RuntimeError for + actionable failures (no agents available) and KeyboardInterrupt when the admin aborts a picker; the + CLI maps both. """ - if from_file is not None: - return setup_from_file(from_file) - # Imported here rather than at module scope: `cli` imports this module, so a top-level import # would be circular. - from ucode.cli import _prompt_for_configuration, configure_shared_state + from ucode.cli import configure_shared_state 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() - # `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. 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) - if not keep_going: - return 0 - - # Discover the workspace's models and gateway URLs. This also logs in and persists local state. state = configure_shared_state(workspace, profile=profile, force_login=False) workspace = state.get("workspace") or workspace - profile = state.get("profile") or profile available = [ tool @@ -1496,10 +1375,10 @@ def setup_command( "serves models for at least one agent." ) - # The local draft is the carry-forward source, falling back to what's published on the workspace: + # The local draft is the carry-forward source, falling back to the passed-in published snapshot: # a fresh machine (or one after `ucode revert`) has no draft, and without the fallback the next # publish would silently wipe the workspace's tracing table and budget policy. - previous = load_managed_state(workspace) or published or {} + previous = load_draft_config(workspace) or published or {} previously_enabled = [ tool for tool in (previous.get("enabled_agents") or {}) if tool in available ] @@ -1510,7 +1389,7 @@ def setup_command( ) if not picked: print_note("No coding agents selected — nothing to configure.") - return 0 + return None _step_banner(2, SETUP_STEP_TITLES[1], command_label) enabled_agents: dict[str, dict] = {} @@ -1546,7 +1425,7 @@ def setup_command( # 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 a `ucode setup tracing` command when it is. + # it once the backend is ready. _carry_forward_sections(previous, manifest) errors = validate_manifest(manifest, state) @@ -1558,58 +1437,57 @@ def setup_command( print_note(error) return 1 - save_managed_state(workspace, manifest) + save_draft_config(workspace, manifest) _render_summary(workspace, manifest) console.print() - print_success("Saved to ~/.ucode/managed-state.json") - _print_next_steps(manifest) - _offer_publish() + print_success("Draft saved to ~/.ucode/managed-state.json") 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, so a - mismatch can't have the section saved for one workspace while local state points at 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. + """Resolve the workspace a section command edits, authenticate, and strictly gate on admin. + + Returns ``(workspace, profile, token)``. Takes the workspace strictly from local state (set by + ``ucode configure``) rather than prompting or falling back to the draft file's workspace, and + gates with a **strict** admin check: a section command like ``ucode configure spend-tiers`` is + admin-only up front, so an undetermined admin status is refused here rather than deferred to a + later publish failure. """ 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." + "No workspace is configured. Run `ucode configure` first, then re-run this command 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) + _require_admin(workspace, token, strict=True) 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. + """The draft a section command edits, seeded from the published snapshot when no draft exists. + + Reads the local draft; on a fresh machine (no draft yet) it seeds from the workspace's published + snapshot so an admin editing one section starts from the live config rather than an empty one. + An empty ``enabled_agents`` counts as "no config": a launch records ``{}`` for a workspace with no + managed config (see ``refresh_managed_config``), so a file existing is not proof of authored work. + Requiring agents first also keeps the budget-policy tiers honest — they can only name agents the + manifest enables. """ - manifest = load_managed_state(workspace) + manifest = load_draft_config(workspace) or load_published_config(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." + f"No managed config has been authored for {workspace} yet. Run `ucode configure` 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. + """Validate the edited manifest structurally, save it as the draft, and advise publishing. 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 @@ -1625,17 +1503,16 @@ def _save_section_update(workspace: str, manifest: dict) -> int: print_note(error) return 1 - save_managed_state(workspace, manifest) + save_draft_config(workspace, manifest) _render_summary(workspace, manifest) console.print() - print_success("Saved to ~/.ucode/managed-state.json") - _print_next_steps(manifest) - _offer_publish() + print_success("Draft saved to ~/.ucode/managed-state.json") + print_managed_next_steps(manifest) return 0 -def setup_budget_policy_command() -> int: - """Author the managed config's tiered spend policy (`ucode setup spend-tiers`).""" +def configure_spend_tiers_command() -> int: + """Author the managed config's tiered spend policy (`ucode configure spend-tiers`). Admin-only.""" workspace, _, token = _resolve_admin_workspace() manifest = _manifest_for_edit(workspace) @@ -1664,87 +1541,6 @@ def setup_budget_policy_command() -> int: 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 publish`, 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 and models you chose here. Admins only." - ) - print_note( - "Each command below edits your local draft; nothing reaches the workspace until " - "`ucode publish`." - ) - - 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 `publish` sends", width=width) - ) - console.print(_command_line("ucode publish", "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 - - -def show_command() -> int: - """Print the authored manifest and the proto-JSON `ucode publish` would publish.""" - # Fall back to the workspace the on-disk file was authored for, so `ucode setup --show` still - # works before `ucode configure` has put a workspace in local state. - workspace = load_state().get("workspace") or managed_state_workspace() - manifest = load_managed_state(workspace) - if manifest is None: - print_note("No managed config has been authored yet. Run `ucode setup` to create one.") - return 0 - _render_summary(workspace or "unknown", manifest) - console.print() - print_heading("Payload for `ucode publish`") - console.print(json.dumps(serialize_managed_config(manifest), indent=2)) - return 0 - - # Server-side failures an admin is actually likely to hit, mapped to something they can act on. The # raw reasons are `HTTP : ` strings from the transport, and the body carries the # API's `error_code`, so matching on that is more robust than on status codes alone. @@ -1773,7 +1569,7 @@ def _with_claude_inventory(state: dict, workspace: str, profile: str | None) -> """``state`` plus the full Claude listing, for validating a manifest against the workspace. ``state["claude_models"]`` holds only the newest id per family (the launch path pins one model - per family alias), but `ucode setup` deliberately offers the older versions too — pinning + per family alias), but authoring deliberately offers the older versions too — pinning ``default_opus_model`` to a known-good ``claude-opus-4-8`` is a normal thing for an admin to want. Validating against ``claude_models`` alone therefore rejected a model the wizard itself had just offered: @@ -1781,7 +1577,7 @@ def _with_claude_inventory(state: dict, workspace: str, profile: str | None) -> claude: model 'system.ai.claude-opus-4-8' is not available on this workspace. The wizard stashes the full listing on ``state["all_claude_models"]`` mid-run, but that is never - persisted — `setup` saves the manifest, not the state — so a separate `ucode publish` process + persisted — authoring saves the manifest, not the state — so a separate `ucode publish` process starts from a fresh ``load_state()`` without it. Re-fetching here makes the check independent of what the wizard happened to leave behind, which also covers a hand-edited or ``--from-file`` manifest authored on another machine. @@ -1834,7 +1630,7 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: # listing needs a token. Nothing is written until well below this point. ensure_databricks_auth(workspace, profile) - validation_manifest = load_managed_state(workspace) if file_path is None else manifest + validation_manifest = load_draft_config(workspace) if file_path is None else manifest errors = validate_manifest( validation_manifest or manifest, _with_claude_inventory(state, workspace, profile) ) @@ -1843,7 +1639,7 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: for error in errors: print_note(error) if file_path is None: - print_note("Re-run `ucode setup` to fix it, or edit ~/.ucode/managed-state.json.") + print_note("Re-run `ucode configure` to fix it, or edit ~/.ucode/managed-state.json.") else: print_note("Fix the config file and re-run `ucode publish -f`.") return 1 @@ -1884,6 +1680,7 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: if not changed: print_success(f"{workspace}'s published config already matches this one.") print_note("Nothing to publish.") + save_published_config(workspace, manifest) return 0 console.print() print_warning("This takes effect for every developer on their next `ucode` run.") @@ -1902,6 +1699,8 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: if publish_reason is not None: raise RuntimeError(_explain_publish_failure(publish_reason)) + save_published_config(workspace, manifest) + name = (published or {}).get("name") or existing_name or "coding-agent-configs/?" print_success(f"Published {name} to {workspace}") print_note("Developers get it automatically the next time they run `ucode`.") @@ -1909,10 +1708,9 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: __all__ = [ + "author_managed_config", + "configure_from_file", + "configure_spend_tiers_command", + "print_managed_next_steps", "publish_command", - "setup_budget_policy_command", - "setup_command", - "setup_from_file", - "setup_help_command", - "show_command", ] diff --git a/src/ucode/ui.py b/src/ucode/ui.py index ce92bcbc..3a38db80 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -341,18 +341,58 @@ def normalize_workspace_url(workspace: str) -> str: return workspace.rstrip("/") +def _preselected_row( + choices: list[questionary.Choice | questionary.Separator], + profiles: list[tuple[str, str]], + preselect: tuple[str, str | None] | None, +) -> questionary.Choice | None: + """The profile row the picker should start on, or ``None`` to start at the top. + + Matches on workspace URL, preferring the row whose profile name also matches, so a machine + already configured for a workspace lands on it instead of scrolling a long profile list. + Returns the ``Choice`` itself, which is what questionary validates a default against. + """ + if not preselect or not preselect[0]: + return None + try: + wanted = normalize_workspace_url(preselect[0]) + except ValueError: + return None + rows = [(host, name) for host, name in profiles if _same_workspace(host, wanted)] + if not rows: + return None + row = next((candidate for candidate in rows if candidate[1] == preselect[1]), rows[0]) + return next( + ( + choice + for choice in choices + if isinstance(choice, questionary.Choice) and choice.value == row + ), + None, + ) + + +def _same_workspace(host: str, normalized: str) -> bool: + try: + return normalize_workspace_url(host) == normalized + except ValueError: + return False + + def prompt_for_workspace( description: str, profiles: list[tuple[str, str]] | None = None, + preselect: tuple[str, str | None] | None = None, ) -> tuple[str, str | None]: """Ask the user for a workspace URL, offering profiles as quick-select. `profiles` is a list of (host_url, profile_name) tuples. Caller fetches them — `ui.py` stays Databricks-agnostic. Duplicate hosts (multiple 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. + returns the exact (host, profile_name) the user selected. `preselect` is an + optional (workspace_url, profile_name) the picker starts on, again supplied + by the caller. Returns ``(url, profile_name)``; profile_name is ``None`` + when the user typed a URL manually. """ console.print() @@ -392,7 +432,12 @@ def prompt_for_workspace( ] ) choice = questionary.select( - "Select workspace:", choices=choices, style=style, pointer="›", qmark="" + "Select workspace:", + choices=choices, + style=style, + pointer="›", + qmark="", + default=_preselected_row(choices, profiles, preselect), ).ask() if isinstance(choice, tuple): host, profile_name = choice diff --git a/tests/test_cli.py b/tests/test_cli.py index 62f284a6..95087c7b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -11,6 +11,7 @@ from unittest.mock import MagicMock, patch import pytest +import typer.main from typer.testing import CliRunner import ucode.databricks as db_mod @@ -859,6 +860,23 @@ def test_shows_mcp_servers_configured_by_ucode(self): assert "MCP Server:" not in result.output assert "Configured tools:" not in result.output + def test_next_step_hints_name_commands_that_exist(self): + root = typer.main.get_command(app) + with patch("ucode.cli.load_state", return_value=MINIMAL_STATE): + result = runner.invoke(app, ["status"]) + + assert result.exit_code == 0, result.output + hints = re.findall(r"`ucode ([^`]+)`", " ".join(_strip_ansi(result.output).split())) + assert hints, "status printed no `ucode ...` hints to check" + for hint in hints: + group = root + for token in hint.split(): + if not re.fullmatch(r"[a-z][a-z-]*", token): + break + commands = getattr(group, "commands", {}) + assert token in commands, f"`ucode {hint}` names no command `{token}`" + group = commands[token] + def test_status_treats_available_tools_as_configured_agents(self): state = { **MINIMAL_STATE, @@ -895,7 +913,7 @@ def test_status_shows_managed_config_box_when_present_and_enabled(self, monkeypa } with ( patch("ucode.cli.load_state", return_value=MINIMAL_STATE), - patch("ucode.cli.load_managed_state", return_value=managed), + patch("ucode.cli.load_published_config", return_value=managed), ): result = runner.invoke(app, ["status"]) @@ -909,7 +927,7 @@ def test_status_hides_managed_config_box_when_feature_disabled(self, monkeypatch managed = {"enabled_agents": {"claude": {}}} with ( patch("ucode.cli.load_state", return_value=MINIMAL_STATE), - patch("ucode.cli.load_managed_state", return_value=managed) as load_managed, + patch("ucode.cli.load_published_config", return_value=managed) as load_managed, ): result = runner.invoke(app, ["status"]) @@ -922,7 +940,7 @@ def test_status_hides_managed_config_box_when_none_present(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") with ( patch("ucode.cli.load_state", return_value=MINIMAL_STATE), - patch("ucode.cli.load_managed_state", return_value=None), + patch("ucode.cli.load_published_config", return_value=None), ): result = runner.invoke(app, ["status"]) @@ -1504,10 +1522,12 @@ def test_no_flag_calls_configure_all(self): patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command") as mock_cfg, + patch("ucode.cli.configure_mcp_command") as mock_mcp, ): result = runner.invoke(app, ["configure"]) assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with(prompt_optional_updates=True, offer_optional_setup=True) + mock_mcp.assert_not_called() def test_optional_setup_installs_ai_tools_and_configures_mcp(self): import ucode.cli as cli_mod @@ -1545,14 +1565,15 @@ def test_optional_setup_decline_does_nothing(self): mock_mcp.assert_not_called() def test_agents_flag_skips_mcp_prompt(self): - # Flag-driven (non-interactive) runs must stay scriptable: no MCP prompt. with ( patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command"), + patch("ucode.cli.configure_mcp_command") as mock_mcp, ): result = runner.invoke(app, ["configure", "--agents", "claude,codex"]) assert result.exit_code == 0, result.output + mock_mcp.assert_not_called() def test_agents_flag_calls_configure_with_tools(self): with ( @@ -1699,9 +1720,6 @@ def test_skip_upgrade_flag_disables_optional_update_prompt(self): patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command") as mock_cfg, - # Fully-interactive configure ends by offering the MCP step; decline it. - patch("ucode.cli.prompt_yes_no", return_value=False), - patch("ucode.cli.configure_mcp_command"), ): result = runner.invoke(app, ["configure", "--skip-upgrade"]) assert result.exit_code == 0, result.output @@ -1713,15 +1731,14 @@ def test_disable_databricks_ai_tools_forwards_false_and_skips_prompt(self): patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command") as mock_cfg, - # Fully-interactive configure ends by offering the MCP step; decline it. - patch("ucode.cli.prompt_yes_no", return_value=False), - patch("ucode.cli.configure_mcp_command"), + patch("ucode.cli.prompt_yes_no_default") as mock_prompt, ): result = runner.invoke(app, ["configure", "--disable-databricks-ai-tools"]) assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with( prompt_optional_updates=True, databricks_ai_tools_enabled=False ) + mock_prompt.assert_not_called() def test_enable_databricks_ai_tools_with_agents_forwards_true(self): with ( @@ -3015,7 +3032,7 @@ def test_disabled_reads_nothing_at_all(self, monkeypatch, env_value): monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False) else: monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", env_value) - for name in ("refresh_managed_config", "load_managed_state"): + for name in ("refresh_managed_config", "load_published_config"): monkeypatch.setattr( f"ucode.cli.{name}", lambda *a, called=name, **k: pytest.fail(f"{called} must not run when disabled"), @@ -3048,7 +3065,7 @@ def test_a_removed_model_list_no_longer_skips_discovery(self, monkeypatch): "enabled_agents": {"claude": {"model_config": {"models": {"default_opus_model": "m"}}}} } fresh = {"enabled_agents": {"claude": {"model_config": {}}}} - monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: stale_cache) + monkeypatch.setattr("ucode.cli.load_published_config", lambda ws: stale_cache) monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (fresh, False)) state = dict(MINIMAL_STATE) @@ -3068,286 +3085,423 @@ def test_a_removed_model_list_no_longer_skips_discovery(self, monkeypatch): assert mock_shared.call_args.kwargs["skip_model_discovery"] is False -class TestConfigureDeprecation: - """`ucode configure` resolves the target workspace first, authenticates, then branches on the - caller's role and whether the workspace already publishes a managed config (AIGTWY-4338).""" +class TestRunManagedConfigureFlow: + """`ucode configure` resolves the target workspace, applies any published managed config to this + machine, and offers authoring to admins — `_run_managed_configure_flow` (AIGTWY-4338).""" @pytest.fixture(autouse=True) def _stub_auth(self, monkeypatch): - # The resolver authenticates up front (before the config read and admin check); keep that a + # The flow authenticates up front (before the config fetch and admin check); keep that a # no-op so these tests never touch a real Databricks login. monkeypatch.setattr("ucode.cli.ensure_databricks_auth", lambda *a, **k: None) @staticmethod - def _resolve(entries=None): + def _run(entries=None, **kwargs): import ucode.cli as cli_mod - return cli_mod._resolve_workspace_then_maybe_reject(entries) + return cli_mod._run_managed_configure_flow(entries, **kwargs) @staticmethod - def _stub_admin(monkeypatch, value): - """Stub the best-effort admin check the has-config branch runs (token + SCIM).""" + def _stub_fetch(monkeypatch, published, *, feature_disabled=False): + """Stub the token fetch and `fetch_published_config` (published, reason, feature_disabled).""" + # A disabled feature always arrives as a read reason, never as a clean empty read. + reason = "FEATURE_DISABLED: managed agent config is off" if feature_disabled else None monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") - monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: value) + monkeypatch.setattr( + "ucode.cli.fetch_published_config", + lambda ws, tok: (published, reason, feature_disabled), + ) + + @staticmethod + def _stub_cached_published(monkeypatch, cached): + """Stub the on-disk published slot the flow falls back to when the fetch fails.""" + monkeypatch.setattr("ucode.cli.load_published_config", lambda ws: cached) @staticmethod - def _stub_setup(monkeypatch): - """Replace ``setup_command`` with a recorder returning exit code 0.""" - setup_calls: list[dict] = [] + def _record_apply(monkeypatch, applied=True): + """Replace `apply_managed_config_locally` with a recorder of its (managed, ws, profile). + + Returns `applied`, standing in for the real function's "did any agent apply" answer. + """ + calls: list[tuple] = [] monkeypatch.setattr( - "ucode.cli.setup_command", lambda **kwargs: setup_calls.append(kwargs) or 0 + "ucode.cli.apply_managed_config_locally", + lambda managed, ws, profile: calls.append((managed, ws, profile)) or applied, ) - return setup_calls + return calls - def test_non_admin_with_config_confirms_and_exits(self, monkeypatch, capsys): - # A non-admin on a managed workspace has nothing to configure locally — the launch path - # applies the config every run. Configure just shows what's in force and points at `ucode`; - # it must not route into setup or write anything. + @staticmethod + def _record_author(monkeypatch, code=0): + """Replace `author_managed_config` with a recorder returning `code`.""" + calls: list[dict] = [] + monkeypatch.setattr( + "ucode.cli.author_managed_config", lambda **kwargs: calls.append(kwargs) or code + ) + return calls + + def test_published_config_applies_and_exits_for_non_admin(self, monkeypatch, capsys): import typer monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) - monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: ({"enabled_agents": {"claude": {}}}, False), - ) - self._stub_admin(monkeypatch, False) + self._stub_fetch(monkeypatch, {"enabled_agents": {"claude": {}}}) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + applied = self._record_apply(monkeypatch) monkeypatch.setattr( - "ucode.cli.setup_command", - lambda **kwargs: pytest.fail("a non-admin must not be routed into setup"), + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("a non-admin must not be routed into authoring"), ) with pytest.raises(typer.Exit) as exc: - self._resolve([("https://w", None)]) + self._run([("https://w", None)]) assert exc.value.exit_code == 0 + assert applied == [({"enabled_agents": {"claude": {}}}, "https://w", None)] + assert "Configuration complete" in capsys.readouterr().out + + def test_no_agent_applied_warns_instead_of_reporting_completion(self, monkeypatch, capsys): + """OS-managed settings can block every enabled agent; the run must not then claim success.""" + import typer + + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + self._stub_fetch(monkeypatch, {"enabled_agents": {"claude": {}}}) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + self._record_apply(monkeypatch, applied=False) + with pytest.raises(typer.Exit): + self._run([("https://w", None)]) out = capsys.readouterr().out - assert "you're all set" in out - assert "Run `ucode`" in out + assert "Could not apply your workspace's managed config" in out + assert "Configuration complete" not in out - def test_fetches_the_config_rather_than_reading_a_cold_cache(self, monkeypatch): - # The gap this guards: on a fresh machine the local cache is empty until the first launch, - # so a cache read would miss a config the workspace does publish. The resolver must fetch. + def test_published_config_admin_declines_update(self, monkeypatch): import typer monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) - # Cold cache — a cache read would wrongly fall through to the local configure flow. - monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: None) + self._stub_fetch(monkeypatch, {"enabled_agents": {"claude": {}}}) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) + monkeypatch.setattr("ucode.cli._prompt_admin_update_workspace", lambda: False) + applied = self._record_apply(monkeypatch) monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: ({"enabled_agents": {"claude": {}}}, False), + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("declining must not author"), ) - self._stub_admin(monkeypatch, False) with pytest.raises(typer.Exit) as exc: - self._resolve([("https://w", None)]) + self._run([("https://w", None)]) assert exc.value.exit_code == 0 + assert len(applied) == 1 - def test_admin_with_config_runs_setup(self, monkeypatch): - # An admin on a workspace that already has a config is dropped into setup — whose - # existing-config menu offers re-author/delete — rather than the confirm-only path. - import typer + def test_admin_update_prompt_defaults_to_no_without_a_tty(self, monkeypatch): + import io + + import ucode.cli as cli_mod + + monkeypatch.setattr("sys.stdin", io.StringIO("")) + assert cli_mod._prompt_admin_update_workspace() is False + def test_scripted_admin_run_on_a_managed_workspace_exits_zero(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda *a, **k: None) monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + self._stub_fetch(monkeypatch, {"enabled_agents": {"claude": {}}}) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) + self._record_apply(monkeypatch) monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: ({"enabled_agents": {"claude": {}}}, False), + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("a non-interactive run must not enter authoring"), ) - self._stub_admin(monkeypatch, True) - setup_calls = self._stub_setup(monkeypatch) + result = runner.invoke(app, ["configure", "--workspaces", "https://w"], input="") + output = _strip_ansi(result.output) + assert result.exit_code == 0, output + assert "This machine now uses your workspace's managed config" in output + assert "Configuration complete" not in output + + def test_published_config_admin_accepts_update_authors_a_draft(self, monkeypatch): + import typer + + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + published = {"enabled_agents": {"claude": {}}} + self._stub_fetch(monkeypatch, published) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) + monkeypatch.setattr("ucode.cli._prompt_admin_update_workspace", lambda: True) + self._record_apply(monkeypatch) + author_calls = self._record_author(monkeypatch) with pytest.raises(typer.Exit) as exc: - self._resolve([("https://w", None)]) + self._run([("https://w", None)]) assert exc.value.exit_code == 0 - assert setup_calls == [ + assert author_calls == [ { "workspace": "https://w", "profile": None, - "command_label": "Configure Unity Gateway", "token": "tok", + "published": published, + "command_label": "ucode configure", } ] - def test_admin_status_unknown_with_config_confirms_without_setup(self, monkeypatch): - # `is_workspace_admin` returns None when the SCIM check fails; treat as non-admin and take - # the confirm path rather than routing an unverifiable caller into the admin-only setup flow. + def test_published_config_admin_update_applies_the_authored_draft(self, monkeypatch, capsys): + """An admin updating a managed workspace must end up on what they just authored. + + Otherwise the run leaves them on the published config with no way to try their own draft + short of publishing it. + """ import typer monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) - monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: ({"enabled_agents": {"claude": {}}}, False), - ) - self._stub_admin(monkeypatch, None) - monkeypatch.setattr( - "ucode.cli.setup_command", - lambda **kwargs: pytest.fail("unknown admin status must not route into setup"), - ) + published = {"enabled_agents": {"claude": {}}} + self._stub_fetch(monkeypatch, published) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) + monkeypatch.setattr("ucode.cli._prompt_admin_update_workspace", lambda: True) + self._record_author(monkeypatch) + draft = {"enabled_agents": {"codex": {}}} + monkeypatch.setattr("ucode.cli.load_draft_config", lambda ws: draft) + applied = self._record_apply(monkeypatch) with pytest.raises(typer.Exit) as exc: - self._resolve([("https://w", None)]) + self._run([("https://w", None)]) assert exc.value.exit_code == 0 + assert applied == [(published, "https://w", None), (draft, "https://w", None)] + out = capsys.readouterr().out + assert "Configuration complete — this machine now uses your authored config" in out + assert out.index("your authored config") < out.index("ucode publish") - def test_token_failure_with_config_confirms_without_checking_admin(self, monkeypatch): - # A token failure must not block the confirm path for a config the workspace does publish. + def test_admin_status_none_with_config_applies_without_authoring(self, monkeypatch): import typer monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) + self._stub_fetch(monkeypatch, {"enabled_agents": {"claude": {}}}) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: None) + self._record_apply(monkeypatch) monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: ({"enabled_agents": {"claude": {}}}, False), + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("unverified admin status must not author"), ) + with pytest.raises(typer.Exit) as exc: + self._run([("https://w", None)]) + assert exc.value.exit_code == 0 + + def test_token_failure_falls_through_to_local_configure(self, monkeypatch): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) def _boom(ws, profile=None): raise RuntimeError("no token") monkeypatch.setattr("ucode.cli.get_databricks_token", _boom) monkeypatch.setattr( - "ucode.cli.is_workspace_admin", - lambda ws, tok: pytest.fail("admin check needs a token"), - ) - monkeypatch.setattr( - "ucode.cli.setup_command", lambda **kwargs: pytest.fail("no setup without a token") + "ucode.cli.fetch_published_config", + lambda ws, tok: pytest.fail("no config fetch without a token"), ) - with pytest.raises(typer.Exit) as exc: - self._resolve([("https://w", None)]) - assert exc.value.exit_code == 0 - - @staticmethod - def _stub_not_admin(monkeypatch): - # No managed config -> the resolver now checks admin status; keep the developer case simple. - monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") - monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + entries = self._run([("https://w", None)]) + assert entries == [("https://w", None)] - def test_prompts_for_the_workspace_before_checking_the_config(self, monkeypatch): - # The whole point: even under a managed config the developer can still switch workspaces, - # so the prompt runs (and the picked workspace is made current) before the config check. + def test_prompts_for_the_workspace_before_fetching_the_config(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") picked = [] monkeypatch.setattr( "ucode.cli._prompt_for_configuration", lambda tool=None: ("https://picked", None) ) monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: picked.append(ws)) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, False)) - self._stub_not_admin(monkeypatch) - entries = self._resolve(None) + self._stub_fetch(monkeypatch, None) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + entries = self._run(None) assert picked == ["https://picked"] assert entries == [("https://picked", None)] - def test_returns_flag_entries_and_proceeds_without_a_managed_config(self, monkeypatch): - # Setting up a new workspace still goes through `ucode configure`, so hand the resolved - # workspace back to the caller instead of re-prompting. + def test_no_config_non_admin_falls_through(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, False)) - self._stub_not_admin(monkeypatch) - entries = self._resolve([("https://w", None)]) + self._stub_fetch(monkeypatch, None) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + monkeypatch.setattr( + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("a non-admin must not author"), + ) + entries = self._run([("https://w", None)]) assert entries == [("https://w", None)] - def test_admin_with_no_config_runs_setup_in_place(self, monkeypatch): - # `configure` is replacing `setup`: an admin on a config-less workspace is dropped straight - # into the setup authoring flow (reusing the resolved workspace/profile) and the command - # exits with setup's own status code — no prompt, no fall-through to the manual flow. + def test_no_config_admin_authors_and_applies_the_draft_locally(self, monkeypatch): import typer monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, False)) - monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") + self._stub_fetch(monkeypatch, None) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) + author_calls = self._record_author(monkeypatch) + draft = {"enabled_agents": {"codex": {}}} + monkeypatch.setattr("ucode.cli.load_draft_config", lambda ws: draft) + applied = self._record_apply(monkeypatch) + with pytest.raises(typer.Exit) as exc: + self._run([("https://w", None)]) + assert exc.value.exit_code == 0 + assert author_calls[0]["published"] is None + assert applied == [(draft, "https://w", None)] + + def test_authoring_no_selection_leaves_an_earlier_draft_alone(self, monkeypatch, capsys): + """`None` means nothing was authored, so a pre-existing draft is neither applied nor advised.""" + import typer + + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + self._stub_fetch(monkeypatch, None) monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) - setup_calls: list[dict] = [] + self._record_author(monkeypatch, code=None) monkeypatch.setattr( - "ucode.cli.setup_command", - lambda **kwargs: setup_calls.append(kwargs) or 0, + "ucode.cli.load_draft_config", lambda ws: {"enabled_agents": {"codex": {}}} ) + applied = self._record_apply(monkeypatch) with pytest.raises(typer.Exit) as exc: - self._resolve([("https://w", None)]) + self._run([("https://w", None)]) assert exc.value.exit_code == 0 - assert setup_calls == [ - { - "workspace": "https://w", - "profile": None, - "command_label": "Configure Unity Gateway", - "token": "tok", - } - ] + assert applied == [] + assert "ucode publish" not in capsys.readouterr().out + + def test_no_config_admin_with_local_options_falls_through(self, monkeypatch, capsys): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + self._stub_fetch(monkeypatch, None) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) + monkeypatch.setattr( + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("local options must not enter managed authoring"), + ) + + entries = self._run([("https://w", None)], allow_interactive_authoring=False) + + assert entries == [("https://w", None)] + assert "Using the requested local configuration options" in capsys.readouterr().out + + def test_no_config_admin_reports_completion_before_publish_advice(self, monkeypatch, capsys): + import typer + + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + self._stub_fetch(monkeypatch, None) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) + self._record_author(monkeypatch) + monkeypatch.setattr( + "ucode.cli.load_draft_config", lambda ws: {"enabled_agents": {"codex": {}}} + ) + self._record_apply(monkeypatch) + with pytest.raises(typer.Exit): + self._run([("https://w", None)]) + out = capsys.readouterr().out + assert "Configuration complete" in out and "ucode publish" in out + assert out.index("Configuration complete") < out.index("ucode publish") - def test_setup_failure_maps_to_a_nonzero_exit(self, monkeypatch): - # `setup_command` raises RuntimeError for actionable failures; the resolver surfaces it as a - # clean non-zero exit rather than letting it bubble as an unhandled error. + def test_read_failure_with_cached_config_applies_it(self, monkeypatch, capsys): import typer monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, False)) monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") + monkeypatch.setattr( + "ucode.cli.fetch_published_config", + lambda ws, tok: (None, "HTTP 500 internal error", False), + ) + cached = {"enabled_agents": {"claude": {}}} + monkeypatch.setattr("ucode.cli.load_published_config", lambda ws: cached) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + applied = self._record_apply(monkeypatch) + with pytest.raises(typer.Exit) as exc: + self._run([("https://w", None)]) + assert exc.value.exit_code == 0 + assert applied == [(cached, "https://w", None)] + assert "Could not read your workspace's managed config" in capsys.readouterr().out + + def test_read_failure_without_cache_falls_through_with_warning(self, monkeypatch, capsys): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") + monkeypatch.setattr( + "ucode.cli.fetch_published_config", + lambda ws, tok: (None, "HTTP 403 permission_denied", False), + ) + monkeypatch.setattr("ucode.cli.load_published_config", lambda ws: None) + monkeypatch.setattr( + "ucode.cli.is_workspace_admin", + lambda ws, tok: pytest.fail( + "no admin check when the read failed and nothing is cached" + ), + ) + entries = self._run([("https://w", None)]) + assert entries == [("https://w", None)] + assert "Could not read your workspace's managed config" in capsys.readouterr().out + + def test_authoring_failure_maps_to_a_nonzero_exit(self, monkeypatch): + import typer + + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) + self._stub_fetch(monkeypatch, None) monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) def _boom(**kwargs): raise RuntimeError("no agents available") - monkeypatch.setattr("ucode.cli.setup_command", _boom) + monkeypatch.setattr("ucode.cli.author_managed_config", _boom) with pytest.raises(typer.Exit) as exc: - self._resolve([("https://w", None)]) + self._run([("https://w", None)]) assert exc.value.exit_code == 1 - def test_non_admin_with_no_config_falls_through_without_setup(self, monkeypatch): + def test_feature_disabled_falls_through_without_publish_advice(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, False)) - monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") - monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + self._stub_fetch(monkeypatch, None, feature_disabled=True) + self._stub_cached_published(monkeypatch, None) monkeypatch.setattr( - "ucode.cli.setup_command", - lambda **kwargs: pytest.fail("a non-admin must not be routed into setup"), + "ucode.cli.is_workspace_admin", + lambda ws, tok: pytest.fail("no admin check once the feature is disabled"), + ) + monkeypatch.setattr( + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("nothing to author when the feature is disabled"), ) - entries = self._resolve([("https://w", None)]) + entries = self._run([("https://w", None)]) assert entries == [("https://w", None)] - def test_admin_check_failure_falls_through_without_setup(self, monkeypatch): - # `is_workspace_admin` returns None when the check itself fails; treat as "not an admin" so - # a developer is never blocked behind an authoring flow they can't complete. + def test_feature_disabled_still_applies_the_cached_config(self, monkeypatch, capsys): + """A machine already under a managed config keeps it, as the launch path does.""" + import typer + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, False)) - monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") - monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: None) + self._stub_fetch(monkeypatch, None, feature_disabled=True) + self._stub_cached_published(monkeypatch, {"enabled_agents": {"claude": {}}}) + applied = self._record_apply(monkeypatch) + monkeypatch.setattr( + "ucode.cli.is_workspace_admin", + lambda ws, tok: pytest.fail("no admin check once the feature is disabled"), + ) monkeypatch.setattr( - "ucode.cli.setup_command", - lambda **kwargs: pytest.fail("no setup when admin status is unverifiable"), + "ucode.cli.author_managed_config", + lambda **kwargs: pytest.fail("nothing to author when the feature is disabled"), ) - entries = self._resolve([("https://w", None)]) - assert entries == [("https://w", None)] + with pytest.raises(typer.Exit) as exc: + self._run([("https://w", None)]) + assert exc.value.exit_code == 0 + assert applied == [({"enabled_agents": {"claude": {}}}, "https://w", None)] + assert "ucode publish" not in capsys.readouterr().out - def test_feature_disabled_server_side_does_not_run_setup(self, monkeypatch): - # The coding-agent-configs feature isn't enabled server-side, so `ucode setup` can't publish. - # Even an admin just falls through to the normal configure flow. + def test_several_explicit_workspaces_go_straight_to_local_configure(self, monkeypatch): + """The managed flow resolves one workspace, so it must not silently skip the rest.""" monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") - monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, True)) - monkeypatch.setattr("ucode.cli.get_databricks_token", lambda ws, profile=None: "tok") - monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: True) monkeypatch.setattr( - "ucode.cli.setup_command", - lambda **kwargs: pytest.fail("setup can't publish when the feature is disabled"), + "ucode.cli.get_databricks_token", + lambda ws, profile=None: pytest.fail("must not resolve a workspace for a multi-run"), ) - entries = self._resolve([("https://w", None)]) - assert entries == [("https://w", None)] + entries = [("https://a", None), ("https://b", "prof")] + assert self._run(entries) == entries def test_configure_command_exits_zero_without_erroring(self, monkeypatch): # `typer.Exit(0)` subclasses RuntimeError, so the command's own RuntimeError handler must # not catch the clean exit and print `str(exc)` -> a bare, meaningless "ERROR 0". monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) - monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) - monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: ({"enabled_agents": {"claude": {}}}, False), - ) - self._stub_admin(monkeypatch, False) + self._stub_fetch(monkeypatch, {"enabled_agents": {"claude": {}}}) + monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda ws, tok: False) + self._record_apply(monkeypatch) with ( patch("ucode.cli.install_databricks_cli"), patch("ucode.cli._prompt_for_configuration", return_value=("https://w", None)), @@ -3363,14 +3517,14 @@ def test_passes_entries_through_when_the_env_var_is_off(self, monkeypatch, capsy else: monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", env_value) monkeypatch.setattr( - "ucode.cli.load_managed_state", - lambda ws: pytest.fail("must not read the config when disabled"), + "ucode.cli.fetch_published_config", + lambda ws, tok: pytest.fail("must not fetch the config when disabled"), ) monkeypatch.setattr( "ucode.cli._prompt_for_configuration", lambda tool=None: pytest.fail("must not prompt when disabled"), ) - assert self._resolve(None) is None + assert self._run(None) is None assert capsys.readouterr().out == "" @@ -3435,7 +3589,7 @@ def _run( else: monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (managed, False)) - monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: cached) + monkeypatch.setattr("ucode.cli.load_published_config", lambda ws: cached) monkeypatch.setattr("ucode.cli.get_databricks_token", lambda *a, **k: "tok") monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda *a, **k: is_admin) monkeypatch.setattr( @@ -3484,11 +3638,11 @@ def test_launch_banner_omits_default_agent_when_a_tier_overrides(self, monkeypat assert "launching OpenCode" in result.output assert "as the default agent" not in result.output - def test_admin_without_a_config_is_pointed_at_setup(self, monkeypatch): + def test_admin_without_a_config_is_pointed_at_configure(self, monkeypatch): result, launched = self._run(monkeypatch, managed=None, is_admin=True) assert result.exit_code == 0, result.output assert launched == [] - assert "ucode setup" in result.output + assert "ucode configure" in result.output def test_non_admin_without_a_config_is_told_to_ask(self, monkeypatch): result, launched = self._run(monkeypatch, managed=None, is_admin=False) @@ -3496,21 +3650,21 @@ def test_non_admin_without_a_config_is_told_to_ask(self, monkeypatch): assert launched == [] assert "Ask a workspace admin" in result.output - def test_admin_without_a_config_sees_no_setup_when_feature_disabled(self, monkeypatch): + def test_admin_without_a_config_sees_no_pointer_when_feature_disabled(self, monkeypatch): result, launched = self._run( monkeypatch, managed=None, is_admin=True, coding_agent_config_feature_disabled=True ) assert result.exit_code == 0, result.output assert launched == [] - assert "ucode setup" not in result.output + assert "ucode configure" not in result.output - def test_non_admin_without_a_config_sees_no_setup_when_feature_disabled(self, monkeypatch): + def test_non_admin_without_a_config_sees_no_pointer_when_feature_disabled(self, monkeypatch): result, launched = self._run( monkeypatch, managed=None, is_admin=False, coding_agent_config_feature_disabled=True ) assert result.exit_code == 0, result.output assert launched == [] - assert "ucode setup" not in result.output + assert "ucode configure" not in result.output def test_dry_run_uses_the_cache_and_does_not_fetch(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") @@ -3521,7 +3675,7 @@ def test_dry_run_uses_the_cache_and_does_not_fetch(self, monkeypatch): "ucode.cli.refresh_managed_config", lambda state: pytest.fail("--dry-run must not fetch"), ) - monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: self.MANAGED) + monkeypatch.setattr("ucode.cli.load_published_config", lambda ws: self.MANAGED) launched: list[tuple] = [] monkeypatch.setattr( "ucode.cli._launch_tool", lambda tool, ctx, **kw: launched.append((tool, kw)) @@ -3531,6 +3685,28 @@ def test_dry_run_uses_the_cache_and_does_not_fetch(self, monkeypatch): # The config bare `ucode` already read is handed down, so the launch path does not refetch. assert launched[0][1]["managed"] == self.MANAGED + def test_dry_run_with_an_empty_cache_prints_guidance_instead_of_crashing(self, monkeypatch): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda *a, **k: None) + monkeypatch.setattr("ucode.cli.apply_pat_environment", lambda *a, **k: None) + monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) + monkeypatch.setattr( + "ucode.cli.refresh_managed_config", + lambda state: pytest.fail("--dry-run must not fetch"), + ) + monkeypatch.setattr("ucode.cli.load_published_config", lambda ws: None) + guidance: list[tuple] = [] + monkeypatch.setattr( + "ucode.cli._print_no_managed_config_guidance", + lambda ws, profile: guidance.append((ws, profile)), + ) + monkeypatch.setattr( + "ucode.cli._launch_tool", lambda *a, **k: pytest.fail("nothing to launch") + ) + result = runner.invoke(app, ["--dry-run"]) + assert result.exit_code == 0, result.output + assert guidance == [("https://w", None)] + def test_skip_preflight_still_resolves_an_agent_from_the_managed_config(self, monkeypatch): # --skip-preflight is now only about auth/gateway re-validation, decoupled from managed # config, so bare `ucode --skip-preflight` still fetches the config and picks its agent. diff --git a/tests/test_databricks.py b/tests/test_databricks.py index d1420571..f8ab6f07 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2741,7 +2741,7 @@ def test_buckets_by_family(self, model_id, expected): class TestModelServicesCache: """A successful listing is memoized per workspace: several callers want different views of the - same paginated walk (bucketed families vs the raw Claude ids), so one `ucode setup` run would + same paginated walk (bucketed families vs the raw Claude ids), so one `ucode configure` run would otherwise page the whole catalog twice.""" @staticmethod @@ -2767,7 +2767,7 @@ def test_repeat_listings_hit_the_api_once(self, monkeypatch): assert calls["n"] == 1 def test_the_two_discovery_helpers_share_one_walk(self, monkeypatch): - # The duplicate spinner in `ucode setup`: `discover_model_services` and + # The duplicate spinner in `ucode configure`: `discover_model_services` and # `discover_claude_models_unbucketed` both page the same endpoint. calls: dict = {} db_mod.clear_model_services_cache() @@ -2818,7 +2818,7 @@ def failing(url, token): class TestModelProviderServicesCache: """The MPS listing is workspace-wide and filtered per agent afterwards, so one call serves every - agent — `ucode setup` used to re-list it once per MPS-capable agent.""" + agent — `ucode configure` used to re-list it once per MPS-capable agent.""" @staticmethod def _counting_listing(calls: dict): diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index e1bf9e89..0a4362ea 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -12,15 +12,14 @@ import ucode.databricks as db_mod import ucode.managed_config as mc_mod from ucode.managed_config import ( + fetch_published_config, get_managed_config, load_draft_config, - load_managed_state, load_published_config, managed_state_workspace, normalize_managed_config, refresh_managed_config, save_draft_config, - save_managed_state, save_published_config, ) from ucode.managed_setup import serialize_managed_config @@ -202,51 +201,57 @@ def _managed_path(self, tmp_path, monkeypatch): monkeypatch.setattr(mc_mod, "MANAGED_STATE_PATH", path) return path - def test_save_then_load_round_trips(self, _managed_path): + def test_published_save_then_load_round_trips(self, _managed_path): cfg = normalize_managed_config(RAW_MANIFEST) - save_managed_state("https://ws.example.com", cfg) - loaded = load_managed_state("https://ws.example.com") - assert loaded == cfg + save_published_config("https://ws.example.com", cfg) + assert load_published_config("https://ws.example.com") == cfg + + def test_draft_save_then_load_round_trips(self, _managed_path): + cfg = normalize_managed_config(RAW_MANIFEST) + save_draft_config("https://ws.example.com", cfg) + assert load_draft_config("https://ws.example.com") == cfg def test_saved_file_is_0600(self, _managed_path): - save_managed_state("https://ws.example.com", {"default_agent": "claude"}) + save_published_config("https://ws.example.com", {"default_agent": "claude"}) mode = stat.S_IMODE(os.stat(_managed_path).st_mode) # Owner-only read/write; no group/other bits. assert mode == 0o600 def test_load_ignores_other_workspace(self, _managed_path): - save_managed_state("https://ws-a.example.com", {"default_agent": "claude"}) - assert load_managed_state("https://ws-b.example.com") is None + save_published_config("https://ws-a.example.com", {"default_agent": "claude"}) + assert load_published_config("https://ws-b.example.com") is None def test_load_missing_returns_none(self, _managed_path): - assert load_managed_state("https://ws.example.com") is None + assert load_published_config("https://ws.example.com") is None + assert load_draft_config("https://ws.example.com") is None def test_load_none_workspace_returns_none(self, _managed_path): - assert load_managed_state(None) is None - - def test_empty_config_overwrites_a_previous_one(self, _managed_path): - # Saving an empty config is how "the admin removed it" is recorded: the stored config must - # be replaced, not left behind for the read-failure fallback to reapply. - save_managed_state("https://ws.example.com", {"default_agent": "claude"}) - save_managed_state("https://ws.example.com", {}) - assert load_managed_state("https://ws.example.com") == {} - - def test_workspace_is_stored_alongside_the_config(self, _managed_path): - # `ucode setup --show` reads this when local state carries no workspace yet, so the authored - # file can still be found and attributed on disk. - save_managed_state("https://ws.example.com", {"default_agent": "claude"}) + assert load_published_config(None) is None + assert load_draft_config(None) is None + + def test_empty_published_overwrites_a_previous_one(self, _managed_path): + save_published_config("https://ws.example.com", {"default_agent": "claude"}) + save_published_config("https://ws.example.com", {}) + assert load_published_config("https://ws.example.com") == {} + + def test_workspace_is_reported_when_exactly_one(self, _managed_path): + save_draft_config("https://ws.example.com", {"default_agent": "claude"}) assert managed_state_workspace() == "https://ws.example.com" + def test_workspace_is_none_when_several_recorded(self, _managed_path): + save_published_config("https://ws-a.example.com", {"default_agent": "claude"}) + save_published_config("https://ws-b.example.com", {"default_agent": "codex"}) + assert managed_state_workspace() is None + def test_workspace_is_none_when_absent(self, _managed_path): assert managed_state_workspace() is None def test_dry_run_writes_nothing(self, _managed_path, monkeypatch): # Under --dry-run the config writers print instead of touching disk, so a launch that - # dry-runs an admin's authored draft never overwrites it. - # Patch the flag itself rather than `is_dry_run`: the shared JSON writer reads the module - # global directly. + # dry-runs an admin's authored draft never overwrites it. Patch the flag itself rather than + # `is_dry_run`: the shared JSON writer reads the module global directly. monkeypatch.setattr(config_io_mod, "_dry_run", True) - save_managed_state("https://ws.example.com", {"default_agent": "claude"}) + save_published_config("https://ws.example.com", {"default_agent": "claude"}) assert not _managed_path.exists() def test_corrupt_file_reads_as_absent(self, _managed_path): @@ -254,15 +259,16 @@ def test_corrupt_file_reads_as_absent(self, _managed_path): # launch falls through rather than raising on JSON it can't parse. _managed_path.parent.mkdir(parents=True, exist_ok=True) _managed_path.write_text("{not json", encoding="utf-8") - assert load_managed_state("https://ws.example.com") is None + assert load_published_config("https://ws.example.com") is None + assert load_draft_config("https://ws.example.com") is None assert managed_state_workspace() is None def test_loaded_config_serializes_to_a_json_encodable_payload(self, _managed_path): # `ucode publish` POSTs the serialized config, so a manifest that survives a disk round-trip # must still serialize to something json.dumps accepts with no custom encoder. cfg = normalize_managed_config(RAW_MANIFEST) - save_managed_state("https://ws.example.com", cfg) - loaded = load_managed_state("https://ws.example.com") + save_draft_config("https://ws.example.com", cfg) + loaded = load_draft_config("https://ws.example.com") assert loaded is not None assert json.loads(json.dumps(serialize_managed_config(loaded))) @@ -412,13 +418,15 @@ def _stub_token(self, monkeypatch): def test_persists_and_returns_the_manifest(self, monkeypatch): saved: list[tuple] = [] monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (MANAGED, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: saved.append((ws, cfg))) + monkeypatch.setattr( + mc_mod, "save_published_config", lambda ws, cfg: saved.append((ws, cfg)) + ) assert refresh_managed_config(_state()) == (MANAGED, False) assert saved == [(WORKSPACE, MANAGED)] def test_no_managed_config_returns_none(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setattr(mc_mod, "save_published_config", lambda ws, cfg: None) result, _ = refresh_managed_config(_state()) assert result is None @@ -426,7 +434,7 @@ def test_read_failure_falls_back_to_the_persisted_config(self, monkeypatch): # The admin's last known policy beats no policy, so a failed fetch reuses what we saved. warnings: list[str] = [] monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: MANAGED) monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) assert refresh_managed_config(_state()) == (MANAGED, False) assert "HTTP 500" in warnings[0] @@ -436,7 +444,7 @@ def test_read_failure_without_persisted_config_is_silent(self, monkeypatch): # Nothing persisted means no managed config is in play, so an expired session shouldn't # produce a warning about a feature this developer doesn't use. monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: None) monkeypatch.setattr( mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") ) @@ -450,7 +458,7 @@ def boom(ws, profile): raise RuntimeError("no token") monkeypatch.setattr(mc_mod, "get_databricks_token", boom) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: MANAGED) monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) assert refresh_managed_config(_state()) == (MANAGED, False) assert "no token" in warnings[0] @@ -460,7 +468,7 @@ def boom(ws, profile): raise RuntimeError("no token") monkeypatch.setattr(mc_mod, "get_databricks_token", boom) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: None) monkeypatch.setattr( mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") ) @@ -472,7 +480,7 @@ def test_permission_denied_without_cache_is_silent(self, monkeypatch): # config in play and warning would be a false positive. denied = 'HTTP 403 Forbidden: {"error_code":"PERMISSION_DENIED"}' monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, denied)) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: None) monkeypatch.setattr( mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") ) @@ -485,10 +493,10 @@ def test_permission_denied_warns_and_keeps_the_cached_config(self, monkeypatch): warnings: list[str] = [] denied = 'HTTP 403 Forbidden: {"error_code":"PERMISSION_DENIED"}' monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, denied)) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: MANAGED) monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) monkeypatch.setattr( - mc_mod, "save_managed_state", lambda ws, cfg: pytest.fail("must not clear the cache") + mc_mod, "save_published_config", lambda ws, cfg: pytest.fail("must not clear the cache") ) assert refresh_managed_config(_state()) == (MANAGED, False) assert "not readable by you" in warnings[0] @@ -497,9 +505,9 @@ def test_no_config_on_the_server_does_not_use_a_stale_persisted_file(self, monke # A successful read saying "no config" means the admin removed it — that's authoritative, # so a previously persisted file must not resurrect the old policy. monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setattr(mc_mod, "save_published_config", lambda ws, cfg: None) monkeypatch.setattr( - mc_mod, "load_managed_state", lambda ws: pytest.fail("must not fall back") + mc_mod, "load_published_config", lambda ws: pytest.fail("must not fall back") ) result, _ = refresh_managed_config(_state()) assert result is None @@ -509,8 +517,10 @@ def test_no_config_on_the_server_clears_the_persisted_one(self, monkeypatch): # failed read would put a dead policy back into force. saved: list[tuple] = [] monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: saved.append((ws, cfg))) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr( + mc_mod, "save_published_config", lambda ws, cfg: saved.append((ws, cfg)) + ) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: None) result, _ = refresh_managed_config(_state()) assert result is None assert saved == [(WORKSPACE, {})] @@ -519,7 +529,7 @@ def test_empty_persisted_config_is_not_treated_as_a_fallback(self, monkeypatch): # The empty marker means "no admin policy", so a later failed read falls through to the # developer's own settings rather than reporting a managed config. monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: {}) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: {}) monkeypatch.setattr( mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") ) @@ -534,11 +544,9 @@ def test_no_workspace_is_a_noop(self, monkeypatch): assert result is None def test_feature_disabled_sets_flag_when_there_is_no_fallback(self, monkeypatch): - # The workspace hasn't enabled coding-agent-configs server-side, so `ucode setup` can't - # publish anything yet. The flag lets callers suppress the setup recommendation. reason = 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED"}' monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, reason)) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: None) monkeypatch.setattr(mc_mod, "print_warning", lambda msg: None) state = _state() result, flag = refresh_managed_config(state) @@ -550,7 +558,7 @@ def test_feature_disabled_with_a_fallback_does_not_set_the_flag(self, monkeypatc # feature-off flag is irrelevant and must not be set. reason = 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED"}' monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, reason)) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: MANAGED) monkeypatch.setattr(mc_mod, "print_warning", lambda msg: None) state = _state() result, flag = refresh_managed_config(state) @@ -559,7 +567,7 @@ def test_feature_disabled_with_a_fallback_does_not_set_the_flag(self, monkeypatc def test_transient_failure_does_not_set_the_flag(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) - monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr(mc_mod, "load_published_config", lambda ws: None) monkeypatch.setattr(mc_mod, "print_warning", lambda msg: None) state = _state() result, flag = refresh_managed_config(state) @@ -568,13 +576,49 @@ def test_transient_failure_does_not_set_the_flag(self, monkeypatch): def test_successful_no_config_clears_the_flag(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setattr(mc_mod, "save_published_config", lambda ws, cfg: None) state = _state() result, flag = refresh_managed_config(state) assert result is None assert flag is False +class TestFetchPublishedConfig: + """`ucode configure`'s fetch: persists the published slot and reports feature availability.""" + + @pytest.fixture(autouse=True) + def _managed_path(self, tmp_path, monkeypatch): + path = tmp_path / ".ucode" / "managed-state.json" + monkeypatch.setattr(mc_mod, "MANAGED_STATE_PATH", path) + return path + + def test_success_persists_published_and_reports_available(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (MANAGED, None)) + published, reason, feature_disabled = fetch_published_config(WORKSPACE, "tok") + assert (published, reason, feature_disabled) == (MANAGED, None, False) + assert load_published_config(WORKSPACE) == MANAGED + + def test_no_config_records_emptiness(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) + published, reason, feature_disabled = fetch_published_config(WORKSPACE, "tok") + assert (published, reason, feature_disabled) == (None, None, False) + assert load_published_config(WORKSPACE) == {} + + def test_feature_disabled_reported_even_with_a_stale_snapshot(self, monkeypatch): + save_published_config(WORKSPACE, MANAGED) + reason = 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED"}' + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, reason)) + published, out_reason, feature_disabled = fetch_published_config(WORKSPACE, "tok") + assert published is None + assert feature_disabled is True + + def test_leaves_the_draft_untouched(self, monkeypatch): + save_draft_config(WORKSPACE, {"default_agent": "claude"}) + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (MANAGED, None)) + fetch_published_config(WORKSPACE, "tok") + assert load_draft_config(WORKSPACE) == {"default_agent": "claude"} + + class TestGetModelRecommendation: """The budget recommendation read. Every response field is optional server-side.""" diff --git a/tests/test_managed_export.py b/tests/test_managed_export.py index a066b84d..c960db92 100644 --- a/tests/test_managed_export.py +++ b/tests/test_managed_export.py @@ -53,7 +53,7 @@ def _with_manifest(manifest: dict | None, workspace: str | None = WORKSPACE): """Patch the module's local reads so a test controls the source config without disk or network.""" with ( patch.object(export_mod, "load_state", return_value={"workspace": workspace}), - patch.object(export_mod, "load_managed_state", return_value=manifest), + patch.object(export_mod, "load_draft_config", return_value=manifest), ): yield @@ -90,9 +90,9 @@ def test_config_roundtrips_through_parser_and_validator(self): def test_no_config_is_actionable(self): with ( patch.object(export_mod, "load_state", return_value={}), - patch.object(export_mod, "load_managed_state", return_value=None), + patch.object(export_mod, "load_draft_config", return_value=None), ): - with pytest.raises(RuntimeError, match="No managed coding-agent config found"): + with pytest.raises(RuntimeError, match="No managed config draft found"): export_mod.build_export_payload() def test_invalid_config_is_rejected(self): @@ -102,7 +102,7 @@ def test_invalid_config_is_rejected(self): export_mod.build_export_payload() def test_falls_back_to_managed_state_workspace(self): - managed_config_mod.save_managed_state(WORKSPACE, FULL_MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, FULL_MANIFEST) with patch.object(export_mod, "load_state", return_value={}): payload = export_mod.build_export_payload() assert payload["default_agent"] == "CODING_AGENT_CLAUDE_CODE" @@ -242,7 +242,7 @@ def test_long_and_short_file_flags_both_write_the_file(self, tmp_path): def test_no_config_exits_nonzero(self): with ( patch.object(export_mod, "load_state", return_value={}), - patch.object(export_mod, "load_managed_state", return_value=None), + patch.object(export_mod, "load_draft_config", return_value=None), ): result = runner.invoke(app, ["export"]) assert result.exit_code == 1 diff --git a/tests/test_managed_publish.py b/tests/test_managed_publish.py index 83abc8eb..91f605ff 100644 --- a/tests/test_managed_publish.py +++ b/tests/test_managed_publish.py @@ -50,6 +50,12 @@ def test_expands_user_home(self, tmp_path, monkeypatch): (tmp_path / "config.json").write_text(json.dumps(_payload()), encoding="utf-8") assert load_publish_payload("~/config.json")["workspace"] == WORKSPACE + def test_no_draft_error_names_no_specific_command(self): + """The message is shared with `ucode export`, so it must not tell a publisher to export.""" + with pytest.raises(RuntimeError, match="No managed config draft found") as excinfo: + load_publish_payload(None) + assert "ucode export" not in str(excinfo.value) + def test_missing_file_is_actionable(self, tmp_path): with pytest.raises(RuntimeError, match="No config file"): load_publish_payload(str(tmp_path / "absent.json")) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 90060a35..34d8cf5f 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -1,4 +1,4 @@ -"""Tests for the interactive `ucode setup` flow and its CLI wiring. +"""Tests for the interactive `ucode configure` authoring flow and its CLI wiring. The wizard is mostly orchestration, so these focus on the parts where it can silently produce a wrong manifest: reading tracing/MCP/skills back out of ``state.json``, classifying MCP URLs into @@ -7,6 +7,7 @@ from __future__ import annotations +import contextlib import json from decimal import Decimal from unittest.mock import patch @@ -91,171 +92,13 @@ def test_unverifiable_check_warns_and_continues(self): assert warn.called -class TestExistingConfigHandling: - RICH_CONFIG = { - "name": "coding-agent-configs/abc", - "enabled_agents": {"claude": {}, "opencode": {}, "pi": {}}, - "mcp_servers": [{"name": "a", "type": "sql"}], - "skills": {"names": ["main.default"]}, - "tracing_table": "main.default.traces", - "budget_policy": {"display_name": "lillys_budget", "budget_id": "abc"}, - } - - def test_continue_when_no_config_exists(self): - # Nothing published, so there is no prompt — the wizard just proceeds. - with ( - patch.object(wizard, "get_managed_config", return_value=(None, None)), - patch.object(wizard, "prompt_for_selection") as select, - patch.object(wizard, "print_warning") as warn, - ): - # 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 - - def test_read_failure_continues_with_a_note(self): - # Can't check isn't the same as "there is one"; don't imply data loss or block the wizard. - with ( - patch.object(wizard, "get_managed_config", return_value=(None, "HTTP 403 Forbidden")), - patch.object(wizard, "prompt_for_selection") as select, - patch.object(wizard, "print_note") as note, - ): - assert wizard._handle_existing_config(WORKSPACE, "token") == (True, None) - assert not select.called - assert note.called - - def test_feature_disabled_blocks_setup_with_an_actionable_error(self): - # When the coding-agent-config APIs aren't enabled, the read fails with a FEATURE_DISABLED - # 404. Stop before authoring a draft that can never be published. - reason = ( - 'HTTP 404 Not Found: {"error_code":"FEATURE_DISABLED",' - '"message":"Coding agent config APIs are not enabled for this workspace."}' - ) - with ( - patch.object(wizard, "get_managed_config", return_value=(None, reason)), - patch.object(wizard, "prompt_for_selection") as select, - patch.object(wizard, "print_note") as note, - pytest.raises(RuntimeError) as exc_info, - ): - wizard._handle_existing_config(WORKSPACE, "token") - assert not select.called - message = str(exc_info.value) - assert message == wizard.CODING_AGENT_CONFIGS_DISABLED_MESSAGE - assert "`ucode configure`" in message - # The raw 404 / JSON body must not leak into the message. - assert "404" not in message - assert "FEATURE_DISABLED" not in message - assert not note.called - - def test_choosing_create_continues_authoring(self): - with ( - patch.object( - wizard, - "get_managed_config", - return_value=({"name": "x", "enabled_agents": {}}, None), - ), - patch.object(wizard, "prompt_for_selection", return_value="create"), - ): - # 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 - # admin should do, and `ucode setup show` prints the real thing for comparison. - with ( - patch.object(wizard, "get_managed_config", return_value=(self.RICH_CONFIG, None)), - patch.object(wizard, "prompt_for_selection", return_value="create"), - patch.object(wizard, "print_warning") as warn, - ): - wizard._handle_existing_config(WORKSPACE, "token") - message = warn.call_args[0][0] - assert "one config covers every agent" in message - for leaked in ("Claude Code", "OpenCode", "lillys_budget", "main.default"): - assert leaked not in message, leaked - - def test_choosing_delete_stops_and_deletes(self): - with ( - patch.object( - wizard, - "get_managed_config", - return_value=({"name": "cfg/1", "enabled_agents": {}}, None), - ), - patch.object(wizard, "prompt_for_selection", return_value="delete"), - patch.object(wizard, "prompt_yes_no_default", return_value=True), - patch.object(wizard, "delete_coding_agent_config", return_value=None) as delete, - ): - keep_going, _ = wizard._handle_existing_config(WORKSPACE, "token") - assert keep_going is False - delete.assert_called_once_with(WORKSPACE, "token", "cfg/1") - - def test_choosing_adopt_confirms_and_stops(self): - existing = {"name": "cfg/1", "enabled_agents": {"claude": {}}} - with ( - patch.object(wizard, "get_managed_config", return_value=(existing, None)), - patch.object(wizard, "prompt_for_selection", return_value="adopt"), - patch("ucode.cli._confirm_managed_config_applied") as confirm, - ): - # Adopting just confirms the config is in force (the launch path applies it) and stops - # the wizard — no re-authoring, no local writes. - keep_going, published = wizard._handle_existing_config(WORKSPACE, "token") - assert keep_going is False - assert published == existing - confirm.assert_called_once_with(existing, WORKSPACE) - - def test_delete_declined_leaves_config_intact(self): - with ( - patch.object( - wizard, - "get_managed_config", - return_value=({"name": "cfg/1", "enabled_agents": {}}, None), - ), - patch.object(wizard, "prompt_for_selection", return_value="delete"), - patch.object(wizard, "prompt_yes_no_default", return_value=False), - patch.object(wizard, "delete_coding_agent_config") as delete, - ): - # Still stops the wizard: the admin chose the delete path, not the author path. - keep_going, _ = wizard._handle_existing_config(WORKSPACE, "token") - assert keep_going is False - assert not delete.called - - def test_delete_failure_raises(self): - with ( - patch.object( - wizard, - "get_managed_config", - return_value=({"name": "cfg/1", "enabled_agents": {}}, None), - ), - patch.object(wizard, "prompt_for_selection", return_value="delete"), - patch.object(wizard, "prompt_yes_no_default", return_value=True), - patch.object(wizard, "delete_coding_agent_config", return_value="HTTP 500"), - pytest.raises(RuntimeError, match="Could not delete"), - ): - wizard._handle_existing_config(WORKSPACE, "token") - - def test_cancelling_the_picker_aborts(self): - with ( - patch.object( - wizard, - "get_managed_config", - return_value=({"name": "x", "enabled_agents": {}}, None), - ), - patch.object(wizard, "prompt_for_selection", return_value=None), - pytest.raises(KeyboardInterrupt), - ): - wizard._handle_existing_config(WORKSPACE, "token") - - class TestStepBanner: - """The step headers brand themselves to the invoking command, so a `ucode configure` run - doesn't show `ucode setup` headers.""" + """The step headers brand themselves to the invoking command; the default is `ucode configure`.""" - def test_defaults_to_ucode_setup(self): + def test_defaults_to_ucode_configure(self): with patch.object(wizard, "print_section") as section: wizard._step_banner(1, "Agents") - assert section.call_args.args[0].startswith("ucode setup · step 1 of ") + assert section.call_args.args[0].startswith("ucode configure · step 1 of ") def test_uses_the_command_label_when_given(self): with patch.object(wizard, "print_section") as section: @@ -263,46 +106,69 @@ def test_uses_the_command_label_when_given(self): assert section.call_args.args[0].startswith("ucode configure · step 2 of ") -class TestSetupCommandToken: - """A caller (e.g. `ucode configure`) can hand setup a token so its admin gate uses the same - identity as the routing decision, instead of fetching a second time.""" +class TestAuthorManagedConfig: + """`author_managed_config` is the guided agent/model picker behind `ucode configure`'s admin path. - def test_reuses_a_passed_token_and_skips_a_second_fetch(self): - seen: list[tuple[str, str]] = [] + The caller resolves/authenticates the workspace, determines admin, and fetches the published + snapshot, then hands them in — so this never re-checks admin or re-fetches the config. It only + ever saves the DRAFT; publishing is a separate explicit step. + """ + + def _run(self, *, picked=("codex",), published=None, models=None): + models = models or {"codex": {"default_model": "system.ai.gpt-5-6"}} with ( + patch("ucode.cli.configure_shared_state", return_value=dict(STATE)), + patch.object(wizard, "check_gateway_endpoint", return_value=True), + patch.object(wizard, "prompt_for_tools", return_value=list(picked)), + patch.object(wizard, "_select_provider_service", return_value=None), patch.object( - wizard, - "get_databricks_token", - side_effect=AssertionError("must not fetch a token when one was passed"), - ), - patch.object( - wizard, - "ensure_databricks_auth", - side_effect=AssertionError("must not re-authenticate when a token was passed"), + wizard, "_prompt_models_for_agent", side_effect=lambda tool, *_: models[tool] ), patch.object( - wizard, "_require_admin", side_effect=lambda ws, tok: seen.append((ws, tok)) + wizard, "prompt_for_selection", return_value=picked[0] if picked else None ), - # Stop right after the admin gate so the heavy discovery/picker path doesn't run. - patch.object(wizard, "_handle_existing_config", return_value=(False, None)), + patch.object(wizard, "prompt_yes_no_default", return_value=False), + patch.object(wizard, "publish_command", side_effect=AssertionError("must not publish")), ): - code = wizard.setup_command(workspace="https://w", profile=None, token="tok") + return wizard.author_managed_config( + workspace=WORKSPACE, profile=None, token="tok", published=published + ) + + def test_saves_a_draft_and_does_not_publish(self): + code = self._run() assert code == 0 - assert seen == [("https://w", "tok")] + draft = managed_config_mod.load_draft_config(WORKSPACE) + assert draft["default_agent"] == "codex" + assert draft["enabled_agents"]["codex"]["model_config"]["default_model"] == ( + "system.ai.gpt-5-6" + ) + assert managed_config_mod.load_published_config(WORKSPACE) is None - def test_fetches_its_own_token_when_none_passed(self): - seen: list[tuple[str, str]] = [] - with ( - patch.object(wizard, "ensure_databricks_auth", return_value=None), - patch.object(wizard, "get_databricks_token", return_value="fetched"), - patch.object( - wizard, "_require_admin", side_effect=lambda ws, tok: seen.append((ws, tok)) - ), - patch.object(wizard, "_handle_existing_config", return_value=(False, None)), - ): - code = wizard.setup_command(workspace="https://w", profile=None) + def test_no_agents_selected_saves_nothing(self): + """None, not 0, so the caller does not mistake a pre-existing draft for this run's result.""" + assert self._run(picked=()) is None + assert managed_config_mod.load_draft_config(WORKSPACE) is None + + def test_carries_forward_budget_from_published_when_no_draft(self): + published = { + "default_agent": "codex", + "enabled_agents": {"codex": {"model_config": {"default_model": "system.ai.gpt-5-6"}}}, + "budget_policy": { + "budget_id": BUDGET_ID, + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "codex", + "default_model": "system.ai.gpt-5-6", + } + ], + }, + } + code = self._run(published=published) assert code == 0 - assert seen == [("https://w", "fetched")] + assert managed_config_mod.load_draft_config(WORKSPACE)["budget_policy"]["budget_id"] == ( + BUDGET_ID + ) class TestModelPrompting: @@ -1362,7 +1228,7 @@ def test_malformed_targets_yield_nothing(self): class TestBudgetPolicy: def test_no_up_front_gate(self): - # Running `ucode setup spend-tiers` is the consent, so the flow asks no "set up a policy?" + # Running `ucode configure spend-tiers` 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 ( @@ -1732,7 +1598,10 @@ def test_summary_has_no_settings_scope_choice(self, capsys): assert "ucode-only" not in out -class TestSetupFromFile: +class TestConfigureFromFile: + """`configure_from_file` (behind `ucode configure --from-file`) loads a hand-written manifest as + the local DRAFT. Admin-only, and never publishes.""" + def _write(self, tmp_path, payload): path = tmp_path / "manifest.json" path.write_text(json.dumps(payload), encoding="utf-8") @@ -1746,61 +1615,60 @@ def _valid(self): }, } - def test_valid_manifest_is_saved(self, tmp_path): + @contextlib.contextmanager + def _admin(self, admin=True): + with ( + patch.object(wizard, "load_state", return_value=STATE), + patch.object(wizard, "ensure_databricks_auth", return_value=None), + patch.object(wizard, "get_databricks_token", return_value="tok"), + patch.object(wizard, "is_workspace_admin", return_value=admin), + ): + yield + + def test_valid_manifest_is_saved_as_a_draft(self, tmp_path): + path = self._write(tmp_path, self._valid()) + with self._admin(): + assert wizard.configure_from_file(str(path)) == 0 + assert managed_config_mod.load_draft_config(WORKSPACE) == self._valid() + assert managed_config_mod.load_published_config(WORKSPACE) is None + + def test_non_admin_is_rejected(self, tmp_path): path = self._write(tmp_path, self._valid()) - with patch.object(wizard, "load_state", return_value=STATE): - assert wizard.setup_from_file(str(path)) == 0 - assert managed_config_mod.load_managed_state(WORKSPACE) == self._valid() + with self._admin(admin=False), pytest.raises(RuntimeError, match="not an admin"): + wizard.configure_from_file(str(path)) + assert managed_config_mod.load_draft_config(WORKSPACE) is None + + def test_undetermined_admin_is_rejected(self, tmp_path): + path = self._write(tmp_path, self._valid()) + with self._admin(admin=None), pytest.raises(RuntimeError, match="admin-only"): + wizard.configure_from_file(str(path)) def test_invalid_manifest_returns_1_and_saves_nothing(self, tmp_path): path = self._write(tmp_path, {"enabled_agents": {"claude": {}}}) - with patch.object(wizard, "load_state", return_value=STATE): - assert wizard.setup_from_file(str(path)) == 1 - assert managed_config_mod.load_managed_state(WORKSPACE) is None + with self._admin(): + assert wizard.configure_from_file(str(path)) == 1 + assert managed_config_mod.load_draft_config(WORKSPACE) is None def test_missing_file_is_actionable(self, tmp_path): - with patch.object(wizard, "load_state", return_value=STATE): - with pytest.raises(RuntimeError, match="Could not read manifest file"): - wizard.setup_from_file(str(tmp_path / "nope.json")) + with pytest.raises(RuntimeError, match="Could not read manifest file"): + wizard.configure_from_file(str(tmp_path / "nope.json")) def test_malformed_json_names_the_line(self, tmp_path): path = tmp_path / "bad.json" path.write_text("{oops", encoding="utf-8") - with patch.object(wizard, "load_state", return_value=STATE): - with pytest.raises(RuntimeError, match="not valid JSON"): - wizard.setup_from_file(str(path)) + with pytest.raises(RuntimeError, match="not valid JSON"): + wizard.configure_from_file(str(path)) def test_non_object_json_is_rejected(self, tmp_path): path = self._write(tmp_path, ["not", "an", "object"]) - with patch.object(wizard, "load_state", return_value=STATE): - with pytest.raises(RuntimeError, match="must contain a JSON object"): - wizard.setup_from_file(str(path)) + with pytest.raises(RuntimeError, match="must contain a JSON object"): + wizard.configure_from_file(str(path)) def test_unconfigured_workspace_is_actionable(self, tmp_path): path = self._write(tmp_path, self._valid()) with patch.object(wizard, "load_state", return_value={}): with pytest.raises(RuntimeError, match="No workspace is configured"): - wizard.setup_from_file(str(path)) - - -class TestShowCommand: - def test_reports_nothing_when_unauthored(self): - with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): - assert wizard.show_command() == 0 - - def test_prints_the_publish_payload(self, capsys): - manifest = { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} - }, - } - managed_config_mod.save_managed_state(WORKSPACE, manifest) - with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): - assert wizard.show_command() == 0 - out = capsys.readouterr().out - # The proto enum spelling is what `publish` sends, so it must appear verbatim. - assert "CODING_AGENT_CLAUDE_CODE" in out + wizard.configure_from_file(str(path)) class TestSummaryPanel: @@ -2009,7 +1877,6 @@ 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"}}}, @@ -2017,7 +1884,7 @@ def fake_sel(prompt, options, **kwargs): class TestCarryForwardSections: - def test_all_optional_sections_survive_a_rerun(self): + def test_carried_sections_survive_a_rerun(self): previous = { **AGENTS_ONLY, "tracing_table": "main.default.traces", @@ -2080,27 +1947,35 @@ def test_budget_policy_naming_a_dropped_agent_is_left_out_with_a_warning(self): class TestNextSteps: def test_marks_configured_and_unconfigured_sections(self, capsys): - wizard._print_next_steps(dict(AGENTS_ONLY)) + manifest = { + **AGENTS_ONLY, + "budget_policy": { + "budget_id": BUDGET_ID, + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-8", + } + ], + }, + } + wizard.print_managed_next_steps(manifest) out = capsys.readouterr().out - assert "ucode setup spend-tiers" in out - assert "not configured" in out + assert "ucode configure spend-tiers" in out + assert "ucode setup" not in out assert "ucode publish" in out - wizard._print_next_steps({**AGENTS_ONLY, "budget_policy": {"budget_id": BUDGET_ID}}) - out = capsys.readouterr().out - assert "ucode setup spend-tiers" in out - assert "not configured" not 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) + wizard.print_managed_next_steps(AGENTS_ONLY) out = capsys.readouterr().out assert "Dry run" in out assert "ucode publish" not in out class TestSectionCommands: - """`ucode setup spend-tiers` — the one managed-config section command, strictly admin-only.""" + """`ucode configure spend-tiers` — the one managed-config section command, strictly admin-only.""" @staticmethod def _admin(**overrides): @@ -2128,26 +2003,32 @@ def _run(self, fn, *, admin_overrides=None, **patches): return fn() def test_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_budget_policy_command) + with pytest.raises(RuntimeError, match="ucode configure"): + self._run(wizard.configure_spend_tiers_command) def test_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_budget_policy_command) + managed_config_mod.save_draft_config(WORKSPACE, {}) + with pytest.raises(RuntimeError, match="ucode configure"): + self._run(wizard.configure_spend_tiers_command) def test_not_admin_raises(self): - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + managed_config_mod.save_draft_config(WORKSPACE, AGENTS_ONLY) with pytest.raises(RuntimeError, match="not an admin"): self._run( - wizard.setup_budget_policy_command, + wizard.configure_spend_tiers_command, admin_overrides={"is_workspace_admin": lambda *a, **k: False}, ) - def test_budget_policy_offers_only_the_manifests_agents(self): - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) + def test_undetermined_admin_raises(self): + managed_config_mod.save_draft_config(WORKSPACE, AGENTS_ONLY) + with pytest.raises(RuntimeError, match="admin-only"): + self._run( + wizard.configure_spend_tiers_command, + admin_overrides={"is_workspace_admin": lambda *a, **k: None}, + ) + + def test_seeds_from_published_when_no_draft(self): + managed_config_mod.save_published_config(WORKSPACE, AGENTS_ONLY) captured = {} def fake_prompt(workspace, token, enabled_agents, state, **kwargs): @@ -2155,11 +2036,24 @@ def fake_prompt(workspace, token, enabled_agents, state, **kwargs): 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) + code = self._run(wizard.configure_spend_tiers_command) assert code == 0 assert captured["agents"] == AGENTS_ONLY["enabled_agents"] - def test_budget_policy_none_leaves_existing_untouched(self): + def test_offers_only_the_manifests_agents(self): + managed_config_mod.save_draft_config(WORKSPACE, AGENTS_ONLY) + captured = {} + + def fake_prompt(workspace, token, enabled_agents, state, **kwargs): + captured["agents"] = enabled_agents + return None + + with patch.object(wizard, "_prompt_budget_policy", side_effect=fake_prompt): + code = self._run(wizard.configure_spend_tiers_command) + assert code == 0 + assert captured["agents"] == AGENTS_ONLY["enabled_agents"] + + def test_none_leaves_existing_untouched(self): # A transient budget-listing failure returns None; it must never delete a saved policy. seeded = { **AGENTS_ONLY, @@ -2174,29 +2068,16 @@ def test_budget_policy_none_leaves_existing_untouched(self): ], }, } - managed_config_mod.save_managed_state(WORKSPACE, seeded) + managed_config_mod.save_draft_config(WORKSPACE, seeded) with patch.object(wizard, "_prompt_budget_policy", return_value=None): - code = self._run(wizard.setup_budget_policy_command) + code = self._run(wizard.configure_spend_tiers_command) assert code == 0 assert ( - managed_config_mod.load_managed_state(WORKSPACE)["budget_policy"] + managed_config_mod.load_draft_config(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 spend-tiers", - "ucode setup show", - "ucode publish", - ): - assert command in out - - class TestPublishDiff: def test_lists_added_removed_and_changed(self, capsys): existing = { @@ -2282,11 +2163,11 @@ def _config_file(tmp_path, manifest, *, workspace=WORKSPACE, spec_version=1, **e def test_unauthored_config_is_an_actionable_error(self): with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): - with pytest.raises(RuntimeError, match="ucode setup"): + with pytest.raises(RuntimeError, match="ucode configure"): wizard.publish_command() def test_creates_when_no_config_exists(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) created = {} def fake_create(workspace, token, payload): @@ -2301,7 +2182,7 @@ def fake_create(workspace, token, payload): def test_updates_in_place_when_a_config_exists(self): # Delete-then-create would leave the workspace with no config if the create failed, so an # existing config must be PATCHed rather than replaced. - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) existing = {"name": "coding-agent-configs/abc", "enabled_agents": {"codex": {}}} updated = {} created = {"called": False} @@ -2328,7 +2209,7 @@ def fake_create(*a, **k): 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) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) # What's live is the manifest normalized the same way `publish` will send it. existing = { "name": "coding-agent-configs/abc", @@ -2348,9 +2229,25 @@ def fake_update(*a, **k): == 0 ) assert updated["called"] is False + assert managed_config_mod.load_published_config(WORKSPACE) == self.MANIFEST + + def test_publishing_refreshes_the_cached_published_config(self): + managed_config_mod.save_published_config(WORKSPACE, {"default_agent": "codex"}) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) + + assert self._run() == 0 + assert managed_config_mod.load_published_config(WORKSPACE) == self.MANIFEST + + def test_a_failed_publish_leaves_the_cached_published_config_alone(self): + managed_config_mod.save_published_config(WORKSPACE, {"default_agent": "codex"}) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) + + with pytest.raises(RuntimeError): + self._run(create_coding_agent_config=lambda *a, **k: (None, "HTTP 500 boom")) + assert managed_config_mod.load_published_config(WORKSPACE) == {"default_agent": "codex"} def test_invalid_manifest_is_not_published(self): - managed_config_mod.save_managed_state( + managed_config_mod.save_draft_config( WORKSPACE, {"default_agent": "codex", "enabled_agents": {"claude": {}}} ) created = {"called": False} @@ -2369,7 +2266,7 @@ def test_an_older_family_version_the_wizard_offered_still_publishes(self): # `publish` process used to reject a model it had just offered: # claude: model 'system.ai.claude-opus-4-1' is not available on this workspace. # `publish` re-fetches the full listing rather than trusting what `setup` left in state. - managed_config_mod.save_managed_state( + managed_config_mod.save_draft_config( WORKSPACE, { "default_agent": "claude", @@ -2407,7 +2304,7 @@ def fake_create(workspace, token, payload): def test_a_failed_inventory_fetch_does_not_block_publishing(self): # The re-fetch is best-effort: a transient listing failure must not turn into a refusal to # publish a manifest that validates against what state already knows. - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) published: dict = {} def fake_create(workspace, token, payload): @@ -2424,7 +2321,7 @@ def fake_create(workspace, token, payload): assert published def test_declining_the_prompt_publishes_nothing(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) created = {"called": False} def fake_create(*a, **k): @@ -2438,7 +2335,7 @@ def fake_create(*a, **k): assert created["called"] is False def test_yes_skips_the_prompt(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) def refuse(*a, **k): raise AssertionError("--yes must not prompt") @@ -2446,7 +2343,7 @@ def refuse(*a, **k): assert self._run(yes=True, prompt_yes_no_default=refuse) == 0 def test_non_admin_is_rejected_before_publishing(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) created = {"called": False} def fake_create(*a, **k): @@ -2461,7 +2358,7 @@ def fake_create(*a, **k): def test_unreadable_existing_config_refuses_to_publish(self): # Publishing without knowing whether a config exists risks silently overwriting one. - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) created = {"called": False} def fake_create(*a, **k): @@ -2476,7 +2373,7 @@ def fake_create(*a, **k): assert created["called"] is False def test_feature_disabled_read_uses_the_shared_blocking_message(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) created = {"called": False} def fake_create(*a, **k): @@ -2494,7 +2391,7 @@ def fake_create(*a, **k): assert created["called"] is False def test_existing_config_without_a_resource_name_is_an_error(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) + managed_config_mod.save_draft_config(WORKSPACE, self.MANIFEST) with pytest.raises(RuntimeError, match="resource name"): self._run(get_managed_config=lambda *a, **k: ({"enabled_agents": {}}, None)) @@ -2569,7 +2466,7 @@ def test_missing_file_is_actionable(self, tmp_path): self._run(file_path=str(tmp_path / "absent.json")) def test_no_file_publishes_a_hand_entered_custom_model(self): - managed_config_mod.save_managed_state( + managed_config_mod.save_draft_config( WORKSPACE, { "default_agent": "codex", @@ -2624,25 +2521,9 @@ def test_unknown_failure_still_surfaces_the_reason(self): class TestCliWiring: - def test_setup_is_registered(self): - result = runner.invoke(app, ["--help"]) - assert result.exit_code == 0 - assert "setup" in result.output - - def test_setup_help_lists_from_file(self): - # Assert on the declared option rather than the rendered help text: Rich ellipsizes option - # names to fit the terminal ("--fro…" below ~40 columns), and CI runners report no width, so - # grepping `--from-file` out of the output fails there while passing on a wide local one. - group = typer.main.get_command(app).commands["setup"] # type: ignore[attr-defined] - declared = {opt for param in group.params for opt in param.opts} - assert "--from-file" in declared + def test_setup_group_is_gone(self): result = runner.invoke(app, ["setup", "--help"]) - assert result.exit_code == 0 - - def test_setup_show_is_registered(self): - result = runner.invoke(app, ["setup", "--help"]) - assert result.exit_code == 0 - assert "show" in result.output + assert result.exit_code != 0 def test_publish_is_registered(self): result = runner.invoke(app, ["--help"]) @@ -2652,7 +2533,7 @@ def test_publish_is_registered(self): def test_publish_declares_yes_and_no_dry_run(self): # `--dry-run` was removed: publish always validates before publishing, so a separate # validate-only mode is redundant. Asserted on declared options rather than rendered help, - # which Rich ellipsizes at narrow widths (see test_setup_help_lists_from_file). + # which Rich ellipsizes at narrow widths. command = typer.main.get_command(app).commands["publish"] # type: ignore[attr-defined] declared = {opt for param in command.params for opt in param.opts} assert "--yes" in declared @@ -2684,8 +2565,8 @@ def test_publish_error_exits_nonzero_with_a_message(self): assert result.exit_code == 1 def test_successful_publish_exits_zero(self): - # Same trap as `setup`: `typer.Exit` subclasses RuntimeError, so raising it inside the - # command's try block would report success as "ERROR 0". + # `typer.Exit(0)` subclasses RuntimeError, so raising it inside the command's try block + # would report success as "ERROR 0". with ( patch("ucode.cli.install_databricks_cli"), patch.object(cli_mod, "publish_command", return_value=0), @@ -2694,69 +2575,91 @@ def test_successful_publish_exits_zero(self): assert result.exit_code == 0 assert "ERROR" not in result.output - def test_successful_setup_exits_zero(self): - # `typer.Exit` subclasses RuntimeError, so a success code must not be caught and reported - # as an error by the command's own RuntimeError handler. + def test_spend_tiers_is_registered_under_configure(self): + group = typer.main.get_command(app).commands["configure"] # type: ignore[attr-defined] + assert "spend-tiers" in group.commands # type: ignore[attr-defined] + + def test_spend_tiers_calls_the_command(self): with ( patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", return_value=0) as setup, + patch("ucode.cli.configure_spend_tiers_command", return_value=0) as fn, ): - result = runner.invoke(app, ["setup"]) + result = runner.invoke(app, ["configure", "spend-tiers"]) assert result.exit_code == 0 - assert setup.called + assert fn.called assert "ERROR" not in _out(result) - def test_nonzero_setup_propagates(self): + def test_spend_tiers_runtime_error_exits_1(self): with ( patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", return_value=1), + patch( + "ucode.cli.configure_spend_tiers_command", + side_effect=RuntimeError("not an admin"), + ), ): - result = runner.invoke(app, ["setup"]) + result = runner.invoke(app, ["configure", "spend-tiers"]) assert result.exit_code == 1 + assert "not an admin" in _out(result) - def test_runtime_error_is_reported_and_exits_1(self): + def test_spend_tiers_interrupt_exits_130(self): with ( patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", side_effect=RuntimeError("you are not an admin")), + patch("ucode.cli.configure_spend_tiers_command", side_effect=KeyboardInterrupt), ): - result = runner.invoke(app, ["setup"]) - assert result.exit_code == 1 - assert "not an admin" in _out(result) + result = runner.invoke(app, ["configure", "spend-tiers"]) + assert result.exit_code == 130 - def test_interrupt_exits_130(self): + def test_configure_from_file_is_forwarded(self): with ( patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", side_effect=KeyboardInterrupt), + patch("ucode.cli.configure_from_file", return_value=0) as fn, ): - result = runner.invoke(app, ["setup"]) - assert result.exit_code == 130 + runner.invoke(app, ["configure", "--from-file", "/tmp/x.json"]) + assert fn.call_args.args[0] == "/tmp/x.json" + + def test_configure_from_file_honors_dry_run(self, monkeypatch): + monkeypatch.setattr(config_io_mod, "_dry_run", False) + + def check_dry_run(_path): + assert config_io_mod.is_dry_run() + return 0 - def test_from_file_is_forwarded(self): with ( patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", return_value=0) as setup, + patch("ucode.cli.configure_from_file", side_effect=check_dry_run), ): - runner.invoke(app, ["setup", "--from-file", "/tmp/x.json"]) - assert setup.call_args.kwargs["from_file"] == "/tmp/x.json" + result = runner.invoke(app, ["configure", "--from-file", "/tmp/x.json", "--dry-run"]) - def test_show_exits_zero(self): - with patch("ucode.cli.show_command", return_value=0): - result = runner.invoke(app, ["setup", "show"]) - assert result.exit_code == 0 + assert result.exit_code == 0, _out(result) + + def test_configure_from_file_runtime_error_exits_1(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.configure_from_file", side_effect=RuntimeError("not an admin")), + ): + result = runner.invoke(app, ["configure", "--from-file", "/tmp/x.json"]) + assert result.exit_code == 1 + assert "not an admin" in _out(result) @pytest.mark.parametrize( - ("command", "target"), - [("spend-tiers", "setup_budget_policy_command")], + "extra", + [ + ["--workspaces", "https://other.example.com"], + ["--profiles", "OTHER"], + ["--agent", "claude"], + ["--tracing"], + ], ) - def test_section_subcommands_are_registered_and_called(self, command, target): + def test_configure_from_file_rejects_options_it_cannot_honor(self, extra): with ( - patch("ucode.cli.install_databricks_cli"), - patch(f"ucode.cli.{target}", return_value=0) as fn, + patch("ucode.cli.install_databricks_cli") as install, + patch("ucode.cli.configure_from_file") as fn, ): - result = runner.invoke(app, ["setup", command]) - assert result.exit_code == 0 - assert fn.called - assert "ERROR" not in _out(result) + result = runner.invoke(app, ["configure", "--from-file", "/tmp/x.json", *extra]) + assert result.exit_code == 2, _out(result) + assert "--from-file can't be combined with" in _out(result) + assert fn.call_count == 0 + assert install.call_count == 0 def test_skills_is_a_top_level_command(self): commands = typer.main.get_command(app).commands # type: ignore[attr-defined] @@ -2766,37 +2669,6 @@ def test_mcp_is_a_top_level_group(self): group = typer.main.get_command(app).commands["mcp"] # type: ignore[attr-defined] assert {"add", "remove", "web-search"} <= set(group.commands) # type: ignore[attr-defined] - 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_budget_policy_command", - side_effect=RuntimeError("run `ucode setup` first"), - ), - ): - result = runner.invoke(app, ["setup", "spend-tiers"]) - 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", "spend-tiers"]) - assert result.exit_code == 130 - def _out(result) -> str: """CliRunner output with stderr folded in, since print_err writes to a stderr console.""" diff --git a/tests/test_ui.py b/tests/test_ui.py index c1ef9980..7bb2d8ac 100644 --- a/tests/test_ui.py +++ b/tests/test_ui.py @@ -116,7 +116,7 @@ class TestClosedStdinAborts: """Ctrl-D must reach the CLI as an abort, not as a traceback.""" def test_percentage_without_a_default_raises_keyboard_interrupt(self): - # `ucode setup`'s tier prompt passes no default. EOFError has no handler above this call — + # `ucode configure spend-tiers`' tier prompt passes no default. EOFError has no handler above this call — # the setup command catches only RuntimeError and KeyboardInterrupt — so a bare EOFError # reached the admin as a raw traceback. with patch("ucode.ui.console.input", side_effect=EOFError): @@ -413,6 +413,41 @@ def test_keeps_duplicate_hosts_as_separate_rows(self, monkeypatch): host_choices = [c for c in choices if isinstance(getattr(c, "value", None), tuple)] assert [c.value for c in host_choices] == profiles + def test_preselects_the_row_matching_workspace_and_profile(self, monkeypatch): + profiles = [ + ("https://a.cloud.databricks.com", "alpha"), + ("https://b.cloud.databricks.com", "beta"), + ] + captured = self._capture_select(monkeypatch, answer=profiles[1]) + prompt_for_workspace( + "setup", profiles, preselect=("https://b.cloud.databricks.com", "beta") + ) + + default = captured["kwargs"]["default"] + assert default.value == profiles[1] + assert default in captured["choices"] + + def test_preselect_matches_on_workspace_when_the_profile_differs(self, monkeypatch): + profiles = [ + ("https://a.cloud.databricks.com", "alpha"), + ("https://b.cloud.databricks.com/", "beta"), + ("https://b.cloud.databricks.com", "beta-too"), + ] + captured = self._capture_select(monkeypatch, answer=profiles[1]) + prompt_for_workspace("setup", profiles, preselect=("https://b.cloud.databricks.com", None)) + + assert captured["kwargs"]["default"].value == profiles[1] + + def test_no_preselect_or_an_unknown_workspace_starts_at_the_top(self, monkeypatch): + profiles = [("https://a.cloud.databricks.com", "alpha")] + captured = self._capture_select(monkeypatch, answer=profiles[0]) + prompt_for_workspace("setup", profiles) + assert captured["kwargs"]["default"] is None + + captured = self._capture_select(monkeypatch, answer=profiles[0]) + prompt_for_workspace("setup", profiles, preselect=("https://gone.databricks.com", "p")) + assert captured["kwargs"]["default"] is None + def test_returns_normalized_url_with_profile(self, monkeypatch): # Picker handed back a URL with a trailing slash — normalize_workspace_url # should strip it before returning.