From 7ea682a084c0c4bde13d3422e8648fce02a3fe75 Mon Sep 17 00:00:00 2001 From: David Liu Date: Tue, 8 Sep 2026 15:07:24 +0000 Subject: [PATCH] Remove `ug setup` and `ug publish`; ug is developer-only Managed coding-agent config authoring moves entirely off the CLI to the AI Gateway API and Unity Gateway UI. Remove the `ug setup` command group (setup, mcps, skills, spend-tiers, help, show) and `ug publish`, and delete their now-dead backing modules managed_wizard.py and managed_publish.py. `ug` is now purely developer-facing: it fetches and applies the published config, and `ug export` stays as a read-only dump. The bare-`ug` no-config guidance no longer tailors advice by admin status or names a removed command; it just notes local settings are used and a workspace admin can configure one (dropping the now-unused admin check). managed_setup.py stays for the developer path and export. Reword docstrings, guidance, and README references off the removed commands. Co-authored-by: Isaac --- README.md | 88 +- src/ucode/cli.py | 203 +-- src/ucode/managed_config.py | 12 +- src/ucode/managed_export.py | 13 +- src/ucode/managed_publish.py | 131 -- src/ucode/managed_resolve.py | 4 +- src/ucode/managed_wizard.py | 2103 ----------------------- tests/test_cli.py | 34 +- tests/test_managed_publish.py | 159 -- tests/test_managed_wizard.py | 2993 --------------------------------- 10 files changed, 43 insertions(+), 5697 deletions(-) delete mode 100644 src/ucode/managed_publish.py delete mode 100644 src/ucode/managed_wizard.py delete mode 100644 tests/test_managed_publish.py delete mode 100644 tests/test_managed_wizard.py diff --git a/README.md b/README.md index d51a9b03..e41eb257 100644 --- a/README.md +++ b/README.md @@ -246,78 +246,6 @@ ucode skill add --location main.default --skills my-skill,other-skill ucode skill add --skills main.default.my-skill,main.default.other-skill ``` -### Managed config for a workspace (admins) - -Author the coding config your developers pick up automatically, instead of asking each of them to -run `ug configure` by hand. Restricted to workspace admins. `ug setup help` prints the whole -sequence; the short version is one command for the agents and models, then a command per optional -section, then publish: - -```bash -ug setup # agents and models (start here) -ug setup mcps # managed MCP servers -ug setup skills # managed skills -ug setup spend-tiers # spend-based routing -ug publish # publish it to the workspace -``` - -`ug setup` walks through the agents to enable and which one bare `ug` 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 Unity Gateway. 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. - -The optional sections each edit their own part of the same config, so you can add an MCP server or -change a spend tier later without walking the whole flow. `ug setup skills --location -main.default,other.schema` skips the prompt. `ug setup spend-tiers` sets a tiered spend policy -that switches the default agent and model as the workspace burns through a budget. Each section -command also offers to publish right away, so you can apply changes incrementally; answering the -section prompts also runs the matching `ug configure` step, which does configure this machine. - -Everything is written to `~/.ucode/managed-state.json` — the one local managed-config file — which -`ug publish` publishes. Re-running `ug setup` keeps the MCP servers, skills, tracing table, and -tiered spend policy already authored, rather than clearing them; to drop one, edit the file and reload -it with `ug setup --from-file`. - -```bash -# Review the manifest and the exact payload `ug publish` would publish. -ug setup show - -# Skip the prompts and load a hand-written config instead (validated before saving). -ug setup --from-file ./managed-config.json -``` - -Once the manifest looks right, publish it: - -```bash -# Validate, show a diff against what's live, and ask before publishing. -ug publish - -# Publish without the confirmation prompt (for CI). -ug publish --yes - -# Publish a config file exported with `ug export` instead of the locally authored one. -ug publish -f ./managed-config.json -ug 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 ug authors is sent — but because `ug setup` carries the -other sections forward, a re-run no longer silently drops them. Developers pick the new config up on -their next ug run. - -With `-f`/`--file`, `publish` reads a config file produced by `ug export` and publishes it through -the same validation, diff, and confirmation flow. The file's `workspace` must match the configured -workspace (it can never redirect publication elsewhere) and its `spec_version` must be a supported -integer; server-owned fields (resource name, workspace ids, timestamps, user ids) and unknown fields -are rejected rather than silently dropped. - ### Exporting the config Any user (not only admins) can print the workspace's managed config as portable JSON with `ug @@ -325,8 +253,8 @@ export`. The output leads with the source `workspace` URL and a `spec_version` ( 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 `ug -publish -f ` consumes. +directory must already exist) while stdout stays empty. The exported file is the portable +`CodingAgentConfig` proto-JSON the AI Gateway API accepts. ```bash # Print the managed config as JSON. @@ -383,16 +311,6 @@ The output looks like: | `ug skill add --location main.default --mcp` | Add schemas to the skills MCP scope, keeping any already configured (additive; never replaces) | | `ug skill add --location main.default` | Download a schema's skills to disk without removing existing downloads | | `ug skill add --skills main.default.my-skill` | Download a named subset of skills (bare names need `--location`; fully-qualified names stand alone) | -| `ug setup` | Author the managed config's agents and models (workspace admins only) | -| `ug setup mcps` | Add or change the managed config's MCP servers | -| `ug setup skills [--location a.b,c.d]` | Add or change the managed config's skills | -| `ug setup spend-tiers` | Set the managed config's tiered spend routing policy | -| `ug setup help` | Walk through the whole setup sequence, marking what's already configured | -| `ug setup show` | Print the authored config and the payload `ug publish` would publish | -| `ug setup --from-file ` | Load a hand-written managed config instead of running the prompts | -| `ug publish` | Publish the authored managed config to the workspace, after a diff and confirmation (admins only) | -| `ug publish -f ` | Publish a config file exported with `ug export` instead of the locally authored one | -| `ug publish --yes` | Publish without the confirmation prompt | Databricks AI Tools are installed only by `ug configure`, never by `ug ` launches. Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ug configure` to @@ -413,7 +331,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 `ug setup` (admins) and refreshed from the workspace on launch | +| `~/.ucode/managed-state.json` | The managed config (published by an admin through the AI Gateway) refreshed from the workspace on launch | | `~/.ucode/managed-backups/` | Baseline backups for OS-managed files changed by ug | Existing files are backed up before being overwritten. `ug revert` restores backups. diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 02d0be74..ef60018b 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -58,7 +58,6 @@ get_databricks_token, install_databricks_cli, is_model_provider_feature_unavailable, - is_workspace_admin, list_profile_entries, list_tool_provider_services, normalize_workspace_url, @@ -90,15 +89,6 @@ recommended_agent, resolve_state, ) -from ucode.managed_wizard import ( - publish_command, - setup_budget_policy_command, - setup_command, - setup_help_command, - setup_mcp_command, - setup_skills_command, - show_command, -) from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -1051,12 +1041,6 @@ def revert() -> int: app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ug.") skill_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(skill_app, name="skill", help="Databricks Skills for 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 `ug setup help`.", -) def _version_callback(value: bool) -> None: @@ -2284,7 +2268,7 @@ def _launch_managed_default( ) return if not managed: - _print_no_managed_config_guidance(current, state.get("profile")) + _print_no_managed_config_guidance() return # The budget tier can move the org to a cheaper agent, so it outranks the config's # default_agent. Fetched here and handed to _launch_tool so it is read once per launch. @@ -2308,22 +2292,12 @@ def _launch_managed_default( ) -def _print_no_managed_config_guidance(workspace: str, profile: str | None) -> None: - """Tell an admin how to publish a config, and everyone else who to ask.""" - print_warning( - "No managed coding agent config was found for this workspace; using your local settings." +def _print_no_managed_config_guidance() -> None: + """Point the developer at per-user configure when no managed config is published.""" + print_note( + "No managed coding agent config is published for this workspace. Run `ug configure` to " + "set up your coding agents, then launch one with `ug ` (for example `ug claude`)." ) - try: - token = get_databricks_token(workspace, profile) - except RuntimeError: - return - 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 `ug setup`.") - else: - # None means the admin check itself failed; point at setup rather than a dead end. - print_note("Run `ug setup` to configure one for your workspace, then `ug publish`.") @app.command( @@ -3007,162 +2981,6 @@ 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 " - "ug's manifest shape) instead. Validated before it is saved.", - ), - ] = None, -) -> None: - """Choose the agents and models for your workspace's managed config (admins only). - - MCP servers, skills, and the tiered spend policy have their own commands — see `ug setup help`. - """ - 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) - 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("mcps") -def setup_mcp_cmd() -> None: - """Choose the MCP servers the managed config gives developers (admins only).""" - # Same `typer.Exit`/RuntimeError ordering trap as the `setup` callback above. - try: - install_databricks_cli() - code = setup_mcp_command() - except RuntimeError as exc: - print_err(str(exc)) - raise typer.Exit(1) from None - except KeyboardInterrupt: - print_err("Interrupted.") - raise typer.Exit(130) from None - if code: - raise typer.Exit(code) - - -@setup_app.command("skills") -def setup_skills_cmd( - location: Annotated[ - str | None, - typer.Option( - "--location", - help="Skill schemas to publish as `.` (comma-separated for several). " - "Skips the prompt.", - ), - ] = None, -) -> None: - """Choose the skills the managed config gives developers (admins only).""" - try: - install_databricks_cli() - # None means "prompt"; an explicit `--location` is parsed to the list to publish. - locations = None if location is None else _parse_skill_locations(location) - code = setup_skills_command(locations) - except RuntimeError as exc: - print_err(str(exc)) - raise typer.Exit(1) from None - except KeyboardInterrupt: - print_err("Interrupted.") - raise typer.Exit(130) from None - if code: - raise typer.Exit(code) - - -@setup_app.command("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 `ug 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[ - str | None, - typer.Option( - "--file", - "-f", - help="Publish a config file exported with `ug export` instead of the locally " - "authored config. Its `workspace` must match the configured workspace.", - ), - ] = None, - yes: Annotated[ - bool, - typer.Option("--yes", "-y", help="Publish without the confirmation prompt."), - ] = False, -) -> None: - """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: `ug setup` only ever writes a valid manifest, 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) - 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) - - @app.command("export") def export_cmd( file_path: Annotated[ @@ -3177,11 +2995,10 @@ def export_cmd( ) -> None: """Export this workspace's managed coding-agent config as portable JSON. - Serializes the local managed config to the external `CodingAgentConfig` format that - `ug publish -f ` consumes, with credentials and server-owned fields (resource name, - workspace id, timestamps, user ids) excluded. Any user can run it; it makes no network calls - and mutates no workspace or local state. Without --file the JSON is printed to stdout; - diagnostics and errors go to stderr. + Serializes the local managed config to the external `CodingAgentConfig` proto-JSON format, + with credentials and server-owned fields (resource name, workspace id, timestamps, user ids) + excluded. Any user can run it; it makes no network calls and mutates no workspace or local + state. Without --file the JSON is printed to stdout; diagnostics and errors go to stderr. """ from ucode.managed_export import export_command diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 74d1daf6..cd2c5912 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -6,15 +6,13 @@ - 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_managed_state` / :func:`load_managed_state`, which the launch path + uses to pull the published copy into the local 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. +The workspace is the source of truth: an admin authors the ``CodingAgentConfig`` through the AI +Gateway API or UI, and each launch pulls the published copy into ``managed-state.json``. ``ucode`` +only reads and applies it; it never authors or publishes. :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 diff --git a/src/ucode/managed_export.py b/src/ucode/managed_export.py index 670ff661..48b60325 100644 --- a/src/ucode/managed_export.py +++ b/src/ucode/managed_export.py @@ -1,10 +1,9 @@ """`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. +Reads the local managed config (the one file :mod:`ucode.managed_config` owns, populated when a +launch refreshes it from the workspace), validates and serializes it to the external proto-JSON +``CodingAgentConfig`` (prefixed with the source ``workspace`` and a ``spec_version`` envelope) 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 @@ -40,8 +39,8 @@ def build_export_payload() -> dict: manifest = load_managed_state(workspace) if not manifest: raise RuntimeError( - "No managed coding-agent config found locally. Run `ug setup` to author one, or run " - "`ug` against a workspace that publishes one, then re-run `ug export`." + "No managed coding-agent config found locally. Run `ug` against a workspace that " + "publishes one, then re-run `ug export`." ) errors = validate_manifest(manifest, None) if errors: diff --git a/src/ucode/managed_publish.py b/src/ucode/managed_publish.py deleted file mode 100644 index 33f5f623..00000000 --- a/src/ucode/managed_publish.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Input handling for ``ucode publish``: turn a config source into a publishable payload. - -``ucode publish`` publishes either the locally authored managed config (no ``-f``) or a config file -produced by ``ucode export`` (``-f ``). Both routes converge here: a source dict is validated -against the configured workspace and the ``spec_version`` envelope, canonicalized through the same -normalize/serialize path the rest of ucode uses, and returned as the internal manifest (for -validation, summaries, and diffs) plus the API payload (the canonical ``CodingAgentConfig`` with -``spec_version`` but without ``workspace``). - -Pure input handling: no auth, no admin check, no network, no publish. Invalid input raises -RuntimeError with an actionable message so ``publish`` fails before it touches the workspace. -""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import cast - -from ucode.managed_config import normalize_managed_config -from ucode.managed_export import EXPORT_SPEC_VERSION, build_export_payload -from ucode.managed_setup import serialize_managed_config -from ucode.ui import normalize_workspace_url - -_ENVELOPE_FIELDS = ("workspace", "spec_version") - - -def load_publish_payload(file_path: str | None) -> dict: - """Return the source config dict for ``ucode publish``. - - With no ``file_path`` the locally authored config is serialized in-process via - :func:`ucode.managed_export.build_export_payload` (no subprocess, no captured stdout). With a - ``file_path`` the file is read as UTF-8 JSON; a missing file, non-UTF-8 bytes, malformed JSON, or - a non-object root each raise RuntimeError with an actionable message. - """ - if file_path is None: - return build_export_payload() - - path = Path(file_path).expanduser() - try: - text = path.read_text(encoding="utf-8") - except FileNotFoundError: - raise RuntimeError( - f"No config file at {path}. Pass an existing file to `ucode publish -f`, or run " - "`ucode publish` with no file to publish the locally authored config." - ) from None - except UnicodeDecodeError as exc: - raise RuntimeError(f"{path} is not valid UTF-8 text: {exc}.") from exc - except OSError as exc: - raise RuntimeError(f"Could not read {path}: {exc}.") from exc - - try: - payload = json.loads(text) - except json.JSONDecodeError as exc: - raise RuntimeError(f"{path} is not valid JSON: {exc}.") from exc - if not isinstance(payload, dict): - raise RuntimeError( - f"{path} must contain a JSON object at the top level, not a {type(payload).__name__}." - ) - return payload - - -def parse_publish_payload(payload: object, workspace: str) -> tuple[dict, dict]: - """Validate and canonicalize a source config into ``(manifest, api_payload)``. - - ``manifest`` is ucode's internal normalized shape (for validation, summaries, and diffs). - ``api_payload`` is the canonical proto-JSON ``CodingAgentConfig`` carrying ``spec_version`` but - not ``workspace``, ready for the configuration API. - - Enforces, before any of it reaches the workspace: an object root; a top-level ``workspace`` that - normalizes to the configured one (the file can never redirect publication elsewhere); a - ``spec_version`` that is a JSON integer (not a boolean or float) equal to the supported version; - no server-owned ``name``; and no unknown or lossy fields (anything normalization would silently - drop is rejected instead). - """ - if not isinstance(payload, dict): - raise RuntimeError(f"The config must be a JSON object, not a {type(payload).__name__}.") - payload = cast("dict[str, object]", payload) - - file_workspace = payload.get("workspace") - if not isinstance(file_workspace, str) or not file_workspace.strip(): - raise RuntimeError( - 'The config is missing a top-level "workspace". Export it with `ucode export` so it ' - "records the workspace it belongs to." - ) - if normalize_workspace_url(file_workspace) != normalize_workspace_url(workspace): - raise RuntimeError( - f"The config is for {file_workspace}, but the configured workspace is {workspace}. " - "`ucode publish` only publishes to the configured workspace; re-run `ucode configure` " - "to switch workspaces." - ) - - spec_version = payload.get("spec_version") - if isinstance(spec_version, bool) or not isinstance(spec_version, int): - raise RuntimeError( - f'The config "spec_version" must be the integer {EXPORT_SPEC_VERSION}. Export it with ' - "`ucode export` to get a config this ucode can publish." - ) - if spec_version != EXPORT_SPEC_VERSION: - raise RuntimeError( - f"The config is spec_version {spec_version}, but this ucode publishes version " - f"{EXPORT_SPEC_VERSION}. Upgrade ucode, or re-export the config." - ) - - config = {key: value for key, value in payload.items() if key not in _ENVELOPE_FIELDS} - if "name" in config: - raise RuntimeError( - 'The config includes a server-owned "name" field. Remove it — the workspace assigns the ' - "resource name when the config is created." - ) - - manifest = normalize_managed_config(config) - canonical = serialize_managed_config(manifest) - canonical.pop("name", None) - if canonical != config: - dropped = sorted(key for key in config if key not in canonical) - altered = sorted( - key for key in config if key in canonical and config[key] != canonical[key] - ) - offending = ", ".join(dropped + altered) - raise RuntimeError( - f"The config has fields ucode does not recognize or cannot publish: {offending}. " - "Server-owned fields (workspace ids, timestamps, user ids) and unknown fields must be " - "removed. Re-export a clean config with `ucode export`." - ) - - api_payload = {"spec_version": spec_version, **canonical} - return manifest, api_payload - - -__all__ = ["load_publish_payload", "parse_publish_payload"] diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 947b174b..4752ff90 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -1,7 +1,7 @@ """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 -from the workspace at launch, both through :mod:`ucode.managed_config`) and the developer's own +The managed config (``~/.ucode/managed-state.json``, published by an admin through the AI Gateway +and refreshed from the workspace at launch 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 anything it leaves unset falls back to diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py deleted file mode 100644 index fa7ca124..00000000 --- a/src/ucode/managed_wizard.py +++ /dev/null @@ -1,2103 +0,0 @@ -"""Interactive `ug 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 ``ug 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: -``ug setup`` picks the agents and models, and ``ug setup mcps`` / ``skills`` / ``spend-tiers`` -each edit their own section of the same manifest. ``ug setup`` carries the other sections forward -untouched (:func:`_carry_forward_sections`), and ``ug setup help`` prints the whole sequence. - -Serialization, validation, and the per-agent model catalogs live in :mod:`ucode.managed_setup`; this -module is the interaction layer on top of them. Sub-flows an admin already knows — MCP, skills — are -delegated to the existing ``ug configure `` commands and their results read back out of -``state.json``, so there is exactly one picker per concern in the codebase. -""" - -from __future__ import annotations - -import json -from collections.abc import Callable -from decimal import Decimal -from pathlib import Path -from typing import cast - -from ucode import config_io -from ucode.agents import TOOL_SPECS, check_gateway_endpoint -from ucode.databricks import ( - 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, - has_cached_model_provider_services, - is_model_provider_feature_unavailable, - is_workspace_admin, - list_model_provider_services, - list_workspace_budgets, - map_claude_family_models, - model_service_exists, - service_usable_for_tool, - update_coding_agent_config, -) -from ucode.managed_config import ( - get_managed_config, - load_managed_state, - managed_state_workspace, - save_managed_state, -) -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, -) -from ucode.state import load_state -from ucode.ui import ( - console, - format_usd, - kv_line, - print_err, - print_heading, - print_note, - print_panel, - print_section, - print_success, - print_warning, - print_warning_panel, - prompt_for_multi_selection, - prompt_for_percentage, - prompt_for_selection, - prompt_for_text, - prompt_for_tools, - prompt_yes_no_default, - spinner, -) - -# Shown whenever the workspace's coding-agent-config APIs return FEATURE_DISABLED. -CODING_AGENT_CONFIGS_DISABLED_MESSAGE = ( - "Workspace-managed coding agent configuration is not available on this workspace. Use " - "`ug configure` to set up agents for individual users instead." -) - -BUDGET_POLICY_BLURB = ( - "As the workspace spends more of a budget, a tiered spend policy automatically switches " - "everyone's default agent and model to a cheaper one — for example Claude Code / Opus normally, " - "Claude Code / Sonnet once spend passes 80%, OpenCode / Kimi past 100%.\n\n" - "It only moves the default. Developers can still pick any model they have access to, and the " - "budget's own hard block is what actually caps spend." -) - -# Agents not offered in `ug setup`'s picker, even when the workspace serves their models. -# `ug gemini` still works as a launch target; it's just not part of the managed config authored -# here. Serialize/validate keep supporting it, so a `--from-file` manifest can still name it. -SETUP_EXCLUDED_AGENTS = frozenset({"gemini"}) - - -def _tracing_table_from_state(state: dict) -> str | None: - """The UC table `ug configure tracing` wired up, or None when tracing is off. - - ``configure tracing`` records the destination as ``uc_destination``; the managed config calls the - same thing ``tracing.table``. - """ - tracing = state.get("tracing") - if not isinstance(tracing, dict) or not tracing.get("enabled"): - return None - destination = tracing.get("uc_destination") - return destination if isinstance(destination, str) and destination else None - - -def _mcp_server_from_url(url: str) -> tuple[str, str] | None: - """Derive a managed-config ``(name, type)`` entry from a registered server's resolved URL. - - ``state.json`` stores each MCP server's resolved URL but not its type, while the managed config - stores ``{name, type}`` and lets the developer's ug rebuild the URL. So map the URL back to the - type *and* the identifier the ai-gateway ``McpServer.name`` field is meant to hold for that type - (a UC name for a UC service, a Genie space id for a genie space, a `.` for - vector-search / uc-functions, a connection name for external). Deriving ``name`` from the URL — - rather than reusing the local display slug — is what lets the developer's ug reconstruct the - URL on launch. Returns None for a URL that matches nothing reconstructable (e.g. an app's - off-workspace host), so those are skipped rather than published unusably. - """ - stripped = url.rstrip("/") - marker = "/ai-gateway/mcp-services/" - if marker in url: - # `.../mcp-services/..` — store the dash form the launch path expects. - service = url.split(marker, 1)[1].split("/", 1)[0] - return service.replace(".", "-"), "mcp-service" - for fragment, tag in ( - ("/api/2.0/mcp/external/", "external"), - ("/api/2.0/mcp/genie/", "genie-space"), - ): - if fragment in url: - # external -> connection name; genie -> space id. Both are the single trailing segment. - return url.split(fragment, 1)[1].split("/", 1)[0], tag - for fragment, tag in ( - ("/api/2.0/mcp/vector-search/", "vector-search"), - ("/api/2.0/mcp/functions/", "uc-functions"), - ): - if fragment in url: - # `...//` — store the `.` the launch path splits back. - rest = url.split(fragment, 1)[1].split("/") - if len(rest) >= 2 and rest[0] and rest[1]: - return f"{rest[0]}.{rest[1]}", tag - return None - if stripped.endswith("/api/2.0/mcp/sql"): - return "databricks-sql", "sql" - # Databricks apps are the residual case: an arbitrary app host with a /mcp suffix. Its host isn't - # reconstructable from the workspace + an id, so it can't be published to the managed config yet. - if stripped.endswith("/mcp"): - return None - return None - - -def _mcp_servers_from_state(state: dict) -> list[dict]: - """The registered MCP servers, as managed-config ``{name, type}`` entries. - - Skips the skills registry connection: skills are published under the manifest's own ``skills`` - field, so including its MCP entry would configure it twice. - """ - from ucode.mcp import SKILLS_MCP_KIND - - servers: list[dict] = [] - seen: set[str] = set() - for entry in state.get("mcp_servers") or []: - if not isinstance(entry, dict) or entry.get("kind") == SKILLS_MCP_KIND: - continue - name = entry.get("name") - url = entry.get("url") - if not isinstance(name, str) or not name or not isinstance(url, str): - continue - resolved = _mcp_server_from_url(url) - if resolved is None: - print_warning( - f"Skipping MCP server '{name}': ug can't publish it to a managed config " - f"(unrecognized or app-hosted URL: {url})." - ) - continue - config_name, tag = resolved - if config_name in seen: - continue - seen.add(config_name) - servers.append({"name": config_name, "type": tag}) - return servers - - -def _skill_names_from_state(state: dict) -> list[str]: - """Skill schemas registered on the skills MCP connection (``catalog.schema`` entries).""" - from ucode.mcp import _skill_mcp_locations - - return [name for name in _skill_mcp_locations(state) if isinstance(name, str) and name] - - -def provider_service_model_options(service: dict) -> list[str]: - """Model ids an admin can pick from a provider service, or [] when they can't be enumerated. - - A service's ``config.targets`` names the provider-side models it exposes, which is exactly the - vocabulary the manifest's ``default_model`` must use when ``model_provider_service`` is set. Two - cases yield nothing to pick from, and the caller falls back to free-text: - - - ``allow_all_targets`` — the service passes through the provider's whole catalog, which ucode - cannot enumerate (there is no list-models call for a provider service). - - no targets at all — e.g. a relayed Anthropic subscription service, which routes by canonical - model name rather than by an explicit target list. - """ - if service.get("allow_all_targets"): - return [] - targets = service.get("targets") - if not isinstance(targets, list): - return [] - return sorted({t for t in targets if isinstance(t, str) and t}) - - -def _select_provider_service(tool: str, workspace: str, token: str) -> dict | None: - """Offer Databricks-hosted vs an external Model Provider Service for ``tool``. - - Returns the chosen service dict (as :func:`list_model_provider_services` shapes it), or None to - stay on Databricks-hosted models. The whole dict is returned rather than just the name so the - model prompt can offer the service's ``targets`` instead of asking the admin to type a model id - from memory. - - Only claude and codex can route through a provider service today; every other agent short-cuts to - Databricks. Mirrors `cli._maybe_select_provider_service`, but returns the choice instead of - persisting it — the wizard is authoring a manifest, not configuring this machine. - """ - if not any( - supports_provider_service(tool, provider_type) - for provider_type in ("anthropic", "amazon_bedrock", "openai") - ): - return None - - display = TOOL_SPECS[tool]["display"] - # The listing is memoized per workspace, so only the first agent's call does any I/O. That one - # takes over a second and deserves a spinner; the rest are instant, and spinning once per agent - # made the wizard look like it re-listed the services every time. - if has_cached_model_provider_services(workspace): - services, reason = list_model_provider_services(workspace, token) - else: - with spinner("Checking for model provider services..."): - services, reason = list_model_provider_services(workspace, token) - if reason is not None: - # A workspace without the feature enabled is the common case and not worth a warning; any - # other failure is worth surfacing, or the admin silently loses the MPS option and has no - # idea why. Mirrors `cli._maybe_select_provider_service`. - if not is_model_provider_feature_unavailable(reason): - print_warning(f"Could not list model provider services: {reason}") - print_note("Falling back to Databricks-hosted models.") - return None - - usable = [service for service in services if service_usable_for_tool(tool, service)] - if tool == "claude": - # Claude subscription relays are not reliable enough for managed configurations yet. - usable = [service for service in usable if not service.get("relayed")] - if not usable: - if services: - # Services exist but none match this agent's dialect — say so, since "no picker appeared" - # is otherwise indistinguishable from the feature being off. - print_note( - f"No model provider service matches {display}'s API dialect " - f"({len(services)} found on this workspace); using Databricks-hosted models." - ) - return None - - choice = prompt_for_selection( - f"How should {display} get its models?", - [ - ("databricks", "Databricks Hosted"), - ("mps", "External Models (Model Provider Service)"), - ], - ) - if choice != "mps": - return None - selected = prompt_for_selection( - f"Select the model provider service for {display}:", - [(service["name"], service["name"]) for service in usable], - searchable=True, - ) - if not selected: - return None - service = next(service for service in usable if service["name"] == selected) - _warn_if_mps_not_broadly_accessible(workspace, token, service["name"]) - return service - - -def _warn_if_mps_not_broadly_accessible(workspace: str, token: str, service_name: str) -> None: - """Warn if the picked MPS's schema isn't granted to all workspace users. - - A developer who pulls a config routing through this MPS needs USE_SCHEMA on its schema, or they - hit "User does not have USE_SCHEMA on Schema ." at launch. This only warns - (never blocks): access may instead come from a team group the check can't see, and an - inconclusive check stays silent. - """ - schema = ".".join(service_name.split(".")[:2]) - if schema.count(".") != 1: - return - with spinner("Checking who can use this service..."): - accessible = all_users_can_use_schema(workspace, token, schema) - if accessible is False: - print_warning( - f"All workspace users don't appear to have USE_SCHEMA on `{schema}`, so developers " - f"who pull this config may not be able to use `{service_name}`. Grant USE_SCHEMA on " - f"`{schema}` to the `account users` group (or the teams that need it) in Unity Catalog." - ) - - -def _prompt_models_for_agent(tool: str, state: dict, provider_service: dict | None) -> dict: - """Build one agent's ``model_config``. Every agent ends up with a ``default_model``. - - Databricks-hosted agents pick from the workspace's discovered models, filtered to the families - that agent can actually serve. Provider-service agents pick from the service's own ``targets``, - falling back to free-text only when those can't be enumerated (``allow_all_targets``, or a - relayed service that routes by canonical name). - - An empty selection is re-prompted rather than accepted: an agent with no ``default_model`` cannot - be the config's ``default_agent`` (the server rejects it) and gives developers nothing to launch, - so "none" is never a useful answer here. Ctrl-C still aborts the whole flow. - - Model ids are stored bare (e.g. ``system.ai.claude-opus-4-8``), not provider-prefixed: each - agent's own writer adds whatever prefix its config format needs (see - ``opencode._resolve_model_selector``), which keeps the manifest agent-neutral. - - Codex takes a single model (the harness selects one); Claude's picks are bucketed into - ``ClaudeDefaultModels`` family slots; the rest keep a flat list plus their chosen default. - """ - display = TOOL_SPECS[tool]["display"] - model_config: dict = {} - if provider_service: - service_name = provider_service["name"] - model_config["model_provider_service"] = service_name - targets = provider_service_model_options(provider_service) - if tool == "claude" and _pins_family_models(targets): - # `targets` (not the raw service) publishes explicit Claude models, and `render_overlay` - # pins each family to a chosen version from them — Bedrock slugs - # (`us.anthropic.claude-opus-4-8-v1:0`) or canonical Anthropic ids (`claude-opus-4-8`). - # So Claude needs a default *per family*, not one overall, mirroring the Databricks-hosted - # path. (A service with no enumerable targets pins nothing and takes a single default — - # handled below.) Keyed on `targets`, the same list the prompt consumes, so the decision - # and the prompt can't disagree — `allow_all_targets` zeroes `targets`, so it correctly - # falls through to the single-default branch even if the raw service also lists Claude. - model_config.update(_prompt_claude_provider_family_models(targets, service_name)) - elif targets: - model_config["default_model"] = _require_selection( - f"Default model for {display} (from {service_name}):", - [(target, target) for target in targets], - ) - else: - # No enumerable target list: the service either passes through the provider's whole - # catalog or routes by canonical model name, so the admin has to name the model. - print_note( - f"{service_name} does not publish an explicit model list, so enter the model id " - "as the provider names it (e.g. claude-sonnet-4-6)." - ) - model_config["default_model"] = _require_text(f"Default model for {display}") - return model_config - - if tool == "claude": - return _prompt_claude_models(state) - - options = model_options_for_agent(tool, state) - if not options: - print_warning(f"No models were discovered for {display} on this workspace.") - return {"default_model": _require_text(f"Default model for {display}")} - - custom: list[str] = [] - if tool in SINGLE_MODEL_AGENTS: - model = _select_hosted_model( - f"Select the default model for {display}:", options, state, custom - ) - single: dict = {"default_model": model} - if custom: - single["custom_models"] = custom - return single - - # 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 - # 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: - model_config["default_model"] = picked[0] - else: - model_config["default_model"] = _require_selection( - f"Default model for {display}:", [(model, model) for model in picked] - ) - - model_config["models"] = picked - if custom: - model_config["custom_models"] = list(dict.fromkeys(custom)) - return model_config - - -# Agents that get a single model rather than a multi-select. Codex's proto has no model list at all. -# Gemini and Copilot do declare `repeated string models`, but their config writers take one model -# (`gemini.write_tool_config(state, model)` / `copilot.write_tool_config(state, model)`) and write a -# single env var — so a published list would be read by nothing. Offering one keeps the manifest -# honest about what ug can apply; widen this when those writers grow a picker. -SINGLE_MODEL_AGENTS = frozenset({"codex", "gemini", "copilot"}) - -# Skip sentinel for a Claude family prompt. Every `ClaudeDefaultModels` slot is optional, and an -# unset one falls back to `default_model`, so leaving a family out is a legitimate choice. -_SKIP_FAMILY = "__skip__" - - -def _confirm_agent(tool: str, agent_config: dict) -> None: - """One consistent closing line per agent in step 2, whatever its model shape. - - Every agent — a single-model codex, a multi-model opencode, a family-slotted claude — ends its - block with the same `✔ configured — · ` line, so the step reads as a - uniform checklist rather than each agent's picker trailing off differently. - """ - display = TOOL_SPECS.get(tool, {}).get("display", tool) - model_config = agent_config.get("model_config") or {} - detail = model_config.get("default_model") or "no model" - provider = model_config.get("model_provider_service") - if provider: - detail = f"{detail} via {provider}" - print_success(f"{display} configured — {detail}") - - -def _render_family_slots(slots: dict[str, str]) -> None: - """Recap the Claude family → model slots just chosen, before the overall-default question. - - Same "form filling in" motif as :func:`_selected_recap`: the per-family answers scrolled by one at - a time, so gathering them into one box makes "which of these is the overall default?" a choice - over something the admin can see rather than recall. - """ - lines = [ - kv_line(slot.removeprefix("default_").removesuffix("_model"), model) - for slot, model in slots.items() - ] - print_panel("Claude Code models", lines) - - -def _prompt_claude_models(state: dict) -> dict: - """Build Claude's ``model_config`` one family slot at a time. - - Claude Code addresses models by family alias, not from a list, so the config is a set of slots: - `default_opus_model`, `default_sonnet_model`, `default_haiku_model`, `default_fable_model`. A flat - multi-select can't express that — and because `state["claude_models"]` holds only the newest id - per family, it could only ever offer one model per family anyway. Asking per family surfaces the - alternatives (six opus versions on a typical workspace, not one) and matches the proto. - - Each family may be skipped; the overall `default_model` is then chosen from the slots that were - filled, so it can never name a model the config doesn't carry. - """ - display = TOOL_SPECS["claude"]["display"] - # No spinner: the model-services listing is already cached by the time the flow reaches here - # (`configure_shared_state` walked it up front), so this is a filter over data in hand, not a - # fetch. Showing "Fetching Claude models..." made the wizard look like it listed the catalog - # twice. - candidates = _claude_candidates(state) - if not candidates: - print_warning(f"No Claude models were discovered for {display} on this workspace.") - return {"default_model": _require_text(f"Default model for {display}")} - - print_note( - "Claude Code picks a model by family, so set a default per family. Skip any family you " - "don't want configured — it falls back to the overall default." - ) - slots: dict[str, str] = {} - custom: list[str] = [] - for family in ANTHROPIC_FAMILIES: - family_models = candidates.get(family) - if not family_models: - continue - rows = [(model, model) for model in family_models] + [ - (_CUSTOM_MODEL, _CUSTOM_MODEL_LABEL), - (_SKIP_FAMILY, f"(skip {family})"), - ] - choice = prompt_for_selection(f"Default {family} model:", rows, searchable=True) - if choice is None: - raise KeyboardInterrupt - if choice == _SKIP_FAMILY: - continue - if choice == _CUSTOM_MODEL: - choice = _prompt_custom_model(state) - custom.append(choice) - slots[CLAUDE_SLOT_FOR_FAMILY[family]] = choice - - if not slots: - # Every slot skipped is a legitimate, minimal config: the proto leaves `models` optional and - # each unset slot falls back to `default_model`, so one model covers every family. Pick it - # from the same candidates rather than asking the admin to type an id. - print_note(f"No families configured, so {display} will use a single model for all of them.") - every_model = list(dict.fromkeys(m for fm in candidates.values() for m in fm)) - fallback_custom: list[str] = [] - model = _select_hosted_model( - f"Which model should {display} use?", every_model, state, fallback_custom - ) - single: dict = {"default_model": model} - if fallback_custom: - single["custom_models"] = fallback_custom - return single - - chosen = list(dict.fromkeys(slots.values())) - model_config: dict = {"models": slots} - if len(chosen) == 1: - # A one-option prompt is a wasted keystroke, but skipping it silently reads as a dropped - # step — say what was inferred so the admin knows the default is set, and to what. - model_config["default_model"] = chosen[0] - print_note(f"Only one model configured, so it's {display}'s overall default.") - else: - _render_family_slots(slots) - model_config["default_model"] = _require_selection( - f"Which of those is {display}'s overall default?", [(m, m) for m in chosen] - ) - if custom: - model_config["custom_models"] = list(dict.fromkeys(custom)) - return model_config - - -def _pins_family_models(targets: list[str]) -> bool: - """True when Claude behind a service is pinned per family at launch. - - Keyed on the *behavior*, not the vendor: when a service publishes explicit Claude targets — a - Bedrock provider-side slug like ``us.anthropic.claude-opus-4-8-v1:0`` *or* a canonical Anthropic - id like ``claude-opus-4-8`` — ``render_overlay`` pins each ``ANTHROPIC_DEFAULT__MODEL`` to - a chosen version, so the wizard prompts one model per family, mirroring the Databricks-hosted - path. No Claude-family targets (``allow_all_targets`` zeroes the list, or a relayed subscription - lists none) means nothing is pinned — Claude Code's canonical names route fine — so it takes a - single default. - - Takes the *enumerated* targets (``provider_service_model_options`` output), the exact list the - per-family prompt consumes, so the decision and the prompt can't disagree. Reading the raw - ``service["targets"]`` here would diverge: an ``allow_all_targets`` service that still lists - Claude models would test True but hand the prompt an empty list, aborting the wizard. - """ - return bool(map_claude_family_models(targets)) - - -def _prompt_claude_provider_family_models(targets: list[str], service_name: str) -> dict: - """Claude family slots (and overall default) chosen from a service's own Claude target ids. - - The Databricks-hosted path (:func:`_prompt_claude_models`) prompts per family because Claude - Code addresses models by family alias; the same holds behind a Model Provider Service, except - the ids are the service's own — Bedrock provider-side slugs or canonical Anthropic ids rather - than ``system.ai.*``. ``render_overlay`` pins each ``ANTHROPIC_DEFAULT__MODEL`` from - them, so a single overall default would leave the other families unpinned. - - Targets are grouped by family via :func:`claude_family_for_model`, which matches the - ``claude--`` segment in any spelling (``anthropic.claude-…`` Bedrock or bare - ``claude-…`` canonical). A target that names no family is offered only as the overall default. - Falls back to a single default when nothing maps to a family at all. - """ - display = TOOL_SPECS["claude"]["display"] - by_family: dict[str, list[str]] = {} - for target in targets: - family = claude_family_for_model(target) - if family: - by_family.setdefault(family, []).append(target) - - if not by_family: - # No target maps to a Claude family (unusual for a Bedrock Claude service); the most this can - # honestly ask for is one overall default. - return { - "default_model": _require_selection( - f"Default model for {display} (from {service_name}):", - [(t, t) for t in targets], - ) - } - - # Quick setup: fill each family with the service's newest id (highest version, broadest region), - # the same pick a developer's own `ug configure` would make. The alternative is choosing a - # specific id per family — e.g. to pin an older, validated version or a particular region. - # map_claude_family_models covers opus/sonnet/haiku but not fable, so a fable-only service has - # nothing to quick-fill — only offer quick setup when it would actually populate a slot. - family_models = map_claude_family_models(targets) - if family_models: - print_note( - "Quick setup fills each Claude family with the newest model this service offers. Answer " - "no to choose a specific model per family instead (pin an older version, a region)." - ) - if family_models and prompt_yes_no_default("Quick setup?", default=True): - slots = {CLAUDE_SLOT_FOR_FAMILY[family]: model for family, model in family_models.items()} - # Overall default = the highest-tier family the service offers, not whichever target happened - # to sort first. Fable is last: it's the premium opt-in model, a poor default. `family_models` - # is non-empty here, so `next` always finds one. - default_family = next( - fam for fam in ("opus", "sonnet", "haiku", "fable") if fam in family_models - ) - model_config = {"models": slots, "default_model": family_models[default_family]} - summary = ", ".join( - f"{fam}={slots[CLAUDE_SLOT_FOR_FAMILY[fam]]}" - for fam in ANTHROPIC_FAMILIES - if CLAUDE_SLOT_FOR_FAMILY[fam] in slots - ) - # A note, not a success line: the loop's `_confirm_agent` prints the single ✔ for the agent. - print_note(f"Quick setup — {summary} (default: {default_family}).") - return model_config - - print_note( - f"Claude Code picks a model by family, so set a default per family from {service_name}. " - "Skip any family you don't want configured — it falls back to the overall default." - ) - slots: dict[str, str] = {} - for family in ANTHROPIC_FAMILIES: - family_targets = by_family.get(family) - if not family_targets: - continue - choice = prompt_for_selection( - f"Default {family} model:", - [(t, t) for t in family_targets] + [(_SKIP_FAMILY, f"(skip {family})")], - searchable=True, - ) - if choice is None: - raise KeyboardInterrupt - if choice != _SKIP_FAMILY: - slots[CLAUDE_SLOT_FOR_FAMILY[family]] = choice - - model_config: dict = {} - if slots: - model_config["models"] = slots - chosen = list(dict.fromkeys(slots.values())) - if len(chosen) == 1: - model_config["default_model"] = chosen[0] - print_note(f"Only one model configured, so it's {display}'s overall default.") - else: - if slots: - _render_family_slots(slots) - # Offered over every target, not just the slots: `default_model` needn't be a family model, - # and a mixed-catalog service may expose one an admin wants as the overall default. - options = chosen or list(targets) - model_config["default_model"] = _require_selection( - f"Which of those is {display}'s overall default?", [(m, m) for m in options] - ) - return model_config - - -def _claude_candidates(state: dict) -> dict[str, list[str]]: - """Claude models grouped by family. Degrades to the per-family picks if the listing fails. - - Caches the full listing on ``state["all_claude_models"]`` so `validate_manifest` recognizes the - older versions these prompts offer — ``claude_models`` alone holds just the newest per family, - and would reject a legitimately-picked ``claude-opus-4-8``. - - INVARIANT: whatever this returns must be recognizable by ``validate_manifest``, which reads - ``all_claude_models`` (falling back to ``claude_models``) via ``_known_models``. The two paths - below both satisfy it, for different reasons: the listing path widens the candidates *and* sets - the cache, while the fallback path sets nothing but also narrows the candidates to - ``claude_models``, which ``_known_models`` already covers. Widening the fallback without also - populating the cache breaks the invariant, and the symptom is a confusing rejection at the very - end of the flow ("claude: model 'system.ai.claude-opus-4-8' is not available on this - workspace") rather than an error at the prompt that offered it. - """ - cached = state.get("all_claude_models") - if isinstance(cached, list) and cached: - return claude_family_candidates([m for m in cached if isinstance(m, str)], state) - - workspace = state.get("workspace") - all_claude: list[str] = [] - if workspace: - try: - token = get_databricks_token(workspace, state.get("profile")) - all_claude, _ = discover_claude_models_unbucketed(workspace, token) - except (RuntimeError, OSError): - # OSError covers a missing `databricks` binary: `get_databricks_token` shells out, so a - # machine without the CLI on PATH raises FileNotFoundError rather than RuntimeError. - # Either way the per-family picks below are a usable fallback. - all_claude = [] - if all_claude: - state["all_claude_models"] = all_claude - return claude_family_candidates(all_claude, state) - - -# Every picker in this flow chooses a model, a provider service, or a budget — lists that on a real -# workspace run to a dozen-plus entries (16 GPT models on the workspace this was built against), so -# they are all filterable by typing. That trades away j/k navigation, which questionary can't offer -# alongside search; arrow keys still work. -def _require_selection(prompt: str, options: list[tuple[str, str]]) -> str: - """Single-select that won't take "nothing" for an answer. - - ``prompt_for_selection`` returns None for both Ctrl-C and an empty submission, and the two are - genuinely indistinguishable here: questionary's ``Question.ask`` catches KeyboardInterrupt - internally and returns None (v2.1.1, question.py), so nothing propagates for a caller to see. - A None is therefore treated as an abort rather than re-asked — re-asking looped forever on - Ctrl-C, printing the error once per keypress and never exiting. - """ - answer = prompt_for_selection(prompt, options, searchable=True) - if not answer: - raise KeyboardInterrupt - return answer - - -def _require_multi_selection( - prompt: str, options: list[tuple[str, str]], preselected: list[str] | None = None -) -> list[str]: - """Multi-select that requires at least one choice. None (Ctrl-C) still aborts.""" - while True: - picked = prompt_for_multi_selection( - prompt, options, preselected=preselected, searchable=True - ) - if picked is None: - raise KeyboardInterrupt - if picked: - return picked - print_err("Select at least one model (space to toggle, enter to confirm).") - - -def _require_text(prompt: str) -> str: - """Free-text prompt that requires a non-empty answer. - - ``required=True`` makes closed stdin abort instead of returning None. Without it a - non-interactive run (piped stdin, CI) spun here forever: ``prompt_for_text`` returns its default - on EOF, the default is None, and the loop re-asked an empty stream. Reachable whenever model - discovery finds nothing, which is exactly when a run is most likely to be scripted. - """ - while True: - answer = prompt_for_text(prompt, required=True) - if answer: - return answer - print_err("Please enter a model id.") - - -# Discovered model lists run long (a dozen-plus ids on a real workspace), so every hosted-model -# picker is searchable and scrolls (see `prompt_for_selection`); all discovered ids are offered, and -# an explicit "type your own" row still covers a custom model service outside `system.ai` that -# discovery never lists at all. -_CUSTOM_MODEL = "__custom_model__" -_CUSTOM_MODEL_LABEL = "✎ Enter a custom model…" - - -def _custom_option_rows(options: list[tuple[str, str]]) -> list[tuple[str, str]]: - """All discovered model rows plus a 'type your own' row, for a hosted-model picker.""" - return list(options) + [(_CUSTOM_MODEL, _CUSTOM_MODEL_LABEL)] - - -def _short_reason(reason: str | None) -> str: - """A one-line reason fit for a prompt, without the raw JSON body the transport appends. - - HTTP failures come back as ``HTTP : `` (the body is a JSON error - blob for gateway/UC errors); an admin at a prompt wants the status, not the payload. Keeps the - ``HTTP `` head and drops a ``{...}`` body, leaving plain reasons (``network - error: ...``) untouched. - """ - if not reason: - return "unknown error" - return reason.split(": {", 1)[0].strip() - - -def _verify_custom_model(state: dict, model: str) -> tuple[bool | None, str | None]: - """Whether ``model`` is a model service on the workspace; None when the check can't run. - - Mirrors how ``_claude_candidates`` reaches the workspace — a token from ``state`` — and turns a - missing workspace or a failed token fetch into an inconclusive result rather than an error. - """ - workspace = state.get("workspace") - if not workspace: - return None, "no workspace in local state" - try: - token = get_databricks_token(workspace, state.get("profile")) - except (RuntimeError, OSError) as exc: - # OSError covers a missing `databricks` binary (get_databricks_token shells out). - return None, str(exc) - return model_service_exists(workspace, token, model) - - -def _prompt_custom_model(state: dict) -> str: - """Prompt for a custom model-service id, re-asking until it exists on the workspace. - - A typo shouldn't get baked into a published config, so the id is checked against the workspace's - model services and a miss is re-prompted. An inconclusive check (no workspace/token, or a - transient API error) is accepted with a warning rather than blocking a possibly-valid model. - """ - while True: - model = _require_text("Custom model (catalog.schema.model)") - exists, reason = _verify_custom_model(state, model) - if exists: - return model - if exists is None: - print_warning( - f"Couldn't verify '{model}' on this workspace ({_short_reason(reason)}); " - "using it as typed." - ) - return model - print_err( - f"'{model}' isn't a model service on this workspace. Check the name and try again " - "(expected catalog.schema.model, e.g. main.default.claude-opus-4-5)." - ) - - -def _select_hosted_model( - prompt: str, options: list[str], state: dict, custom_sink: list[str] -) -> str: - """Single-select over all discovered ``options`` plus a custom-entry row. - - Records any custom id in ``custom_sink`` so the caller can mark it in - ``model_config.custom_models``, which keeps validation from rejecting a model discovery didn't - surface. - """ - choice = _require_selection(prompt, _custom_option_rows([(m, m) for m in options])) - if choice == _CUSTOM_MODEL: - model = _prompt_custom_model(state) - custom_sink.append(model) - return model - return choice - - -def _select_hosted_models_multi( - prompt: str, options: list[str], state: dict, custom_sink: list[str] -) -> list[str]: - """Multi-select over all discovered ``options`` plus a custom-entry row; requires one pick. - - Selecting the custom row prompts for custom model ids — as many as the admin wants, since a - multi-select agent (opencode, pi) can carry a whole list — and folds them into the picks. Records - each custom id in ``custom_sink`` (see :func:`_select_hosted_model`). - """ - rows = _custom_option_rows([(m, m) for m in options]) - picked = _require_multi_selection(prompt, rows) - models = [p for p in picked if p != _CUSTOM_MODEL] - if _CUSTOM_MODEL in picked: - while True: - custom = _prompt_custom_model(state) - if custom not in models: - models.append(custom) - custom_sink.append(custom) - if not prompt_yes_no_default("Add another custom model?", default=False): - break - return models - - -def configured_models_for_agent(agent_config: dict) -> list[str]: - """Models an agent was configured with, in the manifest's own vocabulary. - - ``model_config.models`` is a flat list for most agents but a family-slot dict for claude - (``default_opus_model`` -> id), so both shapes collapse to a list here. The ``default_model`` is - included because codex has no model list at all — it is the only model that agent has. - """ - model_config = agent_config.get("model_config") - if not isinstance(model_config, dict): - return [] - models: list[str] = [] - raw = model_config.get("models") - if isinstance(raw, dict): - models.extend(v for v in raw.values() if isinstance(v, str) and v) - elif isinstance(raw, list): - models.extend(m for m in raw if isinstance(m, str) and m) - default_model = model_config.get("default_model") - if isinstance(default_model, str) and default_model: - models.append(default_model) - # dict.fromkeys de-duplicates while keeping the admin's preference order. - return list(dict.fromkeys(models)) - - -def _render_tier_ladder(tiers: list[dict], threshold: object, *, base_default: str = "") -> None: - """Show the tiers built so far as spend ranges, so the fallback ladder reads at a glance. - - A tier activates once spend passes its percentage and the highest passed tier wins, so each - tier really owns the range from its own percentage up to the next tier's. Rendering those ranges - ("50–90%", "90%+") rather than bare thresholds ("at 50%", "at 90%") is what makes the ladder - legible — the admin sees which agent a developer actually gets at any level of spend. Reprinted - as the ladder grows, so the sequence forms in front of them instead of in their head. - """ - ordered = sorted(tiers, key=lambda t: t["spending_percentage"]) - lines: list[str] = [] - if base_default: - # Below the first tier the manifest's own default applies; naming it anchors the sequence. - first = ordered[0]["spending_percentage"] * 100 - lines.append(kv_line(f"under {first:g}%", f"{base_default} (default)")) - for i, tier in enumerate(ordered): - low = tier["spending_percentage"] * 100 - agent = TOOL_SPECS.get(tier["default_agent"], {}).get("display", tier["default_agent"]) - if i + 1 < len(ordered): - span = f"{low:g}–{ordered[i + 1]['spending_percentage'] * 100:g}%" - else: - span = f"{low:g}%+" - lines.append(kv_line(span, f"{agent} / {tier['default_model']}")) - print_panel("Budget tiers so far", lines) - - -def _prompt_budget_policy( - workspace: str, - token: str, - enabled_agents: dict[str, dict], - state: dict, - *, - base_default: str = "", -) -> dict | None: - """Author a spend-routing ``budget_policy``, or None when the admin backs out or can't. - - Budgets themselves are created in the Databricks console (they're account-level objects), so the - admin picks an existing one here. Tiers are prompted in percent and stored as fractions, which is - what the API validates. - - ``enabled_agents`` is what the manifest gives each agent, so a tier's model choices come from - that rather than the workspace catalog. Offering the catalog would let a tier point an agent at a - model it wasn't given, which neither this validation nor the server's would reject: the tier would - activate and hand the developer a model their agent doesn't have. - - Asks no "set up a tiered spend policy?" gate — running `ug setup spend-tiers` is the answer to - that question, the same way `ug configure ` needs no confirmation. - """ - print_section("Tiered Spend Policy") - - # Check for attachable budgets before anything else: budgets are created in the Databricks - # console, so if there are none (or none that can enforce routing) there is nothing to do here. - # Bail with a boxed warning and skip the explanatory blurb — no point explaining a feature the - # workspace can't use yet. - with spinner("Listing workspace budgets..."): - budgets, reason = list_workspace_budgets(workspace, token) - if reason is not None or not budgets: - print_warning_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 " - "`ug setup spend-tiers`. Currently, only AI Gateway budgets with hard blocks are " - "eligible to be associated with Tiered Spend Policies." - ) - return None - - # Spend routing only works on a budget with a per-user threshold that hard-blocks: without a - # per-user threshold the gateway reports no spend and every tier stays inert, and without a - # BLOCK_USAGE action the policy is never enforced (an email-only alert does not gate spend). The - # listing now exposes each alert's action, so hide the budgets that can't enforce routing. - usable = [budget for budget in budgets if budget.get("has_per_user_block")] - if not usable: - print_warning_panel( - "None of this workspace's AI Gateway budgets have a per-user threshold with a usage " - "block configured, which spend routing enforces. Add a per-user alert threshold with a " - "block action to a budget in the Databricks console, then re-run " - "`ug setup spend-tiers`." - ) - return None - - # Budgets exist — now explain what a policy does, before asking the admin to pick one. Boxed so - # the concept is read as a unit rather than skimmed as one more bullet. - print_panel("What is a Tiered Spend Policy?", [BUDGET_POLICY_BLURB]) - print_note( - "Showing only budgets with a per-user hard block configured, which spend routing enforces." - ) - - budget_id = prompt_for_selection( - "Which budget should this policy track?", - [ - (budget["id"], f"{budget['display_name'] or budget['id']} ({budget['id']})") - for budget in usable - ], - searchable=True, - ) - if not budget_id: - return None - - policy: dict = {"budget_id": budget_id} - # Remember the budget's own name so the summary can show it beside the policy name. It's a local - # display aid only — `_budget_policy_payload` doesn't serialize it, so it never reaches the API. - budget_display_name = next( - (budget["display_name"] for budget in usable if budget["id"] == budget_id), "" - ) - if budget_display_name: - policy["budget_display_name"] = budget_display_name - display_name = prompt_for_text("Policy name", default="coding-agents-tiered-routing") - if display_name: - policy["display_name"] = display_name - - # The per-user monthly cap the budget was created with. Tiers are picked as percentages of it, so - # showing the dollar amount (and what each percentage works out to) tells the admin what the total - # possible per-user spend even is. None when the listing couldn't read it — then we just skip the - # dollar hints and prompt in percent as before. - threshold = next( - (budget.get("per_user_threshold") for budget in usable if budget["id"] == budget_id), None - ) - if threshold is not None: - print_note(f"This budget's per-user limit is {format_usd(threshold)} per month.") - - tiers: list[dict] = [] - seen_percentages: set[float] = set() - seen_combos: set[tuple[str, str]] = set() - print_note( - "Add a tier for each step down: once spend passes the percentage you set, everyone's " - "default switches to the cheaper agent and model you pick." - ) - while True: - index = len(tiers) + 1 - - # Percentage first, in its own retry loop so a duplicate here re-asks only the percentage. - while True: - fraction = prompt_for_percentage( - f"Tier {index}: switch once spend passes what % of the budget? Ex: 50%" - ) - if fraction in seen_percentages: - print_err("That percentage is already used by another tier; pick a different one.") - continue - break - if threshold is not None: - # Echo the dollars this percentage stands for, so the admin can sanity-check the tier - # against the real per-user cap instead of reasoning about percentages in a vacuum. - print_note( - f" {fraction * 100:g}% of {format_usd(threshold)} is " - f"{format_usd(threshold * Decimal(str(fraction)))}." - ) - - # Agent + model in their own retry loop: a duplicate agent/model re-asks just these two, so - # the admin doesn't have to retype the percentage they already entered for this tier. - agent = model = None - while True: - agent = prompt_for_selection( - f"Tier {index}: switch the default to which agent?", - [(tool, TOOL_SPECS[tool]["display"]) for tool in enabled_agents], - ) - if not agent: - break - # Only what this agent was actually configured with; the workspace catalog would offer - # models the agent doesn't have. - options = configured_models_for_agent(enabled_agents.get(agent) or {}) - if not options: - options = model_options_for_agent(agent, state) - if options: - model = prompt_for_selection( - f"Tier {index}: using which model?", - [(m, m) for m in options], - searchable=True, - ) - else: - model = prompt_for_text(f"Tier {index}: using which model?") - if not model: - break - if (agent, model) in seen_combos: - # The highest crossed tier wins, so a second tier on the same agent+model never - # changes what the lower one already selected — a step-down that doesn't step down. - # Reject it rather than build a policy with a silently inert tier; only the agent and - # model are re-asked, the percentage above is kept. - print_err( - f"{TOOL_SPECS[agent]['display']} / {model} is already used by another tier; a " - "repeated agent/model makes this tier do nothing. Pick a different one." - ) - continue - break - # Cancelling the agent or model picker abandons this tier and stops adding more. - if not agent or not model: - break - - seen_percentages.add(fraction) - seen_combos.add((agent, model)) - tiers.append( - { - "spending_percentage": fraction, - "default_agent": agent, - "default_model": model, - } - ) - _render_tier_ladder(tiers, threshold, base_default=base_default) - if not prompt_yes_no_default("Add another tier?", default=False): - break - - if tiers: - policy["tiers"] = tiers - return policy - - -def _render_summary(workspace: str, manifest: dict) -> None: - """Print the authored config in a box so an admin can eyeball it before publishing. - - Boxed rather than printed as loose lines: this is the one block an admin is meant to read as a - whole and check against what they intended, and it lands after a long flow of prompts. - """ - lines: list[str] = [kv_line("Workspace", workspace)] - default_agent = manifest.get("default_agent") - if isinstance(default_agent, str): - lines.append( - kv_line( - "Default agent", TOOL_SPECS.get(default_agent, {}).get("display", default_agent) - ) - ) - - for tool, agent_config in (manifest.get("enabled_agents") or {}).items(): - display = TOOL_SPECS.get(tool, {}).get("display", tool) - model_config = agent_config.get("model_config") or {} - detail = model_config.get("default_model") or "no model" - provider = model_config.get("model_provider_service") - if provider: - detail = f"{detail} via {provider}" - lines.append(kv_line(display, detail)) - # Spell out the per-family slots and model lists: the one-line default alone doesn't show - # which families an admin configured, which is most of what they chose for claude. - models = model_config.get("models") - if isinstance(models, dict): - for slot, model in models.items(): - family = slot.removeprefix("default_").removesuffix("_model") - lines.append(kv_line(f" {family}", str(model))) - elif isinstance(models, list) and len(models) > 1: - lines.append(kv_line(" models", ", ".join(str(m) for m in models))) - - mcp_servers = manifest.get("mcp_servers") or [] - lines.append( - kv_line( - "MCP servers", - ", ".join(str(server.get("name")) for server in mcp_servers) if mcp_servers else "none", - ) - ) - skills = (manifest.get("skills") or {}).get("names") or [] - lines.append(kv_line("Skills", ", ".join(skills) if skills else "none")) - # Managed tracing isn't offered by the flow yet, so a "disabled" line is just noise. Only surface - # it when a `--from-file` config actually set a table. - if manifest.get("tracing_table"): - lines.append(kv_line("Tracing", str(manifest["tracing_table"]))) - - policy = manifest.get("budget_policy") - if isinstance(policy, dict): - tiers = policy.get("tiers") or [] - lines.append( - kv_line("Budget", policy.get("budget_display_name") or policy.get("budget_id") or "set") - ) - lines.append(kv_line("Policy name", policy.get("display_name") or "unnamed")) - for tier in tiers: - agent = tier.get("default_agent") - display = TOOL_SPECS.get(agent, {}).get("display", agent) - percent = float(tier.get("spending_percentage", 0)) * 100 - lines.append(kv_line(f" at {percent:g}%", f"{display} / {tier.get('default_model')}")) - else: - lines.append(kv_line("Tiered Spend Policy", "none")) - - print_panel("Configuration summary", lines) - - -def _config_facts(manifest: dict) -> list[tuple[str, str, str]]: - """Flatten a normalized config into ordered ``(key, label, value)`` facts, for diffing. - - Each fact is one thing an admin would think of as a single setting — the default agent, an agent's - model, its settings scope, an MCP server, a skill, the tracing table, a budget tier. The ``key`` is - a stable identity so the same setting lines up across two configs even when values differ; the - ``label`` is what the admin reads. Deliberately mirrors what :func:`_render_summary` chooses to - show, so the diff and the summary never disagree about what's in a config. - """ - facts: list[tuple[str, str, str]] = [] - - display_name = manifest.get("display_name") - if isinstance(display_name, str) and display_name: - facts.append(("display_name", "Display name", display_name)) - - default_agent = manifest.get("default_agent") - if isinstance(default_agent, str): - display = TOOL_SPECS.get(default_agent, {}).get("display", default_agent) - facts.append(("default_agent", "Default agent", display)) - - for tool, agent_config in (manifest.get("enabled_agents") or {}).items(): - display = TOOL_SPECS.get(tool, {}).get("display", tool) - model_config = agent_config.get("model_config") or {} - detail = model_config.get("default_model") or "no model" - provider = model_config.get("model_provider_service") - if provider: - detail = f"{detail} via {provider}" - facts.append((f"agent:{tool}", display, detail)) - models = model_config.get("models") - if isinstance(models, dict): - for slot, model in models.items(): - family = slot.removeprefix("default_").removesuffix("_model") - facts.append((f"agent:{tool}:model:{family}", f"{display} ({family})", str(model))) - elif isinstance(models, list) and len(models) > 1: - facts.append((f"agent:{tool}:models", f"{display} models", ", ".join(map(str, models)))) - for server in manifest.get("mcp_servers") or []: - name = str(server.get("name")) - facts.append((f"mcp:{name}", f"MCP server {name}", str(server.get("type") or ""))) - - for skill in (manifest.get("skills") or {}).get("names") or []: - facts.append((f"skill:{skill}", f"Skill {skill}", "published")) - - tracing = manifest.get("tracing_table") - if tracing: - facts.append(("tracing_table", "Tracing table", str(tracing))) - - policy = manifest.get("budget_policy") - if isinstance(policy, dict): - facts.append( - ( - "budget:id", - "Budget", - policy.get("budget_display_name") or policy.get("budget_id", ""), - ) - ) - if policy.get("display_name"): - facts.append(("budget:name", "Policy name", str(policy["display_name"]))) - for tier in policy.get("tiers") or []: - agent = tier.get("default_agent") - agent_display = TOOL_SPECS.get(agent, {}).get("display", agent) - percent = float(tier.get("spending_percentage", 0)) * 100 - facts.append( - ( - f"budget:tier:{percent:g}", - f"Budget tier at {percent:g}%", - f"{agent_display} / {tier.get('default_model')}", - ) - ) - - return facts - - -def _render_config_diff(existing: dict | None, incoming: dict, workspace: str) -> bool: - """Show what publishing ``incoming`` changes versus the ``existing`` published config. - - Returns True when there is a difference. Lists only what changes — labelled ADD, DELETE, or - CHANGE (``old → new``) — since the full config was just printed by :func:`_render_summary` above; - repeating the unchanged rows here would bury the actual delta. Both configs are in ug's - normalized shape (the caller round-trips the local manifest through serialize/normalize first), - so the comparison is field-for-field with what the workspace holds. - """ - old = {key: (label, value) for key, label, value in _config_facts(existing or {})} - new = {key: (label, value) for key, label, value in _config_facts(incoming)} - - # Fixed-width verbs so the labels line up in a column and the eye can scan one kind of change. - add = "[green]ADD [/green]" - delete = "[red]DELETE[/red]" - change = "[yellow]CHANGE[/yellow]" - - # Incoming order first (added/changed read top-down like the summary), then removed keys. - ordered = list(new) + [key for key in old if key not in new] - rows: list[str] = [] - for key in ordered: - if key in new and key not in old: - label, value = new[key] - rows.append(f" {add} {label}: {value}") - elif key in old and key not in new: - label, value = old[key] - rows.append(f" {delete} {label}: {value}") - elif old[key][1] != new[key][1]: - label, old_value = old[key] - rows.append(f" {change} {label}: {old_value} → {new[key][1]}") - - if not rows: - return False - print_heading(f"Changes to publish on {workspace}") - for row in rows: - console.print(row) - return True - - -def _require_admin(workspace: str, token: str) -> None: - """Stop unless the caller is a workspace admin. - - 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. - """ - 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}. `ug setup` authors the workspace-wide " - "coding config, so it is restricted to workspace admins." - ) - if admin is None: - print_warning( - "Could not verify workspace admin permissions. Continuing — `ug publish` will fail " - "if you lack them." - ) - else: - 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 MCP servers / skills / - tracing / budget policy forward — the local draft may be missing on a fresh machine or after - ``ug revert``, and without this those sections would be silently dropped on the next publish. - - Deliberately doesn't itemize what the existing config holds. The admin doesn't need an inventory - to act on this, and `ug 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, MCP " - "server, skill, 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 ug 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 " - "ug 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 ug 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 ug's own manifest shape (the same thing the wizard writes), not proto-JSON. - """ - manifest_path = Path(path).expanduser() - try: - raw = manifest_path.read_text(encoding="utf-8") - except OSError as exc: - raise RuntimeError(f"Could not read manifest file: {manifest_path}") from exc - try: - manifest = json.loads(raw) - except json.JSONDecodeError as exc: - raise RuntimeError( - f"{manifest_path} is not valid JSON: {exc.msg} (line {exc.lineno})." - ) from None - if not isinstance(manifest, dict): - raise RuntimeError(f"{manifest_path} must contain a JSON object.") - - state = load_state() - workspace = state.get("workspace") - if not workspace: - raise RuntimeError( - "No workspace is configured. Run `ug configure` first so ug knows which " - "workspace this manifest is for." - ) - - errors = validate_manifest(manifest, state) - if errors: - print_err(f"{manifest_path} is not a valid managed config:") - for error in errors: - print_note(error) - return 1 - - save_managed_state(workspace, manifest) - _render_summary(workspace, manifest) - print_success(f"Saved to {manifest_path.name} -> ~/.ucode/managed-state.json") - _print_next_steps(manifest) - return 0 - - -# The sections that have their own `ug 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]]] = [ - ("ug setup mcps", "MCP servers", lambda m: bool(m.get("mcp_servers"))), - ("ug setup skills", "Skills", lambda m: bool((m.get("skills") or {}).get("names"))), - ( - "ug setup spend-tiers", - "Tiered Spend Policy", - lambda m: isinstance(m.get("budget_policy"), dict), - ), -] - - -def _command_line(command: str, description: str, *, marker: str = " ", width: int = 0) -> str: - """A `` `` row, for command lists that read as a column.""" - return f" {marker} [bold]{command.ljust(width)}[/bold] {description}" - - -# `ug 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 = "ug setup") -> None: - """Announce one phase of the flow as `step N of M`, branded to the invoking command.""" - print_section(f"{command_label} · step {index} of {len(SETUP_STEP_TITLES)} · {title}") - - -def _selected_recap(workspace: str, enabled_agents: dict, default_agent: str | None) -> None: - """A compact panel of what's chosen so far, reprinted as the flow advances. - - Turns the run of prompts into something that reads like a form filling in: each phase reprints - the growing set of decisions before asking the next question. Agents still mid-configuration show - a `…` placeholder for their model. - """ - lines = [kv_line("Workspace", workspace)] - for tool, config in enabled_agents.items(): - model = (config.get("model_config") or {}).get("default_model") or "…" - lines.append(kv_line(TOOL_SPECS.get(tool, {}).get("display", tool), model)) - if default_agent: - lines.append( - kv_line("Default", TOOL_SPECS.get(default_agent, {}).get("display", default_agent)) - ) - print_panel("Selected so far", lines) - - -def _section_status_lines(manifest: dict, width: int = 0) -> list[str]: - """One row per optional section: its command, what it covers, and whether it's configured.""" - width = width or max(len(command) for command, _, _ in SETUP_SECTIONS) - lines: list[str] = [] - for command, label, configured in SETUP_SECTIONS: - if configured(manifest): - marker, state = "[green]✔[/green]", "[green]configured[/green]" - else: - marker, state = "[dim]○[/dim]", "[dim]not configured[/dim]" - lines.append(_command_line(command, f"{label} — {state}", marker=marker, width=width)) - return lines - - -def _print_next_steps(manifest: dict) -> None: - """List the setup commands still worth running, then the publish step. - - Printed rather than prompted: each section is its own command now, so the admin drives the rest of - the setup themselves instead of being walked through a chain they mostly want to skip. Showing - what is already configured keeps a re-run from looking like it lost the other sections — it - didn't; `setup` carries them forward. - """ - console.print() - print_heading("Next steps") - if config_io.is_dry_run(): - # Under --dry-run nothing was written, so the section commands (which read the saved draft) - # and `publish` have nothing to act on. Say so rather than send the admin to commands that - # would report "run `ug setup` first". - print_note("Dry run — nothing was saved. Re-run without --dry-run to author the config.") - return - # These sections aren't required to publish — call them out as optional so an admin doesn't read - # a config with none configured as unfinished. - print_note("[dim]Optional — configure any of these, or skip straight to publishing:[/dim]") - for line in _section_status_lines(manifest): - console.print(line) - print_panel( - "All done?", - ["Publish with [bold]ug publish[/bold] so all developers use this configuration."], - ) - - -def _offer_publish() -> None: - """Offer to publish the saved draft right away, so an admin can apply changes incrementally. - - Each `ug setup` command only writes a local draft. Without this an admin has to remember to run - `ug publish` separately, and a `ug 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 `ug 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 `ug publish`)", default=False - ): - print_note("Draft saved. Run `ug publish` when you're ready to publish.") - return - publish_command(yes=True) - - -# The sections `ug setup` carries forward instead of prompting for, and how to rebuild each one. -CARRIED_SECTIONS: list[tuple[str, str, str]] = [ - ("mcp_servers", "MCP servers", "ug setup mcps"), - ("skills", "Skills", "ug setup skills"), - ("tracing_table", "Tracing table", "ug setup --from-file"), - ("budget_policy", "Tiered Spend Policy", "ug setup spend-tiers"), -] - - -def _carry_forward_sections(previous: dict, manifest: dict) -> None: - """Copy the sections `ug setup` no longer prompts for out of a previously authored config. - - `setup` writes the whole manifest, so without this a re-run would silently clear the MCP servers, - skills, tracing table, and budget policy an admin authored with the other commands — they'd have - to redo every one of them just to change a model. - - Each section is probe-validated before it's carried, and dropped with a warning if it no longer - fits. Otherwise a carried section could make the manifest invalid and block the save outright, - with no way out: the commands that could repair a section read the very manifest that can't be - written. The live case is a budget-policy tier naming an agent the admin just de-selected, but - hand-edited drafts and configs authored by an older ug can trip the others the same way. - """ - # Validating against no inventory keeps this to structural checks, which is all that's at stake - # here: the models were just picked from the workspace's own catalog a few prompts ago. - baseline = validate_manifest(manifest, None) - for key, label, rebuild in CARRIED_SECTIONS: - if key not in previous: - continue - candidate = previous[key] - new_errors = [ - error - for error in validate_manifest({**manifest, key: candidate}, None) - if error not in baseline - ] - if not new_errors: - manifest[key] = candidate - continue - print_warning( - f"{label} from the existing config no longer fits what you just picked, so it was left " - "out:" - ) - for error in new_errors: - print_note(error) - print_note(f"Rebuild it with `{rebuild}`.") - - -def setup_command( - from_file: str | None = None, - *, - workspace: str | None = None, - profile: str | None = None, - command_label: str = "ug 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. MCP servers, skills, and the tiered spend policy each have their - own command (`ug setup mcps` / `skills` / `spend-tiers`), so an admin changing one of them doesn't - have to walk the whole flow again — and this command carries whatever they already authored - forward untouched rather than clearing it (:func:`_carry_forward_sections`). - - ``workspace``/``profile`` let a caller that has already resolved (and authenticated against) a - workspace hand it in so the admin isn't prompted to pick one again — e.g. `ug 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: `ug configure` passes - "Configure unity-gateway CLI" so a user who never typed `ug setup` isn't jarred by it (the - standalone `ug setup` command keeps the default). References to specific sub-commands (`ucode - setup mcps`, `ug apply`, …) stay verbatim — those are real command names, not branding. - - ``token`` lets a caller that already authenticated and admin-checked the workspace (e.g. - `ug 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. - """ - 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 - - 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 ug.") - - 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 - for tool in TOOL_SPECS - if tool not in SETUP_EXCLUDED_AGENTS and check_gateway_endpoint(state, tool) - ] - if not available: - raise RuntimeError( - f"No coding agents are available on {workspace}. Check that the workspace's AI Gateway " - "serves models for at least one agent." - ) - - # The local draft is the carry-forward source, falling back to what's published on the workspace: - # a fresh machine (or one after `ug revert`) has no draft, and without the fallback the next - # publish would silently wipe the workspace's MCP servers, skills, tracing, and budget policy. - previous = load_managed_state(workspace) or published or {} - previously_enabled = [ - tool for tool in (previous.get("enabled_agents") or {}) if tool in available - ] - _step_banner(1, SETUP_STEP_TITLES[0], command_label) - picked = prompt_for_tools( - [(tool, TOOL_SPECS[tool]["display"]) for tool in available], - preselected=previously_enabled or None, - ) - if not picked: - print_note("No coding agents selected — nothing to configure.") - return 0 - - _step_banner(2, SETUP_STEP_TITLES[1], command_label) - enabled_agents: dict[str, dict] = {} - for index, tool in enumerate(picked, start=1): - print_heading(f"{TOOL_SPECS[tool]['display']} ({index} of {len(picked)})") - provider_service = _select_provider_service(tool, workspace, token) - # Always set: `_prompt_models_for_agent` re-prompts rather than returning empty, so every - # enabled agent carries a default_model and any of them can be the default_agent. - agent_config: dict = { - "model_config": _prompt_models_for_agent(tool, state, provider_service) - } - enabled_agents[tool] = agent_config - _confirm_agent(tool, agent_config) - - # Pick the default after configuring each agent, not before: by now the admin has seen every - # agent's models go by, so "which is the default?" is a choice among things they've just set up - # rather than a bare list up front. The recap reprints those picks so the choice is informed. - _step_banner(3, SETUP_STEP_TITLES[2], command_label) - default_agent = picked[0] - if len(picked) > 1: - _selected_recap(workspace, enabled_agents, default_agent=None) - chosen = prompt_for_selection( - "Which coding agent should be the default?", - [(tool, TOOL_SPECS[tool]["display"]) for tool in picked], - ) - if not chosen: - raise KeyboardInterrupt - default_agent = chosen - print_success(f"Default agent set to {TOOL_SPECS[default_agent]['display']}") - - manifest: dict = {"default_agent": default_agent, "enabled_agents": enabled_agents} - - # Tracing is intentionally not prompted here: the managed-tracing path isn't working yet, so - # asking would author a `tracing_table` the workspace can't honor. The manifest field and its - # serialize/validate support stay in place, so a hand-written `--from-file` config can still set - # it once the backend is ready. Re-add a `ug setup tracing` command when it is. - _carry_forward_sections(previous, manifest) - - errors = validate_manifest(manifest, state) - if errors: - # A validation failure here is a wizard bug, not admin error — the pickers only offer valid - # choices. Surface it plainly rather than writing a manifest that `publish` would reject. - print_err("The generated config is not valid:") - for error in errors: - print_note(error) - return 1 - - save_managed_state(workspace, manifest) - _render_summary(workspace, manifest) - console.print() - print_success("Saved to ~/.ucode/managed-state.json") - _print_next_steps(manifest) - _offer_publish() - 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 `ug setup`, this doesn't prompt for a workspace - and takes it strictly from local state rather than falling back to the draft file's workspace: the - MCP and skills pickers re-read ``current_workspace`` themselves (via ``setup_mcp_clients``), so a - mismatch would have them operate against one workspace while the manifest is saved for another. - Requiring ``ug configure`` to have set the current workspace keeps the two in lockstep. It also - skips :func:`_handle_existing_config` — the create-or-delete choice belongs to authoring a config, - not to changing one section of it. - """ - state = load_state() - workspace = state.get("workspace") - if not workspace: - raise RuntimeError( - "No workspace is configured. Run `ug configure` first, then `ug setup` to author " - "this workspace's managed config." - ) - profile = state.get("profile") - ensure_databricks_auth(workspace, profile) - token = get_databricks_token(workspace, profile) - _require_admin(workspace, token) - return workspace, profile, token - - -def _manifest_for_edit(workspace: str) -> dict: - """The authored manifest a section command edits. Raises when `ug setup` hasn't run. - - An empty ``enabled_agents`` counts as "hasn't run": a launch records ``{}`` for a workspace with no - managed config (see ``refresh_managed_config``), so the file existing is not proof an admin - authored anything. Requiring agents first also keeps the budget-policy tiers honest — they can only - name agents the manifest enables. - """ - manifest = load_managed_state(workspace) - if not (manifest or {}).get("enabled_agents"): - raise RuntimeError( - f"No managed config has been authored for {workspace} yet. Run `ug setup` first to " - "pick the agents and models, then re-run this command." - ) - return cast(dict, manifest) - - -def _save_section_update(workspace: str, manifest: dict) -> int: - """Validate the edited manifest structurally, save it, and show what's left to do. - - Validated with no model inventory (``state=None``), so only structure is checked here — not model - availability. That's deliberate: a section command doesn't touch agents or models, so re-checking - them would only reject a legitimately pinned older Claude model (`load_state` keeps just the newest - per family) or, worse, wrongly flag a codex/gemini model whenever the re-fetched inventory happens - to be Claude-only. `ug publish` runs the full model check against the live catalog before - publishing, which is where it belongs. - """ - errors = validate_manifest(manifest, None) - if errors: - print_err("The updated config is not valid:") - for error in errors: - print_note(error) - return 1 - - save_managed_state(workspace, manifest) - _render_summary(workspace, manifest) - console.print() - print_success("Saved to ~/.ucode/managed-state.json") - _print_next_steps(manifest) - _offer_publish() - return 0 - - -def setup_mcp_command() -> int: - """Author the managed config's MCP servers (`ug setup mcps`).""" - workspace, _, _ = _resolve_admin_workspace() - manifest = _manifest_for_edit(workspace) - - print_section("Managed MCP servers") - print_note("Developers get these MCP servers registered automatically when they run ug.") - from ucode.mcp import configure_mcp_command - - # Snapshot the managed-shaped servers before the picker so a cancelled run on an empty local - # state can't delete a section the manifest still has: the picker returns 0 on Esc, and re-reading - # local state would otherwise overwrite the manifest with nothing. - before = _mcp_servers_from_state(load_state()) - # Managed configs can't carry a Databricks app (its host isn't reconstructable from the - # workspace), so hide apps from the picker rather than let an admin pick one that is then - # dropped from the published config. - configure_mcp_command(exclude_sources={"apps"}) - after = _mcp_servers_from_state(load_state()) - - # `after == before` isn't enough to call this a no-op: an admin who ran `ug configure mcp` - # first arrives with those servers already registered, so confirming the picker leaves local state - # unchanged even though the manifest doesn't carry them yet. Also sync when local state already - # holds servers the manifest is missing — but only when servers are actually registered, so an Esc - # on an empty local state still can't wipe a published section. - manifest_servers = manifest.get("mcp_servers") or [] - carries_unsaved = bool(after) and after != manifest_servers - if after == before and not carries_unsaved: - print_note("No changes to the MCP servers — the managed config is unchanged.") - return 0 - - if after: - manifest["mcp_servers"] = after - print_success(f"{len(after)} MCP server(s) in the managed config") - else: - # Deregistering every server locally is how an admin clears the section — there is no - # separate "remove them all" flag. - manifest.pop("mcp_servers", None) - print_note("No MCP servers are registered, so the managed config now carries none.") - return _save_section_update(workspace, manifest) - - -def setup_skills_command(locations: list[str] | None = None) -> int: - """Author the managed config's skills (`ug setup skills`). - - ``locations`` comes from ``--location`` (already parsed to `.` refs); when None - the admin is prompted and the answer is parsed the same way. - """ - workspace, _, _ = _resolve_admin_workspace() - manifest = _manifest_for_edit(workspace) - - print_section("Managed skills") - print_note("Developers get these skills downloaded automatically when they run ug.") - if locations is None: - answer = prompt_for_text( - "Skill schemas to publish, comma-separated `catalog.schema` (blank to leave unchanged)", - default="", - ) - # Route the interactive answer through the same parser as `--location` so `main` (missing the - # schema) is rejected here rather than published as a bogus skill name. - from ucode.cli import _parse_skill_locations - - locations = _parse_skill_locations(answer) - # A blank answer / empty `--location` means "leave the skills alone". Returning before delegating - # matters: `configure_skills_mcp_command([])` is not a no-op — it registers the schema-less skills - # MCP connection into the admin's own agents. - if not locations: - print_note("No skill schemas given — the managed config's skills are unchanged.") - return 0 - - from ucode.mcp import configure_skills_mcp_command - - configure_skills_mcp_command(locations) - skill_names = _skill_names_from_state(load_state()) or locations - manifest["skills"] = {"names": skill_names} - print_success(f"{len(skill_names)} skill schema(s) in the managed config") - return _save_section_update(workspace, manifest) - - -def setup_budget_policy_command() -> int: - """Author the managed config's tiered spend policy (`ug setup spend-tiers`).""" - workspace, _, token = _resolve_admin_workspace() - manifest = _manifest_for_edit(workspace) - - # The manifest's own default, shown as the "under the first rung" row of the fallback ladder so - # the admin sees what developers get before any tier kicks in. - default_agent = manifest.get("default_agent") - default_config = (manifest.get("enabled_agents") or {}).get(default_agent) or {} - default_model = (default_config.get("model_config") or {}).get("default_model") - base_default = ( - f"{TOOL_SPECS.get(default_agent, {}).get('display', default_agent)} / {default_model}" - if default_agent and default_model - else "" - ) - - # `_prompt_budget_policy` returns None on an environmental dead end too (no budgets, no budget with - # a per-user block, or the admin backing out of a picker), not only on an explicit decline. Leave - # any existing policy untouched in every one of those cases — never pop it — so a transient budget - # listing failure can't silently delete a policy the admin already published. - policy = _prompt_budget_policy( - workspace, token, manifest["enabled_agents"], load_state(), base_default=base_default - ) - if not policy: - print_note("The managed config's tiered spend policy is unchanged.") - return 0 - manifest["budget_policy"] = policy - return _save_section_update(workspace, manifest) - - -def setup_help_command() -> int: - """Walk through the whole managed-config setup, marking what this machine has authored. - - Hand-written rather than left to `--help`: the point is the *order* of the commands and the fact - that nothing reaches developers until `ug publish`, neither of which a flag listing conveys. Reads - the manifest but never authenticates, so it works before `ug configure`. - """ - print_section("ug setup") - print_note( - "A managed config is the coding setup your developers pull automatically — they run ug " - "and get the agents, models, MCP servers, and skills you chose here. Admins only." - ) - print_note( - "Each command below edits your local draft; nothing reaches the workspace until " - "`ug 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("ug setup --from-file ")) - - print_heading("1. Start here") - console.print( - _command_line( - "ug 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("ug setup show", "The draft, and the payload `publish` sends", width=width) - ) - console.print(_command_line("ug publish", "Publish it to the workspace", width=width)) - - print_heading("Also") - console.print( - _command_line( - "ug 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 `ug setup` keeps the sections in step 2; to drop one, edit the draft and " - "reload it with `ug setup --from-file`." - ) - return 0 - - -def show_command() -> int: - """Print the authored manifest and the proto-JSON `ug publish` would publish.""" - # Fall back to the workspace the on-disk file was authored for, so `ug setup --show` still - # works before `ug 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 `ug setup` to create one.") - return 0 - _render_summary(workspace or "unknown", manifest) - console.print() - print_heading("Payload for `ug 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. -def _explain_publish_failure(reason: str) -> str: - lowered = reason.lower() - if "feature_disabled" in lowered: - return CODING_AGENT_CONFIGS_DISABLED_MESSAGE - if "permission_denied" in lowered or "http 403" in lowered: - return ( - "Publishing a managed config requires workspace admin. Your account can read the " - "workspace but not author its coding config." - ) - if "already_exists" in lowered: - return ( - "This workspace already has a managed config, but ug couldn't read it to update in " - "place. Run `ug publish` again — if it keeps failing, the existing config may need to " - "be deleted by hand." - ) - if "invalid_parameter_value" in lowered: - # The server names the offending field; passing it through beats paraphrasing. - return f"The workspace rejected the config: {reason}" - return f"Could not publish the managed config: {reason}" - - -def _with_claude_inventory(state: dict, workspace: str, profile: str | None) -> dict: - """``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 `ug setup` 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: - - 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 `ug 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. - - Best-effort: a failed listing returns ``state`` untouched, leaving validation on the narrower - inventory rather than blocking a publish on a transient API error. - """ - if isinstance(state.get("all_claude_models"), list) and state["all_claude_models"]: - return state - try: - token = get_databricks_token(workspace, profile) - all_claude, _ = discover_claude_models_unbucketed(workspace, token) - except (RuntimeError, OSError): - # OSError covers a missing `databricks` binary: `get_databricks_token` shells out, so a - # machine without the CLI on PATH raises FileNotFoundError rather than RuntimeError. - return state - if not all_claude: - return state - return {**state, "all_claude_models": all_claude} - - -def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: - """Publish a managed config to the workspace. - - With no ``file_path`` the locally authored manifest is published; with one, the config file - (produced by ``ug export``) is published instead. Both routes are validated against the - configured workspace and canonicalized before anything is sent. Updates the existing config in - place when there is one, rather than deleting and recreating it: a failed recreate would leave - the workspace with no managed config at all, and every developer would silently fall back to - their own settings. Returns a process exit code. - - Model-availability validation runs against the authored manifest for the no-file case, since the - export round-trip that builds its payload drops the internal ``model_config.custom_models`` - exemption for hand-entered ids that discovery won't surface. - """ - from ucode.cli import _prompt_for_configuration - from ucode.managed_publish import load_publish_payload, parse_publish_payload - - print_section("ug publish") - - state = load_state() - workspace = state.get("workspace") - profile = state.get("profile") - if not workspace: - workspace, profile = _prompt_for_configuration() - - manifest, api_payload = parse_publish_payload(load_publish_payload(file_path), workspace) - - # Auth first: validating a Claude manifest needs the workspace's full model listing, and that - # 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 - errors = validate_manifest( - validation_manifest or manifest, _with_claude_inventory(state, workspace, profile) - ) - if errors: - print_err("The config is not valid, so it was not published:") - for error in errors: - print_note(error) - if file_path is None: - print_note("Re-run `ug setup` to fix it, or edit ~/.ucode/managed-state.json.") - else: - print_note("Fix the config file and re-run `ug publish -f`.") - return 1 - - token = get_databricks_token(workspace, profile) - _require_admin(workspace, token) - - _render_summary(workspace, manifest) - - # Read before writing: the resource name tells us whether to create or update, and shows the - # admin what they are about to overwrite. - 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(): - raise RuntimeError(CODING_AGENT_CONFIGS_DISABLED_MESSAGE) - raise RuntimeError( - f"Could not check whether {workspace} already has a managed config: {reason}. " - "Refusing to publish without knowing, since that could overwrite a config silently." - ) - - existing_name = (existing or {}).get("name") - if existing is not None and not isinstance(existing_name, str): - raise RuntimeError( - "This workspace has a managed config but the API didn't return its resource name, so " - "ug can't update it in place. Delete it in the workspace and re-run `ug publish`." - ) - - console.print() - if existing is None: - print_note(f"This will create a new managed config on {workspace}.") - else: - # Diff against what's live, normalized the same way, so the admin sees exactly what changes - # rather than a bare "this replaces the current config". Comparing the round-tripped payload - # (not the raw manifest) shows the real post-publish state — any field serialization drops - # won't appear as a phantom change. - changed = _render_config_diff(existing, manifest, workspace) - if not changed: - print_success(f"{workspace}'s published config already matches this one.") - print_note("Nothing to publish.") - return 0 - console.print() - print_warning("This takes effect for every developer on their next `ucode` run.") - if not yes and not prompt_yes_no_default("Publish this config?", default=False): - print_note("Nothing was published.") - return 1 - - if existing is None: - with spinner("Publishing the managed config..."): - published, publish_reason = create_coding_agent_config(workspace, token, api_payload) - else: - with spinner("Updating the managed config..."): - published, publish_reason = update_coding_agent_config( - workspace, token, cast("str", existing_name), api_payload - ) - if publish_reason is not None: - raise RuntimeError(_explain_publish_failure(publish_reason)) - - 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`.") - return 0 - - -__all__ = [ - "publish_command", - "setup_budget_policy_command", - "setup_command", - "setup_from_file", - "setup_help_command", - "setup_mcp_command", - "setup_skills_command", - "show_command", -] diff --git a/tests/test_cli.py b/tests/test_cli.py index 8ef41712..387d2bb4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -90,6 +90,14 @@ def test_help_lists_all_agent_subcommands(self): for tool in TOOLS: assert tool in result.output + def test_managed_authoring_commands_are_removed(self): + # Authoring moved to the AI Gateway API/UI, so `ug setup` and `ug publish` no longer exist. + assert runner.invoke(app, ["setup"]).exit_code != 0 + assert runner.invoke(app, ["setup", "mcps"]).exit_code != 0 + assert runner.invoke(app, ["publish"]).exit_code != 0 + # `ug export` (read-only) stays. + assert runner.invoke(app, ["export", "--help"]).exit_code == 0 + @pytest.mark.parametrize("prog_name", ["ug", "ucode"]) def test_help_uses_invoked_name_and_names_ucode_as_an_alias(self, prog_name): result = runner.invoke(app, ["--help"], prog_name=prog_name) @@ -3368,7 +3376,6 @@ def _run( monkeypatch, *, managed, - is_admin=False, args=None, cached=None, coding_agent_config_feature_disabled=False, @@ -3384,8 +3391,6 @@ def _run( 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.get_databricks_token", lambda *a, **k: "tok") - monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda *a, **k: is_admin) monkeypatch.setattr( "ucode.cli._launch_tool", lambda tool, ctx, **kw: launched.append((tool, kw)), @@ -3436,21 +3441,20 @@ 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): - result, launched = self._run(monkeypatch, managed=None, is_admin=True) - assert result.exit_code == 0, result.output - assert launched == [] - assert "ug setup" 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) + def test_no_config_points_the_dev_at_configure(self, monkeypatch): + # With no managed config a developer can still set up locally, so the guidance points at + # `ug configure` (not the removed authoring commands). + result, launched = self._run(monkeypatch, managed=None) assert result.exit_code == 0, result.output assert launched == [] - assert "Ask a workspace admin" in result.output + flat = " ".join(result.output.split()) + assert "ug configure" in flat + assert "ug setup" not in flat + assert "ug publish" not in flat def test_feature_disabled_guides_without_managed_mention(self, monkeypatch): result, launched = self._run( - monkeypatch, managed=None, is_admin=True, coding_agent_config_feature_disabled=True + monkeypatch, managed=None, coding_agent_config_feature_disabled=True ) assert result.exit_code == 0, result.output assert launched == [] @@ -3487,10 +3491,6 @@ def test_dry_run_with_no_cached_config_does_not_crash(self, monkeypatch): lambda state: pytest.fail("--dry-run must not fetch"), ) monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: None) - # The no-config guidance checks admin status; stub the token/admin calls so the test - # doesn't shell out to the `databricks` binary (absent in CI). - monkeypatch.setattr("ucode.cli.get_databricks_token", lambda *a, **k: "tok") - monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda *a, **k: False) monkeypatch.setattr( "ucode.cli._launch_tool", lambda *a, **k: pytest.fail("nothing to launch without a config"), diff --git a/tests/test_managed_publish.py b/tests/test_managed_publish.py deleted file mode 100644 index 37276b60..00000000 --- a/tests/test_managed_publish.py +++ /dev/null @@ -1,159 +0,0 @@ -"""Tests for :mod:`ucode.managed_publish` — the input handling behind `ucode publish`. - -Cover the two config sources (in-process export payload and a `-f` file), the envelope checks -(workspace match, spec_version type and value), and the canonicalization that rejects server-owned, -unknown, and lossy fields before anything reaches the workspace. -""" - -from __future__ import annotations - -import json -from unittest.mock import patch - -import pytest - -import ucode.managed_publish as publish_mod -from ucode.managed_publish import load_publish_payload, parse_publish_payload -from ucode.managed_setup import serialize_managed_config - -WORKSPACE = "https://ws.example.com" - -MANIFEST = { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}, - "codex": {"model_config": {"default_model": "system.ai.gpt-5-6"}}, - }, - "mcp_servers": [{"name": "system.ai.slack", "type": "mcp-service"}], - "skills": {"names": ["main.default"]}, -} - - -def _payload(manifest=MANIFEST, *, workspace=WORKSPACE, spec_version=1, **extra): - config = serialize_managed_config(manifest) - config.pop("name", None) - return {"workspace": workspace, "spec_version": spec_version, **config, **extra} - - -class TestLoadPublishPayload: - def test_no_file_uses_build_export_payload_in_process(self): - sentinel = {"workspace": WORKSPACE, "spec_version": 1} - with patch.object(publish_mod, "build_export_payload", return_value=sentinel) as build: - assert load_publish_payload(None) is sentinel - build.assert_called_once_with() - - def test_reads_json_file(self, tmp_path): - path = tmp_path / "config.json" - path.write_text(json.dumps(_payload()), encoding="utf-8") - assert load_publish_payload(str(path)) == _payload() - - def test_expands_user_home(self, tmp_path, monkeypatch): - monkeypatch.setenv("HOME", str(tmp_path)) - (tmp_path / "config.json").write_text(json.dumps(_payload()), encoding="utf-8") - assert load_publish_payload("~/config.json")["workspace"] == WORKSPACE - - 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")) - - def test_malformed_json_is_actionable(self, tmp_path): - path = tmp_path / "config.json" - path.write_text("{not json", encoding="utf-8") - with pytest.raises(RuntimeError, match="not valid JSON"): - load_publish_payload(str(path)) - - def test_non_object_root_is_rejected(self, tmp_path): - path = tmp_path / "config.json" - path.write_text(json.dumps([1, 2, 3]), encoding="utf-8") - with pytest.raises(RuntimeError, match="JSON object at the top level"): - load_publish_payload(str(path)) - - def test_non_utf8_is_actionable(self, tmp_path): - path = tmp_path / "config.json" - path.write_bytes(b"\xff\xfe not utf-8") - with pytest.raises(RuntimeError, match="UTF-8"): - load_publish_payload(str(path)) - - -class TestParsePublishPayload: - def test_returns_manifest_and_api_payload(self): - manifest, api_payload = parse_publish_payload(_payload(), WORKSPACE) - assert manifest["default_agent"] == "claude" - assert set(manifest["enabled_agents"]) == {"claude", "codex"} - assert api_payload["spec_version"] == 1 - assert api_payload["default_agent"] == "CODING_AGENT_CLAUDE_CODE" - assert "workspace" not in api_payload - assert "name" not in api_payload - - def test_api_payload_matches_serialize_of_manifest(self): - manifest, api_payload = parse_publish_payload(_payload(), WORKSPACE) - canonical = serialize_managed_config(manifest) - assert api_payload == {"spec_version": 1, **canonical} - - def test_non_object_payload_is_rejected(self): - with pytest.raises(RuntimeError, match="JSON object"): - parse_publish_payload([1, 2], WORKSPACE) - - def test_missing_workspace_is_rejected(self): - payload = _payload() - del payload["workspace"] - with pytest.raises(RuntimeError, match="workspace"): - parse_publish_payload(payload, WORKSPACE) - - def test_mismatched_workspace_is_rejected(self): - with pytest.raises(RuntimeError, match="configured workspace"): - parse_publish_payload(_payload(workspace="https://other.example.com"), WORKSPACE) - - def test_workspace_match_ignores_scheme_and_trailing_slash(self): - manifest, _ = parse_publish_payload(_payload(workspace="ws.example.com/"), WORKSPACE) - assert manifest["default_agent"] == "claude" - - def test_missing_spec_version_is_rejected(self): - payload = _payload() - del payload["spec_version"] - with pytest.raises(RuntimeError, match="spec_version"): - parse_publish_payload(payload, WORKSPACE) - - def test_boolean_spec_version_is_rejected(self): - with pytest.raises(RuntimeError, match="spec_version"): - parse_publish_payload(_payload(spec_version=True), WORKSPACE) - - def test_fractional_spec_version_is_rejected(self): - with pytest.raises(RuntimeError, match="spec_version"): - parse_publish_payload(_payload(spec_version=1.0), WORKSPACE) - - def test_unsupported_spec_version_is_rejected(self): - with pytest.raises(RuntimeError, match="spec_version 2"): - parse_publish_payload(_payload(spec_version=2), WORKSPACE) - - def test_server_owned_name_is_rejected(self): - with pytest.raises(RuntimeError, match="server-owned"): - parse_publish_payload(_payload(name="coding-agent-configs/abc"), WORKSPACE) - - def test_server_owned_workspace_id_is_rejected(self): - with pytest.raises(RuntimeError, match="does not recognize"): - parse_publish_payload(_payload(workspace_id="123456"), WORKSPACE) - - def test_unknown_top_level_field_is_rejected(self): - with pytest.raises(RuntimeError, match="does not recognize"): - parse_publish_payload(_payload(bogus="value"), WORKSPACE) - - def test_nested_unknown_field_is_rejected(self): - payload = _payload() - payload["enabled_agents"][0]["config"] = {"unknown_setting": True} - with pytest.raises(RuntimeError, match="enabled_agents"): - parse_publish_payload(payload, WORKSPACE) - - def test_key_order_does_not_matter(self): - config = serialize_managed_config(MANIFEST) - config.pop("name", None) - reordered = {**config, "spec_version": 1, "workspace": WORKSPACE} - manifest, api_payload = parse_publish_payload(reordered, WORKSPACE) - assert api_payload["spec_version"] == 1 - assert manifest["default_agent"] == "claude" - - def test_display_name_round_trips(self): - manifest_with_name = {**MANIFEST, "display_name": "paved-path"} - manifest, api_payload = parse_publish_payload(_payload(manifest_with_name), WORKSPACE) - assert manifest["display_name"] == "paved-path" - assert api_payload["display_name"] == "paved-path" diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py deleted file mode 100644 index a7173414..00000000 --- a/tests/test_managed_wizard.py +++ /dev/null @@ -1,2993 +0,0 @@ -"""Tests for the interactive `ug setup` 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 -managed-config types, the admin gate, and the per-agent model-config shapes. -""" - -from __future__ import annotations - -import json -from decimal import Decimal -from unittest.mock import patch - -import pytest -import typer.main -from typer.testing import CliRunner - -import ucode.cli as cli_mod -import ucode.config_io as config_io_mod -import ucode.managed_config as managed_config_mod -import ucode.managed_wizard as wizard -from ucode.cli import app -from ucode.managed_setup import serialize_managed_config, validate_manifest - -runner = CliRunner() - -WORKSPACE = "https://ws.example.com" - -# `list_workspace_budgets` returns real `budget_configuration_id`s, and validation requires a -# parseable UUID, so the fixtures use one rather than a readable placeholder. -BUDGET_ID = "c6563b45-df9a-4b19-afb2-d42dc2b52576" - -STATE = { - "workspace": WORKSPACE, - "claude_models": { - "opus": "system.ai.claude-opus-4-8", - "sonnet": "system.ai.claude-sonnet-4-6", - }, - "codex_models": ["system.ai.gpt-5-6"], - "gemini_models": ["system.ai.gemini-3-flash"], - "oss_models": ["system.ai.kimi-k2-6"], -} - - -@pytest.fixture(autouse=True) -def _isolate_settings(tmp_path, monkeypatch): - """Point the managed-config file at a tmp dir so no test touches the real ~/.ucode.""" - monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - monkeypatch.setattr(managed_config_mod, "MANAGED_STATE_PATH", tmp_path / "managed-state.json") - monkeypatch.setattr(config_io_mod, "_dry_run", False) - - -class TestTracingReadback: - def test_reads_uc_destination(self): - state = {"tracing": {"enabled": True, "uc_destination": "main.default.ucode-traces"}} - assert wizard._tracing_table_from_state(state) == "main.default.ucode-traces" - - def test_disabled_tracing_yields_none(self): - state = {"tracing": {"enabled": False, "uc_destination": "main.default.t"}} - assert wizard._tracing_table_from_state(state) is None - - def test_enabled_without_destination_yields_none(self): - # A non-UC-backed experiment has no table to publish, so the manifest must omit tracing - # rather than carry an empty value the server would reject. - assert wizard._tracing_table_from_state({"tracing": {"enabled": True}}) is None - - def test_missing_tracing_yields_none(self): - assert wizard._tracing_table_from_state({}) is None - - def test_malformed_tracing_yields_none(self): - assert wizard._tracing_table_from_state({"tracing": "on"}) is None - - -class TestMcpServerFromUrl: - @pytest.mark.parametrize( - ("url", "expected"), - [ - # mcp-service stores the dash form the launch path rebuilds the dotted URL from. - ( - "https://ws.example.com/ai-gateway/mcp-services/system.ai.github", - ("system-ai-github", "mcp-service"), - ), - ("https://ws.example.com/api/2.0/mcp/external/jira-prod", ("jira-prod", "external")), - ("https://ws.example.com/api/2.0/mcp/genie/01ef", ("01ef", "genie-space")), - # vector-search / uc-functions store `.`, not the local display slug. - ( - "https://ws.example.com/api/2.0/mcp/vector-search/my_cat/my_schema", - ("my_cat.my_schema", "vector-search"), - ), - ( - "https://ws.example.com/api/2.0/mcp/functions/dev_cat/dev_fixture", - ("dev_cat.dev_fixture", "uc-functions"), - ), - ("https://ws.example.com/api/2.0/mcp/sql", ("databricks-sql", "sql")), - ], - ) - def test_known_urls(self, url, expected): - assert wizard._mcp_server_from_url(url) == expected - - def test_apps_are_not_publishable(self): - # An app's host isn't reconstructable from the workspace + an id, so it can't be published. - assert ( - wizard._mcp_server_from_url("https://mcp-myapp-123.aws.databricksapps.com/mcp") is None - ) - - def test_unknown_url_yields_none(self): - assert wizard._mcp_server_from_url("https://example.com/something/else") is None - - def test_vector_search_needs_both_catalog_and_schema(self): - assert ( - wizard._mcp_server_from_url("https://ws.example.com/api/2.0/mcp/functions/onlycat") - is None - ) - - -class TestMcpServersFromState: - def test_maps_registered_servers_to_name_and_type(self): - state = { - "mcp_servers": [ - { - "name": "databricks-github", - "url": f"{WORKSPACE}/ai-gateway/mcp-services/system.ai.github", - }, - {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - ] - } - # The published name comes from the URL (the identifier the server field holds), not the - # local display name. - assert wizard._mcp_servers_from_state(state) == [ - {"name": "system-ai-github", "type": "mcp-service"}, - {"name": "databricks-sql", "type": "sql"}, - ] - - def test_publishes_catalog_schema_for_uc_functions(self): - # The lossy local slug is replaced with the dotted catalog.schema the launch path can split. - state = { - "mcp_servers": [ - { - "name": "databricks-functions-dev-cat-dev-fixture", - "url": f"{WORKSPACE}/api/2.0/mcp/functions/dev_cat/dev_fixture", - }, - ] - } - assert wizard._mcp_servers_from_state(state) == [ - {"name": "dev_cat.dev_fixture", "type": "uc-functions"}, - ] - - def test_skips_the_skills_registry_entry(self): - # Skills are published under the manifest's own `skills` field; including the MCP entry too - # would configure them twice. - from ucode.mcp import SKILLS_MCP_KIND - - state = { - "mcp_servers": [ - { - "name": "databricks-skill-registry", - "kind": SKILLS_MCP_KIND, - "url": f"{WORKSPACE}/api/2.0/mcp/sql", - }, - {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - ] - } - assert wizard._mcp_servers_from_state(state) == [{"name": "databricks-sql", "type": "sql"}] - - def test_skips_apps_and_unclassifiable_servers(self): - state = { - "mcp_servers": [ - {"name": "mystery", "url": "https://example.com/nope"}, - {"name": "databricks-app-x", "url": "https://x-1.databricksapps.com/mcp"}, - ] - } - assert wizard._mcp_servers_from_state(state) == [] - - def test_skips_entries_missing_name_or_url(self): - state = { - "mcp_servers": [ - {"url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - {"name": "no-url"}, - "not-a-dict", - ] - } - assert wizard._mcp_servers_from_state(state) == [] - - def test_empty_state_yields_nothing(self): - assert wizard._mcp_servers_from_state({}) == [] - - def test_output_validates_as_a_manifest(self): - state = { - "mcp_servers": [ - {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - ] - } - servers = wizard._mcp_servers_from_state(state) - assert validate_manifest({"mcp_servers": servers}) == [] - - -class TestAdminGate: - def test_non_admin_is_rejected(self): - with patch.object(wizard, "is_workspace_admin", return_value=False): - with pytest.raises(RuntimeError, match="not an admin"): - wizard._require_admin(WORKSPACE, "token") - - def test_admin_passes(self): - with patch.object(wizard, "is_workspace_admin", return_value=True): - wizard._require_admin(WORKSPACE, "token") # must not raise - - def test_unverifiable_check_warns_and_continues(self): - # A failed SCIM call must not block a legitimate admin — the API enforces the same rule. - with ( - patch.object(wizard, "is_workspace_admin", return_value=None), - patch.object(wizard, "print_warning") as warn, - ): - wizard._require_admin(WORKSPACE, "token") - 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 "`ug 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 `ug 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 `ug configure` run - doesn't show `ug setup` headers.""" - - def test_defaults_to_ucode_setup(self): - with patch.object(wizard, "print_section") as section: - wizard._step_banner(1, "Agents") - assert section.call_args.args[0].startswith("ug setup · step 1 of ") - - def test_uses_the_command_label_when_given(self): - with patch.object(wizard, "print_section") as section: - wizard._step_banner(2, "Models", "ug configure") - assert section.call_args.args[0].startswith("ug configure · step 2 of ") - - -class TestSetupCommandToken: - """A caller (e.g. `ug configure`) can hand setup a token so its admin gate uses the same - identity as the routing decision, instead of fetching a second time.""" - - def test_reuses_a_passed_token_and_skips_a_second_fetch(self): - seen: list[tuple[str, str]] = [] - with ( - 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"), - ), - patch.object( - wizard, "_require_admin", side_effect=lambda ws, tok: seen.append((ws, tok)) - ), - # 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)), - ): - code = wizard.setup_command(workspace="https://w", profile=None, token="tok") - assert code == 0 - assert seen == [("https://w", "tok")] - - 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) - assert code == 0 - assert seen == [("https://w", "fetched")] - - -class TestModelPrompting: - def test_codex_takes_a_single_model(self): - with patch.object(wizard, "prompt_for_selection", return_value="system.ai.gpt-5-6"): - config = wizard._prompt_models_for_agent("codex", STATE, None) - # CodexModelConfig has no model list, so the wizard must not build one. - assert config == {"default_model": "system.ai.gpt-5-6"} - - def test_claude_prompts_one_slot_per_family(self): - # Claude Code selects by family alias, so each `ClaudeDefaultModels` slot gets its own - # prompt — and each shows that family's real alternatives, not just the newest. - candidates = { - "opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], - "sonnet": ["system.ai.claude-sonnet-5"], - } - asked: list[str] = [] - - def fake_sel(prompt, options, **kwargs): - asked.append(prompt) - return [v for v, _ in options][0] - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - - assert config["models"] == { - "default_opus_model": "system.ai.claude-opus-5", - "default_sonnet_model": "system.ai.claude-sonnet-5", - } - assert any("opus" in p for p in asked) and any("sonnet" in p for p in asked) - - def test_claude_offers_every_version_in_a_family(self): - candidates = {"opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"]} - offered: list[list[str]] = [] - - def fake_sel(prompt, options, **kwargs): - values = [v for v, _ in options] - offered.append(values) - return values[1] # pick the older opus on purpose - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - - # Pinning a known-good older version has to be expressible. - assert config["models"] == {"default_opus_model": "system.ai.claude-opus-4-8"} - assert "system.ai.claude-opus-4-8" in offered[0] - - def test_claude_families_can_be_skipped(self): - candidates = { - "opus": ["system.ai.claude-opus-5"], - "sonnet": ["system.ai.claude-sonnet-5"], - } - - def fake_sel(prompt, options, **kwargs): - values = [v for v, _ in options] - return wizard._SKIP_FAMILY if "sonnet" in prompt else values[0] - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - - # Every slot is optional in the proto; a skipped one must simply be absent. - assert config["models"] == {"default_opus_model": "system.ai.claude-opus-5"} - - def test_claude_overall_default_comes_from_the_filled_slots(self): - candidates = { - "opus": ["system.ai.claude-opus-5"], - "sonnet": ["system.ai.claude-sonnet-5"], - } - prompts: list[list[str]] = [] - - def fake_sel(prompt, options, **kwargs): - values = [v for v, _ in options] - prompts.append(values) - return values[0] - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - - # The last prompt is the overall default, offering only what the slots hold — so it can - # never name a model the config doesn't carry. - assert set(prompts[-1]) == {"system.ai.claude-opus-5", "system.ai.claude-sonnet-5"} - assert config["default_model"] in config["models"].values() - - def test_claude_single_slot_skips_the_default_prompt(self): - candidates = {"opus": ["system.ai.claude-opus-5"]} - calls = {"n": 0} - - def fake_sel(prompt, options, **kwargs): - calls["n"] += 1 - return [v for v, _ in options][0] - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - - assert calls["n"] == 1 # only the opus prompt; no redundant default question - assert config["default_model"] == "system.ai.claude-opus-5" - - def test_claude_falls_back_to_text_when_nothing_discovered(self): - with ( - patch.object(wizard, "_claude_candidates", return_value={}), - patch.object(wizard, "prompt_for_text", return_value="some-claude"), - patch.object(wizard, "print_warning"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - assert config == {"default_model": "some-claude"} - - def test_single_slot_announces_the_inferred_default(self): - # The one-option prompt is skipped, but silence reads as a dropped step — the admin has to - # learn that the default was inferred. Announced as a note; the per-agent ✔ (`_confirm_agent`) - # in the setup loop carries the final confirmation. - candidates = {"opus": ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"]} - - def fake_sel(prompt, options, **kwargs): - if prompt.startswith("Default opus"): - return "system.ai.claude-opus-4-8" - return wizard._SKIP_FAMILY - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "print_note") as note, - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - - assert config["default_model"] == "system.ai.claude-opus-4-8" - notes = " ".join(str(call.args[0]) for call in note.call_args_list) - assert "overall default" in notes - - def test_claude_all_families_skipped_still_picks_from_the_candidates(self): - # Skipping every slot is a legitimate minimal config — `models` is optional and each unset - # slot falls back to `default_model`, so one model covers every family. The admin shouldn't - # have to type an id we already have. - candidates = { - "opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], - "sonnet": ["system.ai.claude-sonnet-5"], - } - offered: list[list[str]] = [] - - def fake_sel(prompt, options, **kwargs): - values = [v for v, _ in options] - offered.append(values) - if prompt.startswith("Default "): - return wizard._SKIP_FAMILY - return "system.ai.claude-opus-4-8" - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "prompt_for_text") as text, - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - - assert config == {"default_model": "system.ai.claude-opus-4-8"} - assert "models" not in config - assert not text.called, "should pick from candidates, not ask for free text" - # The final prompt offers every candidate across all families (all fit under the picker - # limit here), plus the custom-entry row. - assert set(offered[-1]) == { - "system.ai.claude-opus-5", - "system.ai.claude-opus-4-8", - "system.ai.claude-sonnet-5", - wizard._CUSTOM_MODEL, - } - - def test_older_claude_version_passes_validation(self): - # The picker offers every version in a family, but `claude_models` holds only the newest — - # so validation has to learn about the rest or it rejects a legitimate pick. - candidates = {"opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"]} - state = { - "workspace": "https://ws.example.com", - "claude_models": {"opus": "system.ai.claude-opus-5"}, - } - - def fake_unbucketed(workspace, token): - return ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], None - - with ( - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "discover_claude_models_unbucketed", fake_unbucketed), - patch.object(wizard, "claude_family_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", return_value="system.ai.claude-opus-4-8"), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", state, None) - - manifest = { - "default_agent": "claude", - "enabled_agents": {"claude": {"model_config": config}}, - } - assert validate_manifest(manifest, state) == [] - - def test_claude_cancelled_family_prompt_aborts(self): - with ( - patch.object( - wizard, "_claude_candidates", return_value={"opus": ["system.ai.claude-opus-5"]} - ), - patch.object(wizard, "prompt_for_selection", return_value=None), - patch.object(wizard, "print_note"), - ): - with pytest.raises(KeyboardInterrupt): - wizard._prompt_models_for_agent("claude", STATE, None) - - def test_single_model_agents_get_one_prompt(self): - # Gemini and Copilot declare `repeated string models` in the proto, but their config writers - # take one model and write one env var — a published list would be read by nothing. - for tool in ("codex", "gemini", "copilot"): - options = wizard.model_options_for_agent(tool, STATE) - with ( - patch.object(wizard, "prompt_for_selection", return_value=options[0]) as select, - patch.object(wizard, "prompt_for_multi_selection") as multi, - ): - config = wizard._prompt_models_for_agent(tool, STATE, None) - assert select.called, tool - assert not multi.called, tool - assert config == {"default_model": options[0]}, tool - assert "models" not in config, tool - - def test_claude_catalog_is_fetched_once(self): - # `configure_shared_state` already paged the whole catalog; re-fetching per claude prompt - # pages it again for no new information. - calls = {"n": 0} - - def fake_fetch(workspace, token): - calls["n"] += 1 - return ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], None - - state = {"workspace": "https://ws.example.com", "profile": "p"} - with ( - patch.object(wizard, "discover_claude_models_unbucketed", side_effect=fake_fetch), - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "prompt_for_selection", return_value="system.ai.claude-opus-5"), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - wizard._prompt_models_for_agent("claude", state, None) - wizard._prompt_models_for_agent("claude", state, None) - - assert calls["n"] == 1 - assert state["all_claude_models"] - - def test_nothing_is_prechecked(self): - # 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. - captured: dict = {} - - def fake_multi(prompt, options, preselected=None, **kwargs): - captured["preselected"] = preselected - return [v for v, _ in options][:1] - - with ( - patch.object(wizard, "prompt_for_multi_selection", side_effect=fake_multi), - patch.object(wizard, "prompt_for_selection", return_value="x"), - ): - wizard._prompt_models_for_agent("pi", STATE, None) - - assert not captured["preselected"] - - def test_list_agents_still_multi_select(self): - # OpenCode and Pi really do show a model picker, so their lists are honoured. - for tool in ("opencode", "pi"): - options = wizard.model_options_for_agent(tool, STATE) - with ( - patch.object(wizard, "prompt_for_multi_selection", return_value=options[:2]), - patch.object(wizard, "prompt_for_selection", return_value=options[0]), - ): - config = wizard._prompt_models_for_agent(tool, STATE, None) - assert config["models"] == options[:2], tool - - def test_flat_list_agents_keep_the_picked_list(self): - picked = ["system.ai.claude-opus-4-8", "system.ai.kimi-k2-6"] - with ( - patch.object(wizard, "prompt_for_multi_selection", return_value=picked), - patch.object(wizard, "prompt_for_selection", return_value=picked[0]), - ): - config = wizard._prompt_models_for_agent("opencode", STATE, None) - assert config["models"] == picked - - def test_single_pick_skips_the_default_prompt(self): - with ( - patch.object( - wizard, "prompt_for_multi_selection", return_value=["system.ai.claude-opus-4-8"] - ), - patch.object(wizard, "prompt_for_selection") as select, - ): - config = wizard._prompt_models_for_agent("pi", STATE, None) - assert config["default_model"] == "system.ai.claude-opus-4-8" - assert not select.called - - def test_provider_service_offers_its_targets(self): - # The service's own targets are the model vocabulary the manifest must use, so the admin - # picks from them rather than typing an id from memory. Uses codex here: it takes one model - # from any service (Claude with explicit targets is prompted per family — see the Bedrock/ - # Anthropic per-family tests below). - service = { - "name": "main.default.openai-mps", - "provider_type": "openai", - "targets": ["gpt-5-6", "gpt-5-6-sol"], - "allow_all_targets": False, - } - with ( - patch.object(wizard, "prompt_for_selection", return_value="gpt-5-6") as select, - patch.object(wizard, "prompt_for_text") as text, - ): - config = wizard._prompt_models_for_agent("codex", STATE, service) - assert config == { - "model_provider_service": "main.default.openai-mps", - "default_model": "gpt-5-6", - } - assert not text.called, "should not fall back to free text when targets are known" - # Offered sorted, so the picker order is stable run to run. - assert [value for value, _ in select.call_args[0][1]] == ["gpt-5-6", "gpt-5-6-sol"] - - def test_provider_service_falls_back_to_text_when_targets_unknown(self): - # allow_all_targets passes the provider's whole catalog through; there is nothing to list. - service = { - "name": "main.default.anthropic-mps", - "provider_type": "anthropic", - "targets": [], - "allow_all_targets": True, - } - with ( - patch.object(wizard, "prompt_for_text", return_value="claude-sonnet-5"), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert config == { - "model_provider_service": "main.default.anthropic-mps", - "default_model": "claude-sonnet-5", - } - - def test_relayed_service_falls_back_to_text(self): - # A relayed Anthropic subscription service routes by canonical name, with no target list. - service = { - "name": "main.default.lilly-anthropic", - "provider_type": "anthropic", - "targets": [], - "allow_all_targets": False, - "relayed": True, - } - with ( - patch.object(wizard, "prompt_for_text", return_value="claude-sonnet-4-6"), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert config["default_model"] == "claude-sonnet-4-6" - - def test_falls_back_to_free_text_when_nothing_discovered(self): - with ( - patch.object(wizard, "prompt_for_text", return_value="some-model"), - patch.object(wizard, "print_warning"), - ): - config = wizard._prompt_models_for_agent("pi", {}, None) - assert config == {"default_model": "some-model"} - - def test_bedrock_claude_prompts_per_family(self): - # render_overlay pins ANTHROPIC_DEFAULT__MODEL from a Bedrock service's provider-side - # ids (Claude Code can't route its canonical names there), so Claude needs a default per - # family — a single overall default would leave the other families unpinned. - service = { - "name": "main.default.bedrock", - "provider_type": "amazon_bedrock", - "allow_all_targets": False, - "targets": [ - "anthropic.claude-opus-4-1", - "anthropic.claude-opus-4-8", - "anthropic.claude-sonnet-4-6", - "anthropic.claude-haiku-4-5", - ], - } - asked: list[str] = [] - - def fake_selection(prompt, options, **kwargs): - asked.append(prompt) - return options[0][0] - - with ( - # Decline quick setup to exercise the per-family prompts. - patch.object(wizard, "prompt_yes_no_default", return_value=False), - patch.object(wizard, "prompt_for_selection", side_effect=fake_selection), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert [p for p in asked if p.startswith("Default ")] == [ - "Default opus model:", - "Default sonnet model:", - "Default haiku model:", - ] - assert config["model_provider_service"] == "main.default.bedrock" - assert config["models"] == { - "default_opus_model": "anthropic.claude-opus-4-1", - "default_sonnet_model": "anthropic.claude-sonnet-4-6", - "default_haiku_model": "anthropic.claude-haiku-4-5", - } - # Slots carry the service's own provider-side ids, not workspace `system.ai.*` ids. - assert all(m.startswith("anthropic.") for m in config["models"].values()) - assert config["default_model"] in service["targets"] - - def test_bedrock_claude_skipped_family_is_omitted(self): - service = { - "name": "main.default.bedrock", - "provider_type": "amazon_bedrock", - "allow_all_targets": False, - "targets": ["anthropic.claude-opus-4-8", "anthropic.claude-sonnet-4-6"], - } - # pick=-1 selects the trailing "(skip )" option for every family. - with ( - patch.object(wizard, "prompt_yes_no_default", return_value=False), - patch.object(wizard, "prompt_for_selection", side_effect=lambda p, o, **k: o[-1][0]), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert "models" not in config - assert config["default_model"] in service["targets"] - - def test_bedrock_codex_still_takes_a_single_model(self): - # Only Claude pins per family; codex through any provider takes one model. - service = { - "name": "main.default.bedrock-oai", - "provider_type": "amazon_bedrock", - "allow_all_targets": False, - "targets": ["gpt-5-6", "gpt-5-6-sol"], - } - asked: list[str] = [] - with ( - patch.object( - wizard, - "prompt_for_selection", - side_effect=lambda p, o, **k: (asked.append(p), o[0][0])[1], - ), - ): - config = wizard._prompt_models_for_agent("codex", STATE, service) - assert len(asked) == 1 - assert "models" not in config - - def test_anthropic_claude_with_explicit_targets_prompts_per_family(self): - # An Anthropic service that publishes explicit canonical targets IS pinned per family by - # render_overlay (each ANTHROPIC_DEFAULT__MODEL to a chosen version), so the wizard - # prompts per family — same as Bedrock, just canonical ids instead of provider slugs. - service = { - "name": "main.default.ant", - "provider_type": "anthropic", - "allow_all_targets": False, - "targets": ["claude-opus-4-8", "claude-sonnet-4-6"], - } - asked: list[str] = [] - with ( - patch.object(wizard, "prompt_yes_no_default", return_value=False), - patch.object( - wizard, - "prompt_for_selection", - side_effect=lambda p, o, **k: (asked.append(p), o[0][0])[1], - ), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert [p for p in asked if p.startswith("Default ")] == [ - "Default opus model:", - "Default sonnet model:", - ] - assert config["models"] == { - "default_opus_model": "claude-opus-4-8", - "default_sonnet_model": "claude-sonnet-4-6", - } - - def test_quick_setup_fills_newest_per_family_without_per_family_prompts(self): - # Quick setup (default yes) auto-fills each family with the service's newest id — same pick - # map_claude_family_models makes — and asks no per-family questions. - service = { - "name": "main.default.bedrock", - "provider_type": "amazon_bedrock", - "allow_all_targets": False, - "targets": [ - "anthropic.claude-opus-4-1", - "anthropic.claude-opus-4-8", - "anthropic.claude-sonnet-4-6", - ], - } - selections: list[str] = [] - with ( - patch.object(wizard, "prompt_yes_no_default", return_value=True), - patch.object( - wizard, - "prompt_for_selection", - side_effect=lambda p, o, **k: selections.append(p) or o[0][0], - ), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - # No per-family (or overall-default) picker was shown — quick setup asked nothing. - assert selections == [] - assert config["models"] == { - "default_opus_model": "anthropic.claude-opus-4-8", # newest opus, not 4-1 - "default_sonnet_model": "anthropic.claude-sonnet-4-6", - } - # Overall default is opus (the flagship), not sonnet. - assert config["default_model"] == "anthropic.claude-opus-4-8" - - def test_quick_setup_default_is_flagship_not_target_order(self): - # Regression: the overall default must be the highest-tier family (opus), not whichever - # target sorts first — here haiku leads the list but opus must still win. - service = { - "name": "main.default.bedrock", - "provider_type": "amazon_bedrock", - "allow_all_targets": False, - "targets": [ - "us.anthropic.claude-haiku-4-5", - "us.anthropic.claude-sonnet-5", - "us.anthropic.claude-opus-5", - ], - } - with ( - patch.object(wizard, "prompt_yes_no_default", return_value=True), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert config["default_model"] == "us.anthropic.claude-opus-5" - - def test_quick_setup_default_falls_to_sonnet_without_opus(self): - service = { - "name": "main.default.bedrock", - "provider_type": "amazon_bedrock", - "allow_all_targets": False, - "targets": ["us.anthropic.claude-haiku-4-5", "us.anthropic.claude-sonnet-5"], - } - with ( - patch.object(wizard, "prompt_yes_no_default", return_value=True), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert config["default_model"] == "us.anthropic.claude-sonnet-5" - - def test_anthropic_claude_allow_all_takes_a_single_default(self): - # No explicit targets (allow_all) means nothing is pinned — Claude Code's canonical names - # route fine — so it falls back to a single free-text default, not per-family slots. - service = { - "name": "main.default.ant-all", - "provider_type": "anthropic", - "allow_all_targets": True, - "targets": [], - } - with ( - patch.object(wizard, "prompt_for_text", return_value="claude-sonnet-5"), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - assert "models" not in config - assert config["default_model"] == "claude-sonnet-5" - - def test_allow_all_with_explicit_claude_targets_does_not_abort(self): - # Regression: a service that is allow_all_targets AND lists Claude targets. The family - # decision must key on the enumerated targets (which allow_all zeroes), not the raw service — - # otherwise it takes the per-family branch with an empty list and aborts the wizard. - service = { - "name": "main.default.weird", - "provider_type": "amazon_bedrock", - "allow_all_targets": True, - "targets": ["us.anthropic.claude-opus-5", "us.anthropic.claude-sonnet-5"], - } - with ( - patch.object(wizard, "prompt_for_text", return_value="typed-model"), - patch.object(wizard, "print_note"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, service) - # Falls through to the single free-text default; no per-family slots, no KeyboardInterrupt. - assert "models" not in config - assert config["default_model"] == "typed-model" - - def test_empty_selection_is_re_prompted(self): - # An agent with no default_model can't be the config's default_agent (the server rejects - # it) and gives developers nothing to launch, so "none" is re-asked rather than accepted. - with ( - patch.object( - wizard, - "prompt_for_multi_selection", - side_effect=[[], ["system.ai.claude-opus-4-8"]], - ) as picker, - patch.object(wizard, "print_err") as err, - ): - config = wizard._prompt_models_for_agent("pi", STATE, None) - assert picker.call_count == 2 - assert err.called - assert config["default_model"] == "system.ai.claude-opus-4-8" - - def test_every_agent_always_gets_a_default_model(self): - # The invariant the late-validation bug violated: no agent can come back model-less. - for tool in ("claude", "codex", "gemini", "opencode", "pi", "copilot"): - options = wizard.model_options_for_agent(tool, STATE) - with ( - patch.object(wizard, "prompt_for_multi_selection", return_value=[options[0]]), - patch.object(wizard, "prompt_for_selection", return_value=options[0]), - # Without this the claude pass reaches for the real catalog, which shells out to - # `databricks auth token` and depends on the machine's CLI and credentials. - patch.object(wizard, "discover_claude_models_unbucketed", return_value=([], None)), - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - config = wizard._prompt_models_for_agent(tool, STATE, None) - assert config.get("default_model"), tool - - def test_claude_candidates_survive_a_missing_databricks_cli(self): - # `get_databricks_token` shells out, so a machine without the CLI on PATH raises - # FileNotFoundError, not RuntimeError. That must degrade to the bucketed per-family picks - # rather than aborting the wizard mid-flow. - def no_cli(*args, **kwargs): - raise FileNotFoundError(2, "No such file or directory", "databricks") - - with patch.object(wizard, "get_databricks_token", side_effect=no_cli): - candidates = wizard._claude_candidates(dict(STATE)) - - # STATE's bucketed `claude_models` still supplies one model per family. - assert candidates["opus"] == ["system.ai.claude-opus-4-8"] - assert candidates["sonnet"] == ["system.ai.claude-sonnet-4-6"] - - def test_default_model_is_a_bare_uc_id(self): - # Provider prefixes (e.g. opencode's `databricks-anthropic/`) are added by each agent's own - # writer, so the manifest stays agent-neutral. - with ( - patch.object( - wizard, "prompt_for_multi_selection", return_value=["system.ai.claude-opus-4-8"] - ), - ): - config = wizard._prompt_models_for_agent("opencode", STATE, None) - assert config["default_model"] == "system.ai.claude-opus-4-8" - assert "/" not in config["default_model"] - - def test_dismissed_single_select_aborts_instead_of_re_prompting(self): - # This used to re-prompt. questionary's `ask` swallows Ctrl-C and returns None, so "empty - # submission" and "user aborted" are the same value here — and re-asking spun forever. - # Asserted on the helper directly: the codex path only reaches the picker when discovery - # found models, and falling through to the free-text branch would read real stdin. - with ( - patch.object(wizard, "prompt_for_selection", return_value=None) as picker, - patch.object(wizard, "print_err"), - ): - with pytest.raises(KeyboardInterrupt): - wizard._require_selection("Select the model:", [("a", "A")]) - assert picker.call_count == 1 - - def test_empty_free_text_is_re_prompted(self): - with ( - patch.object(wizard, "prompt_for_text", side_effect=[None, "some-model"]) as text, - patch.object(wizard, "print_err"), - patch.object(wizard, "print_warning"), - ): - config = wizard._prompt_models_for_agent("pi", {}, None) - assert text.call_count == 2 - assert config == {"default_model": "some-model"} - - def test_cancelled_picker_aborts(self): - with patch.object(wizard, "prompt_for_multi_selection", return_value=None): - with pytest.raises(KeyboardInterrupt): - wizard._prompt_models_for_agent("pi", STATE, None) - - -ANTHROPIC_SERVICE = { - "name": "main.default.lilly-anthropic", - "provider_type": "anthropic", - "targets": ["claude-sonnet-4-6"], - "allow_all_targets": False, - "relayed": False, -} -OPENAI_SERVICE = { - "name": "main.default.openai-mps", - "provider_type": "openai", - "targets": ["gpt-5-6"], - "allow_all_targets": False, - "relayed": False, -} - - -class TestCustomModelEntry: - """The hosted-model pickers let an admin type a model service discovery didn't list.""" - - def test_single_agent_can_enter_a_verified_custom_model(self): - # Selecting the custom row prompts for a path, verifies it exists, and records it so - # validation won't reject a model outside the discovered inventory. - with ( - patch.object(wizard, "prompt_for_selection", return_value=wizard._CUSTOM_MODEL), - patch.object(wizard, "prompt_for_text", return_value="main.aarushi.gpt-5-custom"), - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "model_service_exists", return_value=(True, None)) as exists, - ): - config = wizard._prompt_models_for_agent("codex", STATE, None) - assert config == { - "default_model": "main.aarushi.gpt-5-custom", - "custom_models": ["main.aarushi.gpt-5-custom"], - } - exists.assert_called_once_with(WORKSPACE, "tok", "main.aarushi.gpt-5-custom") - - def test_custom_model_reprompts_until_it_exists(self): - # A typo shouldn't get baked into a published config — a miss re-asks. - with ( - patch.object(wizard, "prompt_for_selection", return_value=wizard._CUSTOM_MODEL), - patch.object( - wizard, "prompt_for_text", side_effect=["main.typo.model", "main.real.model"] - ), - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "model_service_exists", side_effect=[(False, None), (True, None)]), - patch.object(wizard, "print_err") as err, - ): - config = wizard._prompt_models_for_agent("codex", STATE, None) - assert config["default_model"] == "main.real.model" - assert err.called - - def test_multi_select_agent_can_add_several_custom_models(self): - # opencode/pi carry a model list, so the custom row keeps prompting until the admin declines. - with ( - patch.object(wizard, "prompt_for_multi_selection", return_value=[wizard._CUSTOM_MODEL]), - patch.object( - wizard, "prompt_for_text", side_effect=["main.default.a", "main.default.b"] - ), - patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), - patch.object(wizard, "prompt_for_selection", return_value="main.default.a"), - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "model_service_exists", return_value=(True, None)), - ): - config = wizard._prompt_models_for_agent("pi", STATE, None) - assert config["models"] == ["main.default.a", "main.default.b"] - assert config["custom_models"] == ["main.default.a", "main.default.b"] - assert config["default_model"] == "main.default.a" - - def test_short_reason_drops_the_json_body(self): - # The warning for an inconclusive check should show the status, not the raw error blob. - assert ( - wizard._short_reason('HTTP 500 Server Error: {"error_code":"INTERNAL"}') - == "HTTP 500 Server Error" - ) - assert wizard._short_reason("network error: timed out") == "network error: timed out" - assert wizard._short_reason(None) == "unknown error" - - def test_unverifiable_custom_model_is_accepted_with_a_warning(self): - # An inconclusive check (transient error / no token) must not block a possibly-valid model. - with ( - patch.object(wizard, "prompt_for_selection", return_value=wizard._CUSTOM_MODEL), - patch.object(wizard, "prompt_for_text", return_value="main.aarushi.maybe"), - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "model_service_exists", return_value=(None, "HTTP 500")), - patch.object(wizard, "print_warning") as warn, - ): - config = wizard._prompt_models_for_agent("codex", STATE, None) - assert config["default_model"] == "main.aarushi.maybe" - assert warn.called - - def test_picker_offers_all_models_plus_custom(self): - # Every discovered model is offered (the searchable picker scrolls), with the custom row last. - models = [f"system.ai.gpt-5-{i}" for i in range(10)] - state = {**STATE, "codex_models": list(models)} - offered: list[list[str]] = [] - - def fake_sel(prompt, options, **kwargs): - offered.append([v for v, _ in options]) - return options[0][0] - - with patch.object(wizard, "prompt_for_selection", side_effect=fake_sel): - wizard._prompt_models_for_agent("codex", state, None) - rows = offered[0] - assert rows[-1] == wizard._CUSTOM_MODEL - # All 10 discovered ids are offered — no truncation to the old top-few cap. - assert set(models) <= set(rows[:-1]) - - def test_claude_family_custom_model_slots_and_validates(self): - # A custom id chosen for a family lands in that slot, is marked custom, and survives - # validation even though discovery never listed it. - candidates = {"opus": ["system.ai.claude-opus-5"]} - - def fake_sel(prompt, options, **kwargs): - values = [v for v, _ in options] - if wizard._CUSTOM_MODEL in values: - return wizard._CUSTOM_MODEL - return values[0] - - with ( - patch.object(wizard, "_claude_candidates", return_value=candidates), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "prompt_for_text", return_value="main.aarushi.claude-opus-4-5"), - patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), - patch.object(wizard, "model_service_exists", return_value=(True, None)), - patch.object(wizard, "print_note"), - patch.object(wizard, "print_success"), - ): - config = wizard._prompt_models_for_agent("claude", STATE, None) - assert config["models"]["default_opus_model"] == "main.aarushi.claude-opus-4-5" - assert config["custom_models"] == ["main.aarushi.claude-opus-4-5"] - assert ( - validate_manifest( - { - "default_agent": "claude", - "enabled_agents": {"claude": {"model_config": config}}, - }, - STATE, - ) - == [] - ) - - -class TestProviderServiceSpinner: - """The MPS listing is cached per workspace, so only the first agent's lookup does any I/O.""" - - SERVICES = [ - { - "name": "main.j.ant", - "provider_type": "anthropic", - "targets": ["claude-opus-5"], - "allow_all_targets": False, - "relayed": False, - } - ] - - def test_spinner_shows_once_not_once_per_agent(self): - # The reported symptom: "Checking for model provider services for ..." appeared for - # every configured agent even though the listing had already been fetched. - spins: list[str] = [] - cached = {"yes": False} - - def fake_list(workspace, token, **kwargs): - cached["yes"] = True - return list(self.SERVICES), None - - def fake_spinner(message): - spins.append(message) - from contextlib import nullcontext - - return nullcontext() - - with ( - patch.object(wizard, "list_model_provider_services", side_effect=fake_list), - patch.object( - wizard, "has_cached_model_provider_services", side_effect=lambda ws: cached["yes"] - ), - patch.object(wizard, "spinner", side_effect=fake_spinner), - patch.object(wizard, "prompt_for_selection", return_value="databricks"), - ): - wizard._select_provider_service("claude", WORKSPACE, "tok") - wizard._select_provider_service("codex", WORKSPACE, "tok") - - listing_spins = [m for m in spins if "provider service" in m] - assert len(listing_spins) == 1, listing_spins - # And it doesn't name an agent, since one lookup covers them all. - assert "Claude Code" not in listing_spins[0] - - -class TestProviderServiceSelection: - def test_agents_without_provider_support_skip_the_prompt(self): - with patch.object(wizard, "list_model_provider_services") as listing: - assert wizard._select_provider_service("opencode", WORKSPACE, "token") is None - assert not listing.called - - def test_feature_disabled_is_silent(self): - # The common case on most workspaces; a warning here would be noise. - with ( - patch.object(wizard, "list_model_provider_services", return_value=([], "HTTP 404")), - patch.object(wizard, "is_model_provider_feature_unavailable", return_value=True), - patch.object(wizard, "print_warning") as warn, - ): - assert wizard._select_provider_service("claude", WORKSPACE, "token") is None - assert not warn.called - - def test_unexpected_listing_failure_warns(self): - # Without this the admin silently loses the MPS option with no idea why. - with ( - patch.object( - wizard, "list_model_provider_services", return_value=([], "HTTP 403 Forbidden") - ), - patch.object(wizard, "is_model_provider_feature_unavailable", return_value=False), - patch.object(wizard, "print_warning") as warn, - patch.object(wizard, "print_note"), - ): - assert wizard._select_provider_service("claude", WORKSPACE, "token") is None - assert warn.called - assert "403" in warn.call_args[0][0] - - def test_services_exist_but_none_match_the_agent_explains_why(self): - # An openai-only workspace offers claude nothing; say so rather than showing no picker. - with ( - patch.object( - wizard, "list_model_provider_services", return_value=([OPENAI_SERVICE], None) - ), - patch.object(wizard, "print_note") as note, - ): - assert wizard._select_provider_service("claude", WORKSPACE, "token") is None - assert note.called - assert "API dialect" in note.call_args[0][0] - - def test_choosing_databricks_returns_none(self): - with ( - patch.object( - wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) - ), - patch.object(wizard, "prompt_for_selection", return_value="databricks"), - ): - assert wizard._select_provider_service("claude", WORKSPACE, "token") is None - - def test_choosing_mps_returns_the_whole_service(self): - # The dict (not just the name) is returned so the model prompt can offer its targets. - with ( - patch.object( - wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) - ), - patch.object( - wizard, - "prompt_for_selection", - side_effect=["mps", "main.default.lilly-anthropic"], - ), - patch.object(wizard, "all_users_can_use_schema", return_value=True), - ): - service = wizard._select_provider_service("claude", WORKSPACE, "token") - assert service == ANTHROPIC_SERVICE - - def test_only_matching_services_are_offered(self): - with ( - patch.object( - wizard, - "list_model_provider_services", - return_value=([ANTHROPIC_SERVICE, OPENAI_SERVICE], None), - ), - patch.object( - wizard, - "prompt_for_selection", - side_effect=["mps", "main.default.lilly-anthropic"], - ) as select, - patch.object(wizard, "all_users_can_use_schema", return_value=True), - ): - wizard._select_provider_service("claude", WORKSPACE, "token") - offered = [value for value, _ in select.call_args_list[1][0][1]] - assert offered == ["main.default.lilly-anthropic"] - - def test_relayed_services_are_not_offered_for_claude(self): - relayed = { - **ANTHROPIC_SERVICE, - "name": "main.default.claude-enterprise", - "targets": [], - "relayed": True, - } - with ( - patch.object( - wizard, - "list_model_provider_services", - return_value=([relayed, ANTHROPIC_SERVICE], None), - ), - patch.object( - wizard, - "prompt_for_selection", - side_effect=["mps", "main.default.lilly-anthropic"], - ) as select, - patch.object(wizard, "all_users_can_use_schema", return_value=True), - ): - service = wizard._select_provider_service("claude", WORKSPACE, "token") - assert service == ANTHROPIC_SERVICE - offered = [value for value, _ in select.call_args_list[1][0][1]] - assert offered == ["main.default.lilly-anthropic"] - - def test_only_relayed_services_falls_back_to_databricks_for_claude(self): - relayed = {**ANTHROPIC_SERVICE, "targets": [], "relayed": True} - with ( - patch.object(wizard, "list_model_provider_services", return_value=([relayed], None)), - patch.object(wizard, "prompt_for_selection") as select, - patch.object(wizard, "print_note"), - ): - assert wizard._select_provider_service("claude", WORKSPACE, "token") is None - assert not select.called - - def test_warns_when_all_users_lack_schema_access(self): - # The picked MPS's schema isn't granted to all workspace users, so developers who pull the - # config may hit "does not have USE_SCHEMA"; warn but still return the service (never block). - with ( - patch.object( - wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) - ), - patch.object( - wizard, "prompt_for_selection", side_effect=["mps", "main.default.lilly-anthropic"] - ), - patch.object(wizard, "all_users_can_use_schema", return_value=False), - patch.object(wizard, "print_warning") as warn, - ): - service = wizard._select_provider_service("claude", WORKSPACE, "token") - assert service == ANTHROPIC_SERVICE - assert warn.called - assert "main.default" in warn.call_args[0][0] - - def test_no_warning_when_access_check_is_inconclusive(self): - # A None result (API unreachable / unexpected shape) must not cry wolf. - with ( - patch.object( - wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) - ), - patch.object( - wizard, "prompt_for_selection", side_effect=["mps", "main.default.lilly-anthropic"] - ), - patch.object(wizard, "all_users_can_use_schema", return_value=None), - patch.object(wizard, "print_warning") as warn, - ): - service = wizard._select_provider_service("claude", WORKSPACE, "token") - assert service == ANTHROPIC_SERVICE - assert not warn.called - - def test_cancelling_the_service_picker_returns_none(self): - with ( - patch.object( - wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) - ), - patch.object(wizard, "prompt_for_selection", side_effect=["mps", None]), - ): - assert wizard._select_provider_service("claude", WORKSPACE, "token") is None - - -class TestProviderServiceModelOptions: - def test_returns_sorted_targets(self): - service = {"targets": ["b-model", "a-model"], "allow_all_targets": False} - assert wizard.provider_service_model_options(service) == ["a-model", "b-model"] - - def test_deduplicates(self): - service = {"targets": ["m", "m"], "allow_all_targets": False} - assert wizard.provider_service_model_options(service) == ["m"] - - def test_allow_all_targets_yields_nothing(self): - service = {"targets": ["m"], "allow_all_targets": True} - assert wizard.provider_service_model_options(service) == [] - - def test_missing_targets_yields_nothing(self): - assert wizard.provider_service_model_options({}) == [] - - def test_malformed_targets_yield_nothing(self): - assert wizard.provider_service_model_options({"targets": "m"}) == [] - - -# Agents as the wizard configures them: the tier picker must offer these, not the workspace catalog. -CLAUDE_ONLY = {"claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}} - - -class TestBudgetPolicy: - def test_no_up_front_gate(self): - # Running `ug setup 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 ( - patch.object(wizard, "prompt_yes_no_default") as ask, - patch.object(wizard, "list_workspace_budgets", return_value=([], "none found")), - patch.object(wizard, "print_warning_panel"), - ): - wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) - assert not ask.called - - def test_no_budgets_warns_in_a_box_and_skips_the_blurb(self): - # No attachable budget: show the dead-end warning as a box and don't explain a feature the - # workspace can't use yet. - with ( - patch.object(wizard, "list_workspace_budgets", return_value=([], "none found")), - patch.object(wizard, "print_warning_panel") as warn_box, - patch.object(wizard, "print_note") as note, - ): - assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None - warning = warn_box.call_args.args[0] - assert "only AI Gateway budgets with hard blocks" in warning - assert "eligible to be associated with Tiered Spend Policies" in warning - assert not note.called # the BUDGET_POLICY_BLURB note is skipped - - def test_no_per_user_block_budgets_warns_and_yields_none(self): - # Spend routing needs a per-user threshold that hard-blocks; a workspace whose only budgets - # lack one has nothing usable to attach a policy to. - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": False}] - with ( - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object(wizard, "print_warning_panel") as warn_box, - ): - assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None - assert warn_box.called - - def test_only_per_user_block_budgets_are_offered(self): - # The picker hides budgets without a per-user hard block rather than letting the admin pick - # one that would leave every tier inert or unenforced. - budgets = [ - {"id": "no-block", "display_name": "email-only", "has_per_user_block": False}, - {"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}, - ] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], - ) as select, - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - ): - policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) - # First selection call is the budget picker; only the per-user budget is offered. - offered = [value for value, _ in select.call_args_list[0][0][1]] - assert offered == [BUDGET_ID] - assert policy is not None and policy["budget_id"] == BUDGET_ID - - def test_percentages_are_stored_as_fractions(self): - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], - ), - patch.object(wizard, "prompt_for_text", return_value="tiered"), - # prompt_for_percentage already converts; it returns the fraction. - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - ): - policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) - assert policy is not None - assert policy["budget_id"] == BUDGET_ID - # The picked budget's name is remembered for the summary. - assert policy["budget_display_name"] == "eng" - assert policy["tiers"] == [ - { - "spending_percentage": 0.8, - "default_agent": "claude", - "default_model": "system.ai.claude-opus-4-8", - } - ] - - def test_shows_per_user_threshold_and_tier_dollars(self): - # The admin picks tiers as percentages, so surface the budget's per-user monthly cap and what - # each percentage works out to in dollars — otherwise a percentage is a number in a vacuum. - budgets = [ - { - "id": BUDGET_ID, - "display_name": "eng", - "has_per_user_block": True, - "per_user_threshold": Decimal("500.00"), - } - ] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], - ), - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - patch.object(wizard, "print_note") as note, - ): - wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) - notes = " ".join(str(call.args[0]) for call in note.call_args_list) - assert "$500" in notes # the per-user monthly cap - assert "$400" in notes # 80% of $500 - - def test_missing_threshold_skips_the_dollar_hints(self): - # A budget whose threshold couldn't be read still works; the prompt just omits the dollars. - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], - ), - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - patch.object(wizard, "print_note") as note, - ): - policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) - notes = " ".join(str(call.args[0]) for call in note.call_args_list) - assert "$" not in notes - assert policy is not None and policy["tiers"][0]["spending_percentage"] == 0.8 - - def test_offers_only_the_models_the_agent_was_configured_with(self): - # Pi's catalog spans every family, so offering the workspace catalog would present four - # models it was never given — and a tier naming one of them silently misroutes developers. - enabled = { - "pi": { - "model_config": { - "default_model": "system.ai.kimi-k2-6", - "models": ["system.ai.kimi-k2-6"], - } - } - } - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "pi", "system.ai.kimi-k2-6"], - ) as select, - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - ): - wizard._prompt_budget_policy(WORKSPACE, "token", enabled, STATE) - # Third call is the model picker. - offered = [value for value, _ in select.call_args_list[2][0][1]] - assert offered == ["system.ai.kimi-k2-6"] - - def test_claude_family_slots_are_flattened_for_the_picker(self): - enabled = { - "claude": { - "model_config": { - "default_model": "system.ai.claude-opus-4-8", - "models": { - "default_opus_model": "system.ai.claude-opus-4-8", - "default_sonnet_model": "system.ai.claude-sonnet-4-6", - }, - } - } - } - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], - ) as select, - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - ): - wizard._prompt_budget_policy(WORKSPACE, "token", enabled, STATE) - offered = [value for value, _ in select.call_args_list[2][0][1]] - assert set(offered) == {"system.ai.claude-opus-4-8", "system.ai.claude-sonnet-4-6"} - - def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): - # An agent configured through a provider service has no enumerable list; better to offer the - # catalog than nothing at all. - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "gemini", "system.ai.gemini-3-flash"], - ) as select, - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - ): - wizard._prompt_budget_policy(WORKSPACE, "token", {"gemini": {}}, STATE) - offered = [value for value, _ in select.call_args_list[2][0][1]] - assert offered == ["system.ai.gemini-3-flash"] - - def test_authored_policy_validates(self): - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], - ), - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - ): - policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) - manifest = { - "default_agent": "claude", - "enabled_agents": CLAUDE_ONLY, - "budget_policy": policy, - } - assert validate_manifest(manifest, STATE) == [] - - def test_a_repeated_agent_model_pair_is_rejected_and_re_prompted(self): - # The highest crossed tier wins, so a second tier on the same agent+model is inert. The loop - # rejects the repeat and re-asks only the agent/model — the percentage already entered for - # this tier is kept, not re-prompted. - two_models = { - "claude": { - "model_config": { - "default_model": "system.ai.claude-opus-4-8", - "models": { - "default_opus_model": "system.ai.claude-opus-4-8", - "default_sonnet_model": "system.ai.claude-sonnet-4-6", - }, - } - } - } - # `has_per_user_block` is required since the budget-threshold gate landed on main: spend - # routing needs a per-user threshold that hard-blocks, so a budget without one is filtered - # out and the policy flow returns before the tier loop this test exercises. - budgets = [{"id": BUDGET_ID, "display_name": "eng", "has_per_user_block": True}] - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object( - wizard, - "prompt_for_selection", - side_effect=[ - BUDGET_ID, - # Tier 1: claude / opus. - "claude", - "system.ai.claude-opus-4-8", - # Tier 2 first attempt: claude / opus again — rejected, so just agent/model re-ask. - "claude", - "system.ai.claude-opus-4-8", - # Tier 2 retry: a genuine step-down. - "claude", - "system.ai.claude-sonnet-4-6", - ], - ), - patch.object(wizard, "prompt_for_text", return_value="tiered"), - # Percentage asked once per tier: 0.5 for tier 1, 0.9 for tier 2. The combo retry does - # not re-ask it. - patch.object(wizard, "prompt_for_percentage", side_effect=[0.5, 0.9]), - patch.object(wizard, "print_err") as err, - ): - policy = wizard._prompt_budget_policy(WORKSPACE, "token", two_models, STATE) - assert [(t["default_agent"], t["default_model"]) for t in policy["tiers"]] == [ - ("claude", "system.ai.claude-opus-4-8"), - ("claude", "system.ai.claude-sonnet-4-6"), - ] - assert any("do nothing" in call.args[0] for call in err.call_args_list) - - -class TestConfiguredModelsForAgent: - def test_flat_list_plus_default(self): - agent = {"model_config": {"default_model": "b", "models": ["a", "b"]}} - assert wizard.configured_models_for_agent(agent) == ["a", "b"] - - def test_claude_slots_are_flattened(self): - agent = { - "model_config": { - "default_model": "opus", - "models": {"default_opus_model": "opus", "default_sonnet_model": "sonnet"}, - } - } - assert set(wizard.configured_models_for_agent(agent)) == {"opus", "sonnet"} - - def test_codex_has_only_a_default(self): - # CodexModelConfig carries no model list, so the default is the whole set. - assert wizard.configured_models_for_agent({"model_config": {"default_model": "gpt-5"}}) == [ - "gpt-5" - ] - - def test_no_model_config_yields_nothing(self): - assert wizard.configured_models_for_agent({}) == [] - - -class TestSummary: - def test_lists_claude_family_slots(self, capsys): - # The one-line default hides which families were configured, which is most of the choice. - manifest = { - "default_agent": "claude", - "enabled_agents": { - "claude": { - "model_config": { - "default_model": "system.ai.claude-opus-4-8", - "models": { - "default_opus_model": "system.ai.claude-opus-4-8", - "default_haiku_model": "system.ai.claude-haiku-4-5", - }, - } - } - }, - } - wizard._render_summary(WORKSPACE, manifest) - out = capsys.readouterr().out - assert "opus" in out and "haiku" in out - assert "system.ai.claude-haiku-4-5" in out - - def test_lists_a_multi_model_agents_models(self, capsys): - manifest = { - "default_agent": "pi", - "enabled_agents": { - "pi": { - "model_config": { - "default_model": "system.ai.kimi-k2-6", - "models": ["system.ai.kimi-k2-6", "system.ai.gpt-5-6"], - } - } - }, - } - wizard._render_summary(WORKSPACE, manifest) - assert "system.ai.gpt-5-6" in capsys.readouterr().out - - def test_single_model_agent_needs_no_extra_line(self, capsys): - manifest = { - "default_agent": "gemini", - "enabled_agents": { - "gemini": {"model_config": {"default_model": "system.ai.gemini-3-flash"}} - }, - } - wizard._render_summary(WORKSPACE, manifest) - out = capsys.readouterr().out - assert "system.ai.gemini-3-flash" in out - assert "models:" not in out - - def test_summary_has_no_settings_scope_choice(self, capsys): - manifest = { - "default_agent": "codex", - "enabled_agents": { - "codex": {"model_config": {"default_model": "system.ai.gpt-5"}}, - "gemini": {"model_config": {"default_model": "system.ai.gemini-3-flash"}}, - }, - } - wizard._render_summary(WORKSPACE, manifest) - out = capsys.readouterr().out - assert "global settings" not in out - assert "ucode-only" not in out - - -class TestSetupFromFile: - def _write(self, tmp_path, payload): - path = tmp_path / "manifest.json" - path.write_text(json.dumps(payload), encoding="utf-8") - return path - - def _valid(self): - return { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} - }, - } - - def test_valid_manifest_is_saved(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() - - 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 - - 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")) - - 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)) - - 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)) - - 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 - - -class TestSummaryPanel: - def test_summary_is_boxed(self, capsys): - # The summary is the one block an admin reads as a whole to check against what they - # intended, and it lands after a long flow of prompts — so it gets a box rather than loose - # lines that blend into the preceding output. - wizard._render_summary( - WORKSPACE, - { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} - }, - }, - ) - out = capsys.readouterr().out - assert "Configuration summary" in out - # Rich box-drawing characters: the panel border. - assert "╭" in out and "╰" in out - assert "system.ai.claude-opus-5" in out - - def test_a_bracketed_policy_name_survives_the_summary(self, capsys): - # Rich reads bracketed text as a style tag and renders nothing for it, so an unescaped - # `[prod] tiered routing` displayed as `tiered routing` — in the block whose whole purpose - # is confirming what the admin is about to publish workspace-wide. - wizard._render_summary( - WORKSPACE, - { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} - }, - "budget_policy": { - "budget_id": "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57", - "display_name": "[prod] tiered routing", - "tiers": [], - }, - }, - ) - assert "[prod] tiered routing" in capsys.readouterr().out - - def test_shows_both_budget_and_policy_names(self, capsys): - # An admin checks the policy against two distinct things: which budget it tracks and what the - # policy itself is called. The summary must surface both, not collapse to one. - wizard._render_summary( - WORKSPACE, - { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} - }, - "budget_policy": { - "budget_id": "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57", - "budget_display_name": "eng-budget", - "display_name": "tiered routing", - "tiers": [], - }, - }, - ) - out = capsys.readouterr().out - assert "eng-budget" in out - assert "tiered routing" in out - - def test_falls_back_to_budget_id_without_a_budget_name(self, capsys): - # `--from-file` and server-read manifests carry no `budget_display_name`, so the budget id is - # all there is to show. - wizard._render_summary( - WORKSPACE, - { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} - }, - "budget_policy": { - "budget_id": "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57", - "display_name": "tiered routing", - "tiers": [], - }, - }, - ) - assert "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57" in capsys.readouterr().out - - -class TestCancelledPromptsAbort: - """A dismissed prompt must abort, not re-ask an input that can't answer.""" - - def test_require_selection_aborts_when_the_picker_is_dismissed(self): - # questionary's Question.ask catches KeyboardInterrupt and returns None (v2.1.1), so Ctrl-C - # is indistinguishable from an empty submission here. Re-asking looped forever. - with patch.object(wizard, "prompt_for_selection", return_value=None) as sel: - with pytest.raises(KeyboardInterrupt): - wizard._require_selection("pick", [("a", "A")]) - assert sel.call_count == 1 - - def test_require_text_asks_for_a_required_answer(self): - # `required=True` is what makes closed stdin raise instead of returning None; without it a - # piped/CI run spins re-asking an exhausted stream. - with patch.object(wizard, "prompt_for_text", return_value="m") as text: - assert wizard._require_text("Default model") == "m" - assert text.call_args.kwargs.get("required") is True - - def test_require_text_aborts_on_closed_stdin(self): - with patch("ucode.ui.console.input", side_effect=EOFError): - with pytest.raises(KeyboardInterrupt): - wizard._require_text("Default model") - - -class TestClaudeCandidatesStayValidatable: - """Whatever the Claude prompts offer, `validate_manifest` must accept.""" - - def _manifest(self, model: str) -> dict: - return { - "default_agent": "claude", - "enabled_agents": {"claude": {"model_config": {"default_model": model}}}, - } - - def test_listing_path_caches_so_older_versions_validate(self): - # The unbucketed listing widens the candidates past `claude_models`, so it must also cache - # them — otherwise picking an older Opus is rejected at the end of the flow. - state = { - "workspace": "https://ws.example.com", - "claude_models": {"opus": "system.ai.claude-opus-5"}, - } - with ( - patch.object(wizard, "get_databricks_token", return_value="t"), - patch.object( - wizard, - "discover_claude_models_unbucketed", - return_value=(["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], None), - ), - ): - candidates = wizard._claude_candidates(state) - offered = [m for models in candidates.values() for m in models] - assert "system.ai.claude-opus-4-8" in offered - for model in offered: - assert validate_manifest(self._manifest(model), state) == [] - - def test_fallback_path_only_offers_what_already_validates(self): - # The fallback caches nothing, so it must not offer anything beyond `claude_models`. - state = { - "workspace": "https://ws.example.com", - "claude_models": {"opus": "system.ai.claude-opus-5"}, - } - with ( - patch.object(wizard, "get_databricks_token", return_value="t"), - patch.object( - wizard, "discover_claude_models_unbucketed", side_effect=RuntimeError("boom") - ), - ): - candidates = wizard._claude_candidates(state) - assert "all_claude_models" not in state - for models in candidates.values(): - for model in models: - assert validate_manifest(self._manifest(model), state) == [] - - -class TestSearchablePickers: - """Long lists (models, provider services, budgets) filter as you type.""" - - def test_model_pickers_are_searchable(self): - seen: list[dict] = [] - - def fake_multi(prompt, options, preselected=None, **kwargs): - seen.append(kwargs) - return [options[0][0]] - - with patch.object(wizard, "prompt_for_multi_selection", side_effect=fake_multi): - wizard._require_multi_selection("pick", [("a", "a"), ("b", "b")]) - assert seen[0].get("searchable") is True - - def test_single_select_pickers_are_searchable(self): - seen: list[dict] = [] - - def fake_sel(prompt, options, **kwargs): - seen.append(kwargs) - return options[0][0] - - with patch.object(wizard, "prompt_for_selection", side_effect=fake_sel): - wizard._require_selection("pick", [("a", "a"), ("b", "b")]) - assert seen[0].get("searchable") is True - - def test_budget_and_tier_pickers_are_searchable(self): - budgets = [{"id": "budget-1", "display_name": "eng", "has_per_user_block": True}] - searchable_prompts: list[str] = [] - - def fake_sel(prompt, options, **kwargs): - if kwargs.get("searchable"): - searchable_prompts.append(prompt) - if "budget" in prompt: - return "budget-1" - if "agent" in prompt: - return "claude" - return "system.ai.claude-opus-4-8" - - with ( - patch.object(wizard, "prompt_yes_no_default", side_effect=[False]), - patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), - patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), - patch.object(wizard, "prompt_for_text", return_value="tiered"), - patch.object(wizard, "prompt_for_percentage", return_value=0.8), - ): - wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) - # Both the budget list and the tier's model list filter as you type. - assert any("budget" in p for p in searchable_prompts), searchable_prompts - assert any("model" in p for p in searchable_prompts), searchable_prompts - - -# A minimal authored manifest (agents + models only), the shape `ug setup` now writes. -AGENTS_ONLY = { - "default_agent": "claude", - "enabled_agents": {"claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}}, -} - - -class TestCarryForwardSections: - def test_all_optional_sections_survive_a_rerun(self): - previous = { - **AGENTS_ONLY, - "mcp_servers": [{"name": "system.ai.github", "type": "mcp-service"}], - "skills": {"names": ["main.default"]}, - "tracing_table": "main.default.traces", - "budget_policy": { - "budget_id": BUDGET_ID, - "tiers": [ - { - "spending_percentage": 0.8, - "default_agent": "claude", - "default_model": "system.ai.claude-opus-4-8", - } - ], - }, - } - manifest = dict(AGENTS_ONLY) - wizard._carry_forward_sections(previous, manifest) - assert manifest["mcp_servers"] == previous["mcp_servers"] - assert manifest["skills"] == previous["skills"] - assert manifest["tracing_table"] == previous["tracing_table"] - assert manifest["budget_policy"] == previous["budget_policy"] - - def test_empty_previous_adds_nothing(self): - manifest = dict(AGENTS_ONLY) - wizard._carry_forward_sections({}, manifest) - assert set(manifest) == set(AGENTS_ONLY) - - def test_budget_policy_naming_a_dropped_agent_is_left_out_with_a_warning(self): - # The admin re-ran `setup` and de-selected codex; the saved policy still routes to it, which - # would fail `validate_manifest` and block the whole save. Drop just the policy, loudly. - previous = { - **AGENTS_ONLY, - "budget_policy": { - "budget_id": BUDGET_ID, - "tiers": [ - { - "spending_percentage": 0.9, - "default_agent": "codex", - "default_model": "system.ai.gpt-5", - } - ], - }, - } - manifest = dict(AGENTS_ONLY) - with patch.object(wizard, "print_warning") as warn: - wizard._carry_forward_sections(previous, manifest) - assert "budget_policy" not in manifest - assert warn.called - # The rest of the manifest is untouched and still valid. - assert validate_manifest(manifest, None) == [] - - -class TestNextSteps: - def test_marks_configured_and_unconfigured_sections(self, capsys): - manifest = {**AGENTS_ONLY, "skills": {"names": ["main.default"]}} - wizard._print_next_steps(manifest) - out = capsys.readouterr().out - assert "ug setup mcps" in out - assert "ug setup skills" in out - assert "ug setup spend-tiers" in out - assert "ug publish" in out - - def test_dry_run_says_nothing_was_saved(self, capsys, monkeypatch): - monkeypatch.setattr(config_io_mod, "_dry_run", True) - wizard._print_next_steps(AGENTS_ONLY) - out = capsys.readouterr().out - assert "Dry run" in out - assert "ug publish" not in out - - -class TestSectionCommands: - """The `ug setup mcps` / `skills` / `spend-tiers` section commands.""" - - @staticmethod - def _admin(**overrides): - """Patch the auth/admin boundary the section commands resolve through.""" - defaults = { - "load_state": lambda: {"workspace": WORKSPACE, "profile": "p", **STATE}, - "ensure_databricks_auth": lambda *a, **k: None, - "get_databricks_token": lambda *a, **k: "tok", - "is_workspace_admin": lambda *a, **k: True, - # Decline the end-of-section "publish now?" offer so a section run saves the draft without - # trying to publish; the publish path is exercised in TestPublishCommand. - "prompt_yes_no_default": lambda *a, **k: False, - } - defaults.update(overrides) - return [patch.object(wizard, name, value) for name, value in defaults.items()] - - def _run(self, fn, *, admin_overrides=None, **patches): - import contextlib - - with contextlib.ExitStack() as stack: - for p in self._admin(**(admin_overrides or {})): - stack.enter_context(p) - for name, value in patches.items(): - stack.enter_context(patch.object(wizard, name, value)) - return fn() - - def test_mcp_requires_an_authored_config(self): - # No manifest on disk → the command can't edit a section that doesn't exist. - with pytest.raises(RuntimeError, match="ug setup"): - self._run(wizard.setup_mcp_command) - - def test_mcp_requires_enabled_agents(self): - # A launch stores `{}` to mean "no managed config"; that must not count as authored. - managed_config_mod.save_managed_state(WORKSPACE, {}) - with pytest.raises(RuntimeError, match="ug setup"): - self._run(wizard.setup_mcp_command) - - def test_mcp_writes_only_its_section(self): - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) - servers = [{"name": "system.ai.github", "type": "mcp-service"}] - # The picker (imported lazily inside the command) registers servers into local state; fake - # that by having the before/after reads bracket a change. - reads = iter([[], servers]) - with ( - patch("ucode.mcp.configure_mcp_command", return_value=0) as picker, - patch.object(wizard, "_mcp_servers_from_state", side_effect=lambda *_: next(reads)), - ): - code = self._run(wizard.setup_mcp_command) - assert code == 0 - assert picker.call_args.kwargs == {"exclude_sources": {"apps"}} - saved = managed_config_mod.load_managed_state(WORKSPACE) - assert saved["mcp_servers"] == servers - assert saved["enabled_agents"] == AGENTS_ONLY["enabled_agents"] - - def test_mcp_cancel_is_a_no_op(self): - # Picker cancelled / nothing changed → the section is left exactly as it was. - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) - with ( - patch("ucode.mcp.configure_mcp_command", return_value=0), - patch.object(wizard, "_mcp_servers_from_state", return_value=[]), - patch.object(wizard, "save_managed_state") as save, - ): - code = self._run(wizard.setup_mcp_command) - assert code == 0 - assert not save.called - - def test_mcp_carries_forward_preregistered_servers(self): - # An admin who ran `ug configure mcp` first arrives with those servers already registered, - # so the picker leaves local state unchanged (before == after). The manifest doesn't carry them - # yet, so `setup mcps` must still save them rather than report "no changes" and drop them. - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) - servers = [{"name": "system.ai.github", "type": "mcp-service"}] - with ( - patch("ucode.mcp.configure_mcp_command", return_value=0), - patch.object(wizard, "_mcp_servers_from_state", return_value=servers), - ): - code = self._run(wizard.setup_mcp_command) - assert code == 0 - assert managed_config_mod.load_managed_state(WORKSPACE)["mcp_servers"] == servers - - def test_mcp_not_admin_raises(self): - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) - with pytest.raises(RuntimeError, match="not an admin"): - self._run( - wizard.setup_mcp_command, - admin_overrides={"is_workspace_admin": lambda *a, **k: False}, - ) - - def test_skills_location_bypasses_the_prompt(self): - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) - with ( - patch("ucode.mcp.configure_skills_mcp_command", return_value=0) as configure, - patch.object(wizard, "_skill_names_from_state", return_value=["main.default"]), - patch.object(wizard, "prompt_for_text") as prompt, - ): - code = self._run(lambda: wizard.setup_skills_command(["main.default"])) - assert code == 0 - assert not prompt.called - configure.assert_called_once_with(["main.default"]) - assert managed_config_mod.load_managed_state(WORKSPACE)["skills"] == { - "names": ["main.default"] - } - - def test_skills_blank_answer_writes_nothing(self): - # A blank answer must not delegate: `configure_skills_mcp_command([])` is not a no-op. - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) - with ( - patch("ucode.mcp.configure_skills_mcp_command") as configure, - patch.object(wizard, "prompt_for_text", return_value=""), - patch.object(wizard, "save_managed_state") as save, - ): - code = self._run(wizard.setup_skills_command) - assert code == 0 - assert not configure.called - assert not save.called - - def test_budget_policy_offers_only_the_manifests_agents(self): - managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) - captured = {} - - def fake_prompt(workspace, token, enabled_agents, state, **kwargs): - captured["agents"] = enabled_agents - return None # decline / dead end → leave the policy unchanged - - with patch.object(wizard, "_prompt_budget_policy", side_effect=fake_prompt): - code = self._run(wizard.setup_budget_policy_command) - assert code == 0 - assert captured["agents"] == AGENTS_ONLY["enabled_agents"] - - def test_budget_policy_none_leaves_existing_untouched(self): - # A transient budget-listing failure returns None; it must never delete a saved policy. - seeded = { - **AGENTS_ONLY, - "budget_policy": { - "budget_id": BUDGET_ID, - "tiers": [ - { - "spending_percentage": 0.8, - "default_agent": "claude", - "default_model": "system.ai.claude-opus-4-8", - } - ], - }, - } - managed_config_mod.save_managed_state(WORKSPACE, seeded) - with patch.object(wizard, "_prompt_budget_policy", return_value=None): - code = self._run(wizard.setup_budget_policy_command) - assert code == 0 - assert ( - managed_config_mod.load_managed_state(WORKSPACE)["budget_policy"] - == seeded["budget_policy"] - ) - - -class TestSetupHelp: - def test_lists_every_setup_command(self, capsys): - wizard.setup_help_command() - out = capsys.readouterr().out - for command in ( - "ug setup", - "ug setup mcps", - "ug setup skills", - "ug setup spend-tiers", - "ug setup show", - "ug publish", - ): - assert command in out - - -class TestPublishDiff: - def test_lists_added_removed_and_changed(self, capsys): - existing = { - "name": "cfg/1", - **AGENTS_ONLY, - "mcp_servers": [{"name": "system.ai.slack", "type": "mcp-service"}], - } - incoming = { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-9"}} - }, - "mcp_servers": [{"name": "system.ai.github", "type": "mcp-service"}], - } - changed = wizard._render_config_diff(existing, incoming, WORKSPACE) - out = capsys.readouterr().out - assert changed is True - assert "CHANGE" in out and "claude-opus-4-8" in out and "claude-opus-4-9" in out - assert "ADD" in out and "system.ai.github" in out # added server - assert "DELETE" in out and "system.ai.slack" in out # removed server - - def test_identical_configs_report_no_change(self, capsys): - assert wizard._render_config_diff(AGENTS_ONLY, AGENTS_ONLY, WORKSPACE) is False - - def test_display_name_change_is_detected(self, capsys): - existing = {"display_name": "old-name", **AGENTS_ONLY} - incoming = {"display_name": "new-name", **AGENTS_ONLY} - changed = wizard._render_config_diff(existing, incoming, WORKSPACE) - out = capsys.readouterr().out - assert changed is True - assert "old-name" in out and "new-name" in out - - -class TestPublishCommand: - MANIFEST = { - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} - }, - } - - @staticmethod - def _patches(**overrides): - """The network/auth boundary `publish_command` sits behind, with per-test overrides.""" - defaults = { - "load_state": lambda: {"workspace": WORKSPACE, "profile": "p", **STATE}, - "ensure_databricks_auth": lambda *a, **k: None, - "get_databricks_token": lambda *a, **k: "tok", - "is_workspace_admin": lambda *a, **k: True, - "get_managed_config": lambda *a, **k: (None, None), - "create_coding_agent_config": lambda *a, **k: ( - {"name": "coding-agent-configs/new"}, - None, - ), - "update_coding_agent_config": lambda *a, **k: ( - {"name": "coding-agent-configs/old"}, - None, - ), - "prompt_yes_no_default": lambda *a, **k: True, - } - defaults.update(overrides) - return [patch.object(wizard, name, value) for name, value in defaults.items()] - - def _run(self, *, yes=False, file_path=None, **overrides): - import contextlib - - with contextlib.ExitStack() as stack: - for p in self._patches(**overrides): - stack.enter_context(p) - return wizard.publish_command(file_path=file_path, yes=yes) - - @staticmethod - def _config_file(tmp_path, manifest, *, workspace=WORKSPACE, spec_version=1, **extra): - config = serialize_managed_config(manifest) - config.pop("name", None) - payload = {"workspace": workspace, "spec_version": spec_version, **config, **extra} - path = tmp_path / "config.json" - path.write_text(json.dumps(payload), encoding="utf-8") - return str(path) - - def test_unauthored_config_is_an_actionable_error(self): - with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): - with pytest.raises(RuntimeError, match="ug setup"): - wizard.publish_command() - - def test_creates_when_no_config_exists(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) - created = {} - - def fake_create(workspace, token, payload): - created.update(workspace=workspace, payload=payload) - return {"name": "coding-agent-configs/new"}, None - - assert self._run(create_coding_agent_config=fake_create) == 0 - assert created["workspace"] == WORKSPACE - # What goes over the wire is proto-JSON, not ug's manifest shape. - assert created["payload"]["default_agent"] == "CODING_AGENT_CLAUDE_CODE" - - 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) - existing = {"name": "coding-agent-configs/abc", "enabled_agents": {"codex": {}}} - updated = {} - created = {"called": False} - - def fake_update(workspace, token, name, payload): - updated.update(name=name, payload=payload) - return {"name": name}, None - - def fake_create(*a, **k): - created["called"] = True - return {}, None - - assert ( - self._run( - get_managed_config=lambda *a, **k: (existing, None), - update_coding_agent_config=fake_update, - create_coding_agent_config=fake_create, - ) - == 0 - ) - assert updated["name"] == "coding-agent-configs/abc" - assert created["called"] is False - - def test_no_publish_when_the_published_config_already_matches(self): - # Publishing a config identical to what's live is a no-op; say so and skip the write rather - # than PATCH the same bytes back. - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) - # What's live is the manifest normalized the same way `publish` will send it. - existing = { - "name": "coding-agent-configs/abc", - **managed_config_mod.normalize_managed_config(serialize_managed_config(self.MANIFEST)), - } - updated = {"called": False} - - def fake_update(*a, **k): - updated["called"] = True - return {}, None - - assert ( - self._run( - get_managed_config=lambda *a, **k: (existing, None), - update_coding_agent_config=fake_update, - ) - == 0 - ) - assert updated["called"] is False - - def test_invalid_manifest_is_not_published(self): - managed_config_mod.save_managed_state( - WORKSPACE, {"default_agent": "codex", "enabled_agents": {"claude": {}}} - ) - created = {"called": False} - - def fake_create(*a, **k): - created["called"] = True - return {}, None - - with pytest.raises(RuntimeError, match="not valid"): - self._run(create_coding_agent_config=fake_create) - assert created["called"] is False - - def test_an_older_family_version_the_wizard_offered_still_publishes(self): - # `setup` offers every version of a Claude family, but `claude_models` keeps only the newest - # per family and the wizard's `all_claude_models` stash is never persisted — so a separate - # `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( - WORKSPACE, - { - "default_agent": "claude", - "enabled_agents": { - "claude": { - "model_config": { - "default_model": "system.ai.claude-opus-4-1", - "models": {"default_opus_model": "system.ai.claude-opus-4-1"}, - } - } - }, - }, - ) - published: dict = {} - - def fake_create(workspace, token, payload): - published["payload"] = payload - return {"name": "coding-agent-configs/new"}, None - - # State carries only the newest Opus, as a fresh `load_state()` would. - narrow = {"workspace": WORKSPACE, "profile": "p", "claude_models": {"opus": "newest"}} - assert ( - self._run( - load_state=lambda: dict(narrow), - discover_claude_models_unbucketed=lambda *a, **k: ( - ["system.ai.claude-opus-4-1", "newest"], - None, - ), - create_coding_agent_config=fake_create, - ) - == 0 - ) - assert published, "the manifest should have been published" - - 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) - published: dict = {} - - def fake_create(workspace, token, payload): - published["payload"] = payload - return {"name": "coding-agent-configs/new"}, None - - assert ( - self._run( - discover_claude_models_unbucketed=lambda *a, **k: ([], "HTTP 500"), - create_coding_agent_config=fake_create, - ) - == 0 - ) - assert published - - def test_declining_the_prompt_publishes_nothing(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) - created = {"called": False} - - def fake_create(*a, **k): - created["called"] = True - return {}, None - - code = self._run( - prompt_yes_no_default=lambda *a, **k: False, create_coding_agent_config=fake_create - ) - assert code == 1 - assert created["called"] is False - - def test_yes_skips_the_prompt(self): - managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) - - def refuse(*a, **k): - raise AssertionError("--yes must not prompt") - - 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) - created = {"called": False} - - def fake_create(*a, **k): - created["called"] = True - return {}, None - - with pytest.raises(RuntimeError, match="not an admin"): - self._run( - is_workspace_admin=lambda *a, **k: False, create_coding_agent_config=fake_create - ) - assert created["called"] is False - - 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) - created = {"called": False} - - def fake_create(*a, **k): - created["called"] = True - return {}, None - - with pytest.raises(RuntimeError, match="Refusing to publish"): - self._run( - get_managed_config=lambda *a, **k: (None, "HTTP 500 Server Error"), - create_coding_agent_config=fake_create, - ) - 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) - created = {"called": False} - - def fake_create(*a, **k): - created["called"] = True - return {}, None - - reason = 'HTTP 404 Not Found: {"error_code":"FEATURE_DISABLED"}' - with pytest.raises(RuntimeError) as exc_info: - self._run( - get_managed_config=lambda *a, **k: (None, reason), - create_coding_agent_config=fake_create, - ) - assert str(exc_info.value) == wizard.CODING_AGENT_CONFIGS_DISABLED_MESSAGE - assert "FEATURE_DISABLED" not in str(exc_info.value) - 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) - with pytest.raises(RuntimeError, match="resource name"): - self._run(get_managed_config=lambda *a, **k: ({"enabled_agents": {}}, None)) - - def test_file_input_creates_and_sends_spec_version_without_workspace(self, tmp_path): - path = self._config_file(tmp_path, self.MANIFEST) - created = {} - - def fake_create(workspace, token, payload): - created.update(workspace=workspace, payload=payload) - return {"name": "coding-agent-configs/new"}, None - - assert self._run(file_path=path, create_coding_agent_config=fake_create) == 0 - assert created["workspace"] == WORKSPACE - assert created["payload"]["spec_version"] == 1 - assert "workspace" not in created["payload"] - assert "name" not in created["payload"] - assert created["payload"]["default_agent"] == "CODING_AGENT_CLAUDE_CODE" - - def test_file_input_updates_in_place_and_sends_spec_version(self, tmp_path): - path = self._config_file(tmp_path, self.MANIFEST) - existing = {"name": "coding-agent-configs/abc", "enabled_agents": {"codex": {}}} - updated = {} - - def fake_update(workspace, token, name, payload): - updated.update(name=name, payload=payload) - return {"name": name}, None - - assert ( - self._run( - file_path=path, - get_managed_config=lambda *a, **k: (existing, None), - update_coding_agent_config=fake_update, - ) - == 0 - ) - assert updated["name"] == "coding-agent-configs/abc" - assert updated["payload"]["spec_version"] == 1 - assert "workspace" not in updated["payload"] - - def test_file_workspace_mismatch_aborts_before_auth_or_mutation(self, tmp_path): - path = self._config_file(tmp_path, self.MANIFEST, workspace="https://other.example.com") - called = {"auth": False, "create": False} - - def fake_auth(*a, **k): - called["auth"] = True - - def fake_create(*a, **k): - called["create"] = True - return {}, None - - with pytest.raises(RuntimeError, match="configured workspace"): - self._run( - file_path=path, - ensure_databricks_auth=fake_auth, - create_coding_agent_config=fake_create, - ) - assert called == {"auth": False, "create": False} - - def test_file_bad_spec_version_aborts_before_auth(self, tmp_path): - path = self._config_file(tmp_path, self.MANIFEST, spec_version=2) - called = {"auth": False} - - def fake_auth(*a, **k): - called["auth"] = True - - with pytest.raises(RuntimeError, match="spec_version"): - self._run(file_path=path, ensure_databricks_auth=fake_auth) - assert called["auth"] is False - - def test_missing_file_is_actionable(self, tmp_path): - with pytest.raises(RuntimeError, match="No config file"): - 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( - WORKSPACE, - { - "default_agent": "codex", - "enabled_agents": { - "codex": { - "model_config": { - "default_model": "main.custom.model", - "custom_models": ["main.custom.model"], - } - } - }, - }, - ) - created = {"called": False} - - def fake_create(*a, **k): - created["called"] = True - return {"name": "coding-agent-configs/new"}, None - - assert self._run(create_coding_agent_config=fake_create) == 0 - assert created["called"] is True - - -class TestPublishFailureMessages: - """The server's error codes, turned into something an admin can act on.""" - - def test_feature_disabled_uses_the_shared_message(self): - message = wizard._explain_publish_failure( - 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED","message":"..."}' - ) - assert message == wizard.CODING_AGENT_CONFIGS_DISABLED_MESSAGE - assert "`ug configure`" in message - - def test_permission_denied_says_admin_is_required(self): - message = wizard._explain_publish_failure( - 'HTTP 403 Forbidden: {"error_code":"PERMISSION_DENIED"}' - ) - assert "workspace admin" in message - - def test_invalid_parameter_value_is_passed_through_verbatim(self): - # The server names the offending field, which is more useful than any paraphrase. - reason = ( - 'HTTP 400 Bad Request: {"error_code":"INVALID_PARAMETER_VALUE",' - '"message":"budget_policy.tiers[0].spending_percentage must be between 0 and 1"}' - ) - message = wizard._explain_publish_failure(reason) - assert "budget_policy.tiers[0].spending_percentage" in message - - def test_unknown_failure_still_surfaces_the_reason(self): - message = wizard._explain_publish_failure("network error: timed out") - assert "timed out" in message - - -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 - 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 - - def test_publish_is_registered(self): - result = runner.invoke(app, ["--help"]) - assert result.exit_code == 0 - assert "publish" in result.output - - 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). - 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 - assert "--dry-run" not in declared - - def test_publish_declares_file_option(self): - 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 "--file" in declared - assert "-f" in declared - - def test_publish_file_flag_is_forwarded(self): - for flag in ("-f", "--file"): - with ( - patch("ucode.cli.install_databricks_cli"), - patch.object(cli_mod, "publish_command", return_value=0) as publish, - ): - runner.invoke(app, ["publish", flag, "/tmp/cfg.json"]) - assert publish.call_args.kwargs["file_path"] == "/tmp/cfg.json" - - def test_publish_error_exits_nonzero_with_a_message(self): - with ( - patch("ucode.cli.install_databricks_cli"), - patch.object( - cli_mod, "publish_command", side_effect=RuntimeError("no config authored") - ), - ): - result = runner.invoke(app, ["publish"]) - 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". - with ( - patch("ucode.cli.install_databricks_cli"), - patch.object(cli_mod, "publish_command", return_value=0), - ): - result = runner.invoke(app, ["publish"]) - 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. - with ( - patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", return_value=0) as setup, - ): - result = runner.invoke(app, ["setup"]) - assert result.exit_code == 0 - assert setup.called - assert "ERROR" not in _out(result) - - def test_nonzero_setup_propagates(self): - with ( - patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", return_value=1), - ): - result = runner.invoke(app, ["setup"]) - assert result.exit_code == 1 - - def test_runtime_error_is_reported_and_exits_1(self): - with ( - patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", side_effect=RuntimeError("you are not an admin")), - ): - result = runner.invoke(app, ["setup"]) - assert result.exit_code == 1 - assert "not an admin" in _out(result) - - def test_interrupt_exits_130(self): - with ( - patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", side_effect=KeyboardInterrupt), - ): - result = runner.invoke(app, ["setup"]) - assert result.exit_code == 130 - - def test_from_file_is_forwarded(self): - with ( - patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_command", return_value=0) as setup, - ): - runner.invoke(app, ["setup", "--from-file", "/tmp/x.json"]) - assert setup.call_args.kwargs["from_file"] == "/tmp/x.json" - - 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 - - @pytest.mark.parametrize( - ("command", "target"), - [ - ("mcps", "setup_mcp_command"), - ("skills", "setup_skills_command"), - ("spend-tiers", "setup_budget_policy_command"), - ], - ) - def test_section_subcommands_are_registered_and_called(self, command, target): - with ( - patch("ucode.cli.install_databricks_cli"), - patch(f"ucode.cli.{target}", return_value=0) as fn, - ): - result = runner.invoke(app, ["setup", command]) - assert result.exit_code == 0 - assert fn.called - assert "ERROR" not in _out(result) - - def test_setup_skills_declares_location(self): - group = typer.main.get_command(app).commands["setup"] # type: ignore[attr-defined] - skills = group.commands["skills"] # type: ignore[attr-defined] - declared = {opt for param in skills.params for opt in param.opts} - assert "--location" in declared - - def test_setup_skills_location_is_parsed_to_a_list(self): - with ( - patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_skills_command", return_value=0) as fn, - ): - runner.invoke(app, ["setup", "skills", "--location", "main.a,main.b"]) - assert fn.call_args.args[0] == ["main.a", "main.b"] - - def test_setup_help_needs_no_auth(self): - # `ug setup help` reads the local draft only — it must not shell out to install the CLI. - with ( - patch("ucode.cli.install_databricks_cli") as install, - patch("ucode.cli.setup_help_command", return_value=0) as fn, - ): - result = runner.invoke(app, ["setup", "help"]) - assert result.exit_code == 0 - assert fn.called - assert not install.called - - def test_section_command_runtime_error_exits_1(self): - with ( - patch("ucode.cli.install_databricks_cli"), - patch("ucode.cli.setup_mcp_command", side_effect=RuntimeError("run `ug setup` first")), - ): - result = runner.invoke(app, ["setup", "mcps"]) - assert result.exit_code == 1 - assert "ug 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.""" - return result.output + (result.stderr if result.stderr_bytes else "")