From 270690886931817aad7197a74aeba61760bd03c7 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Tue, 4 Aug 2026 16:18:11 +0000 Subject: [PATCH 1/2] Add managed-config serializer, validation, and per-agent model catalogs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Admin-write foundation for the managed CodingAgentConfig, mirroring the developer-read side added in the parent commit. No CLI wiring and no network calls: the interactive `ucode setup` wizard and `ucode apply` build on this in follow-up changes. `managed_setup.py`: - `serialize_managed_config` — ucode-native manifest -> proto-JSON `CodingAgentConfig`, the exact inverse of `normalize_managed_config`. Builds each agent's `AgentModelConfig` oneof variant, which the server rejects when it doesn't match its agent: claude gets `ClaudeDefaultModels` family slots, codex gets no model list at all, and opencode/pi/gemini/copilot get a flat list. Output-only fields (workspace_id, timestamps, user ids) are never emitted. - `validate_manifest` — client-side pre-flight mirroring the handler's `validateStoredConfig`, so an admin sees problems before the round-trip rather than as an opaque INVALID_PARAMETER_VALUE. Also cross-checks configured models against the workspace's discovered inventory, skipped for agents routing through a Model Provider Service (those ids come from the provider's catalog). - Per-agent model catalogs: claude -> Claude only, gemini -> Gemini only, codex -> GPT + OSS, opencode/pi/copilot -> everything discovered. - `~/.ucode/managed-settings.json` persistence (0600, workspace-scoped), distinct from `managed-state.json`: this is the authored manifest, that is the pulled one. Note on units: `spending_percentage` is a fraction in [0, 1], not a percent — the server validates that range. `validate_manifest` rejects e.g. 80 with a message naming the expected form, since the spec doc's prose says "80%". `managed_config.py`: makes the two enum maps public so the write side inverts them instead of restating them — a new agent or MCP type is declared once and both directions pick it up. Mechanical rename otherwise. Tests: 77 cases. The load-bearing one is the round-trip property (serialize -> normalize == identity), asserted over the full manifest, the minimal manifest, every known agent, and every MCP type — it pins read and write together so they cannot drift as the proto grows. Co-authored-by: Isaac --- src/ucode/managed_config.py | 16 +- src/ucode/managed_setup.py | 561 ++++++++++++++++++++++++++++++ tests/test_managed_setup.py | 674 ++++++++++++++++++++++++++++++++++++ 3 files changed, 1244 insertions(+), 7 deletions(-) create mode 100644 src/ucode/managed_setup.py create mode 100644 tests/test_managed_setup.py diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index c128d3c..0d6d7ed 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..d2a1638 --- /dev/null +++ b/src/ucode/managed_setup.py @@ -0,0 +1,561 @@ +"""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, + 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. +_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_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". + """ + 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"): + 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 _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") + 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." + ) + if not tier.get("default_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/tests/test_managed_setup.py b/tests/test_managed_setup.py new file mode 100644 index 0000000..3570a7f --- /dev/null +++ b/tests/test_managed_setup.py @@ -0,0 +1,674 @@ +"""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 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_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_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 From 868ccce408f7f45fc754a08db1a9e4a023709809 Mon Sep 17 00:00:00 2001 From: Tien Le <44335153+tt-le@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:12:59 -0400 Subject: [PATCH 2/2] Add `ucode setup`: interactive managed-config authoring (#268) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add `ucode setup`: interactive managed-config authoring The admin-facing wizard on top of the serializer added in the parent commit. Walks a workspace 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 is a separate `ucode apply` (next change) so an admin can review the file first — and `ucode setup show` prints both the summary and the exact proto-JSON that apply will POST. `managed_wizard.py` is the interaction layer only; catalogs, validation, and serialization stay in `managed_setup.py`. Sub-flows an admin already knows are delegated to the existing commands (`configure_tracing_command`, `configure_mcp_command`, `configure_skills_mcp_command`) and their results read back out of `state.json`, so there is still exactly one picker per concern. Claude Code is prompted per model family, not as a flat list. `ClaudeModelConfig.models` is a `ClaudeDefaultModels` message — four optional slots (`default_opus_model` / `default_sonnet_model` / `default_haiku_model` / `default_fable_model`), since Claude Code addresses models by family alias rather than from a list, and a flat multi-select cannot express that. Each family is skippable (the slots are `optional`, and an unset one falls back to `default_model`); 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. With one slot filled there is nothing to choose, so it reports what it inferred rather than asking a single-option question. That required looking past `state["claude_models"]`, which is `{family: newest}` because the launch path pins one model per alias — the workspace tested against has 12 Claude models and the bucketed view shows 4. `discover_claude_models_unbucketed` returns the full `system.ai.claude-*` listing, `claude_family_candidates` groups it by family, and `_known_models` reads the same set so validation cannot reject a version the prompts offered. `discover_model_services` is untouched — every agent's config writer depends on the shape of `claude_models`. Provider-service agents get a real picker too. `list_model_provider_services` already returns each service's `config.targets` (the provider-side model ids); the wizard now selects the whole service dict and offers those, instead of an empty text prompt that required knowing a provider-side id by heart. Two cases legitimately have nothing to enumerate — `allow_all_targets`, and a relayed Anthropic subscription service that routes by canonical name — and those fall back to free text, but say why and give an example. Gemini and Copilot get one model rather than a multi-select: their writers take a single model and set one env var, so extra picks were silently discarded at apply time. This is a wizard restriction, not a serializer one — both protos do declare `repeated string models`, so `_FLAT_MODEL_LIST_AGENTS` is unchanged and a hand-written `--from-file` manifest still serializes a list. OpenCode and Pi keep their multi-select; they really do show a model picker. Nothing is pre-checked in the model multi-selects. `options[0]` is whatever discovery sorted first, not a recommendation — for Pi a Claude model, for Codex the oldest GPT — so pre-checking it made "hit Enter" produce an arbitrary config and read as an endorsement. It only existed to make Enter work at all, which stopped mattering once every agent required a model: an empty selection is re-prompted (`_require_selection` / `_require_multi_selection` / `_require_text`), which also makes `_prompt_models_for_agent` total. Without that, picking a multi-model agent and hitting Enter produced an agent with no `model_config`, and if it was the `default_agent` the run died at the *end* — after tracing, MCP, skills, and budget had all been answered. Ctrl-C still aborts. Budget tiers offer only the models their agent was configured with, via `configured_models_for_agent`. The picker previously read the workspace catalog and called `load_state()`, so it could not have seen this run's choices even in principle — configuring Pi with one Kimi model still offered every family. Picking one of the extras was not harmless: `validateStoredConfig` checks only that a tier's `default_agent` is in `enabled_agents`, never that the agent has the model, so the tier would activate and hand a developer a model their agent was never given. `validate_manifest` gains that missing rule, which is the durable half — it guards `--from-file` and hand-written manifests too. Tiers are prompted in percent and stored as fractions (`prompt_for_percentage` converts), keeping the spec-vs-API units mismatch in one place. Model ids stay bare (`system.ai.claude-opus-4-8`), not provider-prefixed: each agent's writer adds whatever its config format needs (see `opencode._resolve_model_selector`), so the manifest stays agent-neutral and matches the spec's examples. Two readback details worth noting: - Tracing is stored as `tracing.uc_destination` but the managed config calls it `tracing.table`; an enabled-but-not-UC-backed experiment has no table, so it is omitted rather than published empty. - `state.json` records each MCP server's resolved URL but not its type, while the manifest needs `{name, type}`. `_mcp_type_for_url` maps the URL shape back to a type and skips anything unrecognized instead of guessing. `databricks.py`: `is_workspace_admin` (SCIM `Me` group membership) and `list_workspace_budgets` (workspace-scoped, so no account auth is needed — ucode never creates budgets, an admin picks an existing one). The admin gate errors on a definite non-admin and warns-then-continues when the check itself fails, since the API enforces the same rule server-side. `ui.py`: `preselected`/`prompt` on `prompt_for_tools`, plus `prompt_for_multi_selection`, `prompt_for_text`, and `prompt_for_percentage`. Reuses the existing `prompt_yes_no_default` rather than adding a second defaulting yes/no. README gains a "Managed config for a workspace (admins)" section, `setup` rows in Other Commands, and `~/.ucode/managed-settings.json` in Managed Local Files. It states plainly that the wizard leaves your agent configs alone *except* when you accept the tracing / MCP / skills sub-steps, which do configure this machine — the one thing about the flow that isn't obvious. Verified interactively against eng-ml-inference.staging: six opus versions and four sonnet offered, a skipped family absent from `models`, and the result round-trips through `normalize_managed_config` to the proto shape. Tests: 116 cases. Several are mutation-checked — reverting the claude flat multi-select fails eight, reverting the budget-tier picker fails `test_offers_only_the_models_the_agent_was_configured_with`, and dropping the tier-model validation rule fails `test_tier_model_must_be_one_the_agent_has`. Note `typer.Exit` subclasses RuntimeError, so raising it inside the command's try block made a successful run print "ERROR 0"; fixed by exiting after the handler, with `test_successful_setup_exits_zero` failing if that regresses. Known gap: `discover_model_services` and `discover_claude_models_unbucketed` both page `/api/2.1/unity-catalog/model-services`, so a run that configures Claude lists it twice. `_claude_candidates` caches on `state["all_claude_models"]`, which bounds it to once per run, but removing the second listing means widening `discover_model_services`' 5-tuple return — asserted in eight tests and mocked in sixteen — so it is left for a follow-up. Co-authored-by: Isaac * setup: don't require the databricks CLI to build the Claude catalog Two CI failures, both passing locally for environment reasons. `_claude_candidates` guarded its catalog fetch with `except RuntimeError`, but `get_databricks_token` shells out to `databricks auth token` — so a machine without the CLI on PATH raises `FileNotFoundError`, which is an `OSError` and sailed straight through. This is a real bug, not just a test artifact: any user without the Databricks CLI installed hit an uncaught traceback mid-wizard instead of degrading to the bucketed per-family picks. Now catches `OSError` too. `test_claude_candidates_survive_a_missing_databricks_cli` covers it, and `test_every_agent_always_gets_a_default_model` — which reached the network on its claude pass — now stubs the fetch rather than depending on the host. `test_setup_help_lists_from_file` grepped `--from-file` out of rendered `--help` output. Rich ellipsizes option names to fit the terminal, so below roughly 40 columns it renders `--fro…`; hosted runners report no width and get the narrow fallback. It now asserts on the declared Click option, which no amount of wrapping can change, and still checks the help renders successfully. Verified by reverting each fix in turn: the first fails with the same `FileNotFoundError: 'databricks'` CI reported, the second with the same missing `--from-file` assertion at `COLUMNS=30`. Co-authored-by: Isaac * setup: box the summary, filter long pickers, and page the catalog once Review feedback on #268. The configuration summary is boxed. It is the one block an admin is meant to read as a whole and check against what they intended, and it arrives after a long flow of prompts, where loose key/value lines blend into everything printed before them. `ui.print_panel` boxes a body (as opposed to `print_section`, which boxes a bare title) and `ui.kv_line` returns a `print_kv`-styled line for collecting into one. Every picker in the flow now filters as you type. On the workspace this was built against, the model lists run to 16 GPT entries and 12 Claude ones, which is more than is comfortable to arrow through. questionary supports this directly, with one constraint worth naming: it refuses `use_search_filter` together with `use_jk_keys` (j and k are search characters), so search costs j/k navigation. Arrow keys still work, and it is opt-in via a `searchable` flag so short pickers — "Databricks Hosted vs External Models" — keep j/k. The existing-config warning now says what actually happens. It claimed "there is no partial update yet", which was both vague and, once `ucode apply` PATCHes, no longer true. Publishing is a full replace of the workspace's config, so the warning says that plainly and names the agents already published, which is what the admin stands to lose. Field masks and PATCH stay out of it — an admin doesn't need to reason about the transport. The existing-config warning also says what it is warning about. It named the enabled agents in a parenthetical — "already exists for this workspace (claude)" — which read as though the workspace itself were called claude, and listing the contents implied the config was per-agent when there is exactly one per workspace covering agents, MCP servers, skills, tracing, and the budget policy. It now says that scope plainly and itemizes nothing: an inventory doesn't change what the admin should do, and `ucode setup show` prints the real thing for anyone comparing. The duplicate "Fetching available models" spinner is gone. `discover_model_services` and `discover_claude_models_unbucketed` both page the whole metastore catalog, so a run that configured Claude walked it twice. Rather than widen `discover_model_services`' 5-tuple return (asserted in 8 tests, mocked in 19, which is what made this look expensive last time), the memo goes one layer down in `list_model_services`, where both callers already converge — so no signature changes and no stub churn. Cached per process and per workspace; failures are never cached, so a transient error still retries; `use_cache=False` forces a fresh walk. `conftest` clears it between tests so a cached listing can't leak into a test that stubs the endpoint. The Model Provider Service listing gets the same treatment, for the same reason: it is workspace-wide and filtered per agent afterwards, so the wizard re-listed it once per MPS-capable agent. One call now serves them all. Cached entries are copied on both store and read, since the wizard treats the list as its own. Its spinner is now conditional rather than removed. The cold listing takes ~1.4s, so silence there would look like a hang; the cached ones are instant, and spinning "Checking for model provider services for ..." once per agent was what made the wizard look like it re-listed them every time. The message no longer names an agent either, since one lookup covers all of them. Measured against eng-ml-inference.staging: 12 HTTP pages for the model catalog across both callers, down from 24; and one MPS listing for claude + codex, down from two. Tests: +15. Mutation-verified — dropping either cache read fails a cache test, dropping `searchable` fails the picker test, un-boxing the summary fails the panel test, and re-adding the config inventory fails the warning test. The MPS aliasing guard is split in two on purpose: the store-side and read-side copies each need their own case, since a single test passed with either one still present. Two more redundant lines removed while in here. The "Next steps" block no longer suggests `ucode configure --dry-run`: the manifest describes what *developers* should get, while that command previews *this machine's* agent configs, so it implied a local test it does not perform. Publishing is the only next step there is today. And the wizard's own `ensure_databricks_auth` is now quiet, because `configure_shared_state` authenticates a moment later and prints the same "auth already available" success — the wizard's call still has to run first, since the admin gate and the existing-config check both need a token before discovery. Deferred, with Aarushi's agreement: a root `ucode --dry-run` that launches the default agent against the authored config, so an admin can try it locally before publishing. That needs both the flag and bare-`ucode` launch behavior, neither of which exists yet. Co-authored-by: Isaac * setup: show the default a prompt will actually use `prompt_for_text`'s default hint was invisible for exactly the values that have one. Rich parses `[coding-agents-tiered-routing]` as a style tag and renders nothing for it, so the budget policy-name prompt read `Policy name ›` with no sign that enter would pick a name. Numeric defaults like `[80]` are not valid tag names and survived, which is why this held up wherever it was checked. The hint is now escaped whole, brackets included, and says what enter does: `Policy name [coding-agents-tiered-routing] (enter to accept)`. Bracketed text on its own reads as a format example as easily as a value that will be used, so it invited retyping what enter already picks. Also fixes the skills-location prompt, the other call site with a word-like default. `prompt_for_percentage` gets the same formatting so the two cannot drift. No caller passes it a default: tier thresholds deliberately have none, since a threshold decides when developers get downgraded and should be typed rather than accepted by accident. That is now recorded in the docstring. Tests: +6, rendering through a real Rich console rather than asserting on the raw markup string — the latter is what let this ship, since `[tiered]` is present in the markup and absent from the output. Mutation-verified: dropping either `escape` fails three cases, and the pre-fix assertions pass against the bug. Co-authored-by: Isaac * Stop tracking .isaac/, a local tool's state dir `.isaac/config.json` was committed by accident in the review-feedback commit. It holds one field, `sync_reminder_last_shown` — per-machine state for isaac's own reminder, not project config — so every developer running isaac would see it as a spurious modification and could commit their own date over someone else's. Untracked and added to .gitignore so it stays local. The file itself is left on disk; only the tracking is removed. Co-authored-by: Isaac * setup: don't lose bracketed values, or crash on closed stdin Two bugs found reviewing #268. `kv_line` interpolated its value into Rich markup unescaped, and Rich reads bracketed text as a style tag and renders nothing for it. A budget policy named `[prod] tiered routing` displayed as ` tiered routing` in the configuration summary — the one block an admin reads to confirm what they are about to publish workspace-wide, so a silently altered value defeats its whole purpose. The saved manifest was always correct; this was display only. Values reaching it include admin-typed free text: the policy name, skills locations, and the tracing table. `prompt_for_percentage` re-raised `EOFError` when it had no default to fall back on, and nothing above it catches that — the setup command handles `RuntimeError` and `KeyboardInterrupt` only. Pressing Ctrl-D at `Tier 1: activates at what percent of budget?` (the tier prompt passes no default, deliberately) printed a traceback instead of exiting cleanly. It now raises `KeyboardInterrupt`, which the existing handler turns into "Interrupted." and exit 130 — the same thing Ctrl-C does, which is what closed stdin means here. A default, when one is passed, still short-circuits first: `prompt_for_text` and `prompt_yes_no_default` already treat EOF as "take the default", and this was the odd one out only because it can have no default. Tests: +3. Mutation-verified: un-escaping `kv_line` fails the summary test, and re-raising the bare `EOFError` fails the abort test. Co-authored-by: Isaac * setup: abort dismissed prompts instead of re-asking forever Review feedback on #268. Two infinite loops, both confirmed by driving the helpers directly. `_require_selection` re-asked on a None answer, and its docstring justified that by claiming questionary raises KeyboardInterrupt on Ctrl-C before returning. It does not: `Question.ask` catches KeyboardInterrupt internally and returns None (questionary 2.1.1, question.py), so Ctrl-C and an empty submission are the same value and nothing propagates for the caller to distinguish. Ctrl-C therefore printed "Please choose one of the options." and re-prompted, forever. A None is now an abort, which `_require_multi_selection` right below it already did. `_require_text` had the same shape on closed stdin: `prompt_for_text` returns its default on EOF, the default is None, and the loop re-asked an exhausted stream — so a piped or CI run spun printing "Please enter a model id." It now passes `required=True`, a new `prompt_for_text` flag that raises KeyboardInterrupt on EOF instead of returning None, matching what `prompt_for_percentage` does when it has no default to fall back on. Reachable whenever model discovery finds nothing, which is exactly when a run is most likely to be scripted. Also documents the `all_claude_models` invariant that made the two Claude candidate paths safe for non-obvious reasons: the listing path widens the candidates and caches them, while the fallback narrows them to what `_known_models` already covers. Widening the fallback without caching would surface as a rejection at the end of the flow rather than at the prompt that offered the model, so both paths are now pinned by tests that validate everything they offer. Tests: +5. `test_empty_single_select_is_re_prompted` asserted the old looping behavior and is replaced. Co-authored-by: Isaac --- .gitignore | 1 + README.md | 38 + src/ucode/cli.py | 52 ++ src/ucode/databricks.py | 162 +++- src/ucode/managed_setup.py | 83 +- src/ucode/managed_wizard.py | 900 ++++++++++++++++++++++ src/ucode/ui.py | 178 ++++- tests/conftest.py | 4 + tests/test_databricks.py | 165 ++++ tests/test_managed_setup.py | 201 +++++ tests/test_managed_wizard.py | 1402 ++++++++++++++++++++++++++++++++++ tests/test_ui.py | 71 ++ 12 files changed, 3241 insertions(+), 16 deletions(-) create mode 100644 src/ucode/managed_wizard.py create mode 100644 tests/test_managed_wizard.py diff --git a/.gitignore b/.gitignore index 4d1e17c..007899c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ build/ dist/ .venv/ .DS_Store +.isaac/ diff --git a/README.md b/README.md index 0352757..70a343b 100644 --- a/README.md +++ b/README.md @@ -163,6 +163,40 @@ ucode configure skills --location main.default,ml.prod --mcp Each run prints the registered server, its URL, the configured agents, and its tools, and reminds you to run `ucode ` (existing agent sessions need a restart before the MCP tools load). +### Managed config for a workspace (admins) + +```bash +ucode setup +``` + +Author the coding config your developers pick up automatically, instead of asking each of them to +run `ucode configure` by hand. Restricted to workspace admins. + +The flow walks through the agents to enable and which one bare `ucode` launches, then per agent: +Databricks-hosted models or an external Model Provider Service, the models to expose, and whether +the config applies machine-wide or per user. Claude Code is asked one model per family +(opus/sonnet/haiku/fable), since Claude Code selects models by family alias; any family can be +skipped. It then offers tracing, managed MCP servers, skills, and a spend-based budget policy that +switches the default agent and model as the workspace burns through a budget. + +The result is written to `~/.ucode/managed-settings.json`, which `ucode apply` publishes to the +workspace. Your own agent configs are left alone, with one exception: answering yes to tracing, MCP +servers, or skills runs the matching `ucode configure` step, which does configure this machine. + +```bash +# Review the manifest and the exact payload `ucode apply` would publish. +ucode setup show + +# Walk the flow without writing anything. +ucode setup --dry-run + +# Skip the prompts and load a hand-written config instead (validated before saving). +ucode setup --from-file ./managed-settings.json +``` + +Publishing replaces the workspace's config outright — there is no partial update yet, so anything +skipped in a re-run is dropped. + --- ## Other Commands @@ -185,6 +219,9 @@ you to run `ucode ` (existing agent sessions need a restart before the MC | `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | | `ucode configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | | `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | +| `ucode setup` | Author the workspace's managed coding config (workspace admins only) | +| `ucode setup show` | Print the authored config and the payload `ucode apply` would publish | +| `ucode setup --from-file ` | Load a hand-written managed config instead of running the prompts | ## Managed Local Files @@ -199,6 +236,7 @@ you to run `ucode ` (existing agent sessions need a restart before the MC | `~/.copilot/.env` | GitHub Copilot CLI | | `~/.pi/agent/models.json` | Pi | | `~/.cursor/mcp.json` | Cursor Agent (MCP servers only) | +| `~/.ucode/managed-settings.json` | The managed config authored by `ucode setup` (admins) | Existing files are backed up before being overwritten. `ucode revert` restores backups. diff --git a/src/ucode/cli.py b/src/ucode/cli.py index f541079..7873d4c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -68,6 +68,7 @@ managed_unservable_models, resolve_state, ) +from ucode.managed_wizard import setup_command, show_command from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -873,6 +874,10 @@ def revert() -> int: app.add_typer(configure_app, name="configure", help="Configure workspace and tool settings.") mcp_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") +setup_app = typer.Typer(add_completion=False, no_args_is_help=False) +app.add_typer( + setup_app, name="setup", help="Author the workspace's managed coding config (admins only)." +) @mcp_app.command("web-search") @@ -1974,6 +1979,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 b9e48bc..126533a 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -331,6 +331,76 @@ def _http_get_bytes(url: str, token: str, *, timeout: int = 10) -> tuple[bytes | return None, f"network error: {exc.reason}" +# Workspace group whose members are workspace admins. `ucode setup` / `ucode apply` are restricted +# to this group because the coding-agent-config CRUD API enforces the same check server-side. +WORKSPACE_ADMIN_GROUP = "admins" + + +def is_workspace_admin(workspace: str, token: str) -> bool | None: + """Whether the caller is a workspace admin, via their SCIM `Me` group membership. + + Returns True/False, or None when the check itself could not be made (SCIM unreachable or a + malformed response). Callers should treat None as "unknown" and proceed optimistically rather + than blocking: the API enforces the same check server-side, so a false negative here would + needlessly stop a legitimate admin, while a false positive just surfaces the server's + PERMISSION_DENIED later. + """ + hostname = workspace_hostname(workspace) + payload, _ = _http_get_json(f"https://{hostname}/api/2.0/preview/scim/v2/Me", token) + if not isinstance(payload, dict): + return None + groups = payload.get("groups") + if not isinstance(groups, list): + # A well-formed `Me` for a user in no groups omits `groups` entirely, so this is a + # definitive "not an admin" rather than a failed check. + return False + return any( + isinstance(group, dict) and group.get("display") == WORKSPACE_ADMIN_GROUP + for group in groups + ) + + +# Workspace-scoped budget listing. Account-level budget APIs need account auth, which ucode does not +# have; this endpoint resolves the workspace server-side from the caller's token. +_WORKSPACE_BUDGETS_API_PATH = "/api/ai-gateway/v2/workspace-metrics/budgets" + + +def list_workspace_budgets(workspace: str, token: str) -> tuple[list[dict], str | None]: + """List the AI Gateway budgets that apply to this workspace. + + Returns ``(budgets, reason)`` where each budget is ``{"id": ..., "display_name": ...}``. + ``reason`` is None on success, otherwise it explains why the list is empty. ucode never creates + budgets — an admin picks an existing one to attach a spend-routing policy to. + """ + hostname = workspace_hostname(workspace) + url = f"https://{hostname}{_WORKSPACE_BUDGETS_API_PATH}" + payload, reason = _http_get_json(url, token, timeout=30) + if reason is not None: + return [], reason + if not isinstance(payload, dict): + return [], "workspace budget listing returned an unexpected response shape" + raw = payload.get("workspace_ai_gateway_budgets") + if not isinstance(raw, list): + return [], "workspace budget listing returned no budgets" + budgets: list[dict] = [] + for entry in raw: + if not isinstance(entry, dict): + continue + budget_id = entry.get("budget_configuration_id") + if not isinstance(budget_id, str) or not budget_id: + continue + display_name = entry.get("display_name") + budgets.append( + { + "id": budget_id, + "display_name": display_name if isinstance(display_name, str) else "", + } + ) + if not budgets: + return [], "workspace budget listing returned no budgets" + return budgets, None + + def get_current_user_name(workspace: str, token: str) -> str | None: """Return the current user's login (email) via SCIM `Me`, or None on failure. @@ -843,12 +913,20 @@ def run_databricks_login(workspace: str, profile: str | None = None) -> None: print_success("Databricks authentication complete") -def ensure_databricks_auth(workspace: str, profile: str | None = None) -> None: - """Check auth and login only if needed (used by launch path).""" +def ensure_databricks_auth( + workspace: str, profile: str | None = None, *, quiet: bool = False +) -> None: + """Check auth and login only if needed (used by launch path). + + ``quiet`` suppresses the "already available" line for a caller that only needs a token before + some later step re-authenticates and reports it — otherwise the same success prints twice. A + login that actually runs is never silent. + """ with spinner("Checking Databricks auth..."): auth_is_valid = has_valid_databricks_auth(workspace, profile) if auth_is_valid: - print_success(f"Databricks auth already available for {workspace}") + if not quiet: + print_success(f"Databricks auth already available for {workspace}") return run_databricks_login(workspace, profile) @@ -1293,12 +1371,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. @@ -1307,7 +1417,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 @@ -1340,10 +1458,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]: @@ -1523,7 +1657,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. @@ -1539,7 +1673,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 @@ -1576,6 +1728,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_setup.py b/src/ucode/managed_setup.py index d2a1638..2ef8373 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -31,6 +31,7 @@ import ucode.config_io as config_io from ucode.databricks import ( ANTHROPIC_FAMILIES, + model_version_sort_key, tool_supports_provider_type, ) from ucode.managed_config import ( @@ -61,8 +62,9 @@ # (`ClaudeDefaultModels`), and Codex has no model list at all — it selects exactly one model. _FLAT_MODEL_LIST_AGENTS = frozenset({"opencode", "pi", "gemini", "copilot"}) -# Claude family slot names in `ClaudeDefaultModels`, keyed by ucode's family name. -_CLAUDE_SLOT_FOR_FAMILY: dict[str, str] = { +# Claude family slot names in `ClaudeDefaultModels`, keyed by ucode's family name. Public because +# the wizard prompts one slot at a time. +CLAUDE_SLOT_FOR_FAMILY: dict[str, str] = { family: f"default_{family}_model" for family in ANTHROPIC_FAMILIES } @@ -144,6 +146,39 @@ def claude_family_for_model(model: str) -> str | None: return next((family for family in ANTHROPIC_FAMILIES if f"claude-{family}-" in lowered), None) +def claude_family_candidates( + all_claude_models: list[str], state: dict | None = None +) -> dict[str, list[str]]: + """Group Claude model ids by family, newest first. + + ``state["claude_models"]`` holds only one id per family — the newest, chosen by + ``discover_model_services`` for the launch path, which pins exactly one model per family alias. + An admin authoring a managed config needs the alternatives too: pinning ``default_opus_model`` + to a known-good ``claude-opus-4-8`` rather than whatever happens to be newest is a normal thing + to want, and impossible if only the newest is offered. + + ``all_claude_models`` is the unbucketed listing (see + :func:`ucode.databricks.discover_claude_family_candidates`). When it is empty, falls back to the + per-family picks already in ``state`` so the per-slot prompts still work — with one candidate + each. Families with no models are omitted. + """ + models = list(all_claude_models) + if not models and state: + claude_models = state.get("claude_models") + if isinstance(claude_models, dict): + models = [m for m in claude_models.values() if isinstance(m, str) and m] + + candidates: dict[str, list[str]] = {} + for model in models: + family = claude_family_for_model(model) + if family: + candidates.setdefault(family, []).append(model) + for family, found in candidates.items(): + # model_version_sort_key negates version components, so plain ascending is newest-first. + candidates[family] = sorted(set(found), key=model_version_sort_key) + return candidates + + def claude_model_slots(models: list[str]) -> dict[str, str]: """Group picked Claude model ids into ``ClaudeDefaultModels`` slots. @@ -157,7 +192,7 @@ def claude_model_slots(models: list[str]) -> dict[str, str]: family = claude_family_for_model(model) if family is None: continue - slot = _CLAUDE_SLOT_FOR_FAMILY[family] + slot = CLAUDE_SLOT_FOR_FAMILY[family] slots.setdefault(slot, model) return slots @@ -342,12 +377,16 @@ def _known_models(state: dict) -> set[str]: Empty when discovery found nothing (or no state was passed), which callers treat as "can't check" rather than "nothing is valid". + + ``claude_models`` holds only the newest id per family (the launch path pins one model per family + alias), so on its own it would reject the older versions the per-family prompts legitimately + offer. ``all_claude_models`` carries the full listing when the caller has it. """ known: set[str] = set() claude_models = state.get("claude_models") if isinstance(claude_models, dict): known.update(m for m in claude_models.values() if isinstance(m, str) and m) - for key in ("codex_models", "gemini_models", "oss_models"): + for key in ("codex_models", "gemini_models", "oss_models", "all_claude_models"): models = state.get(key) if isinstance(models, list): known.update(m for m in models if isinstance(m, str) and m) @@ -470,6 +509,28 @@ def validate_manifest(manifest: dict, state: dict | None = None) -> list[str]: return errors +def _agent_model_ids(agent_config: dict) -> set[str]: + """Every model id an agent is configured with — its list plus its default. + + Claude's ``models`` is a family-slot dict and the others' a flat list; codex has no list at all, + only ``default_model``. Returns an empty set when nothing is configured, which callers treat as + "can't check" rather than "nothing is allowed". + """ + model_config = agent_config.get("model_config") + if not isinstance(model_config, dict): + return set() + ids: set[str] = set() + raw = model_config.get("models") + if isinstance(raw, dict): + ids.update(v for v in raw.values() if isinstance(v, str) and v) + elif isinstance(raw, list): + ids.update(m for m in raw if isinstance(m, str) and m) + default_model = model_config.get("default_model") + if isinstance(default_model, str) and default_model: + ids.add(default_model) + return ids + + def _validate_budget_policy(budget_policy: dict, enabled_agents: dict[str, dict]) -> list[str]: """Validate a ``budget_policy`` against the agents the manifest enables.""" errors: list[str] = [] @@ -493,6 +554,7 @@ def _validate_budget_policy(budget_policy: dict, enabled_agents: dict[str, dict] else: percentages.append(float(pct)) tier_agent = tier.get("default_agent") + tier_model = tier.get("default_model") if not tier_agent: errors.append(f"budget_policy.tiers[{index}]: default_agent is required.") elif tier_agent not in enabled_agents: @@ -500,7 +562,18 @@ def _validate_budget_policy(budget_policy: dict, enabled_agents: dict[str, dict] f"budget_policy.tiers[{index}]: default_agent '{tier_agent}' must appear " "in enabled_agents." ) - if not tier.get("default_model"): + elif tier_model: + # The server only checks that the tier's agent is enabled, not that it has the model — + # so without this a tier can activate and hand the developer a model their agent was + # never configured with. Skipped when the agent lists no models (it then has only a + # default, or routes through a provider service whose catalog isn't enumerable). + available = _agent_model_ids(enabled_agents[tier_agent]) + if available and tier_model not in available: + errors.append( + f"budget_policy.tiers[{index}]: default_model '{tier_model}' is not one of the " + f"models configured for '{tier_agent}' ({', '.join(sorted(available))})." + ) + if not tier_model: errors.append(f"budget_policy.tiers[{index}]: default_model is required.") if len(set(percentages)) != len(percentages): diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py new file mode 100644 index 0000000..9cf2655 --- /dev/null +++ b/src/ucode/managed_wizard.py @@ -0,0 +1,900 @@ +"""Interactive `ucode setup`: author the workspace's managed coding-agent config. + +Workspace admins run this to build the ``CodingAgentConfig`` their developers will pull. It walks +the admin through agents, per-agent models, tracing, MCP servers, skills, and a spend-routing budget +policy, then writes the manifest to ``~/.ucode/managed-settings.json``. Publishing it to the +workspace is ``ucode apply`` (a separate command, so an admin can review the file first). + +Serialization, validation, and the per-agent model catalogs live in :mod:`ucode.managed_setup`; this +module is the interaction layer on top of them. Sub-flows an admin already knows — tracing, MCP, +skills — are delegated to the existing ``ucode configure `` commands and their results read +back out of ``state.json``, so there is exactly one picker per concern in the codebase. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from ucode.agents import TOOL_SPECS, check_gateway_endpoint +from ucode.databricks import ( + ANTHROPIC_FAMILIES, + discover_claude_models_unbucketed, + ensure_databricks_auth, + get_databricks_token, + has_cached_model_provider_services, + is_model_provider_feature_unavailable, + is_workspace_admin, + list_model_provider_services, + list_workspace_budgets, + service_usable_for_tool, +) +from ucode.managed_config import get_managed_config +from ucode.managed_setup import ( + CLAUDE_SLOT_FOR_FAMILY, + claude_family_candidates, + load_managed_settings, + model_options_for_agent, + save_managed_settings, + serialize_managed_config, + supports_provider_service, + validate_manifest, +) +from ucode.state import load_state +from ucode.ui import ( + console, + kv_line, + print_err, + print_heading, + print_note, + print_panel, + print_section, + print_success, + print_warning, + prompt_for_multi_selection, + prompt_for_percentage, + prompt_for_selection, + prompt_for_text, + prompt_for_tools, + prompt_yes_no_default, + spinner, +) + +# What `use_as_global_settings` actually does, in plain terms. Admins are choosing between a +# machine-wide managed settings file and a per-user one, which is not obvious from the field name. +GLOBAL_SETTINGS_BLURB = ( + "Write this agent's config to the machine's managed settings file, which applies to every " + "user on the machine and cannot be overridden locally. Answer no to write the per-user " + "settings file instead, which developers can still change." +) + +BUDGET_POLICY_BLURB = ( + "A budget policy moves developers onto cheaper agents and models as the workspace spends " + "against a budget — for example Claude Code on Opus by default, then Sonnet at 80%, then " + "OpenCode on Kimi at 100%. It only changes the default; developers can still pick anything " + "they have access to. Hard caps stay with the budget's own blocking threshold." +) + + +def _tracing_table_from_state(state: dict) -> str | None: + """The UC table `ucode configure tracing` wired up, or None when tracing is off. + + ``configure tracing`` records the destination as ``uc_destination``; the managed config calls the + same thing ``tracing.table``. + """ + tracing = state.get("tracing") + if not isinstance(tracing, dict) or not tracing.get("enabled"): + return None + destination = tracing.get("uc_destination") + return destination if isinstance(destination, str) and destination else None + + +def _mcp_type_for_url(url: str) -> str | None: + """Classify a registered MCP server's URL into a managed-config type tag. + + ``state.json`` stores each MCP server's resolved URL but not its type, while the managed config + stores ``{name, type}`` and lets the developer's ucode rebuild the URL. The URL shape is the only + signal available, so map it back. Returns None for a URL that matches nothing known, so unknown + servers are skipped rather than published with a guessed type. + """ + if "/ai-gateway/mcp-services/" in url: + return "mcp-service" + for fragment, tag in ( + ("/api/2.0/mcp/external/", "external"), + ("/api/2.0/mcp/genie/", "genie-space"), + ("/api/2.0/mcp/vector-search/", "vector-search"), + ("/api/2.0/mcp/functions/", "uc-functions"), + ): + if fragment in url: + return tag + if url.rstrip("/").endswith("/api/2.0/mcp/sql"): + return "sql" + # Databricks apps are the residual case: an arbitrary app host with a /mcp suffix. + if url.rstrip("/").endswith("/mcp"): + return "app" + return None + + +def _mcp_servers_from_state(state: dict) -> list[dict]: + """The registered MCP servers, as managed-config ``{name, type}`` entries. + + Skips the skills registry connection: skills are published under the manifest's own ``skills`` + field, so including its MCP entry would configure it twice. + """ + from ucode.mcp import SKILLS_MCP_KIND + + servers: list[dict] = [] + for entry in state.get("mcp_servers") or []: + if not isinstance(entry, dict) or entry.get("kind") == SKILLS_MCP_KIND: + continue + name = entry.get("name") + url = entry.get("url") + if not isinstance(name, str) or not name or not isinstance(url, str): + continue + tag = _mcp_type_for_url(url) + if tag is None: + print_warning(f"Skipping MCP server '{name}': unrecognized URL shape ({url}).") + continue + servers.append({"name": name, "type": tag}) + return servers + + +def _skill_names_from_state(state: dict) -> list[str]: + """Skill schemas registered on the skills MCP connection (``catalog.schema`` entries).""" + from ucode.mcp import _skill_mcp_locations + + return [name for name in _skill_mcp_locations(state) if isinstance(name, str) and name] + + +def provider_service_model_options(service: dict) -> list[str]: + """Model ids an admin can pick from a provider service, or [] when they can't be enumerated. + + A service's ``config.targets`` names the provider-side models it exposes, which is exactly the + vocabulary the manifest's ``default_model`` must use when ``model_provider_service`` is set. Two + cases yield nothing to pick from, and the caller falls back to free-text: + + - ``allow_all_targets`` — the service passes through the provider's whole catalog, which ucode + cannot enumerate (there is no list-models call for a provider service). + - no targets at all — e.g. a relayed Anthropic subscription service, which routes by canonical + model name rather than by an explicit target list. + """ + if service.get("allow_all_targets"): + return [] + targets = service.get("targets") + if not isinstance(targets, list): + return [] + return sorted({t for t in targets if isinstance(t, str) and t}) + + +def _select_provider_service(tool: str, workspace: str, token: str) -> dict | None: + """Offer Databricks-hosted vs an external Model Provider Service for ``tool``. + + Returns the chosen service dict (as :func:`list_model_provider_services` shapes it), or None to + stay on Databricks-hosted models. The whole dict is returned rather than just the name so the + model prompt can offer the service's ``targets`` instead of asking the admin to type a model id + from memory. + + Only claude and codex can route through a provider service today; every other agent short-cuts to + Databricks. Mirrors `cli._maybe_select_provider_service`, but returns the choice instead of + persisting it — the wizard is authoring a manifest, not configuring this machine. + """ + if not any( + supports_provider_service(tool, provider_type) + for provider_type in ("anthropic", "amazon_bedrock", "openai") + ): + return None + + display = TOOL_SPECS[tool]["display"] + # The listing is memoized per workspace, so only the first agent's call does any I/O. That one + # takes over a second and deserves a spinner; the rest are instant, and spinning once per agent + # made the wizard look like it re-listed the services every time. + if has_cached_model_provider_services(workspace): + services, reason = list_model_provider_services(workspace, token) + else: + with spinner("Checking for model provider services..."): + services, reason = list_model_provider_services(workspace, token) + if reason is not None: + # A workspace without the feature enabled is the common case and not worth a warning; any + # other failure is worth surfacing, or the admin silently loses the MPS option and has no + # idea why. Mirrors `cli._maybe_select_provider_service`. + if not is_model_provider_feature_unavailable(reason): + print_warning(f"Could not list model provider services: {reason}") + print_note("Falling back to Databricks-hosted models.") + return None + + usable = [service for service in services if service_usable_for_tool(tool, service)] + if not usable: + if services: + # Services exist but none match this agent's dialect — say so, since "no picker appeared" + # is otherwise indistinguishable from the feature being off. + print_note( + f"No model provider service matches {display}'s API dialect " + f"({len(services)} found on this workspace); using Databricks-hosted models." + ) + return None + + choice = prompt_for_selection( + f"How should {display} get its models?", + [ + ("databricks", "Databricks Hosted"), + ("mps", "External Models (Model Provider Service)"), + ], + ) + if choice != "mps": + return None + selected = prompt_for_selection( + f"Select the model provider service for {display}:", + [(service["name"], service["name"]) for service in usable], + searchable=True, + ) + if not selected: + return None + return next(service for service in usable if service["name"] == selected) + + +def _prompt_models_for_agent(tool: str, state: dict, provider_service: dict | None) -> dict: + """Build one agent's ``model_config``. Every agent ends up with a ``default_model``. + + Databricks-hosted agents pick from the workspace's discovered models, filtered to the families + that agent can actually serve. Provider-service agents pick from the service's own ``targets``, + falling back to free-text only when those can't be enumerated (``allow_all_targets``, or a + relayed service that routes by canonical name). + + An empty selection is re-prompted rather than accepted: an agent with no ``default_model`` cannot + be the config's ``default_agent`` (the server rejects it) and gives developers nothing to launch, + so "none" is never a useful answer here. Ctrl-C still aborts the whole flow. + + Model ids are stored bare (e.g. ``system.ai.claude-opus-4-8``), not provider-prefixed: each + agent's own writer adds whatever prefix its config format needs (see + ``opencode._resolve_model_selector``), which keeps the manifest agent-neutral. + + Codex takes a single model (the harness selects one); Claude's picks are bucketed into + ``ClaudeDefaultModels`` family slots; the rest keep a flat list plus their chosen default. + """ + display = TOOL_SPECS[tool]["display"] + model_config: dict = {} + if provider_service: + service_name = provider_service["name"] + model_config["model_provider_service"] = service_name + targets = provider_service_model_options(provider_service) + if targets: + model_config["default_model"] = _require_selection( + f"Default model for {display} (from {service_name}):", + [(target, target) for target in targets], + ) + else: + # No enumerable target list: the service either passes through the provider's whole + # catalog or routes by canonical model name, so the admin has to name the model. + print_note( + f"{service_name} does not publish an explicit model list, so enter the model id " + "as the provider names it (e.g. claude-sonnet-4-6)." + ) + model_config["default_model"] = _require_text(f"Default model for {display}") + return model_config + + if tool == "claude": + return _prompt_claude_models(state) + + options = model_options_for_agent(tool, state) + if not options: + print_warning(f"No models were discovered for {display} on this workspace.") + return {"default_model": _require_text(f"Default model for {display}")} + + if tool in SINGLE_MODEL_AGENTS: + return { + "default_model": _require_selection( + f"Select the model for {display}:", [(model, model) for model in options] + ) + } + + # Nothing pre-checked: the first option is whatever discovery sorted first, not a + # recommendation — for pi it is a Claude model, for codex the oldest GPT. Pre-checking it made + # "hit Enter" produce an arbitrary config. (A worthwhile follow-up is to pre-check the models + # this workspace was configured with last time, which `load_managed_settings` already loads for + # the agent picker, so a re-run becomes an edit rather than a re-entry.) + picked = _require_multi_selection( + f"Select models for {display}:", + [(model, model) for model in options], + ) + if len(picked) == 1: + model_config["default_model"] = picked[0] + else: + model_config["default_model"] = _require_selection( + f"Default model for {display}:", [(model, model) for model in picked] + ) + + model_config["models"] = picked + return model_config + + +# Agents that get a single model rather than a multi-select. Codex's proto has no model list at all. +# Gemini and Copilot do declare `repeated string models`, but their config writers take one model +# (`gemini.write_tool_config(state, model)` / `copilot.write_tool_config(state, model)`) and write a +# single env var — so a published list would be read by nothing. Offering one keeps the manifest +# honest about what ucode can apply; widen this when those writers grow a picker. +SINGLE_MODEL_AGENTS = frozenset({"codex", "gemini", "copilot"}) + +# Skip sentinel for a Claude family prompt. Every `ClaudeDefaultModels` slot is optional, and an +# unset one falls back to `default_model`, so leaving a family out is a legitimate choice. +_SKIP_FAMILY = "__skip__" + + +def _prompt_claude_models(state: dict) -> dict: + """Build Claude's ``model_config`` one family slot at a time. + + Claude Code addresses models by family alias, not from a list, so the config is a set of slots: + `default_opus_model`, `default_sonnet_model`, `default_haiku_model`, `default_fable_model`. A flat + multi-select can't express that — and because `state["claude_models"]` holds only the newest id + per family, it could only ever offer one model per family anyway. Asking per family surfaces the + alternatives (six opus versions on a typical workspace, not one) and matches the proto. + + Each family may be skipped; the overall `default_model` is then chosen from the slots that were + filled, so it can never name a model the config doesn't carry. + """ + display = TOOL_SPECS["claude"]["display"] + # No spinner: the model-services listing is already cached by the time the flow reaches here + # (`configure_shared_state` walked it up front), so this is a filter over data in hand, not a + # fetch. Showing "Fetching Claude models..." made the wizard look like it listed the catalog + # twice. + candidates = _claude_candidates(state) + if not candidates: + print_warning(f"No Claude models were discovered for {display} on this workspace.") + return {"default_model": _require_text(f"Default model for {display}")} + + print_note( + "Claude Code picks a model by family, so set a default per family. Skip any family you " + "don't want configured — it falls back to the overall default." + ) + slots: dict[str, str] = {} + for family in ANTHROPIC_FAMILIES: + family_models = candidates.get(family) + if not family_models: + continue + choice = prompt_for_selection( + f"Default {family} model:", + [(model, model) for model in family_models] + [(_SKIP_FAMILY, f"(skip {family})")], + searchable=True, + ) + if choice is None: + raise KeyboardInterrupt + if choice != _SKIP_FAMILY: + slots[CLAUDE_SLOT_FOR_FAMILY[family]] = choice + + if not slots: + # Every slot skipped is a legitimate, minimal config: the proto leaves `models` optional and + # each unset slot falls back to `default_model`, so one model covers every family. Pick it + # from the same candidates rather than asking the admin to type an id. + print_note(f"No families configured, so {display} will use a single model for all of them.") + every_model = [m for family_models in candidates.values() for m in family_models] + return { + "default_model": _require_selection( + f"Which model should {display} use?", + [(m, m) for m in dict.fromkeys(every_model)], + ) + } + + chosen = list(dict.fromkeys(slots.values())) + model_config: dict = {"models": slots} + if len(chosen) == 1: + # A one-option prompt is a wasted keystroke, but skipping it silently reads as a dropped + # step — say what was inferred so the admin knows the default is set, and to what. + model_config["default_model"] = chosen[0] + print_success(f"Overall default for {display}: {chosen[0]} (the only model configured)") + else: + model_config["default_model"] = _require_selection( + f"Which of those is {display}'s overall default?", [(m, m) for m in chosen] + ) + return model_config + + +def _claude_candidates(state: dict) -> dict[str, list[str]]: + """Claude models grouped by family. Degrades to the per-family picks if the listing fails. + + Caches the full listing on ``state["all_claude_models"]`` so `validate_manifest` recognizes the + older versions these prompts offer — ``claude_models`` alone holds just the newest per family, + and would reject a legitimately-picked ``claude-opus-4-8``. + + INVARIANT: whatever this returns must be recognizable by ``validate_manifest``, which reads + ``all_claude_models`` (falling back to ``claude_models``) via ``_known_models``. The two paths + below both satisfy it, for different reasons: the listing path widens the candidates *and* sets + the cache, while the fallback path sets nothing but also narrows the candidates to + ``claude_models``, which ``_known_models`` already covers. Widening the fallback without also + populating the cache breaks the invariant, and the symptom is a confusing rejection at the very + end of the flow ("claude: model 'system.ai.claude-opus-4-8' is not available on this + workspace") rather than an error at the prompt that offered it. + """ + cached = state.get("all_claude_models") + if isinstance(cached, list) and cached: + return claude_family_candidates([m for m in cached if isinstance(m, str)], state) + + workspace = state.get("workspace") + all_claude: list[str] = [] + if workspace: + try: + token = get_databricks_token(workspace, state.get("profile")) + all_claude, _ = discover_claude_models_unbucketed(workspace, token) + except (RuntimeError, OSError): + # OSError covers a missing `databricks` binary: `get_databricks_token` shells out, so a + # machine without the CLI on PATH raises FileNotFoundError rather than RuntimeError. + # Either way the per-family picks below are a usable fallback. + all_claude = [] + if all_claude: + state["all_claude_models"] = all_claude + return claude_family_candidates(all_claude, state) + + +# Every picker in this flow chooses a model, a provider service, or a budget — lists that on a real +# workspace run to a dozen-plus entries (16 GPT models on the workspace this was built against), so +# they are all filterable by typing. That trades away j/k navigation, which questionary can't offer +# alongside search; arrow keys still work. +def _require_selection(prompt: str, options: list[tuple[str, str]]) -> str: + """Single-select that won't take "nothing" for an answer. + + ``prompt_for_selection`` returns None for both Ctrl-C and an empty submission, and the two are + genuinely indistinguishable here: questionary's ``Question.ask`` catches KeyboardInterrupt + internally and returns None (v2.1.1, question.py), so nothing propagates for a caller to see. + A None is therefore treated as an abort rather than re-asked — re-asking looped forever on + Ctrl-C, printing the error once per keypress and never exiting. + """ + answer = prompt_for_selection(prompt, options, searchable=True) + if not answer: + raise KeyboardInterrupt + return answer + + +def _require_multi_selection( + prompt: str, options: list[tuple[str, str]], preselected: list[str] | None = None +) -> list[str]: + """Multi-select that requires at least one choice. None (Ctrl-C) still aborts.""" + while True: + picked = prompt_for_multi_selection( + prompt, options, preselected=preselected, searchable=True + ) + if picked is None: + raise KeyboardInterrupt + if picked: + return picked + print_err("Select at least one model (space to toggle, enter to confirm).") + + +def _require_text(prompt: str) -> str: + """Free-text prompt that requires a non-empty answer. + + ``required=True`` makes closed stdin abort instead of returning None. Without it a + non-interactive run (piped stdin, CI) spun here forever: ``prompt_for_text`` returns its default + on EOF, the default is None, and the loop re-asked an empty stream. Reachable whenever model + discovery finds nothing, which is exactly when a run is most likely to be scripted. + """ + while True: + answer = prompt_for_text(prompt, required=True) + if answer: + return answer + print_err("Please enter a model id.") + + +def configured_models_for_agent(agent_config: dict) -> list[str]: + """Models an agent was configured with, in the manifest's own vocabulary. + + ``model_config.models`` is a flat list for most agents but a family-slot dict for claude + (``default_opus_model`` -> id), so both shapes collapse to a list here. The ``default_model`` is + included because codex has no model list at all — it is the only model that agent has. + """ + model_config = agent_config.get("model_config") + if not isinstance(model_config, dict): + return [] + models: list[str] = [] + raw = model_config.get("models") + if isinstance(raw, dict): + models.extend(v for v in raw.values() if isinstance(v, str) and v) + elif isinstance(raw, list): + models.extend(m for m in raw if isinstance(m, str) and m) + default_model = model_config.get("default_model") + if isinstance(default_model, str) and default_model: + models.append(default_model) + # dict.fromkeys de-duplicates while keeping the admin's preference order. + return list(dict.fromkeys(models)) + + +def _prompt_budget_policy( + workspace: str, token: str, enabled_agents: dict[str, dict], state: dict +) -> dict | None: + """Author a spend-routing ``budget_policy``, or None when the admin declines or can't. + + Budgets themselves are created in the Databricks console (they're account-level objects), so the + admin picks an existing one here. Tiers are prompted in percent and stored as fractions, which is + what the API validates. + + A tier's model choices come from what the admin configured for that agent earlier in this run — + not the workspace catalog. Offering the catalog would let a tier point an agent at a model it + wasn't given, which neither this validation nor the server's would reject: the tier would + activate and hand the developer a model their agent doesn't have. + """ + print_section("Budget policy") + print_note(BUDGET_POLICY_BLURB) + if not prompt_yes_no_default("Set up a budget policy for this workspace?", default=False): + return None + + with spinner("Listing workspace budgets..."): + budgets, reason = list_workspace_budgets(workspace, token) + if reason is not None or not budgets: + print_warning( + "No AI Gateway budgets are visible for this workspace, so there is nothing to attach a " + "policy to. Create a budget in the Databricks console first, then re-run `ucode setup`." + ) + return None + + budget_id = prompt_for_selection( + "Which budget should this policy track?", + [ + (budget["id"], f"{budget['display_name'] or budget['id']} ({budget['id']})") + for budget in budgets + ], + searchable=True, + ) + if not budget_id: + return None + + policy: dict = {"budget_id": budget_id} + display_name = prompt_for_text("Policy name", default="coding-agents-tiered-routing") + if display_name: + policy["display_name"] = display_name + + tiers: list[dict] = [] + seen_percentages: set[float] = set() + print_note( + "Add one tier per step-down. Each tier activates once spend reaches its percentage, and " + "the highest activated tier wins." + ) + while True: + index = len(tiers) + 1 + fraction = prompt_for_percentage(f"Tier {index}: activates at what percent of budget?") + if fraction in seen_percentages: + print_err("That percentage is already used by another tier; pick a different one.") + continue + agent = prompt_for_selection( + f"Tier {index}: which agent becomes the default?", + [(tool, TOOL_SPECS[tool]["display"]) for tool in enabled_agents], + ) + if not agent: + break + # Only what this agent was actually configured with; the workspace catalog would offer + # models the agent doesn't have. + options = configured_models_for_agent(enabled_agents.get(agent) or {}) + if not options: + options = model_options_for_agent(agent, state) + if options: + model = prompt_for_selection( + f"Tier {index}: which model?", [(m, m) for m in options], searchable=True + ) + else: + model = prompt_for_text(f"Tier {index}: which model?") + if not model: + break + seen_percentages.add(fraction) + tiers.append( + { + "spending_percentage": fraction, + "default_agent": agent, + "default_model": model, + } + ) + if not prompt_yes_no_default("Add another tier?", default=False): + break + + if tiers: + policy["tiers"] = tiers + return policy + + +def _render_summary(workspace: str, manifest: dict) -> None: + """Print the authored config in a box so an admin can eyeball it before publishing. + + Boxed rather than printed as loose lines: this is the one block an admin is meant to read as a + whole and check against what they intended, and it lands after a long flow of prompts. + """ + lines: list[str] = [kv_line("Workspace", workspace)] + default_agent = manifest.get("default_agent") + if isinstance(default_agent, str): + lines.append( + kv_line( + "Default agent", TOOL_SPECS.get(default_agent, {}).get("display", default_agent) + ) + ) + + for tool, agent_config in (manifest.get("enabled_agents") or {}).items(): + display = TOOL_SPECS.get(tool, {}).get("display", tool) + model_config = agent_config.get("model_config") or {} + detail = model_config.get("default_model") or "no model" + provider = model_config.get("model_provider_service") + if provider: + detail = f"{detail} via {provider}" + scope = "machine-wide" if agent_config.get("use_as_global_settings") else "per-user" + lines.append(kv_line(display, f"{detail} ({scope})")) + # Spell out the per-family slots and model lists: the one-line default alone doesn't show + # which families an admin configured, which is most of what they chose for claude. + models = model_config.get("models") + if isinstance(models, dict): + for slot, model in models.items(): + family = slot.removeprefix("default_").removesuffix("_model") + lines.append(kv_line(f" {family}", str(model))) + elif isinstance(models, list) and len(models) > 1: + lines.append(kv_line(" models", ", ".join(str(m) for m in models))) + + mcp_servers = manifest.get("mcp_servers") or [] + lines.append( + kv_line( + "MCP servers", + ", ".join(str(server.get("name")) for server in mcp_servers) if mcp_servers else "none", + ) + ) + skills = (manifest.get("skills") or {}).get("names") or [] + lines.append(kv_line("Skills", ", ".join(skills) if skills else "none")) + lines.append(kv_line("Tracing", manifest.get("tracing_table") or "disabled")) + + policy = manifest.get("budget_policy") + if isinstance(policy, dict): + tiers = policy.get("tiers") or [] + lines.append( + kv_line("Budget policy", policy.get("display_name") or policy.get("budget_id") or "set") + ) + for tier in tiers: + agent = tier.get("default_agent") + display = TOOL_SPECS.get(agent, {}).get("display", agent) + percent = float(tier.get("spending_percentage", 0)) * 100 + lines.append(kv_line(f" at {percent:g}%", f"{display} / {tier.get('default_model')}")) + else: + lines.append(kv_line("Budget policy", "none")) + + print_panel("Configuration summary", lines) + + +def _require_admin(workspace: str, token: str) -> None: + """Stop unless the caller is a workspace admin. + + An unverifiable check (SCIM unreachable) warns and continues: the API enforces the same rule, so + the worst case is a clear PERMISSION_DENIED at publish time rather than a false block here. + """ + with spinner("Checking workspace admin permissions..."): + admin = is_workspace_admin(workspace, token) + if admin is False: + raise RuntimeError( + f"You are not an admin of {workspace}. `ucode setup` authors the workspace-wide " + "coding config, so it is restricted to workspace admins." + ) + if admin is None: + print_warning( + "Could not verify workspace admin permissions. Continuing — `ucode apply` will fail " + "if you lack them." + ) + else: + print_success("Admin permissions verified") + + +def _warn_on_existing_config(workspace: str, token: str) -> None: + """Warn when the workspace already has a published config that `ucode apply` would replace. + + Deliberately doesn't itemize what the existing config holds. The admin doesn't need an inventory + to act on this — the instruction is the same either way ("include everything you want to keep") + — and `ucode setup show` prints the real thing for anyone who wants to compare. + """ + with spinner("Checking for an existing managed config..."): + existing, reason = get_managed_config(workspace, token) + if reason is not None: + print_note(f"Could not check for an existing config: {reason}") + return + if existing is None: + return + print_warning( + "This workspace already has a managed configuration — one config covers every agent, MCP " + "server, skill, tracing table, and budget policy for the whole workspace. Publishing " + "replaces all of it, so make sure this run includes everything you want to keep." + ) + + +def setup_from_file(path: str) -> int: + """Validate an admin-written manifest and save it, skipping the interactive flow. + + The non-interactive path for CI and for admins who'd rather keep the JSON in version control. + Reads ucode's own manifest shape (the same thing the wizard writes), not proto-JSON. + """ + manifest_path = Path(path).expanduser() + try: + raw = manifest_path.read_text(encoding="utf-8") + except OSError as exc: + raise RuntimeError(f"Could not read manifest file: {manifest_path}") from exc + try: + manifest = json.loads(raw) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"{manifest_path} is not valid JSON: {exc.msg} (line {exc.lineno})." + ) from None + if not isinstance(manifest, dict): + raise RuntimeError(f"{manifest_path} must contain a JSON object.") + + state = load_state() + workspace = state.get("workspace") + if not workspace: + raise RuntimeError( + "No workspace is configured. Run `ucode configure` first so ucode knows which " + "workspace this manifest is for." + ) + + errors = validate_manifest(manifest, state) + if errors: + print_err(f"{manifest_path} is not a valid managed config:") + for error in errors: + print_note(error) + return 1 + + save_managed_settings(workspace, manifest) + _render_summary(workspace, manifest) + print_success(f"Saved to {manifest_path.name} -> ~/.ucode/managed-settings.json") + _print_next_steps() + return 0 + + +def _print_next_steps() -> None: + console.print() + print_heading("Next steps") + # Deliberately only `apply`. There is no way yet to try the authored config locally: the + # manifest describes what developers should get, while `ucode configure --dry-run` previews + # this machine's own agent configs, so pointing at it implied a local test it doesn't perform. + print_note("Publish it to the workspace: ucode apply") + + +def setup_command(from_file: str | None = None) -> int: + """Author the workspace's managed coding-agent config interactively. + + Returns a process exit code. Raises RuntimeError for actionable failures (not an admin, no + agents available) and KeyboardInterrupt when the admin aborts a picker; the CLI maps both. + """ + if from_file is not None: + return setup_from_file(from_file) + + # Imported here rather than at module scope: `cli` imports this module, so a top-level import + # would be circular. + from ucode.cli import _prompt_for_configuration, configure_shared_state + + print_section("ucode setup") + print_note("Author the managed coding config for this workspace.") + print_note("Developers pull it automatically when they run ucode.") + + workspace, profile = _prompt_for_configuration() + # `configure_shared_state` below authenticates too and prints its own success line, so this one + # stays quiet rather than reporting the same thing twice. It still has to run first: the admin + # gate and the existing-config check both need a token before discovery. + ensure_databricks_auth(workspace, profile, quiet=True) + token = get_databricks_token(workspace, profile) + + _require_admin(workspace, token) + _warn_on_existing_config(workspace, token) + + # Discover the workspace's models and gateway URLs. This also logs in and persists local state, + # which is what lets the admin dry-run the config on their own machine afterwards. + state = configure_shared_state(workspace, profile=profile, force_login=False) + workspace = state.get("workspace") or workspace + profile = state.get("profile") or profile + + available = [tool for tool in TOOL_SPECS if check_gateway_endpoint(state, tool)] + if not available: + raise RuntimeError( + f"No coding agents are available on {workspace}. Check that the workspace's AI Gateway " + "serves models for at least one agent." + ) + + previous = load_managed_settings(workspace) or {} + previously_enabled = [ + tool for tool in (previous.get("enabled_agents") or {}) if tool in TOOL_SPECS + ] + picked = prompt_for_tools( + [(tool, TOOL_SPECS[tool]["display"]) for tool in available], + preselected=previously_enabled or None, + ) + if not picked: + print_note("No coding agents selected — nothing to configure.") + return 0 + + default_agent = picked[0] + if len(picked) > 1: + chosen = prompt_for_selection( + "Which agent should launch when a developer runs `ucode`?", + [(tool, TOOL_SPECS[tool]["display"]) for tool in picked], + ) + if not chosen: + raise KeyboardInterrupt + default_agent = chosen + print_success(f"Default agent set to {TOOL_SPECS[default_agent]['display']}") + + enabled_agents: dict[str, dict] = {} + for tool in picked: + print_heading(TOOL_SPECS[tool]["display"]) + provider_service = _select_provider_service(tool, workspace, token) + # Always set: `_prompt_models_for_agent` re-prompts rather than returning empty, so every + # enabled agent carries a default_model and any of them can be the default_agent. + agent_config: dict = { + "model_config": _prompt_models_for_agent(tool, state, provider_service) + } + agent_config["use_as_global_settings"] = prompt_yes_no_default( + f"Apply {TOOL_SPECS[tool]['display']} config machine-wide? ({GLOBAL_SETTINGS_BLURB})", + default=False, + ) + enabled_agents[tool] = agent_config + + manifest: dict = {"default_agent": default_agent, "enabled_agents": enabled_agents} + + print_section("Tracing") + if prompt_yes_no_default( + "Send coding-session traces to an MLflow experiment in this workspace?", + default=bool(_tracing_table_from_state(state)), + ): + from ucode.tracing import configure_tracing_command + + configure_tracing_command(workspaces=[(workspace, profile)]) + tracing_table = _tracing_table_from_state(load_state()) + if tracing_table: + manifest["tracing_table"] = tracing_table + print_success(f"Tracing configured ({tracing_table})") + else: + print_warning("Tracing was not enabled, so it is left out of the managed config.") + + print_section("MCP servers") + if prompt_yes_no_default("Set up managed MCP servers for this workspace?", default=False): + from ucode.mcp import configure_mcp_command + + configure_mcp_command() + mcp_servers = _mcp_servers_from_state(load_state()) + if mcp_servers: + manifest["mcp_servers"] = mcp_servers + print_success(f"{len(mcp_servers)} MCP server(s) added to the managed config") + + print_section("Skills") + if prompt_yes_no_default("Set up managed skills for this workspace?", default=False): + locations = prompt_for_text( + "Skill schemas to publish, comma-separated `catalog.schema` (blank to skip)", + default="", + ) + parsed: list[str] = [item.strip() for item in (locations or "").split(",") if item.strip()] + if parsed: + from ucode.mcp import configure_skills_mcp_command + + configure_skills_mcp_command(parsed) + skill_names = _skill_names_from_state(load_state()) or parsed + manifest["skills"] = {"names": skill_names} + print_success(f"{len(skill_names)} skill schema(s) added to the managed config") + + budget_policy = _prompt_budget_policy(workspace, token, enabled_agents, state) + if budget_policy: + manifest["budget_policy"] = budget_policy + + errors = validate_manifest(manifest, state) + if errors: + # A validation failure here is a wizard bug, not admin error — the pickers only offer valid + # choices. Surface it plainly rather than writing a manifest that `apply` would reject. + print_err("The generated config is not valid:") + for error in errors: + print_note(error) + return 1 + + save_managed_settings(workspace, manifest) + _render_summary(workspace, manifest) + console.print() + print_success("Saved to ~/.ucode/managed-settings.json") + _print_next_steps() + return 0 + + +def show_command() -> int: + """Print the authored manifest and the proto-JSON `ucode apply` would publish.""" + workspace = load_state().get("workspace") + manifest = load_managed_settings(workspace) + if manifest is None: + print_note("No managed config has been authored yet. Run `ucode setup` to create one.") + return 0 + _render_summary(workspace or "unknown", manifest) + console.print() + print_heading("Payload for `ucode apply`") + console.print(json.dumps(serialize_managed_config(manifest), indent=2)) + return 0 + + +__all__ = ["setup_command", "setup_from_file", "show_command"] diff --git a/src/ucode/ui.py b/src/ucode/ui.py index 81ffc2e..b3ff6d1 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -13,6 +13,7 @@ import questionary from rich.console import Console +from rich.markup import escape from rich.panel import Panel from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn @@ -51,6 +52,28 @@ def print_kv(key: str, val: str) -> None: console.print(f" [bold]{key}:[/bold] [cyan]{val}[/cyan]") +def kv_line(key: str, val: str) -> str: + """A `print_kv`-styled line, returned instead of printed, for collecting into a panel. + + The value is markup-escaped. Rich reads bracketed text as a style tag and renders nothing for + it, so a policy name of ``[prod] tiered routing`` displayed as ``tiered routing`` in the config + summary — the one block an admin reads to confirm what they are about to publish workspace-wide. + Values here include admin-typed free text (policy name, skills locations, tracing table). + """ + return f"[bold]{escape(key)}:[/bold] [cyan]{escape(val)}[/cyan]" + + +def print_panel(title: str, lines: list[str]) -> None: + """Render `lines` inside a titled box. + + Unlike :func:`print_section`, which boxes a bare title, this boxes the body — so a block that + should be read as one unit (a config summary an admin is about to publish) reads as one, rather + than as loose lines that blend into whatever the flow printed before it. + """ + console.print() + console.print(Panel("\n".join(lines), title=title, style="blue", expand=False)) + + def print_note(text: str) -> None: console.print(f"[dim]•[/dim] {text}") @@ -268,12 +291,18 @@ def prompt_for_workspace( print_err(str(exc)) -def prompt_for_tools(available: list[tuple[str, str]]) -> list[str]: +def prompt_for_tools( + available: list[tuple[str, str]], + preselected: list[str] | set[str] | None = None, + prompt: str = "Select coding agents to configure:", +) -> list[str]: """Multi-select picker for coding agents. `available` is [(tool_id, display_name), ...]. Returns the chosen tool_ids. - All options are checked by default so hitting Enter selects everything. - Returns [] if the user submits an empty selection. + When ``preselected`` is None every option is checked by default, so hitting + Enter selects everything; pass a subset to pre-check only those (e.g. the + agents an existing managed config already enables). Returns [] if the user + submits an empty selection. """ style = questionary.Style( [ @@ -287,12 +316,17 @@ def prompt_for_tools(available: list[tuple[str, str]]) -> list[str]: ("answer", "fg:cyan"), ] ) + preselected_set = {str(item) for item in preselected} if preselected is not None else None choices = [ - questionary.Choice(title=display, value=tool_id, checked=True) + questionary.Choice( + title=display, + value=tool_id, + checked=(preselected_set is None or tool_id in preselected_set), + ) for tool_id, display in available ] answer = questionary.checkbox( - "Select coding agents to configure:", + prompt, choices=choices, style=style, pointer="›", @@ -302,11 +336,139 @@ def prompt_for_tools(available: list[tuple[str, str]]) -> list[str]: return list(answer) if answer else [] -def prompt_for_selection(prompt: str, options: list[tuple[str, str]]) -> str | None: +def prompt_for_multi_selection( + prompt: str, + options: list[tuple[str, str]], + preselected: list[str] | set[str] | None = None, + *, + searchable: bool = False, +) -> list[str] | None: + """Multi-select picker over arbitrary `(value, label)` options. + + Distinct from :func:`prompt_for_tools`, which is agent-specific and defaults to + everything checked: here nothing is checked unless ``preselected`` says so, since + an admin picking models wants an explicit choice rather than "all of them". + Returns the chosen values, [] on an empty submission, or None if cancelled + (Ctrl-C) so callers can distinguish "chose nothing" from "aborted". + + ``searchable`` lets the user narrow a long list by typing; see + :func:`prompt_for_selection` for why it trades away j/k navigation. + """ + style = questionary.Style( + [ + ("pointer", "fg:cyan bold"), + ("highlighted", "noinherit"), + ("selected", "noinherit"), + ("answer", "fg:cyan"), + ] + ) + preselected_set = {str(item) for item in preselected} if preselected is not None else set() + choices = [ + questionary.Choice(title=option_label, value=value, checked=value in preselected_set) + for value, option_label in options + ] + instruction = "(space to toggle, enter to confirm)" + if searchable: + instruction = "(type to filter, space to toggle, enter to confirm)" + answer = questionary.checkbox( + prompt, + choices=choices, + style=style, + pointer="›", + qmark="", + instruction=instruction, + use_search_filter=searchable, + use_jk_keys=not searchable, + ).ask() + return None if answer is None else list(answer) + + +def prompt_for_text( + prompt: str, *, default: str | None = None, required: bool = False +) -> str | None: + """Free-text prompt, used when model discovery found nothing to pick from. + + Returns the trimmed input, ``default`` on an empty answer, or None when there is no + default and the user submits nothing (or closes stdin). + + ``required=True`` raises ``KeyboardInterrupt`` on closed stdin instead of returning None, for + callers that loop until they get a value: returning None to such a caller spins forever on a + piped or exhausted stdin. Matches :func:`prompt_for_percentage`, which has no default and does + the same. + + A default is shown as ``[value] (enter to accept)`` rather than the bare ``[value]``: bracketed + text alone reads as a format example as easily as a value that will be used, so it invited + retyping what pressing enter would already pick. + + The whole bracketed hint is markup-escaped, brackets included. Rich reads + ``[coding-agents-tiered-routing]`` as a style tag and prints nothing for it, so an unescaped + word-like default vanished from the prompt entirely — numeric ones like ``[80]`` are not valid + tags and survived, which is why this looked fine wherever it was checked. + """ + hint = f" {escape(f'[{default}]')} (enter to accept)" if default else "" + while True: + try: + raw_value = console.input(f"{label(prompt)}{muted(hint)} {muted('›')} ").strip() + except EOFError as exc: + if required: + raise KeyboardInterrupt from exc + return default + if raw_value: + return raw_value + if default is not None: + return default + print_err("Please enter a value.") + + +def prompt_for_percentage(prompt: str, *, default: float | None = None) -> float: + """Prompt for a percentage (0-100) and return it as a fraction in [0, 1]. + + Budget tiers are fractions in the API (the server validates 0..1), but admins think in + percent — and the spec's own prose says "80%". Prompting in percent and converting here + keeps that mismatch in one place instead of at every call site. + + No caller passes ``default`` today, and tier thresholds deliberately have none: a threshold + decides when developers get downgraded, so it should be typed rather than accepted by accident. + The hint is still formatted (and escaped) the same way :func:`prompt_for_text` formats its own, + so the two cannot drift if a default is ever introduced. + + Raises ``KeyboardInterrupt`` on closed stdin when there is no default — see the handler below. + """ + hint = f" {escape(f'[{default * 100:g}]')} (enter to accept)" if default is not None else "" + while True: + try: + raw_value = console.input(f"{label(prompt)}{muted(hint)} {muted('› ')}").strip() + except EOFError as exc: + if default is not None: + return default + # Closed stdin with no default to fall back on is the admin abandoning the prompt, which + # is what Ctrl-C means here too. Raised as KeyboardInterrupt so the CLI's existing + # handler prints "Interrupted." and exits 130; a bare EOFError has no handler anywhere + # above this and reached the admin as a traceback. + raise KeyboardInterrupt from exc + if not raw_value and default is not None: + return default + try: + percent = float(raw_value.rstrip("%")) + except ValueError: + print_err("Please enter a number between 0 and 100.") + continue + if 0 <= percent <= 100: + return percent / 100 + print_err("Please enter a number between 0 and 100.") + + +def prompt_for_selection( + prompt: str, options: list[tuple[str, str]], *, searchable: bool = False +) -> str | None: """Single-select arrow-key picker. `options` is [(value, label), ...]. The prompt renders above the choices (questionary convention). Returns the chosen value, or None if the user cancels (Ctrl-C / empty). + + ``searchable`` lets the user narrow a long list by typing. It costs j/k navigation — questionary + rejects both at once, since j and k are also search characters — so it is opt-in for the pickers + that are actually long (model and budget lists), leaving short ones on plain arrow keys. """ style = questionary.Style( [ @@ -323,7 +485,9 @@ def prompt_for_selection(prompt: str, options: list[tuple[str, str]]) -> str | N style=style, pointer="›", qmark="", - instruction="(use arrow keys)", + instruction="(type to filter, arrow keys to move)" if searchable else "(use arrow keys)", + use_search_filter=searchable, + use_jk_keys=not searchable, ).ask() return answer diff --git a/tests/conftest.py b/tests/conftest.py index 0cc7932..a60d07d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,12 +24,16 @@ def _isolate_ucode_state(tmp_path, monkeypatch): it can never touch the developer's real ~/.ucode/state.json. """ import ucode.config_io as config_io_mod + import ucode.databricks as databricks_mod import ucode.state as state_mod state_dir = tmp_path / ".ucode" state_dir.mkdir() monkeypatch.setattr(state_mod, "STATE_PATH", state_dir / "state.json") monkeypatch.setattr(config_io_mod, "APP_DIR", state_dir) + # The model-services listing is memoized for the life of the process, so without this a cached + # result would leak into the next test and make a stubbed listing look like it was never called. + databricks_mod.clear_model_services_cache() def _workspace() -> str: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index cc7b027..2272025 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2093,3 +2093,168 @@ class TestClassifyModelFamily: ) 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"] diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index 3570a7f..8a08200 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -421,6 +421,73 @@ def test_slots_serialize_into_the_claude_variant(self): assert variant["models"] == {"default_opus_model": "system.ai.claude-opus-4-8"} +class TestClaudeFamilyCandidates: + """Discovery keeps one id per family for the launch path; authoring needs the alternatives.""" + + ALL = [ + "system.ai.claude-opus-4-1", + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-sonnet-5", + "system.ai.claude-haiku-4-5", + "system.ai.gpt-5-6", + ] + + def test_groups_by_family(self): + from ucode.managed_setup import claude_family_candidates + + got = claude_family_candidates(self.ALL) + assert set(got) == {"opus", "sonnet", "haiku"} + assert got["haiku"] == ["system.ai.claude-haiku-4-5"] + + def test_newest_first_within_a_family(self): + from ucode.managed_setup import claude_family_candidates + + assert claude_family_candidates(self.ALL)["opus"] == [ + "system.ai.claude-opus-5", + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-4-1", + ] + + def test_non_claude_models_are_ignored(self): + from ucode.managed_setup import claude_family_candidates + + assert not any( + "gpt" in m for models in claude_family_candidates(self.ALL).values() for m in models + ) + + def test_deduplicates(self): + from ucode.managed_setup import claude_family_candidates + + got = claude_family_candidates(["system.ai.claude-opus-5", "system.ai.claude-opus-5"]) + assert got["opus"] == ["system.ai.claude-opus-5"] + + def test_falls_back_to_the_per_family_picks(self): + # Without the full listing, the bucketed state is all that's available — one per family, but + # enough for the per-slot prompts to work. + from ucode.managed_setup import claude_family_candidates + + got = claude_family_candidates([], {"claude_models": {"opus": "system.ai.claude-opus-5"}}) + assert got == {"opus": ["system.ai.claude-opus-5"]} + + def test_empty_everything_yields_nothing(self): + from ucode.managed_setup import claude_family_candidates + + assert claude_family_candidates([], {}) == {} + + def test_slot_names_match_the_proto(self): + # ClaudeDefaultModels fields, verified against ai-gateway-api service.proto. + from ucode.managed_setup import CLAUDE_SLOT_FOR_FAMILY + + assert set(CLAUDE_SLOT_FOR_FAMILY.values()) == { + "default_fable_model", + "default_opus_model", + "default_sonnet_model", + "default_haiku_model", + } + + class TestValidate: def test_full_manifest_is_valid(self): assert validate_manifest(_full_manifest(), STATE) == [] @@ -472,6 +539,40 @@ def test_unknown_model_is_rejected(self): errors = validate_manifest(manifest, STATE) assert any("not available on this workspace" in e for e in errors) + def test_older_claude_version_is_recognized(self): + # `claude_models` holds only the newest per family, so without the full listing an older + # version the per-family prompts offered would be wrongly rejected. + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-4-8", + "models": {"default_opus_model": "system.ai.claude-opus-4-8"}, + } + } + }, + } + state = { + "claude_models": {"opus": "system.ai.claude-opus-5"}, + "all_claude_models": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], + } + assert validate_manifest(manifest, state) == [] + + def test_unknown_claude_version_is_still_rejected(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-9-9"}} + }, + } + state = { + "claude_models": {"opus": "system.ai.claude-opus-5"}, + "all_claude_models": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], + } + errors = validate_manifest(manifest, state) + assert any("not available on this workspace" in e for e in errors) + def test_model_check_skipped_without_state(self): manifest = { "default_agent": "claude", @@ -586,6 +687,106 @@ def test_tier_needs_a_default_model(self): errors = validate_manifest(manifest, STATE) assert any("default_model is required" in e for e in errors) + def test_tier_model_must_be_one_the_agent_has(self): + # The server only checks that the tier's agent is enabled, so without this a tier activates + # and hands the developer a model their agent was never configured with. + manifest = { + "default_agent": "pi", + "enabled_agents": { + "pi": { + "model_config": { + "default_model": "system.ai.kimi-k2-6", + "models": ["system.ai.kimi-k2-6"], + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "pi", + "default_model": "system.ai.gpt-5-6", + } + ], + }, + } + errors = validate_manifest(manifest, STATE) + assert any("is not one of the models configured for 'pi'" in e for e in errors), errors + + def test_tier_model_from_the_agents_list_is_accepted(self): + manifest = { + "default_agent": "pi", + "enabled_agents": { + "pi": { + "model_config": { + "default_model": "system.ai.kimi-k2-6", + "models": ["system.ai.kimi-k2-6", "system.ai.gpt-5-6"], + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "pi", + "default_model": "system.ai.gpt-5-6", + } + ], + }, + } + assert validate_manifest(manifest, STATE) == [] + + def test_tier_model_matching_a_claude_family_slot_is_accepted(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-4-8", + "models": {"default_sonnet_model": "system.ai.claude-sonnet-4-6"}, + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-sonnet-4-6", + } + ], + }, + } + assert validate_manifest(manifest, STATE) == [] + + def test_tier_model_check_skipped_when_the_agent_lists_nothing(self): + # A provider-service agent has no enumerable catalog, so there is nothing to check against. + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "model_provider_service": "main.default.anthropic-mps", + "default_model": "claude-sonnet-5", + } + } + }, + "budget_policy": { + "budget_id": "b", + "tiers": [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "claude-sonnet-5", + } + ], + }, + } + assert validate_manifest(manifest, STATE) == [] + def test_budget_policy_alone_still_requires_a_default_agent(self): errors = validate_manifest({"budget_policy": {"budget_id": "b"}}) assert any("default_agent is required" in e for e in errors) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py new file mode 100644 index 0000000..126d630 --- /dev/null +++ b/tests/test_managed_wizard.py @@ -0,0 +1,1402 @@ +"""Tests for the interactive `ucode setup` flow and its CLI wiring. + +The wizard is mostly orchestration, so these focus on the parts where it can silently produce a +wrong manifest: reading tracing/MCP/skills back out of ``state.json``, classifying MCP URLs into +managed-config types, the admin gate, and the per-agent model-config shapes. +""" + +from __future__ import annotations + +import json +from unittest.mock import patch + +import pytest +import typer.main +from typer.testing import CliRunner + +import ucode.config_io as config_io_mod +import ucode.managed_setup as managed_setup_mod +import ucode.managed_wizard as wizard +from ucode.cli import app +from ucode.managed_setup import validate_manifest + +runner = CliRunner() + +WORKSPACE = "https://ws.example.com" + +STATE = { + "workspace": WORKSPACE, + "claude_models": { + "opus": "system.ai.claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-4-6", + }, + "codex_models": ["system.ai.gpt-5-6"], + "gemini_models": ["system.ai.gemini-3-flash"], + "oss_models": ["system.ai.kimi-k2-6"], +} + + +@pytest.fixture(autouse=True) +def _isolate_settings(tmp_path, monkeypatch): + """Point the manifest path at a tmp dir so no test touches the real ~/.ucode.""" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr( + managed_setup_mod, "MANAGED_SETTINGS_PATH", tmp_path / "managed-settings.json" + ) + monkeypatch.setattr(config_io_mod, "_dry_run", False) + + +class TestTracingReadback: + def test_reads_uc_destination(self): + state = {"tracing": {"enabled": True, "uc_destination": "main.default.ucode-traces"}} + assert wizard._tracing_table_from_state(state) == "main.default.ucode-traces" + + def test_disabled_tracing_yields_none(self): + state = {"tracing": {"enabled": False, "uc_destination": "main.default.t"}} + assert wizard._tracing_table_from_state(state) is None + + def test_enabled_without_destination_yields_none(self): + # A non-UC-backed experiment has no table to publish, so the manifest must omit tracing + # rather than carry an empty value the server would reject. + assert wizard._tracing_table_from_state({"tracing": {"enabled": True}}) is None + + def test_missing_tracing_yields_none(self): + assert wizard._tracing_table_from_state({}) is None + + def test_malformed_tracing_yields_none(self): + assert wizard._tracing_table_from_state({"tracing": "on"}) is None + + +class TestMcpUrlClassification: + @pytest.mark.parametrize( + ("url", "expected"), + [ + ("https://ws.example.com/ai-gateway/mcp-services/system.ai.github", "mcp-service"), + ("https://ws.example.com/api/2.0/mcp/external/jira-prod", "external"), + ("https://ws.example.com/api/2.0/mcp/genie/01ef", "genie-space"), + ("https://ws.example.com/api/2.0/mcp/vector-search/main/default", "vector-search"), + ("https://ws.example.com/api/2.0/mcp/functions/main/default", "uc-functions"), + ("https://ws.example.com/api/2.0/mcp/sql", "sql"), + ("https://mcp-myapp-123.aws.databricksapps.com/mcp", "app"), + ], + ) + def test_known_urls(self, url, expected): + assert wizard._mcp_type_for_url(url) == expected + + def test_trailing_slash_is_tolerated(self): + assert wizard._mcp_type_for_url("https://ws.example.com/api/2.0/mcp/sql/") == "sql" + + def test_unknown_url_yields_none(self): + # Better to skip a server than publish it with a guessed type. + assert wizard._mcp_type_for_url("https://example.com/something/else") is None + + def test_sql_is_not_confused_for_app(self): + # Both end in a fixed segment; sql must win since it is checked first. + assert wizard._mcp_type_for_url("https://ws.example.com/api/2.0/mcp/sql") == "sql" + + +class TestMcpServersFromState: + def test_maps_registered_servers_to_name_and_type(self): + state = { + "mcp_servers": [ + { + "name": "databricks-github", + "url": f"{WORKSPACE}/ai-gateway/mcp-services/system.ai.github", + }, + {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, + ] + } + assert wizard._mcp_servers_from_state(state) == [ + {"name": "databricks-github", "type": "mcp-service"}, + {"name": "databricks-sql", "type": "sql"}, + ] + + def test_skips_the_skills_registry_entry(self): + # Skills are published under the manifest's own `skills` field; including the MCP entry too + # would configure them twice. + from ucode.mcp import SKILLS_MCP_KIND + + state = { + "mcp_servers": [ + { + "name": "databricks-skill-registry", + "kind": SKILLS_MCP_KIND, + "url": f"{WORKSPACE}/api/2.0/mcp/sql", + }, + {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, + ] + } + assert wizard._mcp_servers_from_state(state) == [{"name": "databricks-sql", "type": "sql"}] + + def test_skips_unclassifiable_servers(self): + state = {"mcp_servers": [{"name": "mystery", "url": "https://example.com/nope"}]} + assert wizard._mcp_servers_from_state(state) == [] + + def test_skips_entries_missing_name_or_url(self): + state = { + "mcp_servers": [ + {"url": f"{WORKSPACE}/api/2.0/mcp/sql"}, + {"name": "no-url"}, + "not-a-dict", + ] + } + assert wizard._mcp_servers_from_state(state) == [] + + def test_empty_state_yields_nothing(self): + assert wizard._mcp_servers_from_state({}) == [] + + def test_output_validates_as_a_manifest(self): + state = { + "mcp_servers": [ + {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, + ] + } + servers = wizard._mcp_servers_from_state(state) + assert validate_manifest({"mcp_servers": servers}) == [] + + +class TestAdminGate: + def test_non_admin_is_rejected(self): + with patch.object(wizard, "is_workspace_admin", return_value=False): + with pytest.raises(RuntimeError, match="not an admin"): + wizard._require_admin(WORKSPACE, "token") + + def test_admin_passes(self): + with patch.object(wizard, "is_workspace_admin", return_value=True): + wizard._require_admin(WORKSPACE, "token") # must not raise + + def test_unverifiable_check_warns_and_continues(self): + # A failed SCIM call must not block a legitimate admin — the API enforces the same rule. + with ( + patch.object(wizard, "is_workspace_admin", return_value=None), + patch.object(wizard, "print_warning") as warn, + ): + wizard._require_admin(WORKSPACE, "token") + assert warn.called + + +class TestExistingConfigWarning: + @staticmethod + def _warn(existing: dict) -> str: + with ( + patch.object(wizard, "get_managed_config", return_value=(existing, None)), + patch.object(wizard, "print_warning") as warn, + ): + wizard._warn_on_existing_config(WORKSPACE, "token") + assert warn.called + return warn.call_args[0][0] + + def test_warns_that_publishing_replaces_the_whole_config(self): + message = self._warn({"enabled_agents": {"claude": {}, "codex": {}}}) + # There is one config per workspace covering everything, so the warning says that rather + # than reading like a per-agent notice. + assert "one config covers every agent" in message + assert "replaces all of it" in message + assert "everything you want to keep" in message + + def test_warning_does_not_itemize_the_existing_config(self): + # The message is the same whatever the config holds: an inventory doesn't change what the + # admin should do, and `ucode setup show` prints the real thing for comparison. + rich = self._warn( + { + "enabled_agents": {"claude": {}, "opencode": {}, "pi": {}}, + "mcp_servers": [{"name": "a", "type": "sql"}], + "skills": {"names": ["main.default"]}, + "tracing_table": "main.default.traces", + "budget_policy": {"display_name": "lillys_budget", "budget_id": "abc"}, + } + ) + assert rich == self._warn({"enabled_agents": {}}) + for leaked in ("Claude Code", "OpenCode", "lillys_budget", "main.default"): + assert leaked not in rich, leaked + + def test_silent_when_no_config_exists(self): + with ( + patch.object(wizard, "get_managed_config", return_value=(None, None)), + patch.object(wizard, "print_warning") as warn, + ): + wizard._warn_on_existing_config(WORKSPACE, "token") + assert not warn.called + + def test_read_failure_is_a_note_not_a_warning(self): + # Can't check isn't the same as "there is one"; don't imply data loss. + with ( + patch.object(wizard, "get_managed_config", return_value=(None, "HTTP 403 Forbidden")), + patch.object(wizard, "print_warning") as warn, + patch.object(wizard, "print_note") as note, + ): + wizard._warn_on_existing_config(WORKSPACE, "token") + assert not warn.called + assert note.called + + +class TestModelPrompting: + def test_codex_takes_a_single_model(self): + with patch.object(wizard, "prompt_for_selection", return_value="system.ai.gpt-5-6"): + config = wizard._prompt_models_for_agent("codex", STATE, None) + # CodexModelConfig has no model list, so the wizard must not build one. + assert config == {"default_model": "system.ai.gpt-5-6"} + + def test_claude_prompts_one_slot_per_family(self): + # Claude Code selects by family alias, so each `ClaudeDefaultModels` slot gets its own + # prompt — and each shows that family's real alternatives, not just the newest. + candidates = { + "opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], + "sonnet": ["system.ai.claude-sonnet-5"], + } + asked: list[str] = [] + + def fake_sel(prompt, options, **kwargs): + asked.append(prompt) + return [v for v, _ in options][0] + + with ( + patch.object(wizard, "_claude_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + assert config["models"] == { + "default_opus_model": "system.ai.claude-opus-5", + "default_sonnet_model": "system.ai.claude-sonnet-5", + } + assert any("opus" in p for p in asked) and any("sonnet" in p for p in asked) + + def test_claude_offers_every_version_in_a_family(self): + candidates = {"opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"]} + offered: list[list[str]] = [] + + def fake_sel(prompt, options, **kwargs): + values = [v for v, _ in options] + offered.append(values) + return values[1] # pick the older opus on purpose + + with ( + patch.object(wizard, "_claude_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + # Pinning a known-good older version has to be expressible. + assert config["models"] == {"default_opus_model": "system.ai.claude-opus-4-8"} + assert "system.ai.claude-opus-4-8" in offered[0] + + def test_claude_families_can_be_skipped(self): + candidates = { + "opus": ["system.ai.claude-opus-5"], + "sonnet": ["system.ai.claude-sonnet-5"], + } + + def fake_sel(prompt, options, **kwargs): + values = [v for v, _ in options] + return wizard._SKIP_FAMILY if "sonnet" in prompt else values[0] + + with ( + patch.object(wizard, "_claude_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + # Every slot is optional in the proto; a skipped one must simply be absent. + assert config["models"] == {"default_opus_model": "system.ai.claude-opus-5"} + + def test_claude_overall_default_comes_from_the_filled_slots(self): + candidates = { + "opus": ["system.ai.claude-opus-5"], + "sonnet": ["system.ai.claude-sonnet-5"], + } + prompts: list[list[str]] = [] + + def fake_sel(prompt, options, **kwargs): + values = [v for v, _ in options] + prompts.append(values) + return values[0] + + with ( + patch.object(wizard, "_claude_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + # The last prompt is the overall default, offering only what the slots hold — so it can + # never name a model the config doesn't carry. + assert set(prompts[-1]) == {"system.ai.claude-opus-5", "system.ai.claude-sonnet-5"} + assert config["default_model"] in config["models"].values() + + def test_claude_single_slot_skips_the_default_prompt(self): + candidates = {"opus": ["system.ai.claude-opus-5"]} + calls = {"n": 0} + + def fake_sel(prompt, options, **kwargs): + calls["n"] += 1 + return [v for v, _ in options][0] + + with ( + patch.object(wizard, "_claude_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + assert calls["n"] == 1 # only the opus prompt; no redundant default question + assert config["default_model"] == "system.ai.claude-opus-5" + + def test_claude_falls_back_to_text_when_nothing_discovered(self): + with ( + patch.object(wizard, "_claude_candidates", return_value={}), + patch.object(wizard, "prompt_for_text", return_value="some-claude"), + patch.object(wizard, "print_warning"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + assert config == {"default_model": "some-claude"} + + def test_single_slot_announces_the_inferred_default(self): + # The one-option prompt is skipped, but silence reads as a dropped step — the admin has to + # learn that the default is set, and to what. + candidates = {"opus": ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"]} + + def fake_sel(prompt, options, **kwargs): + if prompt.startswith("Default opus"): + return "system.ai.claude-opus-4-8" + return wizard._SKIP_FAMILY + + with ( + patch.object(wizard, "_claude_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "print_note"), + patch.object(wizard, "print_success") as success, + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + assert config["default_model"] == "system.ai.claude-opus-4-8" + assert success.called + assert "system.ai.claude-opus-4-8" in success.call_args[0][0] + + def test_claude_all_families_skipped_still_picks_from_the_candidates(self): + # Skipping every slot is a legitimate minimal config — `models` is optional and each unset + # slot falls back to `default_model`, so one model covers every family. The admin shouldn't + # have to type an id we already have. + candidates = { + "opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], + "sonnet": ["system.ai.claude-sonnet-5"], + } + offered: list[list[str]] = [] + + def fake_sel(prompt, options, **kwargs): + values = [v for v, _ in options] + offered.append(values) + if prompt.startswith("Default "): + return wizard._SKIP_FAMILY + return "system.ai.claude-opus-4-8" + + with ( + patch.object(wizard, "_claude_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "prompt_for_text") as text, + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, None) + + assert config == {"default_model": "system.ai.claude-opus-4-8"} + assert "models" not in config + assert not text.called, "should pick from candidates, not ask for free text" + # The final prompt offers every candidate across all families. + assert set(offered[-1]) == { + "system.ai.claude-opus-5", + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-5", + } + + def test_older_claude_version_passes_validation(self): + # The picker offers every version in a family, but `claude_models` holds only the newest — + # so validation has to learn about the rest or it rejects a legitimate pick. + candidates = {"opus": ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"]} + state = { + "workspace": "https://ws.example.com", + "claude_models": {"opus": "system.ai.claude-opus-5"}, + } + + def fake_unbucketed(workspace, token): + return ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], None + + with ( + patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), + patch.object(wizard, "discover_claude_models_unbucketed", fake_unbucketed), + patch.object(wizard, "claude_family_candidates", return_value=candidates), + patch.object(wizard, "prompt_for_selection", return_value="system.ai.claude-opus-4-8"), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", state, None) + + manifest = { + "default_agent": "claude", + "enabled_agents": {"claude": {"model_config": config}}, + } + assert validate_manifest(manifest, state) == [] + + def test_claude_cancelled_family_prompt_aborts(self): + with ( + patch.object( + wizard, "_claude_candidates", return_value={"opus": ["system.ai.claude-opus-5"]} + ), + patch.object(wizard, "prompt_for_selection", return_value=None), + patch.object(wizard, "print_note"), + ): + with pytest.raises(KeyboardInterrupt): + wizard._prompt_models_for_agent("claude", STATE, None) + + def test_single_model_agents_get_one_prompt(self): + # Gemini and Copilot declare `repeated string models` in the proto, but their config writers + # take one model and write one env var — a published list would be read by nothing. + for tool in ("codex", "gemini", "copilot"): + options = wizard.model_options_for_agent(tool, STATE) + with ( + patch.object(wizard, "prompt_for_selection", return_value=options[0]) as select, + patch.object(wizard, "prompt_for_multi_selection") as multi, + ): + config = wizard._prompt_models_for_agent(tool, STATE, None) + assert select.called, tool + assert not multi.called, tool + assert config == {"default_model": options[0]}, tool + assert "models" not in config, tool + + def test_claude_catalog_is_fetched_once(self): + # `configure_shared_state` already paged the whole catalog; re-fetching per claude prompt + # pages it again for no new information. + calls = {"n": 0} + + def fake_fetch(workspace, token): + calls["n"] += 1 + return ["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], None + + state = {"workspace": "https://ws.example.com", "profile": "p"} + with ( + patch.object(wizard, "discover_claude_models_unbucketed", side_effect=fake_fetch), + patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), + patch.object(wizard, "prompt_for_selection", return_value="system.ai.claude-opus-5"), + patch.object(wizard, "print_note"), + patch.object(wizard, "print_success"), + ): + wizard._prompt_models_for_agent("claude", state, None) + wizard._prompt_models_for_agent("claude", state, None) + + assert calls["n"] == 1 + assert state["all_claude_models"] + + def test_nothing_is_prechecked(self): + # The first option is whatever discovery sorted first, not a recommendation — for pi it is a + # Claude model, for codex the oldest GPT. Pre-checking it made "hit Enter" produce an + # arbitrary config. + captured: dict = {} + + def fake_multi(prompt, options, preselected=None, **kwargs): + captured["preselected"] = preselected + return [v for v, _ in options][:1] + + with ( + patch.object(wizard, "prompt_for_multi_selection", side_effect=fake_multi), + patch.object(wizard, "prompt_for_selection", return_value="x"), + ): + wizard._prompt_models_for_agent("pi", STATE, None) + + assert not captured["preselected"] + + def test_list_agents_still_multi_select(self): + # OpenCode and Pi really do show a model picker, so their lists are honoured. + for tool in ("opencode", "pi"): + options = wizard.model_options_for_agent(tool, STATE) + with ( + patch.object(wizard, "prompt_for_multi_selection", return_value=options[:2]), + patch.object(wizard, "prompt_for_selection", return_value=options[0]), + ): + config = wizard._prompt_models_for_agent(tool, STATE, None) + assert config["models"] == options[:2], tool + + def test_flat_list_agents_keep_the_picked_list(self): + picked = ["system.ai.claude-opus-4-8", "system.ai.kimi-k2-6"] + with ( + patch.object(wizard, "prompt_for_multi_selection", return_value=picked), + patch.object(wizard, "prompt_for_selection", return_value=picked[0]), + ): + config = wizard._prompt_models_for_agent("opencode", STATE, None) + assert config["models"] == picked + + def test_single_pick_skips_the_default_prompt(self): + with ( + patch.object( + wizard, "prompt_for_multi_selection", return_value=["system.ai.claude-opus-4-8"] + ), + patch.object(wizard, "prompt_for_selection") as select, + ): + config = wizard._prompt_models_for_agent("pi", STATE, None) + assert config["default_model"] == "system.ai.claude-opus-4-8" + assert not select.called + + def test_provider_service_offers_its_targets(self): + # The service's own targets are the model vocabulary the manifest must use, so the admin + # picks from them rather than typing an id from memory. + service = { + "name": "main.default.anthropic-mps", + "provider_type": "anthropic", + "targets": ["claude-sonnet-4-6", "claude-opus-4-8"], + "allow_all_targets": False, + } + with ( + patch.object(wizard, "prompt_for_selection", return_value="claude-opus-4-8") as select, + patch.object(wizard, "prompt_for_text") as text, + ): + config = wizard._prompt_models_for_agent("claude", STATE, service) + assert config == { + "model_provider_service": "main.default.anthropic-mps", + "default_model": "claude-opus-4-8", + } + assert not text.called, "should not fall back to free text when targets are known" + # Offered sorted, so the picker order is stable run to run. + assert [value for value, _ in select.call_args[0][1]] == [ + "claude-opus-4-8", + "claude-sonnet-4-6", + ] + + def test_provider_service_falls_back_to_text_when_targets_unknown(self): + # allow_all_targets passes the provider's whole catalog through; there is nothing to list. + service = { + "name": "main.default.anthropic-mps", + "provider_type": "anthropic", + "targets": [], + "allow_all_targets": True, + } + with ( + patch.object(wizard, "prompt_for_text", return_value="claude-sonnet-5"), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, service) + assert config == { + "model_provider_service": "main.default.anthropic-mps", + "default_model": "claude-sonnet-5", + } + + def test_relayed_service_falls_back_to_text(self): + # A relayed Anthropic subscription service routes by canonical name, with no target list. + service = { + "name": "main.default.lilly-anthropic", + "provider_type": "anthropic", + "targets": [], + "allow_all_targets": False, + "relayed": True, + } + with ( + patch.object(wizard, "prompt_for_text", return_value="claude-sonnet-4-6"), + patch.object(wizard, "print_note"), + ): + config = wizard._prompt_models_for_agent("claude", STATE, service) + assert config["default_model"] == "claude-sonnet-4-6" + + def test_falls_back_to_free_text_when_nothing_discovered(self): + with ( + patch.object(wizard, "prompt_for_text", return_value="some-model"), + patch.object(wizard, "print_warning"), + ): + config = wizard._prompt_models_for_agent("pi", {}, None) + assert config == {"default_model": "some-model"} + + def test_empty_selection_is_re_prompted(self): + # An agent with no default_model can't be the config's default_agent (the server rejects + # it) and gives developers nothing to launch, so "none" is re-asked rather than accepted. + with ( + patch.object( + wizard, + "prompt_for_multi_selection", + side_effect=[[], ["system.ai.claude-opus-4-8"]], + ) as picker, + patch.object(wizard, "print_err") as err, + ): + config = wizard._prompt_models_for_agent("pi", STATE, None) + assert picker.call_count == 2 + assert err.called + assert config["default_model"] == "system.ai.claude-opus-4-8" + + def test_every_agent_always_gets_a_default_model(self): + # The invariant the late-validation bug violated: no agent can come back model-less. + for tool in ("claude", "codex", "gemini", "opencode", "pi", "copilot"): + options = wizard.model_options_for_agent(tool, STATE) + with ( + patch.object(wizard, "prompt_for_multi_selection", return_value=[options[0]]), + patch.object(wizard, "prompt_for_selection", return_value=options[0]), + # Without this the claude pass reaches for the real catalog, which shells out to + # `databricks auth token` and depends on the machine's CLI and credentials. + patch.object(wizard, "discover_claude_models_unbucketed", return_value=([], None)), + patch.object(wizard, "get_databricks_token", lambda *a, **k: "tok"), + patch.object(wizard, "print_note"), + patch.object(wizard, "print_success"), + ): + config = wizard._prompt_models_for_agent(tool, STATE, None) + assert config.get("default_model"), tool + + def test_claude_candidates_survive_a_missing_databricks_cli(self): + # `get_databricks_token` shells out, so a machine without the CLI on PATH raises + # FileNotFoundError, not RuntimeError. That must degrade to the bucketed per-family picks + # rather than aborting the wizard mid-flow. + def no_cli(*args, **kwargs): + raise FileNotFoundError(2, "No such file or directory", "databricks") + + with patch.object(wizard, "get_databricks_token", side_effect=no_cli): + candidates = wizard._claude_candidates(dict(STATE)) + + # STATE's bucketed `claude_models` still supplies one model per family. + assert candidates["opus"] == ["system.ai.claude-opus-4-8"] + assert candidates["sonnet"] == ["system.ai.claude-sonnet-4-6"] + + def test_default_model_is_a_bare_uc_id(self): + # Provider prefixes (e.g. opencode's `databricks-anthropic/`) are added by each agent's own + # writer, so the manifest stays agent-neutral. + with ( + patch.object( + wizard, "prompt_for_multi_selection", return_value=["system.ai.claude-opus-4-8"] + ), + ): + config = wizard._prompt_models_for_agent("opencode", STATE, None) + assert config["default_model"] == "system.ai.claude-opus-4-8" + assert "/" not in config["default_model"] + + def test_dismissed_single_select_aborts_instead_of_re_prompting(self): + # This used to re-prompt. questionary's `ask` swallows Ctrl-C and returns None, so "empty + # submission" and "user aborted" are the same value here — and re-asking spun forever. + # Asserted on the helper directly: the codex path only reaches the picker when discovery + # found models, and falling through to the free-text branch would read real stdin. + with ( + patch.object(wizard, "prompt_for_selection", return_value=None) as picker, + patch.object(wizard, "print_err"), + ): + with pytest.raises(KeyboardInterrupt): + wizard._require_selection("Select the model:", [("a", "A")]) + assert picker.call_count == 1 + + def test_empty_free_text_is_re_prompted(self): + with ( + patch.object(wizard, "prompt_for_text", side_effect=[None, "some-model"]) as text, + patch.object(wizard, "print_err"), + patch.object(wizard, "print_warning"), + ): + config = wizard._prompt_models_for_agent("pi", {}, None) + assert text.call_count == 2 + assert config == {"default_model": "some-model"} + + def test_cancelled_picker_aborts(self): + with patch.object(wizard, "prompt_for_multi_selection", return_value=None): + with pytest.raises(KeyboardInterrupt): + wizard._prompt_models_for_agent("pi", STATE, None) + + +ANTHROPIC_SERVICE = { + "name": "main.default.lilly-anthropic", + "provider_type": "anthropic", + "targets": ["claude-sonnet-4-6"], + "allow_all_targets": False, + "relayed": False, +} +OPENAI_SERVICE = { + "name": "main.default.openai-mps", + "provider_type": "openai", + "targets": ["gpt-5-6"], + "allow_all_targets": False, + "relayed": False, +} + + +class TestProviderServiceSpinner: + """The MPS listing is cached per workspace, so only the first agent's lookup does any I/O.""" + + SERVICES = [ + { + "name": "main.j.ant", + "provider_type": "anthropic", + "targets": ["claude-opus-5"], + "allow_all_targets": False, + "relayed": False, + } + ] + + def test_spinner_shows_once_not_once_per_agent(self): + # The reported symptom: "Checking for model provider services for ..." appeared for + # every configured agent even though the listing had already been fetched. + spins: list[str] = [] + cached = {"yes": False} + + def fake_list(workspace, token, **kwargs): + cached["yes"] = True + return list(self.SERVICES), None + + def fake_spinner(message): + spins.append(message) + from contextlib import nullcontext + + return nullcontext() + + with ( + patch.object(wizard, "list_model_provider_services", side_effect=fake_list), + patch.object( + wizard, "has_cached_model_provider_services", side_effect=lambda ws: cached["yes"] + ), + patch.object(wizard, "spinner", side_effect=fake_spinner), + patch.object(wizard, "prompt_for_selection", return_value="databricks"), + ): + wizard._select_provider_service("claude", WORKSPACE, "tok") + wizard._select_provider_service("codex", WORKSPACE, "tok") + + listing_spins = [m for m in spins if "provider service" in m] + assert len(listing_spins) == 1, listing_spins + # And it doesn't name an agent, since one lookup covers them all. + assert "Claude Code" not in listing_spins[0] + + +class TestProviderServiceSelection: + def test_agents_without_provider_support_skip_the_prompt(self): + with patch.object(wizard, "list_model_provider_services") as listing: + assert wizard._select_provider_service("opencode", WORKSPACE, "token") is None + assert not listing.called + + def test_feature_disabled_is_silent(self): + # The common case on most workspaces; a warning here would be noise. + with ( + patch.object(wizard, "list_model_provider_services", return_value=([], "HTTP 404")), + patch.object(wizard, "is_model_provider_feature_unavailable", return_value=True), + patch.object(wizard, "print_warning") as warn, + ): + assert wizard._select_provider_service("claude", WORKSPACE, "token") is None + assert not warn.called + + def test_unexpected_listing_failure_warns(self): + # Without this the admin silently loses the MPS option with no idea why. + with ( + patch.object( + wizard, "list_model_provider_services", return_value=([], "HTTP 403 Forbidden") + ), + patch.object(wizard, "is_model_provider_feature_unavailable", return_value=False), + patch.object(wizard, "print_warning") as warn, + patch.object(wizard, "print_note"), + ): + assert wizard._select_provider_service("claude", WORKSPACE, "token") is None + assert warn.called + assert "403" in warn.call_args[0][0] + + def test_services_exist_but_none_match_the_agent_explains_why(self): + # An openai-only workspace offers claude nothing; say so rather than showing no picker. + with ( + patch.object( + wizard, "list_model_provider_services", return_value=([OPENAI_SERVICE], None) + ), + patch.object(wizard, "print_note") as note, + ): + assert wizard._select_provider_service("claude", WORKSPACE, "token") is None + assert note.called + assert "API dialect" in note.call_args[0][0] + + def test_choosing_databricks_returns_none(self): + with ( + patch.object( + wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) + ), + patch.object(wizard, "prompt_for_selection", return_value="databricks"), + ): + assert wizard._select_provider_service("claude", WORKSPACE, "token") is None + + def test_choosing_mps_returns_the_whole_service(self): + # The dict (not just the name) is returned so the model prompt can offer its targets. + with ( + patch.object( + wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) + ), + patch.object( + wizard, + "prompt_for_selection", + side_effect=["mps", "main.default.lilly-anthropic"], + ), + ): + service = wizard._select_provider_service("claude", WORKSPACE, "token") + assert service == ANTHROPIC_SERVICE + + def test_only_matching_services_are_offered(self): + with ( + patch.object( + wizard, + "list_model_provider_services", + return_value=([ANTHROPIC_SERVICE, OPENAI_SERVICE], None), + ), + patch.object( + wizard, + "prompt_for_selection", + side_effect=["mps", "main.default.lilly-anthropic"], + ) as select, + ): + wizard._select_provider_service("claude", WORKSPACE, "token") + offered = [value for value, _ in select.call_args_list[1][0][1]] + assert offered == ["main.default.lilly-anthropic"] + + def test_cancelling_the_service_picker_returns_none(self): + with ( + patch.object( + wizard, "list_model_provider_services", return_value=([ANTHROPIC_SERVICE], None) + ), + patch.object(wizard, "prompt_for_selection", side_effect=["mps", None]), + ): + assert wizard._select_provider_service("claude", WORKSPACE, "token") is None + + +class TestProviderServiceModelOptions: + def test_returns_sorted_targets(self): + service = {"targets": ["b-model", "a-model"], "allow_all_targets": False} + assert wizard.provider_service_model_options(service) == ["a-model", "b-model"] + + def test_deduplicates(self): + service = {"targets": ["m", "m"], "allow_all_targets": False} + assert wizard.provider_service_model_options(service) == ["m"] + + def test_allow_all_targets_yields_nothing(self): + service = {"targets": ["m"], "allow_all_targets": True} + assert wizard.provider_service_model_options(service) == [] + + def test_missing_targets_yields_nothing(self): + assert wizard.provider_service_model_options({}) == [] + + def test_malformed_targets_yield_nothing(self): + assert wizard.provider_service_model_options({"targets": "m"}) == [] + + +# Agents as the wizard configures them: the tier picker must offer these, not the workspace catalog. +CLAUDE_ONLY = {"claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}} + + +class TestBudgetPolicy: + def test_declining_yields_none(self): + with patch.object(wizard, "prompt_yes_no_default", return_value=False): + assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None + + def test_no_budgets_warns_and_yields_none(self): + with ( + patch.object(wizard, "prompt_yes_no_default", return_value=True), + patch.object(wizard, "list_workspace_budgets", return_value=([], "none found")), + patch.object(wizard, "print_warning") as warn, + ): + assert wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) is None + assert warn.called + + def test_percentages_are_stored_as_fractions(self): + budgets = [{"id": "budget-1", "display_name": "eng"}] + with ( + patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), + patch.object( + wizard, + "prompt_for_selection", + side_effect=["budget-1", "claude", "system.ai.claude-opus-4-8"], + ), + patch.object(wizard, "prompt_for_text", return_value="tiered"), + # prompt_for_percentage already converts; it returns the fraction. + patch.object(wizard, "prompt_for_percentage", return_value=0.8), + ): + policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) + assert policy is not None + assert policy["budget_id"] == "budget-1" + assert policy["tiers"] == [ + { + "spending_percentage": 0.8, + "default_agent": "claude", + "default_model": "system.ai.claude-opus-4-8", + } + ] + + def test_offers_only_the_models_the_agent_was_configured_with(self): + # Pi's catalog spans every family, so offering the workspace catalog would present four + # models it was never given — and a tier naming one of them silently misroutes developers. + enabled = { + "pi": { + "model_config": { + "default_model": "system.ai.kimi-k2-6", + "models": ["system.ai.kimi-k2-6"], + } + } + } + budgets = [{"id": "budget-1", "display_name": "eng"}] + with ( + patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), + patch.object( + wizard, + "prompt_for_selection", + side_effect=["budget-1", "pi", "system.ai.kimi-k2-6"], + ) as select, + patch.object(wizard, "prompt_for_text", return_value="tiered"), + patch.object(wizard, "prompt_for_percentage", return_value=0.8), + ): + wizard._prompt_budget_policy(WORKSPACE, "token", enabled, STATE) + # Third call is the model picker. + offered = [value for value, _ in select.call_args_list[2][0][1]] + assert offered == ["system.ai.kimi-k2-6"] + + def test_claude_family_slots_are_flattened_for_the_picker(self): + enabled = { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-4-8", + "models": { + "default_opus_model": "system.ai.claude-opus-4-8", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + }, + } + } + } + budgets = [{"id": "budget-1", "display_name": "eng"}] + with ( + patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), + patch.object( + wizard, + "prompt_for_selection", + side_effect=["budget-1", "claude", "system.ai.claude-opus-4-8"], + ) as select, + patch.object(wizard, "prompt_for_text", return_value="tiered"), + patch.object(wizard, "prompt_for_percentage", return_value=0.8), + ): + wizard._prompt_budget_policy(WORKSPACE, "token", enabled, STATE) + offered = [value for value, _ in select.call_args_list[2][0][1]] + assert set(offered) == {"system.ai.claude-opus-4-8", "system.ai.claude-sonnet-4-6"} + + def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): + # An agent configured through a provider service has no enumerable list; better to offer the + # catalog than nothing at all. + budgets = [{"id": "budget-1", "display_name": "eng"}] + with ( + patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), + patch.object( + wizard, + "prompt_for_selection", + side_effect=["budget-1", "gemini", "system.ai.gemini-3-flash"], + ) as select, + patch.object(wizard, "prompt_for_text", return_value="tiered"), + patch.object(wizard, "prompt_for_percentage", return_value=0.8), + ): + wizard._prompt_budget_policy(WORKSPACE, "token", {"gemini": {}}, STATE) + offered = [value for value, _ in select.call_args_list[2][0][1]] + assert offered == ["system.ai.gemini-3-flash"] + + def test_authored_policy_validates(self): + budgets = [{"id": "budget-1", "display_name": "eng"}] + with ( + patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), + patch.object( + wizard, + "prompt_for_selection", + side_effect=["budget-1", "claude", "system.ai.claude-opus-4-8"], + ), + patch.object(wizard, "prompt_for_text", return_value="tiered"), + patch.object(wizard, "prompt_for_percentage", return_value=0.8), + ): + policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) + manifest = { + "default_agent": "claude", + "enabled_agents": CLAUDE_ONLY, + "budget_policy": policy, + } + assert validate_manifest(manifest, STATE) == [] + + +class TestConfiguredModelsForAgent: + def test_flat_list_plus_default(self): + agent = {"model_config": {"default_model": "b", "models": ["a", "b"]}} + assert wizard.configured_models_for_agent(agent) == ["a", "b"] + + def test_claude_slots_are_flattened(self): + agent = { + "model_config": { + "default_model": "opus", + "models": {"default_opus_model": "opus", "default_sonnet_model": "sonnet"}, + } + } + assert set(wizard.configured_models_for_agent(agent)) == {"opus", "sonnet"} + + def test_codex_has_only_a_default(self): + # CodexModelConfig carries no model list, so the default is the whole set. + assert wizard.configured_models_for_agent({"model_config": {"default_model": "gpt-5"}}) == [ + "gpt-5" + ] + + def test_no_model_config_yields_nothing(self): + assert wizard.configured_models_for_agent({}) == [] + + +class TestSummary: + def test_lists_claude_family_slots(self, capsys): + # The one-line default hides which families were configured, which is most of the choice. + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-4-8", + "models": { + "default_opus_model": "system.ai.claude-opus-4-8", + "default_haiku_model": "system.ai.claude-haiku-4-5", + }, + } + } + }, + } + wizard._render_summary(WORKSPACE, manifest) + out = capsys.readouterr().out + assert "opus" in out and "haiku" in out + assert "system.ai.claude-haiku-4-5" in out + + def test_lists_a_multi_model_agents_models(self, capsys): + manifest = { + "default_agent": "pi", + "enabled_agents": { + "pi": { + "model_config": { + "default_model": "system.ai.kimi-k2-6", + "models": ["system.ai.kimi-k2-6", "system.ai.gpt-5-6"], + } + } + }, + } + wizard._render_summary(WORKSPACE, manifest) + assert "system.ai.gpt-5-6" in capsys.readouterr().out + + def test_single_model_agent_needs_no_extra_line(self, capsys): + manifest = { + "default_agent": "gemini", + "enabled_agents": { + "gemini": {"model_config": {"default_model": "system.ai.gemini-3-flash"}} + }, + } + wizard._render_summary(WORKSPACE, manifest) + out = capsys.readouterr().out + assert "system.ai.gemini-3-flash" in out + assert "models:" not in out + + +class TestSetupFromFile: + def _write(self, tmp_path, payload): + path = tmp_path / "manifest.json" + path.write_text(json.dumps(payload), encoding="utf-8") + return path + + def _valid(self): + return { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} + }, + } + + def test_valid_manifest_is_saved(self, tmp_path): + path = self._write(tmp_path, self._valid()) + with patch.object(wizard, "load_state", return_value=STATE): + assert wizard.setup_from_file(str(path)) == 0 + assert managed_setup_mod.load_managed_settings(WORKSPACE) == self._valid() + + def test_invalid_manifest_returns_1_and_saves_nothing(self, tmp_path): + path = self._write(tmp_path, {"enabled_agents": {"claude": {}}}) + with patch.object(wizard, "load_state", return_value=STATE): + assert wizard.setup_from_file(str(path)) == 1 + assert managed_setup_mod.load_managed_settings(WORKSPACE) is None + + def test_missing_file_is_actionable(self, tmp_path): + with patch.object(wizard, "load_state", return_value=STATE): + with pytest.raises(RuntimeError, match="Could not read manifest file"): + wizard.setup_from_file(str(tmp_path / "nope.json")) + + def test_malformed_json_names_the_line(self, tmp_path): + path = tmp_path / "bad.json" + path.write_text("{oops", encoding="utf-8") + with patch.object(wizard, "load_state", return_value=STATE): + with pytest.raises(RuntimeError, match="not valid JSON"): + wizard.setup_from_file(str(path)) + + def test_non_object_json_is_rejected(self, tmp_path): + path = self._write(tmp_path, ["not", "an", "object"]) + with patch.object(wizard, "load_state", return_value=STATE): + with pytest.raises(RuntimeError, match="must contain a JSON object"): + wizard.setup_from_file(str(path)) + + def test_unconfigured_workspace_is_actionable(self, tmp_path): + path = self._write(tmp_path, self._valid()) + with patch.object(wizard, "load_state", return_value={}): + with pytest.raises(RuntimeError, match="No workspace is configured"): + wizard.setup_from_file(str(path)) + + +class TestShowCommand: + def test_reports_nothing_when_unauthored(self): + with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): + assert wizard.show_command() == 0 + + def test_prints_the_apply_payload(self, capsys): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} + }, + } + managed_setup_mod.save_managed_settings(WORKSPACE, manifest) + with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): + assert wizard.show_command() == 0 + out = capsys.readouterr().out + # The proto enum spelling is what `apply` sends, so it must appear verbatim. + assert "CODING_AGENT_CLAUDE_CODE" in out + + +class TestSummaryPanel: + def test_summary_is_boxed(self, capsys): + # The summary is the one block an admin reads as a whole to check against what they + # intended, and it lands after a long flow of prompts — so it gets a box rather than loose + # lines that blend into the preceding output. + wizard._render_summary( + WORKSPACE, + { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} + }, + }, + ) + out = capsys.readouterr().out + assert "Configuration summary" in out + # Rich box-drawing characters: the panel border. + assert "╭" in out and "╰" in out + assert "system.ai.claude-opus-5" in out + + def test_a_bracketed_policy_name_survives_the_summary(self, capsys): + # Rich reads bracketed text as a style tag and renders nothing for it, so an unescaped + # `[prod] tiered routing` displayed as `tiered routing` — in the block whose whole purpose + # is confirming what the admin is about to publish workspace-wide. + wizard._render_summary( + WORKSPACE, + { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} + }, + "budget_policy": { + "budget_id": "19165ea4-ff8d-4fbb-b6ce-fc5abe7e1c57", + "display_name": "[prod] tiered routing", + "tiers": [], + }, + }, + ) + assert "[prod] tiered routing" in capsys.readouterr().out + + +class TestCancelledPromptsAbort: + """A dismissed prompt must abort, not re-ask an input that can't answer.""" + + def test_require_selection_aborts_when_the_picker_is_dismissed(self): + # questionary's Question.ask catches KeyboardInterrupt and returns None (v2.1.1), so Ctrl-C + # is indistinguishable from an empty submission here. Re-asking looped forever. + with patch.object(wizard, "prompt_for_selection", return_value=None) as sel: + with pytest.raises(KeyboardInterrupt): + wizard._require_selection("pick", [("a", "A")]) + assert sel.call_count == 1 + + def test_require_text_asks_for_a_required_answer(self): + # `required=True` is what makes closed stdin raise instead of returning None; without it a + # piped/CI run spins re-asking an exhausted stream. + with patch.object(wizard, "prompt_for_text", return_value="m") as text: + assert wizard._require_text("Default model") == "m" + assert text.call_args.kwargs.get("required") is True + + def test_require_text_aborts_on_closed_stdin(self): + with patch("ucode.ui.console.input", side_effect=EOFError): + with pytest.raises(KeyboardInterrupt): + wizard._require_text("Default model") + + +class TestClaudeCandidatesStayValidatable: + """Whatever the Claude prompts offer, `validate_manifest` must accept.""" + + def _manifest(self, model: str) -> dict: + return { + "default_agent": "claude", + "enabled_agents": {"claude": {"model_config": {"default_model": model}}}, + } + + def test_listing_path_caches_so_older_versions_validate(self): + # The unbucketed listing widens the candidates past `claude_models`, so it must also cache + # them — otherwise picking an older Opus is rejected at the end of the flow. + state = { + "workspace": "https://ws.example.com", + "claude_models": {"opus": "system.ai.claude-opus-5"}, + } + with ( + patch.object(wizard, "get_databricks_token", return_value="t"), + patch.object( + wizard, + "discover_claude_models_unbucketed", + return_value=(["system.ai.claude-opus-5", "system.ai.claude-opus-4-8"], None), + ), + ): + candidates = wizard._claude_candidates(state) + offered = [m for models in candidates.values() for m in models] + assert "system.ai.claude-opus-4-8" in offered + for model in offered: + assert validate_manifest(self._manifest(model), state) == [] + + def test_fallback_path_only_offers_what_already_validates(self): + # The fallback caches nothing, so it must not offer anything beyond `claude_models`. + state = { + "workspace": "https://ws.example.com", + "claude_models": {"opus": "system.ai.claude-opus-5"}, + } + with ( + patch.object(wizard, "get_databricks_token", return_value="t"), + patch.object( + wizard, "discover_claude_models_unbucketed", side_effect=RuntimeError("boom") + ), + ): + candidates = wizard._claude_candidates(state) + assert "all_claude_models" not in state + for models in candidates.values(): + for model in models: + assert validate_manifest(self._manifest(model), state) == [] + + +class TestSearchablePickers: + """Long lists (models, provider services, budgets) filter as you type.""" + + def test_model_pickers_are_searchable(self): + seen: list[dict] = [] + + def fake_multi(prompt, options, preselected=None, **kwargs): + seen.append(kwargs) + return [options[0][0]] + + with patch.object(wizard, "prompt_for_multi_selection", side_effect=fake_multi): + wizard._require_multi_selection("pick", [("a", "a"), ("b", "b")]) + assert seen[0].get("searchable") is True + + def test_single_select_pickers_are_searchable(self): + seen: list[dict] = [] + + def fake_sel(prompt, options, **kwargs): + seen.append(kwargs) + return options[0][0] + + with patch.object(wizard, "prompt_for_selection", side_effect=fake_sel): + wizard._require_selection("pick", [("a", "a"), ("b", "b")]) + assert seen[0].get("searchable") is True + + def test_budget_and_tier_pickers_are_searchable(self): + budgets = [{"id": "budget-1", "display_name": "eng"}] + searchable_prompts: list[str] = [] + + def fake_sel(prompt, options, **kwargs): + if kwargs.get("searchable"): + searchable_prompts.append(prompt) + if "budget" in prompt: + return "budget-1" + if "agent" in prompt: + return "claude" + return "system.ai.claude-opus-4-8" + + with ( + patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), + patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), + patch.object(wizard, "prompt_for_selection", side_effect=fake_sel), + patch.object(wizard, "prompt_for_text", return_value="tiered"), + patch.object(wizard, "prompt_for_percentage", return_value=0.8), + ): + wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) + # Both the budget list and the tier's model list filter as you type. + assert any("budget" in p for p in searchable_prompts), searchable_prompts + assert any("model" in p for p in searchable_prompts), searchable_prompts + + +class TestCliWiring: + def test_setup_is_registered(self): + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0 + assert "setup" in result.output + + def test_setup_help_lists_from_file(self): + # Assert on the declared option rather than the rendered help text: Rich ellipsizes option + # names to fit the terminal ("--fro…" below ~40 columns), and CI runners report no width, so + # grepping `--from-file` out of the output fails there while passing on a wide local one. + group = typer.main.get_command(app).commands["setup"] # type: ignore[attr-defined] + declared = {opt for param in group.params for opt in param.opts} + assert "--from-file" in declared + result = runner.invoke(app, ["setup", "--help"]) + assert result.exit_code == 0 + + def test_setup_show_is_registered(self): + result = runner.invoke(app, ["setup", "--help"]) + assert result.exit_code == 0 + assert "show" in result.output + + def test_successful_setup_exits_zero(self): + # `typer.Exit` subclasses RuntimeError, so a success code must not be caught and reported + # as an error by the command's own RuntimeError handler. + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_command", return_value=0) as setup, + ): + result = runner.invoke(app, ["setup"]) + assert result.exit_code == 0 + assert setup.called + assert "ERROR" not in _out(result) + + def test_nonzero_setup_propagates(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_command", return_value=1), + ): + result = runner.invoke(app, ["setup"]) + assert result.exit_code == 1 + + def test_runtime_error_is_reported_and_exits_1(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_command", side_effect=RuntimeError("you are not an admin")), + ): + result = runner.invoke(app, ["setup"]) + assert result.exit_code == 1 + assert "not an admin" in _out(result) + + def test_interrupt_exits_130(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_command", side_effect=KeyboardInterrupt), + ): + result = runner.invoke(app, ["setup"]) + assert result.exit_code == 130 + + def test_from_file_is_forwarded(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_command", return_value=0) as setup, + ): + runner.invoke(app, ["setup", "--from-file", "/tmp/x.json"]) + assert setup.call_args.kwargs["from_file"] == "/tmp/x.json" + + def test_dry_run_sets_the_flag(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.setup_command", return_value=0), + patch("ucode.cli.set_dry_run") as set_flag, + ): + runner.invoke(app, ["setup", "--dry-run"]) + set_flag.assert_called_once_with(True) + + def test_show_exits_zero(self): + with patch("ucode.cli.show_command", return_value=0): + result = runner.invoke(app, ["setup", "show"]) + assert result.exit_code == 0 + + +def _out(result) -> str: + """CliRunner output with stderr folded in, since print_err writes to a stderr console.""" + return result.output + (result.stderr if result.stderr_bytes else "") diff --git a/tests/test_ui.py b/tests/test_ui.py index 3c0fbc8..bb2c9c5 100644 --- a/tests/test_ui.py +++ b/tests/test_ui.py @@ -2,15 +2,19 @@ from __future__ import annotations +import io from datetime import timedelta from unittest.mock import patch import pytest +from rich.console import Console from ucode.ui import ( format_duration, format_token_count, normalize_workspace_url, + prompt_for_percentage, + prompt_for_text, prompt_for_workspace, prompt_yes_no_default, render_box_table, @@ -50,6 +54,73 @@ def test_explicit_yes_overrides_default_false(self, monkeypatch): assert prompt_yes_no_default("go?", default=False) is True +def _visible(markup: str) -> str: + """What the user actually sees, with Rich markup resolved. + + Asserting on the raw markup string is what let a swallowed default ship: `[tiered]` is present + in the markup and absent from the output, because Rich reads it as a style tag. + """ + console = Console(file=io.StringIO(), force_terminal=False, width=200) + console.print(markup) + return console.file.getvalue().rstrip() + + +class TestDefaultsAreLabelledAsAcceptable: + """A shown default must say that enter takes it, or it reads as a format example.""" + + def test_text_default_says_enter_accepts_it(self): + with patch("ucode.ui.console.input", return_value="") as inp: + assert prompt_for_text("Policy name", default="tiered") == "tiered" + rendered = _visible(inp.call_args[0][0]) + assert "[tiered]" in rendered + assert "enter to accept" in rendered + + def test_a_word_like_default_is_not_eaten_as_markup(self): + # Rich treats `[coding-agents-tiered-routing]` as a style tag and renders nothing for it, so + # the real wizard default vanished from the prompt while `[80]` survived. + with patch("ucode.ui.console.input", return_value="") as inp: + prompt_for_text("Policy name", default="coding-agents-tiered-routing") + assert "[coding-agents-tiered-routing]" in _visible(inp.call_args[0][0]) + + def test_a_dotted_default_is_not_eaten_as_markup(self): + with patch("ucode.ui.console.input", return_value="") as inp: + prompt_for_text("Skills location", default="main.default") + assert "[main.default]" in _visible(inp.call_args[0][0]) + + def test_percentage_default_says_enter_accepts_it(self): + with patch("ucode.ui.console.input", return_value="") as inp: + assert prompt_for_percentage("at what percent?", default=0.8) == 0.8 + rendered = _visible(inp.call_args[0][0]) + # Prompted in percent even though the API takes a fraction. + assert "[80]" in rendered + assert "enter to accept" in rendered + + def test_no_default_shows_no_hint(self): + with patch("ucode.ui.console.input", return_value="typed") as inp: + assert prompt_for_text("Model") == "typed" + assert "enter to accept" not in _visible(inp.call_args[0][0]) + + def test_typing_still_overrides_the_default(self): + with patch("ucode.ui.console.input", return_value="mine"): + assert prompt_for_text("Policy name", default="tiered") == "mine" + + +class TestClosedStdinAborts: + """Ctrl-D must reach the CLI as an abort, not as a traceback.""" + + def test_percentage_without_a_default_raises_keyboard_interrupt(self): + # `ucode setup`'s tier prompt passes no default. EOFError has no handler above this call — + # the setup command catches only RuntimeError and KeyboardInterrupt — so a bare EOFError + # reached the admin as a raw traceback. + with patch("ucode.ui.console.input", side_effect=EOFError): + with pytest.raises(KeyboardInterrupt): + prompt_for_percentage("Tier 1: activates at what percent of budget?") + + def test_percentage_with_a_default_still_takes_it(self): + with patch("ucode.ui.console.input", side_effect=EOFError): + assert prompt_for_percentage("at what percent?", default=0.8) == 0.8 + + class TestNormalizeWorkspaceUrl: def test_adds_https_when_missing(self): assert normalize_workspace_url("example.databricks.com") == "https://example.databricks.com"