diff --git a/.gitignore b/.gitignore index 4d1e17c..007899c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ build/ dist/ .venv/ .DS_Store +.isaac/ diff --git a/README.md b/README.md index 0352757..70a343b 100644 --- a/README.md +++ b/README.md @@ -163,6 +163,40 @@ ucode configure skills --location main.default,ml.prod --mcp Each run prints the registered server, its URL, the configured agents, and its tools, and reminds you to run `ucode ` (existing agent sessions need a restart before the MCP tools load). +### Managed config for a workspace (admins) + +```bash +ucode setup +``` + +Author the coding config your developers pick up automatically, instead of asking each of them to +run `ucode configure` by hand. Restricted to workspace admins. + +The flow walks through the agents to enable and which one bare `ucode` launches, then per agent: +Databricks-hosted models or an external Model Provider Service, the models to expose, and whether +the config applies machine-wide or per user. Claude Code is asked one model per family +(opus/sonnet/haiku/fable), since Claude Code selects models by family alias; any family can be +skipped. It then offers tracing, managed MCP servers, skills, and a spend-based budget policy that +switches the default agent and model as the workspace burns through a budget. + +The result is written to `~/.ucode/managed-settings.json`, which `ucode apply` publishes to the +workspace. Your own agent configs are left alone, with one exception: answering yes to tracing, MCP +servers, or skills runs the matching `ucode configure` step, which does configure this machine. + +```bash +# Review the manifest and the exact payload `ucode apply` would publish. +ucode setup show + +# Walk the flow without writing anything. +ucode setup --dry-run + +# Skip the prompts and load a hand-written config instead (validated before saving). +ucode setup --from-file ./managed-settings.json +``` + +Publishing replaces the workspace's config outright — there is no partial update yet, so anything +skipped in a re-run is dropped. + --- ## Other Commands @@ -185,6 +219,9 @@ you to run `ucode ` (existing agent sessions need a restart before the MC | `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | | `ucode configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | | `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | +| `ucode setup` | Author the workspace's managed coding config (workspace admins only) | +| `ucode setup show` | Print the authored config and the payload `ucode apply` would publish | +| `ucode setup --from-file ` | Load a hand-written managed config instead of running the prompts | ## Managed Local Files @@ -199,6 +236,7 @@ you to run `ucode ` (existing agent sessions need a restart before the MC | `~/.copilot/.env` | GitHub Copilot CLI | | `~/.pi/agent/models.json` | Pi | | `~/.cursor/mcp.json` | Cursor Agent (MCP servers only) | +| `~/.ucode/managed-settings.json` | The managed config authored by `ucode setup` (admins) | Existing files are backed up before being overwritten. `ucode revert` restores backups. diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 57c081e..f05e70c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -57,6 +57,7 @@ ) from ucode.managed_config import managed_agent_config_enabled, managed_launch_state from ucode.managed_resolve import managed_default_model, managed_provider_service +from ucode.managed_wizard import setup_command, show_command from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -862,6 +863,10 @@ def revert() -> int: app.add_typer(configure_app, name="configure", help="Configure workspace and tool settings.") mcp_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") +setup_app = typer.Typer(add_completion=False, no_args_is_help=False) +app.add_typer( + setup_app, name="setup", help="Author the workspace's managed coding config (admins only)." +) @mcp_app.command("web-search") @@ -1915,6 +1920,53 @@ def configure_tracing( raise typer.Exit(130) from None +@setup_app.callback(invoke_without_command=True) +def setup( + ctx: typer.Context, + from_file: Annotated[ + str | None, + typer.Option( + "--from-file", + help="Skip the interactive flow and load a hand-written managed config (JSON, in " + "ucode's manifest shape) instead. Validated before it is saved.", + ), + ] = None, + dry_run: Annotated[ + bool, + typer.Option("--dry-run", help="Walk the flow without writing any files."), + ] = False, +) -> None: + """Author the managed coding config for your workspace (workspace admins only).""" + if ctx.invoked_subcommand is not None: + return + set_dry_run(dry_run) + # `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("show") +def setup_show_cmd() -> None: + """Print the authored managed config and the payload `ucode apply` 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("status") def status_cmd() -> None: """Show current workspace, tool configs, and saved model selections.""" diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d9cf166..758bb25 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -331,6 +331,76 @@ def _http_get_bytes(url: str, token: str, *, timeout: int = 10) -> tuple[bytes | return None, f"network error: {exc.reason}" +# Workspace group whose members are workspace admins. `ucode setup` / `ucode apply` are restricted +# to this group because the coding-agent-config CRUD API enforces the same check server-side. +WORKSPACE_ADMIN_GROUP = "admins" + + +def is_workspace_admin(workspace: str, token: str) -> bool | None: + """Whether the caller is a workspace admin, via their SCIM `Me` group membership. + + Returns True/False, or None when the check itself could not be made (SCIM unreachable or a + malformed response). Callers should treat None as "unknown" and proceed optimistically rather + than blocking: the API enforces the same check server-side, so a false negative here would + needlessly stop a legitimate admin, while a false positive just surfaces the server's + PERMISSION_DENIED later. + """ + hostname = workspace_hostname(workspace) + payload, _ = _http_get_json(f"https://{hostname}/api/2.0/preview/scim/v2/Me", token) + if not isinstance(payload, dict): + return None + groups = payload.get("groups") + if not isinstance(groups, list): + # A well-formed `Me` for a user in no groups omits `groups` entirely, so this is a + # definitive "not an admin" rather than a failed check. + return False + return any( + isinstance(group, dict) and group.get("display") == WORKSPACE_ADMIN_GROUP + for group in groups + ) + + +# Workspace-scoped budget listing. Account-level budget APIs need account auth, which ucode does not +# have; this endpoint resolves the workspace server-side from the caller's token. +_WORKSPACE_BUDGETS_API_PATH = "/api/ai-gateway/v2/workspace-metrics/budgets" + + +def list_workspace_budgets(workspace: str, token: str) -> tuple[list[dict], str | None]: + """List the AI Gateway budgets that apply to this workspace. + + Returns ``(budgets, reason)`` where each budget is ``{"id": ..., "display_name": ...}``. + ``reason`` is None on success, otherwise it explains why the list is empty. ucode never creates + budgets — an admin picks an existing one to attach a spend-routing policy to. + """ + hostname = workspace_hostname(workspace) + url = f"https://{hostname}{_WORKSPACE_BUDGETS_API_PATH}" + payload, reason = _http_get_json(url, token, timeout=30) + if reason is not None: + return [], reason + if not isinstance(payload, dict): + return [], "workspace budget listing returned an unexpected response shape" + raw = payload.get("workspace_ai_gateway_budgets") + if not isinstance(raw, list): + return [], "workspace budget listing returned no budgets" + budgets: list[dict] = [] + for entry in raw: + if not isinstance(entry, dict): + continue + budget_id = entry.get("budget_configuration_id") + if not isinstance(budget_id, str) or not budget_id: + continue + display_name = entry.get("display_name") + budgets.append( + { + "id": budget_id, + "display_name": display_name if isinstance(display_name, str) else "", + } + ) + if not budgets: + return [], "workspace budget listing returned no budgets" + return budgets, None + + def get_current_user_name(workspace: str, token: str) -> str | None: """Return the current user's login (email) via SCIM `Me`, or None on failure. @@ -843,12 +913,20 @@ def run_databricks_login(workspace: str, profile: str | None = None) -> None: print_success("Databricks authentication complete") -def ensure_databricks_auth(workspace: str, profile: str | None = None) -> None: - """Check auth and login only if needed (used by launch path).""" +def ensure_databricks_auth( + workspace: str, profile: str | None = None, *, quiet: bool = False +) -> None: + """Check auth and login only if needed (used by launch path). + + ``quiet`` suppresses the "already available" line for a caller that only needs a token before + some later step re-authenticates and reports it — otherwise the same success prints twice. A + login that actually runs is never silent. + """ with spinner("Checking Databricks auth..."): auth_is_valid = has_valid_databricks_auth(workspace, profile) if auth_is_valid: - print_success(f"Databricks auth already available for {workspace}") + if not quiet: + print_success(f"Databricks auth already available for {workspace}") return run_databricks_login(workspace, profile) @@ -1272,12 +1350,42 @@ def _get_model_services_page( return payload, reason +# Successful model-service listings for this process, keyed by workspace. The listing is a paginated +# walk of the whole metastore catalog, and several callers want different views of the same result +# (`discover_model_services` buckets it per family, `discover_claude_models_unbucketed` keeps the raw +# Claude ids), so a single `ucode setup` run would otherwise page it twice. Cached per process, not +# persisted: a long-lived process is not a thing here, and a new model appearing mid-command is not +# worth a second walk. Failures are never cached, so a transient error still retries. +_MODEL_SERVICES_CACHE: dict[str, list[str]] = {} + +# Same idea for the Model Provider Service listing (a different endpoint). It is workspace-wide and +# filtered per agent afterwards, so `ucode setup` would otherwise re-list it once per MPS-capable +# agent. +_MODEL_PROVIDER_SERVICES_CACHE: dict[str, list[dict]] = {} + + +def clear_model_services_cache() -> None: + """Forget cached model-service listings (used by tests, and after a workspace switch).""" + _MODEL_SERVICES_CACHE.clear() + _MODEL_PROVIDER_SERVICES_CACHE.clear() + + +def has_cached_model_provider_services(workspace: str) -> bool: + """True when :func:`list_model_provider_services` will answer from cache. + + Lets a caller skip a progress spinner it doesn't need: the cold listing takes over a second, so + it deserves one, but repeating it per agent on an instant cache hit is just noise. + """ + return workspace in _MODEL_PROVIDER_SERVICES_CACHE + + def list_model_services( workspace: str, token: str, *, page_size: int = _MODEL_SERVICES_PAGE_SIZE, max_pages: int = 100, + use_cache: bool = True, ) -> tuple[list[str], str | None]: """List all `system.ai.*` model ids via the UC model-services API. @@ -1286,7 +1394,15 @@ def list_model_services( de-duplicated, sorted list of ``system.ai.`` ids. Returns (ids, reason); reason is None on success, otherwise it describes why the list is empty (HTTP/network error or no services). + + A successful result is memoized per workspace for the life of the process; pass + ``use_cache=False`` to force a fresh walk. """ + if use_cache: + cached = _MODEL_SERVICES_CACHE.get(workspace) + if cached is not None: + return list(cached), None + hostname = workspace_hostname(workspace) ids: list[str] = [] page_token: str | None = None @@ -1319,10 +1435,26 @@ def list_model_services( deduped = sorted(set(ids)) if deduped: + if use_cache: + _MODEL_SERVICES_CACHE[workspace] = list(deduped) return deduped, None return [], last_reason or "model-services listing returned no models" +def discover_claude_models_unbucketed(workspace: str, token: str) -> tuple[list[str], str | None]: + """Every `system.ai.claude-*` id on the workspace, unbucketed. + + `discover_model_services` keeps only the newest id per family because the launch path pins one + model per Claude family alias. An admin authoring a managed config needs the alternatives too + (see `managed_setup.claude_family_candidates`), so this returns the full set without disturbing + that shape. + """ + ids, reason = list_model_services(workspace, token) + if not ids: + return [], reason + return [m for m in ids if "claude-" in m.lower()], None + + def discover_model_services( workspace: str, token: str ) -> tuple[dict[str, str], list[str], list[str], list[str], str | None]: @@ -1495,7 +1627,9 @@ def _provider_type_tag(provider_type: str | None) -> str: return tag.lower() -def list_model_provider_services(workspace: str, token: str) -> tuple[list[dict], str | None]: +def list_model_provider_services( + workspace: str, token: str, *, use_cache: bool = True +) -> tuple[list[dict], str | None]: """List Unity Catalog Model Provider Services on the workspace. Returns ``(services, reason)`` where each service is @@ -1505,7 +1639,20 @@ def list_model_provider_services(workspace: str, token: str) -> tuple[list[dict] Bedrock model names). ``relayed`` is True for a credential-less Anthropic service (Claude Max/Team/Enterprise subscription relay). A non-None ``reason`` means the listing call itself failed. + + The listing is workspace-wide — it is filtered per agent afterwards by + :func:`service_usable_for_tool` — so a successful result is memoized per workspace for the life + of the process, like the model-services listing. Without that, `ucode setup` re-lists it once + per MPS-capable agent. Pass ``use_cache=False`` to force a fresh call. """ + if use_cache: + cached = _MODEL_PROVIDER_SERVICES_CACHE.get(workspace) + if cached is not None: + # A fresh list of fresh dicts each time: callers treat the result as theirs (the wizard + # filters it per agent), so handing out the cached objects would let one caller's edit + # reach the next. + return [dict(service) for service in cached], None + hostname = workspace_hostname(workspace) url = f"https://{hostname}/api/2.1/unity-catalog/model-provider-services" payload, reason = _http_get_json(url, token, timeout=30) @@ -1542,6 +1689,8 @@ def list_model_provider_services(workspace: str, token: str) -> tuple[list[dict] } ) services.sort(key=lambda s: s["name"]) + if use_cache: + _MODEL_PROVIDER_SERVICES_CACHE[workspace] = [dict(service) for service in services] return services, None diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index d2a1638..2ef8373 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -31,6 +31,7 @@ import ucode.config_io as config_io from ucode.databricks import ( ANTHROPIC_FAMILIES, + model_version_sort_key, tool_supports_provider_type, ) from ucode.managed_config import ( @@ -61,8 +62,9 @@ # (`ClaudeDefaultModels`), and Codex has no model list at all — it selects exactly one model. _FLAT_MODEL_LIST_AGENTS = frozenset({"opencode", "pi", "gemini", "copilot"}) -# Claude family slot names in `ClaudeDefaultModels`, keyed by ucode's family name. -_CLAUDE_SLOT_FOR_FAMILY: dict[str, str] = { +# Claude family slot names in `ClaudeDefaultModels`, keyed by ucode's family name. Public because +# the wizard prompts one slot at a time. +CLAUDE_SLOT_FOR_FAMILY: dict[str, str] = { family: f"default_{family}_model" for family in ANTHROPIC_FAMILIES } @@ -144,6 +146,39 @@ def claude_family_for_model(model: str) -> str | None: return next((family for family in ANTHROPIC_FAMILIES if f"claude-{family}-" in lowered), None) +def claude_family_candidates( + all_claude_models: list[str], state: dict | None = None +) -> dict[str, list[str]]: + """Group Claude model ids by family, newest first. + + ``state["claude_models"]`` holds only one id per family — the newest, chosen by + ``discover_model_services`` for the launch path, which pins exactly one model per family alias. + An admin authoring a managed config needs the alternatives too: pinning ``default_opus_model`` + to a known-good ``claude-opus-4-8`` rather than whatever happens to be newest is a normal thing + to want, and impossible if only the newest is offered. + + ``all_claude_models`` is the unbucketed listing (see + :func:`ucode.databricks.discover_claude_family_candidates`). When it is empty, falls back to the + per-family picks already in ``state`` so the per-slot prompts still work — with one candidate + each. Families with no models are omitted. + """ + models = list(all_claude_models) + if not models and state: + claude_models = state.get("claude_models") + if isinstance(claude_models, dict): + models = [m for m in claude_models.values() if isinstance(m, str) and m] + + candidates: dict[str, list[str]] = {} + for model in models: + family = claude_family_for_model(model) + if family: + candidates.setdefault(family, []).append(model) + for family, found in candidates.items(): + # model_version_sort_key negates version components, so plain ascending is newest-first. + candidates[family] = sorted(set(found), key=model_version_sort_key) + return candidates + + def claude_model_slots(models: list[str]) -> dict[str, str]: """Group picked Claude model ids into ``ClaudeDefaultModels`` slots. @@ -157,7 +192,7 @@ def claude_model_slots(models: list[str]) -> dict[str, str]: family = claude_family_for_model(model) if family is None: continue - slot = _CLAUDE_SLOT_FOR_FAMILY[family] + slot = CLAUDE_SLOT_FOR_FAMILY[family] slots.setdefault(slot, model) return slots @@ -342,12 +377,16 @@ def _known_models(state: dict) -> set[str]: Empty when discovery found nothing (or no state was passed), which callers treat as "can't check" rather than "nothing is valid". + + ``claude_models`` holds only the newest id per family (the launch path pins one model per family + alias), so on its own it would reject the older versions the per-family prompts legitimately + offer. ``all_claude_models`` carries the full listing when the caller has it. """ known: set[str] = set() claude_models = state.get("claude_models") if isinstance(claude_models, dict): known.update(m for m in claude_models.values() if isinstance(m, str) and m) - for key in ("codex_models", "gemini_models", "oss_models"): + for key in ("codex_models", "gemini_models", "oss_models", "all_claude_models"): models = state.get(key) if isinstance(models, list): known.update(m for m in models if isinstance(m, str) and m) @@ -470,6 +509,28 @@ def validate_manifest(manifest: dict, state: dict | None = None) -> list[str]: return errors +def _agent_model_ids(agent_config: dict) -> set[str]: + """Every model id an agent is configured with — its list plus its default. + + Claude's ``models`` is a family-slot dict and the others' a flat list; codex has no list at all, + only ``default_model``. Returns an empty set when nothing is configured, which callers treat as + "can't check" rather than "nothing is allowed". + """ + model_config = agent_config.get("model_config") + if not isinstance(model_config, dict): + return set() + ids: set[str] = set() + raw = model_config.get("models") + if isinstance(raw, dict): + ids.update(v for v in raw.values() if isinstance(v, str) and v) + elif isinstance(raw, list): + ids.update(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: + ids.add(default_model) + return ids + + def _validate_budget_policy(budget_policy: dict, enabled_agents: dict[str, dict]) -> list[str]: """Validate a ``budget_policy`` against the agents the manifest enables.""" errors: list[str] = [] @@ -493,6 +554,7 @@ def _validate_budget_policy(budget_policy: dict, enabled_agents: dict[str, dict] else: percentages.append(float(pct)) tier_agent = tier.get("default_agent") + tier_model = tier.get("default_model") if not tier_agent: errors.append(f"budget_policy.tiers[{index}]: default_agent is required.") elif tier_agent not in enabled_agents: @@ -500,7 +562,18 @@ def _validate_budget_policy(budget_policy: dict, enabled_agents: dict[str, dict] f"budget_policy.tiers[{index}]: default_agent '{tier_agent}' must appear " "in enabled_agents." ) - if not tier.get("default_model"): + elif tier_model: + # The server only checks that the tier's agent is enabled, not that it has the model — + # so without this a tier can activate and hand the developer a model their agent was + # never configured with. Skipped when the agent lists no models (it then has only a + # default, or routes through a provider service whose catalog isn't enumerable). + available = _agent_model_ids(enabled_agents[tier_agent]) + if available and tier_model not in available: + errors.append( + f"budget_policy.tiers[{index}]: default_model '{tier_model}' is not one of the " + f"models configured for '{tier_agent}' ({', '.join(sorted(available))})." + ) + if not tier_model: errors.append(f"budget_policy.tiers[{index}]: default_model is required.") if len(set(percentages)) != len(percentages): diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py new file mode 100644 index 0000000..9cf2655 --- /dev/null +++ b/src/ucode/managed_wizard.py @@ -0,0 +1,900 @@ +"""Interactive `ucode setup`: author the workspace's managed coding-agent config. + +Workspace admins run this to build the ``CodingAgentConfig`` their developers will pull. It walks +the admin through agents, per-agent models, tracing, MCP servers, skills, and a spend-routing budget +policy, then writes the manifest to ``~/.ucode/managed-settings.json``. Publishing it to the +workspace is ``ucode apply`` (a separate command, so an admin can review the file first). + +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 — tracing, MCP, +skills — are delegated to the existing ``ucode 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 pathlib import Path + +from ucode.agents import TOOL_SPECS, check_gateway_endpoint +from ucode.databricks import ( + ANTHROPIC_FAMILIES, + 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, + service_usable_for_tool, +) +from ucode.managed_config import get_managed_config +from ucode.managed_setup import ( + CLAUDE_SLOT_FOR_FAMILY, + claude_family_candidates, + load_managed_settings, + model_options_for_agent, + save_managed_settings, + serialize_managed_config, + supports_provider_service, + validate_manifest, +) +from ucode.state import load_state +from ucode.ui import ( + console, + kv_line, + print_err, + print_heading, + print_note, + print_panel, + print_section, + print_success, + print_warning, + prompt_for_multi_selection, + prompt_for_percentage, + prompt_for_selection, + prompt_for_text, + prompt_for_tools, + prompt_yes_no_default, + spinner, +) + +# What `use_as_global_settings` actually does, in plain terms. Admins are choosing between a +# machine-wide managed settings file and a per-user one, which is not obvious from the field name. +GLOBAL_SETTINGS_BLURB = ( + "Write this agent's config to the machine's managed settings file, which applies to every " + "user on the machine and cannot be overridden locally. Answer no to write the per-user " + "settings file instead, which developers can still change." +) + +BUDGET_POLICY_BLURB = ( + "A budget policy moves developers onto cheaper agents and models as the workspace spends " + "against a budget — for example Claude Code on Opus by default, then Sonnet at 80%, then " + "OpenCode on Kimi at 100%. It only changes the default; developers can still pick anything " + "they have access to. Hard caps stay with the budget's own blocking threshold." +) + + +def _tracing_table_from_state(state: dict) -> str | None: + """The UC table `ucode 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_type_for_url(url: str) -> str | None: + """Classify a registered MCP server's URL into a managed-config type tag. + + ``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 ucode rebuild the URL. The URL shape is the only + signal available, so map it back. Returns None for a URL that matches nothing known, so unknown + servers are skipped rather than published with a guessed type. + """ + if "/ai-gateway/mcp-services/" in url: + return "mcp-service" + for fragment, tag in ( + ("/api/2.0/mcp/external/", "external"), + ("/api/2.0/mcp/genie/", "genie-space"), + ("/api/2.0/mcp/vector-search/", "vector-search"), + ("/api/2.0/mcp/functions/", "uc-functions"), + ): + if fragment in url: + return tag + if url.rstrip("/").endswith("/api/2.0/mcp/sql"): + return "sql" + # Databricks apps are the residual case: an arbitrary app host with a /mcp suffix. + if url.rstrip("/").endswith("/mcp"): + return "app" + 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] = [] + 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 + tag = _mcp_type_for_url(url) + if tag is None: + print_warning(f"Skipping MCP server '{name}': unrecognized URL shape ({url}).") + continue + servers.append({"name": 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 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 + return next(service for service in usable if service["name"] == selected) + + +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 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}")} + + if tool in SINGLE_MODEL_AGENTS: + return { + "default_model": _require_selection( + f"Select the model for {display}:", [(model, model) for model in options] + ) + } + + # 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_settings` already loads for + # the agent picker, so a re-run becomes an edit rather than a re-entry.) + picked = _require_multi_selection( + f"Select models for {display}:", + [(model, model) for model in options], + ) + 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 + 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 ucode 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 _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] = {} + for family in ANTHROPIC_FAMILIES: + family_models = candidates.get(family) + if not family_models: + continue + choice = prompt_for_selection( + f"Default {family} model:", + [(model, model) for model in family_models] + [(_SKIP_FAMILY, f"(skip {family})")], + searchable=True, + ) + if choice is None: + raise KeyboardInterrupt + if choice != _SKIP_FAMILY: + 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 = [m for family_models in candidates.values() for m in family_models] + return { + "default_model": _require_selection( + f"Which model should {display} use?", + [(m, m) for m in dict.fromkeys(every_model)], + ) + } + + 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_success(f"Overall default for {display}: {chosen[0]} (the only model configured)") + else: + model_config["default_model"] = _require_selection( + f"Which of those is {display}'s overall default?", [(m, m) for m in chosen] + ) + 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.") + + +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 _prompt_budget_policy( + workspace: str, token: str, enabled_agents: dict[str, dict], state: dict +) -> dict | None: + """Author a spend-routing ``budget_policy``, or None when the admin declines 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. + + A tier's model choices come from what the admin configured for that agent earlier in this run — + not 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. + """ + print_section("Budget policy") + print_note(BUDGET_POLICY_BLURB) + if not prompt_yes_no_default("Set up a budget policy for this workspace?", default=False): + return None + + with spinner("Listing workspace budgets..."): + budgets, reason = list_workspace_budgets(workspace, token) + if reason is not None or not budgets: + print_warning( + "No AI Gateway budgets are visible for this workspace, so there is nothing to attach a " + "policy to. Create a budget in the Databricks console first, then re-run `ucode setup`." + ) + return None + + 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 budgets + ], + searchable=True, + ) + if not budget_id: + return None + + policy: dict = {"budget_id": budget_id} + display_name = prompt_for_text("Policy name", default="coding-agents-tiered-routing") + if display_name: + policy["display_name"] = display_name + + tiers: list[dict] = [] + seen_percentages: set[float] = set() + print_note( + "Add one tier per step-down. Each tier activates once spend reaches its percentage, and " + "the highest activated tier wins." + ) + while True: + index = len(tiers) + 1 + fraction = prompt_for_percentage(f"Tier {index}: activates at what percent of budget?") + if fraction in seen_percentages: + print_err("That percentage is already used by another tier; pick a different one.") + continue + agent = prompt_for_selection( + f"Tier {index}: which agent becomes the default?", + [(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}: which model?", [(m, m) for m in options], searchable=True + ) + else: + model = prompt_for_text(f"Tier {index}: which model?") + if not model: + break + seen_percentages.add(fraction) + tiers.append( + { + "spending_percentage": fraction, + "default_agent": agent, + "default_model": model, + } + ) + 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}" + scope = "machine-wide" if agent_config.get("use_as_global_settings") else "per-user" + lines.append(kv_line(display, f"{detail} ({scope})")) + # 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")) + lines.append(kv_line("Tracing", manifest.get("tracing_table") or "disabled")) + + policy = manifest.get("budget_policy") + if isinstance(policy, dict): + tiers = policy.get("tiers") or [] + lines.append( + kv_line("Budget policy", policy.get("display_name") or policy.get("budget_id") or "set") + ) + 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("Budget policy", "none")) + + print_panel("Configuration summary", lines) + + +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}. `ucode 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 — `ucode apply` will fail " + "if you lack them." + ) + else: + print_success("Admin permissions verified") + + +def _warn_on_existing_config(workspace: str, token: str) -> None: + """Warn when the workspace already has a published config that `ucode apply` would replace. + + Deliberately doesn't itemize what the existing config holds. The admin doesn't need an inventory + to act on this — the instruction is the same either way ("include everything you want to keep") + — and `ucode setup show` prints the real thing for anyone who wants to compare. + """ + with spinner("Checking for an existing managed config..."): + existing, reason = get_managed_config(workspace, token) + if reason is not None: + print_note(f"Could not check for an existing config: {reason}") + return + if existing is None: + return + 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. Publishing " + "replaces all of it, so make sure this run includes everything you want to keep." + ) + + +def setup_from_file(path: str) -> int: + """Validate an admin-written manifest and save it, skipping the interactive flow. + + The non-interactive path for CI and for admins who'd rather keep the JSON in version control. + Reads ucode's own manifest shape (the same thing the wizard writes), not proto-JSON. + """ + 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 `ucode configure` first so ucode 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_settings(workspace, manifest) + _render_summary(workspace, manifest) + print_success(f"Saved to {manifest_path.name} -> ~/.ucode/managed-settings.json") + _print_next_steps() + return 0 + + +def _print_next_steps() -> None: + console.print() + print_heading("Next steps") + # Deliberately only `apply`. There is no way yet to try the authored config locally: the + # manifest describes what developers should get, while `ucode configure --dry-run` previews + # this machine's own agent configs, so pointing at it implied a local test it doesn't perform. + print_note("Publish it to the workspace: ucode apply") + + +def setup_command(from_file: str | None = None) -> int: + """Author the workspace's managed coding-agent config interactively. + + 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("ucode setup") + print_note("Author the managed coding config for this workspace.") + print_note("Developers pull it automatically when they run ucode.") + + 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. + ensure_databricks_auth(workspace, profile, quiet=True) + token = get_databricks_token(workspace, profile) + + _require_admin(workspace, token) + _warn_on_existing_config(workspace, token) + + # Discover the workspace's models and gateway URLs. This also logs in and persists local state, + # which is what lets the admin dry-run the config on their own machine afterwards. + 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 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." + ) + + previous = load_managed_settings(workspace) or {} + previously_enabled = [ + tool for tool in (previous.get("enabled_agents") or {}) if tool in TOOL_SPECS + ] + 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 + + default_agent = picked[0] + if len(picked) > 1: + chosen = prompt_for_selection( + "Which agent should launch when a developer runs `ucode`?", + [(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']}") + + enabled_agents: dict[str, dict] = {} + for tool in picked: + print_heading(TOOL_SPECS[tool]["display"]) + 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) + } + agent_config["use_as_global_settings"] = prompt_yes_no_default( + f"Apply {TOOL_SPECS[tool]['display']} config machine-wide? ({GLOBAL_SETTINGS_BLURB})", + default=False, + ) + enabled_agents[tool] = agent_config + + manifest: dict = {"default_agent": default_agent, "enabled_agents": enabled_agents} + + print_section("Tracing") + if prompt_yes_no_default( + "Send coding-session traces to an MLflow experiment in this workspace?", + default=bool(_tracing_table_from_state(state)), + ): + from ucode.tracing import configure_tracing_command + + configure_tracing_command(workspaces=[(workspace, profile)]) + tracing_table = _tracing_table_from_state(load_state()) + if tracing_table: + manifest["tracing_table"] = tracing_table + print_success(f"Tracing configured ({tracing_table})") + else: + print_warning("Tracing was not enabled, so it is left out of the managed config.") + + print_section("MCP servers") + if prompt_yes_no_default("Set up managed MCP servers for this workspace?", default=False): + from ucode.mcp import configure_mcp_command + + configure_mcp_command() + mcp_servers = _mcp_servers_from_state(load_state()) + if mcp_servers: + manifest["mcp_servers"] = mcp_servers + print_success(f"{len(mcp_servers)} MCP server(s) added to the managed config") + + print_section("Skills") + if prompt_yes_no_default("Set up managed skills for this workspace?", default=False): + locations = prompt_for_text( + "Skill schemas to publish, comma-separated `catalog.schema` (blank to skip)", + default="", + ) + parsed: list[str] = [item.strip() for item in (locations or "").split(",") if item.strip()] + if parsed: + from ucode.mcp import configure_skills_mcp_command + + configure_skills_mcp_command(parsed) + skill_names = _skill_names_from_state(load_state()) or parsed + manifest["skills"] = {"names": skill_names} + print_success(f"{len(skill_names)} skill schema(s) added to the managed config") + + budget_policy = _prompt_budget_policy(workspace, token, enabled_agents, state) + if budget_policy: + manifest["budget_policy"] = budget_policy + + 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 `apply` would reject. + print_err("The generated config is not valid:") + for error in errors: + print_note(error) + return 1 + + save_managed_settings(workspace, manifest) + _render_summary(workspace, manifest) + console.print() + print_success("Saved to ~/.ucode/managed-settings.json") + _print_next_steps() + return 0 + + +def show_command() -> int: + """Print the authored manifest and the proto-JSON `ucode apply` would publish.""" + workspace = load_state().get("workspace") + manifest = load_managed_settings(workspace) + if manifest is None: + print_note("No managed config has been authored yet. Run `ucode setup` to create one.") + return 0 + _render_summary(workspace or "unknown", manifest) + console.print() + print_heading("Payload for `ucode apply`") + console.print(json.dumps(serialize_managed_config(manifest), indent=2)) + return 0 + + +__all__ = ["setup_command", "setup_from_file", "show_command"] diff --git a/src/ucode/ui.py b/src/ucode/ui.py index 81ffc2e..b3ff6d1 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -13,6 +13,7 @@ import questionary from rich.console import Console +from rich.markup import escape from rich.panel import Panel from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn @@ -51,6 +52,28 @@ def print_kv(key: str, val: str) -> None: console.print(f" [bold]{key}:[/bold] [cyan]{val}[/cyan]") +def kv_line(key: str, val: str) -> str: + """A `print_kv`-styled line, returned instead of printed, for collecting into a panel. + + The value is markup-escaped. Rich reads bracketed text as a style tag and renders nothing for + it, so a policy name of ``[prod] tiered routing`` displayed as ``tiered routing`` in the config + summary — the one block an admin reads to confirm what they are about to publish workspace-wide. + Values here include admin-typed free text (policy name, skills locations, tracing table). + """ + return f"[bold]{escape(key)}:[/bold] [cyan]{escape(val)}[/cyan]" + + +def print_panel(title: str, lines: list[str]) -> None: + """Render `lines` inside a titled box. + + Unlike :func:`print_section`, which boxes a bare title, this boxes the body — so a block that + should be read as one unit (a config summary an admin is about to publish) reads as one, rather + than as loose lines that blend into whatever the flow printed before it. + """ + console.print() + console.print(Panel("\n".join(lines), title=title, style="blue", expand=False)) + + def print_note(text: str) -> None: console.print(f"[dim]•[/dim] {text}") @@ -268,12 +291,18 @@ def prompt_for_workspace( print_err(str(exc)) -def prompt_for_tools(available: list[tuple[str, str]]) -> list[str]: +def prompt_for_tools( + available: list[tuple[str, str]], + preselected: list[str] | set[str] | None = None, + prompt: str = "Select coding agents to configure:", +) -> list[str]: """Multi-select picker for coding agents. `available` is [(tool_id, display_name), ...]. Returns the chosen tool_ids. - All options are checked by default so hitting Enter selects everything. - Returns [] if the user submits an empty selection. + When ``preselected`` is None every option is checked by default, so hitting + Enter selects everything; pass a subset to pre-check only those (e.g. the + agents an existing managed config already enables). Returns [] if the user + submits an empty selection. """ style = questionary.Style( [ @@ -287,12 +316,17 @@ def prompt_for_tools(available: list[tuple[str, str]]) -> list[str]: ("answer", "fg:cyan"), ] ) + preselected_set = {str(item) for item in preselected} if preselected is not None else None choices = [ - questionary.Choice(title=display, value=tool_id, checked=True) + questionary.Choice( + title=display, + value=tool_id, + checked=(preselected_set is None or tool_id in preselected_set), + ) for tool_id, display in available ] answer = questionary.checkbox( - "Select coding agents to configure:", + prompt, choices=choices, style=style, pointer="›", @@ -302,11 +336,139 @@ def prompt_for_tools(available: list[tuple[str, str]]) -> list[str]: return list(answer) if answer else [] -def prompt_for_selection(prompt: str, options: list[tuple[str, str]]) -> str | None: +def prompt_for_multi_selection( + prompt: str, + options: list[tuple[str, str]], + preselected: list[str] | set[str] | None = None, + *, + searchable: bool = False, +) -> list[str] | None: + """Multi-select picker over arbitrary `(value, label)` options. + + Distinct from :func:`prompt_for_tools`, which is agent-specific and defaults to + everything checked: here nothing is checked unless ``preselected`` says so, since + an admin picking models wants an explicit choice rather than "all of them". + Returns the chosen values, [] on an empty submission, or None if cancelled + (Ctrl-C) so callers can distinguish "chose nothing" from "aborted". + + ``searchable`` lets the user narrow a long list by typing; see + :func:`prompt_for_selection` for why it trades away j/k navigation. + """ + style = questionary.Style( + [ + ("pointer", "fg:cyan bold"), + ("highlighted", "noinherit"), + ("selected", "noinherit"), + ("answer", "fg:cyan"), + ] + ) + preselected_set = {str(item) for item in preselected} if preselected is not None else set() + choices = [ + questionary.Choice(title=option_label, value=value, checked=value in preselected_set) + for value, option_label in options + ] + instruction = "(space to toggle, enter to confirm)" + if searchable: + instruction = "(type to filter, space to toggle, enter to confirm)" + answer = questionary.checkbox( + prompt, + choices=choices, + style=style, + pointer="›", + qmark="", + instruction=instruction, + use_search_filter=searchable, + use_jk_keys=not searchable, + ).ask() + return None if answer is None else list(answer) + + +def prompt_for_text( + prompt: str, *, default: str | None = None, required: bool = False +) -> str | None: + """Free-text prompt, used when model discovery found nothing to pick from. + + Returns the trimmed input, ``default`` on an empty answer, or None when there is no + default and the user submits nothing (or closes stdin). + + ``required=True`` raises ``KeyboardInterrupt`` on closed stdin instead of returning None, for + callers that loop until they get a value: returning None to such a caller spins forever on a + piped or exhausted stdin. Matches :func:`prompt_for_percentage`, which has no default and does + the same. + + A default is shown as ``[value] (enter to accept)`` rather than the bare ``[value]``: bracketed + text alone reads as a format example as easily as a value that will be used, so it invited + retyping what pressing enter would already pick. + + The whole bracketed hint is markup-escaped, brackets included. Rich reads + ``[coding-agents-tiered-routing]`` as a style tag and prints nothing for it, so an unescaped + word-like default vanished from the prompt entirely — numeric ones like ``[80]`` are not valid + tags and survived, which is why this looked fine wherever it was checked. + """ + hint = f" {escape(f'[{default}]')} (enter to accept)" if default else "" + while True: + try: + raw_value = console.input(f"{label(prompt)}{muted(hint)} {muted('›')} ").strip() + except EOFError as exc: + if required: + raise KeyboardInterrupt from exc + return default + if raw_value: + return raw_value + if default is not None: + return default + print_err("Please enter a value.") + + +def prompt_for_percentage(prompt: str, *, default: float | None = None) -> float: + """Prompt for a percentage (0-100) and return it as a fraction in [0, 1]. + + Budget tiers are fractions in the API (the server validates 0..1), but admins think in + percent — and the spec's own prose says "80%". Prompting in percent and converting here + keeps that mismatch in one place instead of at every call site. + + No caller passes ``default`` today, and tier thresholds deliberately have none: a threshold + decides when developers get downgraded, so it should be typed rather than accepted by accident. + The hint is still formatted (and escaped) the same way :func:`prompt_for_text` formats its own, + so the two cannot drift if a default is ever introduced. + + Raises ``KeyboardInterrupt`` on closed stdin when there is no default — see the handler below. + """ + hint = f" {escape(f'[{default * 100:g}]')} (enter to accept)" if default is not None else "" + while True: + try: + raw_value = console.input(f"{label(prompt)}{muted(hint)} {muted('› ')}").strip() + except EOFError as exc: + if default is not None: + return default + # Closed stdin with no default to fall back on is the admin abandoning the prompt, which + # is what Ctrl-C means here too. Raised as KeyboardInterrupt so the CLI's existing + # handler prints "Interrupted." and exits 130; a bare EOFError has no handler anywhere + # above this and reached the admin as a traceback. + raise KeyboardInterrupt from exc + if not raw_value and default is not None: + return default + try: + percent = float(raw_value.rstrip("%")) + except ValueError: + print_err("Please enter a number between 0 and 100.") + continue + if 0 <= percent <= 100: + return percent / 100 + print_err("Please enter a number between 0 and 100.") + + +def prompt_for_selection( + prompt: str, options: list[tuple[str, str]], *, searchable: bool = False +) -> str | None: """Single-select arrow-key picker. `options` is [(value, label), ...]. The prompt renders above the choices (questionary convention). Returns the chosen value, or None if the user cancels (Ctrl-C / empty). + + ``searchable`` lets the user narrow a long list by typing. It costs j/k navigation — questionary + rejects both at once, since j and k are also search characters — so it is opt-in for the pickers + that are actually long (model and budget lists), leaving short ones on plain arrow keys. """ style = questionary.Style( [ @@ -323,7 +485,9 @@ def prompt_for_selection(prompt: str, options: list[tuple[str, str]]) -> str | N style=style, pointer="›", qmark="", - instruction="(use arrow keys)", + instruction="(type to filter, arrow keys to move)" if searchable else "(use arrow keys)", + use_search_filter=searchable, + use_jk_keys=not searchable, ).ask() return answer diff --git a/tests/conftest.py b/tests/conftest.py index 0cc7932..a60d07d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,12 +24,16 @@ def _isolate_ucode_state(tmp_path, monkeypatch): it can never touch the developer's real ~/.ucode/state.json. """ import ucode.config_io as config_io_mod + import ucode.databricks as databricks_mod import ucode.state as state_mod state_dir = tmp_path / ".ucode" state_dir.mkdir() monkeypatch.setattr(state_mod, "STATE_PATH", state_dir / "state.json") monkeypatch.setattr(config_io_mod, "APP_DIR", state_dir) + # The model-services listing is memoized for the life of the process, so without this a cached + # result would leak into the next test and make a stubbed listing look like it was never called. + databricks_mod.clear_model_services_cache() def _workspace() -> str: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7e1a73a..1b142cf 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1940,3 +1940,168 @@ def test_failure_surfaces_cli_stderr(self, monkeypatch): install_ai_tools(["copilot"]) assert len(warnings) == 1 assert "copilot: cli-not-on-path: could not resolve copilot" in warnings[0] + + +class TestModelServicesCache: + """A successful listing is memoized per workspace: several callers want different views of the + same paginated walk (bucketed families vs the raw Claude ids), so one `ucode setup` run would + otherwise page the whole catalog twice.""" + + @staticmethod + def _counting_page(calls: dict): + def page(url, token): + calls["n"] = calls.get("n", 0) + 1 + return { + "model_services": [ + {"name": "model-services/system.ai.claude-opus-5"}, + {"name": "model-services/system.ai.claude-opus-4-8"}, + ] + }, None + + return page + + def test_repeat_listings_hit_the_api_once(self, monkeypatch): + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_get_model_services_page", self._counting_page(calls)) + first, _ = db_mod.list_model_services(WS, "tok") + second, _ = db_mod.list_model_services(WS, "tok") + assert first == second + assert calls["n"] == 1 + + def test_the_two_discovery_helpers_share_one_walk(self, monkeypatch): + # The duplicate spinner in `ucode setup`: `discover_model_services` and + # `discover_claude_models_unbucketed` both page the same endpoint. + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_get_model_services_page", self._counting_page(calls)) + claude, _codex, _gemini, _oss, _reason = db_mod.discover_model_services(WS, "tok") + unbucketed, _ = db_mod.discover_claude_models_unbucketed(WS, "tok") + assert calls["n"] == 1 + # Both views still come back intact: newest-per-family, and the full list. + assert claude["opus"] == "system.ai.claude-opus-5" + assert unbucketed == ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"] + + def test_use_cache_false_forces_a_fresh_walk(self, monkeypatch): + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_get_model_services_page", self._counting_page(calls)) + db_mod.list_model_services(WS, "tok") + db_mod.list_model_services(WS, "tok", use_cache=False) + assert calls["n"] == 2 + + def test_each_workspace_is_cached_separately(self, monkeypatch): + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_get_model_services_page", self._counting_page(calls)) + db_mod.list_model_services(WS, "tok") + db_mod.list_model_services("https://other.databricks.com", "tok") + assert calls["n"] == 2 + + def test_failures_are_not_cached(self, monkeypatch): + # A transient error must not poison the rest of the process into believing there are no + # models on the workspace. + calls: dict = {} + + def failing(url, token): + calls["n"] = calls.get("n", 0) + 1 + return None, "HTTP 500 Server Error" + + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_get_model_services_page", failing) + ids, reason = db_mod.list_model_services(WS, "tok") + assert ids == [] and reason is not None + + monkeypatch.setattr(db_mod, "_get_model_services_page", self._counting_page(calls)) + ids, reason = db_mod.list_model_services(WS, "tok") + assert reason is None + assert ids + + +class TestModelProviderServicesCache: + """The MPS listing is workspace-wide and filtered per agent afterwards, so one call serves every + agent — `ucode setup` used to re-list it once per MPS-capable agent.""" + + @staticmethod + def _counting_listing(calls: dict): + def get_json(url, token, timeout=10): + calls["n"] = calls.get("n", 0) + 1 + return { + "model_provider_services": [ + { + "name": "model-provider-services/main.j.ant", + "config": { + "provider_type": "ANTHROPIC", + "targets": [{"model": "claude-opus-5"}], + }, + }, + { + "name": "model-provider-services/main.j.oai", + "config": {"provider_type": "OPENAI", "targets": [{"model": "gpt-5"}]}, + }, + ] + }, None + + return get_json + + def test_one_call_serves_every_agent(self, monkeypatch): + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_http_get_json", self._counting_listing(calls)) + claude, _ = db_mod.list_tool_provider_services("claude", WS, "tok") + codex, _ = db_mod.list_tool_provider_services("codex", WS, "tok") + assert calls["n"] == 1 + # Each agent still gets only the services matching its API dialect. + assert claude == ["main.j.ant"] + assert codex == ["main.j.oai"] + + def test_use_cache_false_forces_a_fresh_call(self, monkeypatch): + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_http_get_json", self._counting_listing(calls)) + db_mod.list_model_provider_services(WS, "tok") + db_mod.list_model_provider_services(WS, "tok", use_cache=False) + assert calls["n"] == 2 + + def test_each_workspace_is_cached_separately(self, monkeypatch): + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_http_get_json", self._counting_listing(calls)) + db_mod.list_model_provider_services(WS, "tok") + db_mod.list_model_provider_services("https://other.databricks.com", "tok") + assert calls["n"] == 2 + + def test_failures_are_not_cached(self, monkeypatch): + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_http_get_json", lambda *a, **k: (None, "HTTP 500")) + services, reason = db_mod.list_model_provider_services(WS, "tok") + assert services == [] and reason is not None + monkeypatch.setattr(db_mod, "_http_get_json", self._counting_listing(calls)) + services, reason = db_mod.list_model_provider_services(WS, "tok") + assert reason is None and services + + def test_the_first_caller_cannot_corrupt_the_cache(self, monkeypatch): + # The caller that populates the cache gets the same list that was stored, so mutating it + # would poison every later reader. + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_http_get_json", self._counting_listing(calls)) + first, _ = db_mod.list_model_provider_services(WS, "tok") + first[0]["name"] = "clobbered" + first.pop() + second, _ = db_mod.list_model_provider_services(WS, "tok") + assert [s["name"] for s in second] == ["main.j.ant", "main.j.oai"] + + def test_a_later_caller_cannot_corrupt_the_cache(self, monkeypatch): + # And so does every cache *hit* — the wizard filters this list per agent, so the second + # agent's read must not see what the first one did to it. + calls: dict = {} + db_mod.clear_model_services_cache() + monkeypatch.setattr(db_mod, "_http_get_json", self._counting_listing(calls)) + db_mod.list_model_provider_services(WS, "tok") # populate + hit, _ = db_mod.list_model_provider_services(WS, "tok") + hit[0]["name"] = "clobbered" + hit.pop() + again, _ = db_mod.list_model_provider_services(WS, "tok") + assert [s["name"] for s in again] == ["main.j.ant", "main.j.oai"] diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index 3570a7f..8a08200 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -421,6 +421,73 @@ def test_slots_serialize_into_the_claude_variant(self): assert variant["models"] == {"default_opus_model": "system.ai.claude-opus-4-8"} +class TestClaudeFamilyCandidates: + """Discovery keeps one id per family for the launch path; authoring needs the alternatives.""" + + ALL = [ + "system.ai.claude-opus-4-1", + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-sonnet-5", + "system.ai.claude-haiku-4-5", + "system.ai.gpt-5-6", + ] + + def test_groups_by_family(self): + from ucode.managed_setup import claude_family_candidates + + got = claude_family_candidates(self.ALL) + assert set(got) == {"opus", "sonnet", "haiku"} + assert got["haiku"] == ["system.ai.claude-haiku-4-5"] + + def test_newest_first_within_a_family(self): + from ucode.managed_setup import claude_family_candidates + + assert claude_family_candidates(self.ALL)["opus"] == [ + "system.ai.claude-opus-5", + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-4-1", + ] + + def test_non_claude_models_are_ignored(self): + from ucode.managed_setup import claude_family_candidates + + assert not any( + "gpt" in m for models in claude_family_candidates(self.ALL).values() for m in models + ) + + def test_deduplicates(self): + from ucode.managed_setup import claude_family_candidates + + got = claude_family_candidates(["system.ai.claude-opus-5", "system.ai.claude-opus-5"]) + assert got["opus"] == ["system.ai.claude-opus-5"] + + def test_falls_back_to_the_per_family_picks(self): + # Without the full listing, the bucketed state is all that's available — one per family, but + # enough for the per-slot prompts to work. + from ucode.managed_setup import claude_family_candidates + + got = claude_family_candidates([], {"claude_models": {"opus": "system.ai.claude-opus-5"}}) + assert got == {"opus": ["system.ai.claude-opus-5"]} + + def test_empty_everything_yields_nothing(self): + from ucode.managed_setup import claude_family_candidates + + assert claude_family_candidates([], {}) == {} + + def test_slot_names_match_the_proto(self): + # ClaudeDefaultModels fields, verified against ai-gateway-api service.proto. + from ucode.managed_setup import CLAUDE_SLOT_FOR_FAMILY + + assert set(CLAUDE_SLOT_FOR_FAMILY.values()) == { + "default_fable_model", + "default_opus_model", + "default_sonnet_model", + "default_haiku_model", + } + + class TestValidate: def test_full_manifest_is_valid(self): assert validate_manifest(_full_manifest(), STATE) == [] @@ -472,6 +539,40 @@ def test_unknown_model_is_rejected(self): errors = validate_manifest(manifest, STATE) assert any("not available on this workspace" in e for e in errors) + def test_older_claude_version_is_recognized(self): + # `claude_models` holds only the newest per family, so without the full listing an older + # version the per-family prompts offered would be wrongly rejected. + 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"}, + } + } + }, + } + state = { + "claude_models": {"opus": "system.ai.claude-opus-5"}, + "all_claude_models": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], + } + assert validate_manifest(manifest, state) == [] + + def test_unknown_claude_version_is_still_rejected(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-9-9"}} + }, + } + state = { + "claude_models": {"opus": "system.ai.claude-opus-5"}, + "all_claude_models": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], + } + errors = validate_manifest(manifest, state) + assert any("not available on this workspace" in e for e in errors) + def test_model_check_skipped_without_state(self): manifest = { "default_agent": "claude", @@ -586,6 +687,106 @@ def test_tier_needs_a_default_model(self): errors = validate_manifest(manifest, STATE) assert any("default_model is required" in e for e in errors) + def test_tier_model_must_be_one_the_agent_has(self): + # The server only checks that the tier's agent is enabled, so without this a tier activates + # and hands the developer a model their agent was never configured with. + manifest = { + "default_agent": "pi", + "enabled_agents": { + "pi": { + "model_config": { + "default_model": "system.ai.kimi-k2-6", + "models": ["system.ai.kimi-k2-6"], + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "pi", + "default_model": "system.ai.gpt-5-6", + } + ], + }, + } + errors = validate_manifest(manifest, STATE) + assert any("is not one of the models configured for 'pi'" in e for e in errors), errors + + def test_tier_model_from_the_agents_list_is_accepted(self): + 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"], + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "pi", + "default_model": "system.ai.gpt-5-6", + } + ], + }, + } + assert validate_manifest(manifest, STATE) == [] + + def test_tier_model_matching_a_claude_family_slot_is_accepted(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-4-8", + "models": {"default_sonnet_model": "system.ai.claude-sonnet-4-6"}, + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-sonnet-4-6", + } + ], + }, + } + assert validate_manifest(manifest, STATE) == [] + + def test_tier_model_check_skipped_when_the_agent_lists_nothing(self): + # A provider-service agent has no enumerable catalog, so there is nothing to check against. + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "model_provider_service": "main.default.anthropic-mps", + "default_model": "claude-sonnet-5", + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "claude-sonnet-5", + } + ], + }, + } + assert validate_manifest(manifest, STATE) == [] + def test_budget_policy_alone_still_requires_a_default_agent(self): errors = validate_manifest({"budget_policy": {"budget_id": "b"}}) assert any("default_agent is required" in e for e in errors) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py new file mode 100644 index 0000000..126d630 --- /dev/null +++ b/tests/test_managed_wizard.py @@ -0,0 +1,1402 @@ +"""Tests for the interactive `ucode 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 unittest.mock import patch + +import pytest +import typer.main +from typer.testing import CliRunner + +import ucode.config_io as config_io_mod +import ucode.managed_setup as managed_setup_mod +import ucode.managed_wizard as wizard +from ucode.cli import app +from ucode.managed_setup import validate_manifest + +runner = CliRunner() + +WORKSPACE = "https://ws.example.com" + +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 manifest path at a tmp dir so no test touches the real ~/.ucode.""" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr( + managed_setup_mod, "MANAGED_SETTINGS_PATH", tmp_path / "managed-settings.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 TestMcpUrlClassification: + @pytest.mark.parametrize( + ("url", "expected"), + [ + ("https://ws.example.com/ai-gateway/mcp-services/system.ai.github", "mcp-service"), + ("https://ws.example.com/api/2.0/mcp/external/jira-prod", "external"), + ("https://ws.example.com/api/2.0/mcp/genie/01ef", "genie-space"), + ("https://ws.example.com/api/2.0/mcp/vector-search/main/default", "vector-search"), + ("https://ws.example.com/api/2.0/mcp/functions/main/default", "uc-functions"), + ("https://ws.example.com/api/2.0/mcp/sql", "sql"), + ("https://mcp-myapp-123.aws.databricksapps.com/mcp", "app"), + ], + ) + def test_known_urls(self, url, expected): + assert wizard._mcp_type_for_url(url) == expected + + def test_trailing_slash_is_tolerated(self): + assert wizard._mcp_type_for_url("https://ws.example.com/api/2.0/mcp/sql/") == "sql" + + def test_unknown_url_yields_none(self): + # Better to skip a server than publish it with a guessed type. + assert wizard._mcp_type_for_url("https://example.com/something/else") is None + + def test_sql_is_not_confused_for_app(self): + # Both end in a fixed segment; sql must win since it is checked first. + assert wizard._mcp_type_for_url("https://ws.example.com/api/2.0/mcp/sql") == "sql" + + +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"}, + ] + } + assert wizard._mcp_servers_from_state(state) == [ + {"name": "databricks-github", "type": "mcp-service"}, + {"name": "databricks-sql", "type": "sql"}, + ] + + 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_unclassifiable_servers(self): + state = {"mcp_servers": [{"name": "mystery", "url": "https://example.com/nope"}]} + 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 TestExistingConfigWarning: + @staticmethod + def _warn(existing: dict) -> str: + with ( + patch.object(wizard, "get_managed_config", return_value=(existing, None)), + patch.object(wizard, "print_warning") as warn, + ): + wizard._warn_on_existing_config(WORKSPACE, "token") + assert warn.called + return warn.call_args[0][0] + + def test_warns_that_publishing_replaces_the_whole_config(self): + message = self._warn({"enabled_agents": {"claude": {}, "codex": {}}}) + # There is one config per workspace covering everything, so the warning says that rather + # than reading like a per-agent notice. + assert "one config covers every agent" in message + assert "replaces all of it" in message + assert "everything you want to keep" in message + + def test_warning_does_not_itemize_the_existing_config(self): + # The message is the same whatever the config holds: an inventory doesn't change what the + # admin should do, and `ucode setup show` prints the real thing for comparison. + rich = self._warn( + { + "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"}, + } + ) + assert rich == self._warn({"enabled_agents": {}}) + for leaked in ("Claude Code", "OpenCode", "lillys_budget", "main.default"): + assert leaked not in rich, leaked + + def test_silent_when_no_config_exists(self): + with ( + patch.object(wizard, "get_managed_config", return_value=(None, None)), + patch.object(wizard, "print_warning") as warn, + ): + wizard._warn_on_existing_config(WORKSPACE, "token") + assert not warn.called + + def test_read_failure_is_a_note_not_a_warning(self): + # Can't check isn't the same as "there is one"; don't imply data loss. + with ( + patch.object(wizard, "get_managed_config", return_value=(None, "HTTP 403 Forbidden")), + patch.object(wizard, "print_warning") as warn, + patch.object(wizard, "print_note") as note, + ): + wizard._warn_on_existing_config(WORKSPACE, "token") + assert not warn.called + assert note.called + + +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 is set, and to what. + 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"), + patch.object(wizard, "print_success") as success, + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + assert config["default_model"] == "system.ai.claude-opus-4-8" + assert success.called + assert "system.ai.claude-opus-4-8" in success.call_args[0][0] + + 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. + assert set(offered[-1]) == { + "system.ai.claude-opus-5", + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-5", + } + + 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. + service = { + "name": "main.default.anthropic-mps", + "provider_type": "anthropic", + "targets": ["claude-sonnet-4-6", "claude-opus-4-8"], + "allow_all_targets": False, + } + with ( + patch.object(wizard, "prompt_for_selection", return_value="claude-opus-4-8") as select, + patch.object(wizard, "prompt_for_text") as text, + ): + config = wizard._prompt_models_for_agent("claude", STATE, service) + assert config == { + "model_provider_service": "main.default.anthropic-mps", + "default_model": "claude-opus-4-8", + } + 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]] == [ + "claude-opus-4-8", + "claude-sonnet-4-6", + ] + + 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_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 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"], + ), + ): + 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, + ): + 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_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_declining_yields_none(self): + with patch.object(wizard, "prompt_yes_no_default", return_value=False): + assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None + + def test_no_budgets_warns_and_yields_none(self): + with ( + patch.object(wizard, "prompt_yes_no_default", return_value=True), + patch.object(wizard, "list_workspace_budgets", return_value=([], "none found")), + patch.object(wizard, "print_warning") as warn, + ): + assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None + assert warn.called + + def test_percentages_are_stored_as_fractions(self): + budgets = [{"id": "budget-1", "display_name": "eng"}] + 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-1", "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-1" + assert policy["tiers"] == [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-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-1", "display_name": "eng"}] + 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-1", "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-1", "display_name": "eng"}] + 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-1", "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-1", "display_name": "eng"}] + 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-1", "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-1", "display_name": "eng"}] + 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-1", "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) == [] + + +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 + + +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_setup_mod.load_managed_settings(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_setup_mod.load_managed_settings(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_apply_payload(self, capsys): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} + }, + } + managed_setup_mod.save_managed_settings(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 `apply` 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 + + +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"}] + 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=[True, 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 + + +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_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_dry_run_sets_the_flag(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_command", return_value=0), + patch("ucode.cli.set_dry_run") as set_flag, + ): + runner.invoke(app, ["setup", "--dry-run"]) + set_flag.assert_called_once_with(True) + + 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 + + +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 "") diff --git a/tests/test_ui.py b/tests/test_ui.py index 3c0fbc8..bb2c9c5 100644 --- a/tests/test_ui.py +++ b/tests/test_ui.py @@ -2,15 +2,19 @@ from __future__ import annotations +import io from datetime import timedelta from unittest.mock import patch import pytest +from rich.console import Console from ucode.ui import ( format_duration, format_token_count, normalize_workspace_url, + prompt_for_percentage, + prompt_for_text, prompt_for_workspace, prompt_yes_no_default, render_box_table, @@ -50,6 +54,73 @@ def test_explicit_yes_overrides_default_false(self, monkeypatch): assert prompt_yes_no_default("go?", default=False) is True +def _visible(markup: str) -> str: + """What the user actually sees, with Rich markup resolved. + + Asserting on the raw markup string is what let a swallowed default ship: `[tiered]` is present + in the markup and absent from the output, because Rich reads it as a style tag. + """ + console = Console(file=io.StringIO(), force_terminal=False, width=200) + console.print(markup) + return console.file.getvalue().rstrip() + + +class TestDefaultsAreLabelledAsAcceptable: + """A shown default must say that enter takes it, or it reads as a format example.""" + + def test_text_default_says_enter_accepts_it(self): + with patch("ucode.ui.console.input", return_value="") as inp: + assert prompt_for_text("Policy name", default="tiered") == "tiered" + rendered = _visible(inp.call_args[0][0]) + assert "[tiered]" in rendered + assert "enter to accept" in rendered + + def test_a_word_like_default_is_not_eaten_as_markup(self): + # Rich treats `[coding-agents-tiered-routing]` as a style tag and renders nothing for it, so + # the real wizard default vanished from the prompt while `[80]` survived. + with patch("ucode.ui.console.input", return_value="") as inp: + prompt_for_text("Policy name", default="coding-agents-tiered-routing") + assert "[coding-agents-tiered-routing]" in _visible(inp.call_args[0][0]) + + def test_a_dotted_default_is_not_eaten_as_markup(self): + with patch("ucode.ui.console.input", return_value="") as inp: + prompt_for_text("Skills location", default="main.default") + assert "[main.default]" in _visible(inp.call_args[0][0]) + + def test_percentage_default_says_enter_accepts_it(self): + with patch("ucode.ui.console.input", return_value="") as inp: + assert prompt_for_percentage("at what percent?", default=0.8) == 0.8 + rendered = _visible(inp.call_args[0][0]) + # Prompted in percent even though the API takes a fraction. + assert "[80]" in rendered + assert "enter to accept" in rendered + + def test_no_default_shows_no_hint(self): + with patch("ucode.ui.console.input", return_value="typed") as inp: + assert prompt_for_text("Model") == "typed" + assert "enter to accept" not in _visible(inp.call_args[0][0]) + + def test_typing_still_overrides_the_default(self): + with patch("ucode.ui.console.input", return_value="mine"): + assert prompt_for_text("Policy name", default="tiered") == "mine" + + +class TestClosedStdinAborts: + """Ctrl-D must reach the CLI as an abort, not as a traceback.""" + + def test_percentage_without_a_default_raises_keyboard_interrupt(self): + # `ucode setup`'s tier prompt passes no default. EOFError has no handler above this call — + # the setup command catches only RuntimeError and KeyboardInterrupt — so a bare EOFError + # reached the admin as a raw traceback. + with patch("ucode.ui.console.input", side_effect=EOFError): + with pytest.raises(KeyboardInterrupt): + prompt_for_percentage("Tier 1: activates at what percent of budget?") + + def test_percentage_with_a_default_still_takes_it(self): + with patch("ucode.ui.console.input", side_effect=EOFError): + assert prompt_for_percentage("at what percent?", default=0.8) == 0.8 + + class TestNormalizeWorkspaceUrl: def test_adds_https_when_missing(self): assert normalize_workspace_url("example.databricks.com") == "https://example.databricks.com"