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 338fe29..b9dd630 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,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 @@ -193,6 +227,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 @@ -207,6 +244,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 baf1ad0..c4eedfb 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -69,6 +69,7 @@ managed_unservable_models, resolve_state, ) +from ucode.managed_wizard import setup_command, show_command from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -955,6 +956,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)." +) def _version_callback(value: bool) -> None: @@ -2192,6 +2197,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 c912e2c..73bbf04 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -331,6 +331,8 @@ 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" @@ -345,16 +347,18 @@ 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), so callers can say "unknown" rather than misreport an admin as a - non-admin. + 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. """ payload = _scim_me(workspace, token) if payload is None: return None groups = payload.get("groups") if not isinstance(groups, list): - # A well-formed `Me` for a user in no groups omits `groups`, so this is a definitive - # "not an admin" rather than a failed check. + # 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 @@ -362,6 +366,47 @@ def is_workspace_admin(workspace: str, token: str) -> bool | None: ) +# 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. @@ -873,12 +918,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) @@ -1323,12 +1376,44 @@ 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. Keyed by ``(workspace, parent)`` — a schema-scoped listing is a different result set than +# the metastore-wide one, so they must not share an entry. +_MODEL_PROVIDER_SERVICES_CACHE: dict[tuple[str, str], list[dict]] = {} + + +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, parent: str | None = None) -> 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. Takes + ``parent`` for the same reason the cache is keyed on it — a scoped listing is a separate entry. + """ + return (workspace, parent or "") 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. @@ -1337,7 +1422,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 @@ -1370,10 +1463,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]: @@ -1553,7 +1662,7 @@ def _provider_type_tag(provider_type: str | None) -> str: def list_model_provider_services( - workspace: str, token: str, *, parent: str | None = None + workspace: str, token: str, *, parent: str | None = None, use_cache: bool = True ) -> tuple[list[dict], str | None]: """List Unity Catalog Model Provider Services on the workspace. @@ -1569,7 +1678,25 @@ def list_model_provider_services( remainder silently dropped, so a service that plainly existed looked absent. ``parent`` scopes the listing to one ``catalog.schema`` — the metastore-wide default is documented as an internal, likely-to-be-deprecated scope, so prefer passing it when the schema is known. + + A successful result is memoized per workspace for the life of the process, like the + model-services listing: the listing is workspace-wide (filtered per agent afterwards by + :func:`service_usable_for_tool`), so without the memo `ucode setup` re-lists it once per + MPS-capable agent. Pass ``use_cache=False`` to force a fresh call. """ + # Keyed by workspace *and* parent: a `parent`-scoped listing holds only that schema's services, + # so caching it under the workspace alone would serve a partial list to an unscoped caller (and + # vice versa) — a service that plainly exists would look absent, the same failure pagination was + # added to fix. + cache_key = (workspace, parent or "") + if use_cache: + cached = _MODEL_PROVIDER_SERVICES_CACHE.get(cache_key) + 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) services: list[dict] = [] page_token: str | None = None @@ -1606,6 +1733,8 @@ def list_model_provider_services( if not services and last_reason is not None: return [], last_reason services.sort(key=lambda s: s["name"]) + if use_cache: + _MODEL_PROVIDER_SERVICES_CACHE[cache_key] = [dict(service) for service in services] return services, None diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 90d4d58..aaaf6ed 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -37,8 +37,10 @@ NO_MANAGED_CONFIG_MESSAGE = "No coding-agent config has been set up by your workspace admin yet." # CodingAgent proto enum -> ucode tool name. Anything unrecognized (e.g. a newer agent this ucode -# build doesn't know) is dropped during normalization rather than guessed at. -_AGENT_ENUM_TO_TOOL: dict[str, str] = { +# build doesn't know) is dropped during normalization rather than guessed at. Public because the +# admin-write side (``managed_setup``) inverts these maps to serialize, so a new agent or MCP type +# only has to be declared once. +AGENT_ENUM_TO_TOOL: dict[str, str] = { "CODING_AGENT_CLAUDE_CODE": "claude", "CODING_AGENT_CODEX": "codex", "CODING_AGENT_GEMINI": "gemini", @@ -49,7 +51,7 @@ # McpServerType proto enum -> ucode's short type tag. Mirrors the selection prefixes in ``mcp.py``; # the actual name->URL resolution happens there when the manifest is applied (a later change). -_MCP_TYPE_ENUM_TO_TAG: dict[str, str] = { +MCP_TYPE_ENUM_TO_TAG: dict[str, str] = { "MCP_SERVER_TYPE_UC_SERVICE": "mcp-service", "MCP_SERVER_TYPE_EXTERNAL": "external", "MCP_SERVER_TYPE_GENIE": "genie-space", @@ -131,7 +133,7 @@ def _normalize_enabled_agent(entry: object) -> tuple[str, dict] | None: entry_dict = _as_dict(entry) if not entry_dict: return None - tool = _AGENT_ENUM_TO_TOOL.get(_str(entry_dict.get("agent")) or "") + tool = AGENT_ENUM_TO_TOOL.get(_str(entry_dict.get("agent")) or "") if tool is None: return None config_in = _as_dict(entry_dict.get("config")) @@ -166,7 +168,7 @@ def _normalize_mcp_servers(value: object) -> list[dict]: for entry in value: entry_dict = _as_dict(entry) name = _str(entry_dict.get("name")) - tag = _MCP_TYPE_ENUM_TO_TAG.get(_str(entry_dict.get("type")) or "") + tag = MCP_TYPE_ENUM_TO_TAG.get(_str(entry_dict.get("type")) or "") if name and tag: out.append({"name": name, "type": tag}) return out @@ -191,7 +193,7 @@ def _normalize_budget_policy(value: object) -> dict | None: if not isinstance(pct, (int, float)) or isinstance(pct, bool): continue tier_out: dict = {"spending_percentage": float(pct)} - agent = _AGENT_ENUM_TO_TOOL.get(_str(tier_dict.get("default_agent")) or "") + agent = AGENT_ENUM_TO_TOOL.get(_str(tier_dict.get("default_agent")) or "") if agent: tier_out["default_agent"] = agent model = _str(tier_dict.get("default_model")) @@ -214,7 +216,7 @@ def normalize_managed_config(raw: dict) -> dict: name = _str(raw.get("name")) if name: result["name"] = name - default_agent = _AGENT_ENUM_TO_TOOL.get(_str(raw.get("default_agent")) or "") + default_agent = AGENT_ENUM_TO_TOOL.get(_str(raw.get("default_agent")) or "") if default_agent: result["default_agent"] = default_agent enabled_agents: dict[str, dict] = {} diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py new file mode 100644 index 0000000..2ef8373 --- /dev/null +++ b/src/ucode/managed_setup.py @@ -0,0 +1,634 @@ +"""Admin-authored managed coding-agent config: model catalogs, validation, and serialization. + +This module owns the admin-write half of the managed config, mirroring the developer-read half in +:mod:`ucode.managed_config`: + +- which models an admin may pick for each agent (:func:`model_options_for_agent`), +- validating a manifest before it is published (:func:`validate_manifest`), +- serializing ucode's internal manifest shape into proto-JSON ``CodingAgentConfig`` + (:func:`serialize_managed_config`), and +- persisting the authored manifest to ``~/.ucode/managed-settings.json``. + +The manifest shape here is exactly the one :func:`ucode.managed_config.normalize_managed_config` +produces, so ``serialize`` then ``normalize`` round-trips to the input. The enum maps are derived by +inverting that module's maps rather than restated, so a new agent or MCP type only has to be added +once. + +``managed-settings.json`` (authored by an admin, published by ``ucode apply``) is distinct from +``managed-state.json`` (pulled from the workspace by a developer, owned by ``managed_config``). + +The interactive wizard that calls these helpers, and the publish step, live in later changes; this +module deliberately stops at "catalogs + validate + serialize + persist". +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import cast + +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 ( + AGENT_ENUM_TO_TOOL, + MCP_TYPE_ENUM_TO_TAG, +) + +MANAGED_SETTINGS_PATH = config_io.APP_DIR / "managed-settings.json" + +# ucode tool name -> CodingAgent proto enum, and ucode MCP type tag -> McpServerType proto enum. +# Inverted from the read side's maps so the two directions cannot drift: adding an agent to +# `managed_config._AGENT_ENUM_TO_TOOL` makes it serializable here automatically. +AGENT_TOOL_TO_ENUM: dict[str, str] = {tool: enum for enum, tool in AGENT_ENUM_TO_TOOL.items()} +MCP_TAG_TO_TYPE_ENUM: dict[str, str] = {tag: enum for enum, tag in MCP_TYPE_ENUM_TO_TAG.items()} + +# `AgentModelConfig` oneof variant key per agent. The server rejects a config whose variant doesn't +# match its agent (`validateAgentModelConfig`), so this mapping is not cosmetic. +_AGENT_MODEL_CONFIG_VARIANT: dict[str, str] = { + "claude": "claude", + "codex": "codex", + "opencode": "opencode", + "pi": "pi", + "gemini": "gemini", + "copilot": "copilot", +} + +# Agents whose model config carries a flat `models` list. Claude instead uses per-family slots +# (`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. 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 +} + +# Which discovered model families each agent may be configured with, on the Databricks-hosted path. +# Claude Code only speaks the Anthropic dialect, Gemini CLI only Gemini; Codex's `/v1/models` route +# serves GPT plus the OSS models; the multi-provider harnesses can use anything discovered. +_AGENT_MODEL_FAMILIES: dict[str, tuple[str, ...]] = { + "claude": ("claude",), + "gemini": ("gemini",), + "codex": ("codex", "oss"), + "opencode": ("claude", "codex", "gemini", "oss"), + "pi": ("claude", "codex", "gemini", "oss"), + "copilot": ("claude", "codex", "gemini", "oss"), +} + + +def _as_dict(value: object) -> dict[str, object]: + """Return ``value`` as a ``dict[str, object]`` when it is a dict, else an empty dict. + + Mirrors the read side's helper: a bare ``isinstance(x, dict)`` narrows to ``dict[Never, Never]``, + which rejects string keys, so ``.get("name")`` on an untyped manifest fails type-checking without + this. + """ + return cast("dict[str, object]", value) if isinstance(value, dict) else {} + + +def model_families_for_agent(tool: str) -> tuple[str, ...]: + """The discovered model families ``tool`` can be configured with.""" + return _AGENT_MODEL_FAMILIES.get(tool, ()) + + +def supports_provider_service(tool: str, provider_type: str) -> bool: + """True when ``tool`` can route through a ``provider_type`` Model Provider Service. + + Thin pass-through to :func:`ucode.databricks.tool_supports_provider_type` so the wizard has one + obvious place to ask. Only claude (anthropic / amazon_bedrock) and codex (openai) have MPS + support today; the other harnesses are Databricks-hosted only. + """ + return tool_supports_provider_type(tool, provider_type) + + +def model_options_for_agent(tool: str, state: dict) -> list[str]: + """Models an admin may pick for ``tool``, drawn from the workspace's discovered inventory. + + ``state`` is a hydrated workspace state (as ``configure_shared_state`` produces): ``claude_models`` + is a family->id dict, while ``codex_models`` / ``gemini_models`` / ``oss_models`` are lists. + Returns a de-duplicated list in a stable order (Claude newest-family first, then the family + order in :data:`_AGENT_MODEL_FAMILIES`) — empty when nothing was discovered for those families, + in which case the caller should fall back to free-text entry. + """ + options: list[str] = [] + for family in model_families_for_agent(tool): + if family == "claude": + claude_models = state.get("claude_models") + if isinstance(claude_models, dict): + # ANTHROPIC_FAMILIES is newest-tier-first, so slot order is meaningful. + for name in ANTHROPIC_FAMILIES: + model = claude_models.get(name) + if isinstance(model, str) and model: + options.append(model) + continue + key = {"codex": "codex_models", "gemini": "gemini_models", "oss": "oss_models"}[family] + models = state.get(key) + if isinstance(models, list): + options.extend(m for m in models if isinstance(m, str) and m) + # dict.fromkeys de-duplicates while preserving first-seen order (a model can match more than + # one family bucket, e.g. an OSS id that also appears in the codex listing). + return list(dict.fromkeys(options)) + + +def claude_family_for_model(model: str) -> str | None: + """The Claude family (``opus``/``sonnet``/``haiku``/``fable``) a model id belongs to, or None. + + Used to place picked Claude models into their `ClaudeDefaultModels` slots. Matches on the + family segment so both discovery spellings work (``system.ai.claude-opus-4-8`` and + ``databricks-claude-opus-4-8``). + """ + lowered = model.lower() + 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. + + Claude Code addresses models by family alias rather than by list, so the wizard's multi-select + has to be bucketed into ``default_opus_model`` / ``default_sonnet_model`` / etc. Ids whose family + can't be identified are skipped; when two ids share a family the first wins (the caller's list + order is the admin's preference order). + """ + slots: dict[str, str] = {} + for model in models: + family = claude_family_for_model(model) + if family is None: + continue + slot = CLAUDE_SLOT_FOR_FAMILY[family] + slots.setdefault(slot, model) + return slots + + +def _model_config_payload(tool: str, model_config: dict) -> dict: + """Build one ``AgentModelConfig`` oneof variant body for ``tool``. + + Shapes per the proto: claude gets `models` as a `ClaudeDefaultModels` slot object, codex gets + no model list at all, and the rest get a flat repeated `models`. + """ + body: dict = {} + mps = model_config.get("model_provider_service") + if isinstance(mps, str) and mps: + body["model_provider_service"] = mps + default_model = model_config.get("default_model") + if isinstance(default_model, str) and default_model: + body["default_model"] = default_model + + models = model_config.get("models") + if tool == "claude": + if isinstance(models, dict): + slots = { + slot: value + for slot, value in models.items() + if isinstance(slot, str) and isinstance(value, str) and value + } + if slots: + body["models"] = slots + elif tool in _FLAT_MODEL_LIST_AGENTS: + if isinstance(models, list): + model_list = [m for m in models if isinstance(m, str) and m] + if model_list: + body["models"] = model_list + # codex intentionally carries no model list — CodexModelConfig has only + # model_provider_service + default_model. + return body + + +def _enabled_agent_payload(tool: str, agent_config: dict) -> dict: + """Build one ``EnabledAgent`` entry (agent enum + its ``AgentConfig``).""" + config: dict = {} + use_as_global = agent_config.get("use_as_global_settings") + if isinstance(use_as_global, bool): + config["use_as_global_settings"] = use_as_global + headers = agent_config.get("custom_headers") + if isinstance(headers, dict): + clean = {k: v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)} + if clean: + config["custom_headers"] = clean + tracing_table = agent_config.get("tracing_table") + if isinstance(tracing_table, str) and tracing_table: + config["tracing_config"] = {"table": tracing_table} + model_config = agent_config.get("model_config") + if isinstance(model_config, dict): + body = _model_config_payload(tool, model_config) + if body: + variant = _AGENT_MODEL_CONFIG_VARIANT[tool] + config["model_config"] = {variant: body} + + entry: dict = {"agent": AGENT_TOOL_TO_ENUM[tool]} + if config: + entry["config"] = config + return entry + + +def _budget_policy_payload(budget_policy: dict) -> dict: + """Build the ``BudgetPolicy`` body, dropping tiers that name an unknown agent. + + ``spending_percentage`` is passed through as-is: it is a fraction in [0, 1] both in ucode's + manifest and in the proto (the server validates that range). Callers prompting an admin in + percent must divide before building the manifest. + """ + payload: dict = {} + display_name = budget_policy.get("display_name") + if isinstance(display_name, str) and display_name: + payload["display_name"] = display_name + budget_id = budget_policy.get("budget_id") + if isinstance(budget_id, str) and budget_id: + payload["budget_id"] = budget_id + + tiers: list[dict] = [] + raw_tiers = budget_policy.get("tiers") + for tier in raw_tiers if isinstance(raw_tiers, list) else []: + if not isinstance(tier, dict): + continue + pct = tier.get("spending_percentage") + if not isinstance(pct, (int, float)) or isinstance(pct, bool): + continue + tier_payload: dict = {"spending_percentage": float(pct)} + agent_enum = AGENT_TOOL_TO_ENUM.get(str(tier.get("default_agent") or "")) + if agent_enum: + tier_payload["default_agent"] = agent_enum + default_model = tier.get("default_model") + if isinstance(default_model, str) and default_model: + tier_payload["default_model"] = default_model + tiers.append(tier_payload) + if tiers: + payload["tiers"] = tiers + return payload + + +def serialize_managed_config(manifest: dict) -> dict: + """Serialize ucode's internal manifest into a proto-JSON ``CodingAgentConfig``. + + The exact inverse of :func:`ucode.managed_config.normalize_managed_config`: tool names become + ``CODING_AGENT_*`` enums, MCP type tags become ``MCP_SERVER_TYPE_*``, and each agent's model + config is wrapped in its matching ``AgentModelConfig`` oneof variant. Agents and MCP types this + build doesn't recognize are dropped, mirroring the read side. + + Output-only proto fields (``workspace_id``, timestamps, user ids) are never emitted. ``name`` is + carried through when present so an update path can address an existing resource; ``ucode apply`` + omits it on create and lets the server assign one. + """ + payload: dict = {} + + name = manifest.get("name") + if isinstance(name, str) and name: + payload["name"] = name + display_name = manifest.get("display_name") + if isinstance(display_name, str) and display_name: + payload["display_name"] = display_name + + default_agent = AGENT_TOOL_TO_ENUM.get(str(manifest.get("default_agent") or "")) + if default_agent: + payload["default_agent"] = default_agent + + enabled_agents = manifest.get("enabled_agents") + if isinstance(enabled_agents, dict): + entries = [ + _enabled_agent_payload(tool, agent_config) + for tool, agent_config in enabled_agents.items() + if tool in AGENT_TOOL_TO_ENUM and isinstance(agent_config, dict) + ] + if entries: + payload["enabled_agents"] = entries + + mcp_servers = manifest.get("mcp_servers") + if isinstance(mcp_servers, list): + servers: list[dict] = [] + for server in mcp_servers: + if not isinstance(server, dict): + continue + server_name = server.get("name") + type_enum = MCP_TAG_TO_TYPE_ENUM.get(str(server.get("type") or "")) + if isinstance(server_name, str) and server_name and type_enum: + servers.append({"name": server_name, "type": type_enum}) + if servers: + payload["mcp_servers"] = servers + + skills = manifest.get("skills") + if isinstance(skills, dict): + names = skills.get("names") + if isinstance(names, list): + skill_names = [n for n in names if isinstance(n, str) and n] + if skill_names: + payload["skills"] = {"names": skill_names} + + tracing_table = manifest.get("tracing_table") + if isinstance(tracing_table, str) and tracing_table: + payload["tracing"] = {"table": tracing_table} + + budget_policy = manifest.get("budget_policy") + if isinstance(budget_policy, dict): + policy = _budget_policy_payload(budget_policy) + if policy: + payload["budget_policy"] = policy + + return payload + + +def _manifest_default_model(agent_config: dict) -> str | None: + """The ``default_model`` on an agent's model config, or None when unset/empty.""" + model_config = agent_config.get("model_config") + if not isinstance(model_config, dict): + return None + default_model = model_config.get("default_model") + return default_model if isinstance(default_model, str) and default_model else None + + +def _known_models(state: dict) -> set[str]: + """Every model id discovered on the workspace, across all families. + + 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", "all_claude_models"): + models = state.get(key) + if isinstance(models, list): + known.update(m for m in models if isinstance(m, str) and m) + return known + + +def _validate_agent_models(tool: str, agent_config: dict, known: set[str]) -> list[str]: + """Check one agent's configured models against the workspace inventory. + + Skipped entirely when the agent routes through a Model Provider Service: those model ids come + from the provider's own catalog, not from UC model services, so the workspace inventory says + nothing about them. + """ + model_config = agent_config.get("model_config") + if not isinstance(model_config, dict): + return [] + if model_config.get("model_provider_service"): + return [] + + referenced: list[str] = [] + default_model = model_config.get("default_model") + if isinstance(default_model, str) and default_model: + referenced.append(default_model) + models = model_config.get("models") + if isinstance(models, dict): + referenced.extend(m for m in models.values() if isinstance(m, str) and m) + elif isinstance(models, list): + referenced.extend(m for m in models if isinstance(m, str) and m) + + return [ + f"{tool}: model '{model}' is not available on this workspace." + for model in dict.fromkeys(referenced) + if model not in known + ] + + +def validate_manifest(manifest: dict, state: dict | None = None) -> list[str]: + """Validate a manifest before publishing; returns human-readable errors (empty when valid). + + Mirrors the server's ``validateStoredConfig`` so an admin sees problems locally instead of via + an ``INVALID_PARAMETER_VALUE`` round-trip: + + - ``default_agent`` is required once any agent configuration is present, must appear in + ``enabled_agents``, and that agent must have a non-empty ``default_model``; + - every ``enabled_agents`` key must be an agent this ucode build knows; + - each MCP server needs a name and a recognized type; skill names must be non-empty; + - ``tracing_table`` must be non-empty when the key is present; + - a ``budget_policy`` needs a ``budget_id``, and each tier needs a ``spending_percentage`` in + [0, 1] (unique across tiers), a ``default_agent`` that appears in ``enabled_agents``, and a + ``default_model``. + + When ``state`` is provided, configured models are additionally checked against the workspace's + discovered inventory — skipped for agents routing through a Model Provider Service, and skipped + entirely when discovery returned nothing. + """ + errors: list[str] = [] + + enabled_agents_raw = manifest.get("enabled_agents") + enabled_agents: dict[str, dict] = {} + if isinstance(enabled_agents_raw, dict): + for tool, agent_config in enabled_agents_raw.items(): + if tool not in AGENT_TOOL_TO_ENUM: + valid = ", ".join(sorted(AGENT_TOOL_TO_ENUM)) + errors.append( + f"enabled_agents: '{tool}' is not a supported agent (valid: {valid})." + ) + continue + if not isinstance(agent_config, dict): + errors.append(f"enabled_agents: '{tool}' config must be an object.") + continue + enabled_agents[tool] = agent_config + + default_agent = manifest.get("default_agent") + budget_policy = manifest.get("budget_policy") + has_agent_selection = bool(default_agent or enabled_agents_raw or budget_policy) + if has_agent_selection: + if not default_agent: + errors.append("default_agent is required when agent configuration is present.") + elif default_agent not in enabled_agents: + errors.append(f"default_agent '{default_agent}' must appear in enabled_agents.") + elif not _manifest_default_model(enabled_agents[default_agent]): + errors.append( + f"default_agent '{default_agent}' must have a non-empty model_config.default_model." + ) + + known = _known_models(state or {}) + if known: + for tool, agent_config in enabled_agents.items(): + errors.extend(_validate_agent_models(tool, agent_config, known)) + + mcp_servers = manifest.get("mcp_servers") + if isinstance(mcp_servers, list): + for index, raw_server in enumerate(mcp_servers, start=1): + if not isinstance(raw_server, dict): + errors.append(f"mcp_servers[{index}] must be an object.") + continue + server = _as_dict(raw_server) + if not server.get("name"): + errors.append(f"mcp_servers[{index}]: name is required.") + server_type = str(server.get("type") or "") + if server_type not in MCP_TAG_TO_TYPE_ENUM: + valid = ", ".join(sorted(MCP_TAG_TO_TYPE_ENUM)) + errors.append( + f"mcp_servers[{index}]: type '{server_type}' is not recognized " + f"(valid: {valid})." + ) + + skills = manifest.get("skills") + if isinstance(skills, dict): + names = skills.get("names") + if isinstance(names, list) and any(not isinstance(name, str) or not name for name in names): + errors.append("skills.names must not contain empty names.") + + if "tracing_table" in manifest and not manifest.get("tracing_table"): + errors.append("tracing_table must not be empty.") + + if isinstance(budget_policy, dict): + errors.extend(_validate_budget_policy(budget_policy, enabled_agents)) + + 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] = [] + if not budget_policy.get("budget_id"): + errors.append("budget_policy.budget_id is required.") + + percentages: list[float] = [] + tiers = budget_policy.get("tiers") + for index, tier in enumerate(tiers if isinstance(tiers, list) else [], start=1): + if not isinstance(tier, dict): + errors.append(f"budget_policy.tiers[{index}] must be an object.") + continue + pct = tier.get("spending_percentage") + if not isinstance(pct, (int, float)) or isinstance(pct, bool): + errors.append(f"budget_policy.tiers[{index}]: spending_percentage is required.") + elif not 0 <= float(pct) <= 1: + errors.append( + f"budget_policy.tiers[{index}]: spending_percentage must be a fraction " + f"between 0 and 1 (got {pct})." + ) + 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: + errors.append( + f"budget_policy.tiers[{index}]: default_agent '{tier_agent}' must appear " + "in enabled_agents." + ) + 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): + errors.append("budget_policy tier spending_percentage values must be unique.") + return errors + + +def save_managed_settings(workspace: str, manifest: dict) -> None: + """Persist the authored manifest to ``~/.ucode/managed-settings.json``. No-op in dry-run. + + Stored alongside its workspace so ``ucode apply`` can refuse to publish a manifest that was + authored against a different workspace. + """ + if config_io.is_dry_run(): + return + payload = {"workspace": workspace, "config": manifest} + config_io.ensure_parent_dir(MANAGED_SETTINGS_PATH) + try: + MANAGED_SETTINGS_PATH.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + except OSError as exc: + raise RuntimeError( + f"Failed to write managed settings file: {MANAGED_SETTINGS_PATH}" + ) from exc + _restrict_permissions(MANAGED_SETTINGS_PATH) + + +def _restrict_permissions(path: Path) -> None: + """Best-effort chmod 0600, matching how ``managed_config`` protects ``managed-state.json``. + + An unpublished manifest can name internal catalogs, budgets, and MCP servers, so it should not + be group- or world-readable. No-op where unsupported (e.g. Windows). + """ + try: + os.chmod(path, 0o600) + except (OSError, NotImplementedError): + pass + + +def load_managed_settings(workspace: str | None = None) -> dict | None: + """Load the authored manifest, or None when absent (or authored for another workspace). + + Passing ``workspace`` scopes the read the way :func:`ucode.managed_config.load_managed_state` + does, so a manifest left over from a different workspace is ignored rather than published to the + wrong place. Omit it to read whatever is on disk. + """ + data = config_io.read_json_safe(MANAGED_SETTINGS_PATH) + if not data: + return None + if workspace is not None and data.get("workspace") != workspace: + return None + manifest = data.get("config") + return manifest if isinstance(manifest, dict) else None + + +def managed_settings_workspace() -> str | None: + """The workspace the on-disk manifest was authored for, or None when there is no manifest.""" + workspace = config_io.read_json_safe(MANAGED_SETTINGS_PATH).get("workspace") + return workspace if isinstance(workspace, str) and workspace else None 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 a7afd47..bfdf3e3 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2095,6 +2095,171 @@ def test_buckets_by_family(self, model_id, expected): assert classify_model_family(model_id) == expected +class TestModelServicesCache: + """A successful listing is memoized per workspace: several callers want different views of the + same paginated walk (bucketed families vs the raw Claude ids), so one `ucode setup` run would + 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"] + + class TestIsWorkspaceAdmin: """Admin detection reuses the SCIM `Me` payload, which carries group membership.""" diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py new file mode 100644 index 0000000..8a08200 --- /dev/null +++ b/tests/test_managed_setup.py @@ -0,0 +1,875 @@ +"""Tests for the admin-write half of the managed coding-agent config. + +The most valuable case here is the round-trip: ``serialize_managed_config`` followed by +``managed_config.normalize_managed_config`` must return the manifest it started from. That single +property pins the write side to the read side, so the two cannot drift as the proto grows. +""" + +from __future__ import annotations + +import json +import stat + +import pytest + +import ucode.config_io as config_io_mod +import ucode.managed_setup as managed_setup_mod +from ucode.managed_config import ( + AGENT_ENUM_TO_TOOL, + MCP_TYPE_ENUM_TO_TAG, + normalize_managed_config, +) +from ucode.managed_setup import ( + AGENT_TOOL_TO_ENUM, + MCP_TAG_TO_TYPE_ENUM, + claude_family_for_model, + claude_model_slots, + load_managed_settings, + managed_settings_workspace, + model_families_for_agent, + model_options_for_agent, + save_managed_settings, + serialize_managed_config, + supports_provider_service, + validate_manifest, +) + +WORKSPACE = "https://ws.example.com" + +# A workspace state shaped like `configure_shared_state` produces. +STATE = { + "workspace": WORKSPACE, + "claude_models": { + "opus": "system.ai.claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-4-6", + "haiku": "system.ai.claude-haiku-4-5", + }, + "codex_models": ["system.ai.gpt-5-6"], + "gemini_models": ["system.ai.gemini-3-flash"], + "oss_models": ["system.ai.kimi-k2-6"], +} + + +def _minimal_manifest() -> dict: + """The smallest manifest that passes validation: one agent, which is the default.""" + return { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": {"default_model": "system.ai.claude-opus-4-8"}, + } + }, + } + + +def _full_manifest() -> dict: + """A manifest exercising every field the read side normalizes.""" + return { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "use_as_global_settings": True, + "custom_headers": {"x-databricks-workspace": "eng-ml-inference"}, + "tracing_table": "main.default.claude-traces", + "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", + }, + }, + }, + "codex": { + "use_as_global_settings": False, + "model_config": {"default_model": "system.ai.gpt-5-6"}, + }, + "opencode": { + "model_config": { + "default_model": "system.ai.claude-opus-4-8", + "models": ["system.ai.claude-opus-4-8", "system.ai.kimi-k2-6"], + }, + }, + }, + "mcp_servers": [ + {"name": "system.ai.github", "type": "mcp-service"}, + {"name": "genie-space-id", "type": "genie-space"}, + ], + "skills": {"names": ["system.ai.pdf-extraction"]}, + "tracing_table": "main.default.ucode-traces", + "budget_policy": { + "display_name": "eng-tiered-routing", + "budget_id": "c6563b45-df9a-4b19-afb2-d42dc2b52576", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-sonnet-4-6", + }, + { + "spending_percentage": 1.0, + "default_agent": "opencode", + "default_model": "system.ai.kimi-k2-6", + }, + ], + }, + } + + +class TestEnumMaps: + def test_agent_map_is_the_inverse_of_the_read_side(self): + assert AGENT_TOOL_TO_ENUM == {tool: enum for enum, tool in AGENT_ENUM_TO_TOOL.items()} + + def test_mcp_map_is_the_inverse_of_the_read_side(self): + assert MCP_TAG_TO_TYPE_ENUM == {tag: enum for enum, tag in MCP_TYPE_ENUM_TO_TAG.items()} + + def test_agent_map_round_trips(self): + for tool, enum in AGENT_TOOL_TO_ENUM.items(): + assert AGENT_ENUM_TO_TOOL[enum] == tool + + def test_inversion_is_lossless(self): + # A duplicated tool name on the read side would silently collapse an entry here. + assert len(AGENT_TOOL_TO_ENUM) == len(AGENT_ENUM_TO_TOOL) + assert len(MCP_TAG_TO_TYPE_ENUM) == len(MCP_TYPE_ENUM_TO_TAG) + + +class TestRoundTrip: + """serialize -> normalize must be the identity on a ucode-native manifest.""" + + def test_full_manifest_round_trips(self): + manifest = _full_manifest() + assert normalize_managed_config(serialize_managed_config(manifest)) == manifest + + def test_minimal_manifest_round_trips(self): + manifest = _minimal_manifest() + assert normalize_managed_config(serialize_managed_config(manifest)) == manifest + + def test_every_known_agent_round_trips(self): + # Each agent's oneof variant must survive a round trip, including the flat-list agents and + # codex (which has no model list at all). + for tool in AGENT_TOOL_TO_ENUM: + model_config: dict = {"default_model": "system.ai.some-model"} + if tool == "claude": + model_config["models"] = {"default_opus_model": "system.ai.claude-opus-4-8"} + elif tool != "codex": + model_config["models"] = ["system.ai.some-model"] + manifest = { + "default_agent": tool, + "enabled_agents": {tool: {"model_config": model_config}}, + } + assert normalize_managed_config(serialize_managed_config(manifest)) == manifest, tool + + def test_every_mcp_type_round_trips(self): + for tag in MCP_TAG_TO_TYPE_ENUM: + manifest = {"mcp_servers": [{"name": "some-server", "type": tag}]} + assert normalize_managed_config(serialize_managed_config(manifest)) == manifest, tag + + +class TestSerialize: + def test_maps_tool_names_to_proto_enums(self): + payload = serialize_managed_config(_minimal_manifest()) + assert payload["default_agent"] == "CODING_AGENT_CLAUDE_CODE" + assert payload["enabled_agents"][0]["agent"] == "CODING_AGENT_CLAUDE_CODE" + + def test_claude_model_config_uses_family_slots(self): + payload = serialize_managed_config(_full_manifest()) + claude = next( + entry + for entry in payload["enabled_agents"] + if entry["agent"] == "CODING_AGENT_CLAUDE_CODE" + ) + variant = claude["config"]["model_config"] + assert set(variant) == {"claude"} + assert variant["claude"]["models"] == { + "default_opus_model": "system.ai.claude-opus-4-8", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + } + + def test_codex_model_config_has_no_model_list(self): + # CodexModelConfig carries only model_provider_service + default_model. + manifest = { + "default_agent": "codex", + "enabled_agents": { + "codex": { + "model_config": { + "default_model": "system.ai.gpt-5-6", + # Even if a caller passes a list, it must not be serialized. + "models": ["system.ai.gpt-5-6"], + } + } + }, + } + payload = serialize_managed_config(manifest) + variant = payload["enabled_agents"][0]["config"]["model_config"]["codex"] + assert "models" not in variant + assert variant["default_model"] == "system.ai.gpt-5-6" + + def test_flat_list_agents_use_repeated_models(self): + payload = serialize_managed_config(_full_manifest()) + opencode = next( + entry + for entry in payload["enabled_agents"] + if entry["agent"] == "CODING_AGENT_OPENCODE" + ) + variant = opencode["config"]["model_config"]["opencode"] + assert variant["models"] == ["system.ai.claude-opus-4-8", "system.ai.kimi-k2-6"] + + def test_model_provider_service_is_carried_through(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "model_provider_service": "main.default.anthropic-mps", + "default_model": "claude-sonnet-4-6", + } + } + }, + } + payload = serialize_managed_config(manifest) + variant = payload["enabled_agents"][0]["config"]["model_config"]["claude"] + assert variant["model_provider_service"] == "main.default.anthropic-mps" + + def test_mcp_types_map_to_proto_enums(self): + payload = serialize_managed_config(_full_manifest()) + assert payload["mcp_servers"] == [ + {"name": "system.ai.github", "type": "MCP_SERVER_TYPE_UC_SERVICE"}, + {"name": "genie-space-id", "type": "MCP_SERVER_TYPE_GENIE"}, + ] + + def test_tracing_becomes_a_table_object(self): + payload = serialize_managed_config(_full_manifest()) + assert payload["tracing"] == {"table": "main.default.ucode-traces"} + + def test_per_agent_tracing_override(self): + payload = serialize_managed_config(_full_manifest()) + claude = next( + entry + for entry in payload["enabled_agents"] + if entry["agent"] == "CODING_AGENT_CLAUDE_CODE" + ) + assert claude["config"]["tracing_config"] == {"table": "main.default.claude-traces"} + + def test_budget_tiers_keep_fractions(self): + # The server validates 0 <= spending_percentage <= 1, so these stay fractions. + payload = serialize_managed_config(_full_manifest()) + tiers = payload["budget_policy"]["tiers"] + assert [tier["spending_percentage"] for tier in tiers] == [0.8, 1.0] + assert tiers[1]["default_agent"] == "CODING_AGENT_OPENCODE" + + def test_unknown_agent_is_dropped(self): + payload = serialize_managed_config( + { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "m"}}, + "some-future-agent": {"model_config": {"default_model": "m"}}, + }, + } + ) + assert [entry["agent"] for entry in payload["enabled_agents"]] == [ + "CODING_AGENT_CLAUDE_CODE" + ] + + def test_unknown_mcp_type_is_dropped(self): + payload = serialize_managed_config( + {"mcp_servers": [{"name": "a", "type": "not-a-type"}, {"name": "b", "type": "sql"}]} + ) + assert payload["mcp_servers"] == [{"name": "b", "type": "MCP_SERVER_TYPE_DATABRICKS_SQL"}] + + def test_empty_manifest_serializes_to_empty_payload(self): + assert serialize_managed_config({}) == {} + + def test_output_only_fields_are_never_emitted(self): + # A caller (or a round-tripped GET) may carry server-owned fields; they must not be sent. + payload = serialize_managed_config( + { + **_minimal_manifest(), + "workspace_id": 12345, + "create_time": "2026-01-01T00:00:00Z", + "created_user_id": 42, + } + ) + assert "workspace_id" not in payload + assert "create_time" not in payload + assert "created_user_id" not in payload + + def test_name_is_carried_through_when_present(self): + payload = serialize_managed_config( + {**_minimal_manifest(), "name": "coding-agent-configs/abc"} + ) + assert payload["name"] == "coding-agent-configs/abc" + + def test_use_as_global_settings_false_is_preserved(self): + # `False` is meaningful (write to the user-level file), so it must not be dropped as falsy. + payload = serialize_managed_config( + { + "default_agent": "codex", + "enabled_agents": { + "codex": { + "use_as_global_settings": False, + "model_config": {"default_model": "m"}, + } + }, + } + ) + assert payload["enabled_agents"][0]["config"]["use_as_global_settings"] is False + + +class TestModelOptions: + def test_claude_only_sees_claude_models(self): + options = model_options_for_agent("claude", STATE) + assert options == [ + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5", + ] + + def test_gemini_only_sees_gemini_models(self): + assert model_options_for_agent("gemini", STATE) == ["system.ai.gemini-3-flash"] + + def test_codex_sees_gpt_and_oss(self): + assert model_options_for_agent("codex", STATE) == [ + "system.ai.gpt-5-6", + "system.ai.kimi-k2-6", + ] + + def test_multi_provider_agents_see_everything(self): + for tool in ("opencode", "pi", "copilot"): + options = model_options_for_agent(tool, STATE) + assert "system.ai.claude-opus-4-8" in options, tool + assert "system.ai.gpt-5-6" in options, tool + assert "system.ai.gemini-3-flash" in options, tool + assert "system.ai.kimi-k2-6" in options, tool + + def test_empty_state_yields_no_options(self): + assert model_options_for_agent("claude", {}) == [] + + def test_options_are_deduplicated(self): + # An id can land in two family buckets (e.g. an OSS model also listed under gpt). + state = {"codex_models": ["system.ai.kimi-k2-6"], "oss_models": ["system.ai.kimi-k2-6"]} + assert model_options_for_agent("codex", state) == ["system.ai.kimi-k2-6"] + + def test_unknown_agent_has_no_families(self): + assert model_families_for_agent("not-an-agent") == () + assert model_options_for_agent("not-an-agent", STATE) == [] + + def test_malformed_state_is_ignored(self): + state = {"claude_models": "not-a-dict", "codex_models": {"not": "a list"}} + assert model_options_for_agent("claude", state) == [] + assert model_options_for_agent("codex", state) == [] + + +class TestProviderServiceSupport: + def test_claude_supports_anthropic_and_bedrock(self): + assert supports_provider_service("claude", "anthropic") + assert supports_provider_service("claude", "amazon_bedrock") + + def test_codex_supports_openai(self): + assert supports_provider_service("codex", "openai") + + def test_claude_does_not_support_openai(self): + assert not supports_provider_service("claude", "openai") + + def test_other_agents_have_no_provider_support(self): + for tool in ("gemini", "opencode", "pi", "copilot"): + assert not supports_provider_service(tool, "anthropic"), tool + + +class TestClaudeSlots: + @pytest.mark.parametrize( + ("model", "expected"), + [ + ("system.ai.claude-opus-4-8", "opus"), + ("databricks-claude-sonnet-4-6", "sonnet"), + ("system.ai.claude-haiku-4-5", "haiku"), + ("system.ai.claude-fable-5", "fable"), + ("system.ai.gpt-5-6", None), + ("claude-without-family", None), + ], + ) + def test_family_detection(self, model, expected): + assert claude_family_for_model(model) == expected + + def test_groups_models_into_slots(self): + slots = claude_model_slots(["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-4-6"]) + assert slots == { + "default_opus_model": "system.ai.claude-opus-4-8", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + } + + def test_first_model_wins_within_a_family(self): + slots = claude_model_slots(["system.ai.claude-opus-4-8", "system.ai.claude-opus-4-7"]) + assert slots == {"default_opus_model": "system.ai.claude-opus-4-8"} + + def test_unidentifiable_models_are_skipped(self): + assert claude_model_slots(["system.ai.gpt-5-6"]) == {} + + def test_slots_serialize_into_the_claude_variant(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-4-8", + "models": claude_model_slots(["system.ai.claude-opus-4-8"]), + } + } + }, + } + payload = serialize_managed_config(manifest) + variant = payload["enabled_agents"][0]["config"]["model_config"]["claude"] + 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) == [] + + def test_minimal_manifest_is_valid(self): + assert validate_manifest(_minimal_manifest(), STATE) == [] + + def test_empty_manifest_is_valid(self): + # Nothing configured is not an error here; the wizard decides whether to publish it. + assert validate_manifest({}, STATE) == [] + + def test_default_agent_required_when_agents_present(self): + manifest = {"enabled_agents": {"claude": {"model_config": {"default_model": "m"}}}} + errors = validate_manifest(manifest) + assert any("default_agent is required" in e for e in errors) + + def test_default_agent_must_be_enabled(self): + manifest = { + "default_agent": "codex", + "enabled_agents": {"claude": {"model_config": {"default_model": "m"}}}, + } + errors = validate_manifest(manifest) + assert any("must appear in enabled_agents" in e for e in errors) + + def test_default_agent_needs_a_default_model(self): + manifest = { + "default_agent": "claude", + "enabled_agents": {"claude": {"use_as_global_settings": True}}, + } + errors = validate_manifest(manifest) + assert any("model_config.default_model" in e for e in errors) + + def test_unknown_agent_is_rejected(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}, + "not-an-agent": {}, + }, + } + errors = validate_manifest(manifest, STATE) + assert any("not a supported agent" in e for e in errors) + + def test_unknown_model_is_rejected(self): + manifest = { + "default_agent": "claude", + "enabled_agents": {"claude": {"model_config": {"default_model": "system.ai.nope"}}}, + } + 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", + "enabled_agents": {"claude": {"model_config": {"default_model": "anything"}}}, + } + assert validate_manifest(manifest) == [] + + def test_model_check_skipped_for_provider_service(self): + # MPS model ids come from the provider's catalog, not the UC inventory. + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "model_provider_service": "main.default.anthropic-mps", + "default_model": "claude-sonnet-5", + } + } + }, + } + assert validate_manifest(manifest, STATE) == [] + + def test_mcp_server_needs_a_name(self): + errors = validate_manifest({"mcp_servers": [{"type": "sql"}]}) + assert any("name is required" in e for e in errors) + + def test_mcp_server_needs_a_known_type(self): + errors = validate_manifest({"mcp_servers": [{"name": "a", "type": "bogus"}]}) + assert any("is not recognized" in e for e in errors) + + def test_empty_skill_name_is_rejected(self): + errors = validate_manifest({"skills": {"names": ["ok", ""]}}) + assert any("skills.names" in e for e in errors) + + def test_empty_tracing_table_is_rejected(self): + errors = validate_manifest({"tracing_table": ""}) + assert any("tracing_table" in e for e in errors) + + def test_budget_policy_needs_a_budget_id(self): + manifest = { + **_minimal_manifest(), + "budget_policy": { + "tiers": [ + { + "spending_percentage": 0.5, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-8", + } + ] + }, + } + errors = validate_manifest(manifest, STATE) + assert any("budget_policy.budget_id is required" in e for e in errors) + + @pytest.mark.parametrize("pct", [1.5, -0.1, 80]) + def test_tier_percentage_must_be_a_fraction(self, pct): + # 80 is the classic mistake: the spec doc writes percents, the API wants fractions. + manifest = { + **_minimal_manifest(), + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": pct, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-8", + } + ], + }, + } + errors = validate_manifest(manifest, STATE) + assert any("fraction" in e for e in errors), errors + + def test_tier_percentages_must_be_unique(self): + tier = { + "spending_percentage": 0.5, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-8", + } + manifest = { + **_minimal_manifest(), + "budget_policy": {"budget_id": "b", "tiers": [tier, dict(tier)]}, + } + errors = validate_manifest(manifest, STATE) + assert any("must be unique" in e for e in errors) + + def test_tier_agent_must_be_enabled(self): + manifest = { + **_minimal_manifest(), + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.5, + "default_agent": "opencode", + "default_model": "system.ai.kimi-k2-6", + } + ], + }, + } + errors = validate_manifest(manifest, STATE) + assert any("must appear in enabled_agents" in e for e in errors) + + def test_tier_needs_a_default_model(self): + manifest = { + **_minimal_manifest(), + "budget_policy": { + "budget_id": "b", + "tiers": [{"spending_percentage": 0.5, "default_agent": "claude"}], + }, + } + 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) + + def test_errors_accumulate(self): + manifest = { + "default_agent": "codex", + "enabled_agents": {"claude": {}}, + "mcp_servers": [{"type": "bogus"}], + } + assert len(validate_manifest(manifest, STATE)) >= 3 + + +class TestPersistence: + def test_round_trips_through_disk(self, tmp_path, monkeypatch): + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr( + managed_setup_mod, "MANAGED_SETTINGS_PATH", tmp_path / "managed-settings.json" + ) + manifest = _full_manifest() + save_managed_settings(WORKSPACE, manifest) + assert load_managed_settings(WORKSPACE) == manifest + + def test_stores_the_workspace_alongside_the_manifest(self, tmp_path, monkeypatch): + path = tmp_path / "managed-settings.json" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(managed_setup_mod, "MANAGED_SETTINGS_PATH", path) + save_managed_settings(WORKSPACE, _minimal_manifest()) + assert json.loads(path.read_text())["workspace"] == WORKSPACE + assert managed_settings_workspace() == WORKSPACE + + def test_load_is_workspace_scoped(self, tmp_path, monkeypatch): + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr( + managed_setup_mod, "MANAGED_SETTINGS_PATH", tmp_path / "managed-settings.json" + ) + save_managed_settings(WORKSPACE, _minimal_manifest()) + # A manifest authored for another workspace must not be published to this one. + assert load_managed_settings("https://other.example.com") is None + + def test_load_without_a_workspace_returns_whatever_is_on_disk(self, tmp_path, monkeypatch): + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr( + managed_setup_mod, "MANAGED_SETTINGS_PATH", tmp_path / "managed-settings.json" + ) + save_managed_settings(WORKSPACE, _minimal_manifest()) + assert load_managed_settings() == _minimal_manifest() + + def test_load_returns_none_when_absent(self, tmp_path, monkeypatch): + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(managed_setup_mod, "MANAGED_SETTINGS_PATH", tmp_path / "missing.json") + assert load_managed_settings(WORKSPACE) is None + assert managed_settings_workspace() is None + + def test_dry_run_writes_nothing(self, tmp_path, monkeypatch): + path = tmp_path / "managed-settings.json" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(managed_setup_mod, "MANAGED_SETTINGS_PATH", path) + monkeypatch.setattr(config_io_mod, "is_dry_run", lambda: True) + save_managed_settings(WORKSPACE, _minimal_manifest()) + assert not path.exists() + + def test_corrupt_file_reads_as_absent(self, tmp_path, monkeypatch): + path = tmp_path / "managed-settings.json" + path.write_text("{not json", encoding="utf-8") + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(managed_setup_mod, "MANAGED_SETTINGS_PATH", path) + assert load_managed_settings(WORKSPACE) is None + + def test_serialized_payload_is_json_encodable(self, tmp_path, monkeypatch): + # `ucode apply` POSTs this, so it must survive json.dumps with no custom encoder. + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr( + managed_setup_mod, "MANAGED_SETTINGS_PATH", tmp_path / "managed-settings.json" + ) + save_managed_settings(WORKSPACE, _full_manifest()) + manifest = load_managed_settings(WORKSPACE) + assert manifest is not None + assert json.loads(json.dumps(serialize_managed_config(manifest))) + + def test_settings_file_is_user_only(self, tmp_path, monkeypatch): + path = tmp_path / "managed-settings.json" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(managed_setup_mod, "MANAGED_SETTINGS_PATH", path) + save_managed_settings(WORKSPACE, _minimal_manifest()) + assert stat.S_IMODE(path.stat().st_mode) & 0o077 == 0 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"