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.