Fix model selection for multi-agent configuration

README.md

@@ -45,7 +45,7 @@ To configure all tools at once:
 ucode configure
 ```
 
-To configure specific tools without the picker, pass a comma-separated list:
+To configure specific tools without the agent picker, pass a comma-separated list:
 
 ```bash
 ucode configure --agents claude,codex
@@ -102,8 +102,8 @@ Discovered external MCP connections are listed directly. MCP auth uses a Databri
 | `ucode usage` | Show AI Gateway usage summary |
 | `ucode revert` | Clear saved state and restore backed-up config files |
 | `ucode configure --dry-run` | Preview config files without writing them |
-| `ucode configure --agents claude,codex` | Configure specific agents without the interactive picker |
-| `ucode configure --workspaces https://first.databricks.com,https://second.databricks.com` | Configure workspaces without the interactive picker |
+| `ucode configure --agents claude,codex` | Configure specific agents without the interactive agent picker |
+| `ucode configure --workspaces https://first.databricks.com,https://second.databricks.com` | Configure workspaces without the interactive workspace picker |
 | `ucode configure --profiles DEFAULT` | Configure using existing Databricks CLI profiles (hosts come from `~/.databrickscfg`) |
 | `ucode configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login |
 | `ucode configure --skip-validate` | Write configs without sending a test message through each agent |

src/ucode/agents/__init__.py

@@ -20,6 +20,7 @@
 from ucode.databricks import (
     install_databricks_cli,
 )
+from ucode.model_selection import available_models_for_tool as available_models_for_tool
 from ucode.state import load_state, save_state
 from ucode.telemetry import agent_version
 from ucode.ui import (
@@ -288,7 +289,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool:
     if tool == "opencode":
         return bool(state.get("opencode_models"))
     if tool == "codex":
-        return bool(state.get("codex_models"))
+        return bool(available_models_for_tool("codex", state))
     if tool == "gemini":
         return bool(state.get("gemini_models"))
     if tool == "copilot":

src/ucode/agents/claude.py

@@ -23,6 +23,7 @@
     build_tool_base_url,
     get_databricks_token,
 )
+from ucode.model_selection import selected_model_for_tool
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 from ucode.tracing import tracing_env
@@ -428,6 +429,9 @@ def _ensure_mlflow_cli() -> bool:
 
 
 def default_model(state: dict) -> str | None:
+    selected = selected_model_for_tool("claude", state)
+    if selected:
+        return selected
     claude_models = state.get("claude_models") or {}
     return claude_models.get("opus") or claude_models.get("sonnet") or claude_models.get("haiku")
 

src/ucode/agents/codex.py

@@ -20,6 +20,7 @@
     build_tool_base_url,
     get_databricks_token,
 )
+from ucode.model_selection import selected_model_for_tool
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 
@@ -336,6 +337,10 @@ def default_model(state: dict) -> str | None:
     would be rejected with a Unity Catalog endpoint-name error. When no
     candidate parses as GPT we return None rather than pinning an unroutable id.
     """
+    selected = selected_model_for_tool("codex", state)
+    if selected and _parse_gpt(selected) is not None:
+        return selected
+
     codex_models = state.get("codex_models") or []
     parsed: list[tuple[str, tuple[int, int | None, int | None, str]]] = [
         (mid, gpt) for mid in codex_models if (gpt := _parse_gpt(mid)) is not None

src/ucode/agents/copilot.py

@@ -35,6 +35,7 @@
     build_copilot_base_url,
     get_databricks_token,
 )
+from ucode.model_selection import selected_model_for_tool
 from ucode.state import mark_tool_managed, save_state
 
 COPILOT_CONFIG_DIR = Path.home() / ".copilot"
@@ -72,6 +73,9 @@ def is_update_available() -> tuple[str, str] | None:
 
 def default_model(state: dict) -> str | None:
     """Prefer Claude sonnet, then opus/haiku, then codex."""
+    selected = selected_model_for_tool("copilot", state)
+    if selected:
+        return selected
     claude_models = state.get("claude_models") or {}
     for family in ("sonnet", "opus", "haiku"):
         if claude_models.get(family):

src/ucode/agents/gemini.py

@@ -25,6 +25,7 @@
     build_tool_base_url,
     get_databricks_token,
 )
+from ucode.model_selection import selected_model_for_tool
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 
@@ -184,6 +185,9 @@ def write_tool_config(
 
 
 def default_model(state: dict) -> str | None:
+    selected = selected_model_for_tool("gemini", state)
+    if selected:
+        return selected
     gemini_models = state.get("gemini_models") or []
     return gemini_models[0] if gemini_models else None
 

src/ucode/agents/opencode.py

@@ -21,6 +21,7 @@
     build_opencode_base_urls,
     get_databricks_token,
 )
+from ucode.model_selection import selected_model_for_tool
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 
@@ -192,6 +193,9 @@ def remove_mcp_server_config(name: str) -> bool:
 
 
 def default_model(state: dict) -> str | None:
+    selected = selected_model_for_tool("opencode", state)
+    if selected:
+        return selected
     opencode_models = state.get("opencode_models") or {}
     anthropic = opencode_models.get("anthropic") or []
     if anthropic:

src/ucode/agents/pi.py

@@ -46,6 +46,7 @@
     build_pi_base_urls,
     get_databricks_token,
 )
+from ucode.model_selection import claude_model_options, selected_model_for_tool
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 
@@ -81,17 +82,28 @@ def is_update_available() -> tuple[str, str] | None:
     return available_npm_package_update(SPEC["package"])
 
 
+def _unique_models(models: list[str]) -> list[str]:
+    seen: set[str] = set()
+    unique: list[str] = []
+    for model in models:
+        model = model.strip()
+        if model and model not in seen:
+            seen.add(model)
+            unique.append(model)
+    return unique
+
+
 def _resolve_model_selector(
     model: str,
-    claude_models: dict[str, str],
+    claude_models: list[str],
     codex_models: list[str],
     gemini_models: list[str],
 ) -> str:
     """Return a Pi model selector in `<provider>/<model>` form when possible."""
     for name in PROVIDER_NAMES:
         if model.startswith(f"{name}/"):
             return model
-    if model in claude_models.values():
+    if model in claude_models:
         return f"databricks-claude/{model}"
     if model in codex_models:
         return f"databricks-openai/{model}"
@@ -107,6 +119,7 @@ def render_overlay(
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    selectable_claude_models: list[str] | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json."""
     providers: dict = {}
@@ -115,7 +128,7 @@ def render_overlay(
     # `/` and a space so it can never collide — safe to pass as a literal.
     ua_headers = {"User-Agent": f"ucode/{ucode_version()} pi/{agent_version('pi')}"}
 
-    claude_ids = sorted(set(claude_models.values()))
+    claude_ids = _unique_models(selectable_claude_models or list(claude_models.values()))
     if claude_ids:
         providers["databricks-claude"] = {
             "baseUrl": pi_base_urls["claude"],
@@ -151,7 +164,7 @@ def render_overlay(
         }
         keys.append(["providers", "databricks-gemini"])
     overlay: dict = {
-        "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models),
+        "model": _resolve_model_selector(model, claude_ids, codex_models, gemini_models),
     }
     if providers:
         overlay["providers"] = providers
@@ -178,6 +191,7 @@ def write_tool_config(
         state.get("claude_models") or {},
         state.get("codex_models") or [],
         state.get("gemini_models") or [],
+        claude_model_options(state),
     )
     existing = read_json_safe(PI_CONFIG_PATH)
     providers = existing.get("providers")
@@ -207,6 +221,9 @@ def _write_settings(model_selector: str) -> None:
 
 def default_model(state: dict) -> str | None:
     """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini."""
+    selected = selected_model_for_tool("pi", state)
+    if selected:
+        return selected
     claude_models = state.get("claude_models") or {}
     for family in ("opus", "sonnet", "haiku"):
         if claude_models.get(family):

src/ucode/cli.py

@@ -11,10 +11,12 @@
 
 from ucode.agents import (
     TOOL_SPECS,
+    available_models_for_tool,
     check_gateway_endpoint,
     configure_selected_tools,
     configure_single_tool,
     configure_tool,
+    default_model_for_tool,
     ensure_bootstrap_dependencies,
     ensure_provider_state,
     install_tool_binary,
@@ -32,10 +34,10 @@
 from ucode.databricks import (
     apply_pat_environment,
     build_shared_base_urls,
-    discover_claude_models,
+    discover_claude_models_with_options,
     discover_codex_models,
     discover_gemini_models,
-    discover_model_services,
+    discover_model_services_with_options,
     ensure_ai_gateway_v2,
     ensure_databricks_auth,
     find_profile_name_for_host,
@@ -64,6 +66,7 @@
     print_note,
     print_section,
     print_success,
+    prompt_for_model,
     prompt_for_tools,
     prompt_for_workspace,
     set_verbosity,
@@ -255,18 +258,27 @@ def configure_shared_state(
     gemini_reason: str | None = None
     codex_reason: str | None = None
     claude_models = {}
+    claude_model_options = []
     gemini_models = []
     codex_models = []
     # UC-first, best-effort: one UC model-services call yields all families as
     # `system.ai.<model-name>` ids, bucketed by name. If a family comes back
     # empty (workspace without UC model-services, or the listing failed), fall
     # back to the per-family AI Gateway listing for that family only.
     with spinner("Fetching available models..."):
-        ms_claude, ms_codex, ms_gemini, ms_reason = discover_model_services(workspace, token)
+        ms_claude, ms_claude_options, ms_codex, ms_gemini, ms_reason = (
+            discover_model_services_with_options(workspace, token)
+        )
         if want_claude:
-            claude_models, claude_reason = ms_claude, ms_reason
+            claude_models, claude_model_options, claude_reason = (
+                ms_claude,
+                ms_claude_options,
+                ms_reason,
+            )
             if not claude_models:
-                claude_models, claude_reason = discover_claude_models(workspace, token)
+                claude_models, claude_model_options, claude_reason = (
+                    discover_claude_models_with_options(workspace, token)
+                )
         if want_gemini:
             gemini_models, gemini_reason = ms_gemini, ms_reason
             if not gemini_models:
@@ -276,7 +288,9 @@ def configure_shared_state(
             if not codex_models:
                 codex_models, codex_reason = discover_codex_models(workspace, token)
     opencode_models: dict[str, list[str]] = {}
-    if claude_models:
+    if claude_model_options:
+        opencode_models["anthropic"] = claude_model_options
+    elif claude_models:
         opencode_models["anthropic"] = list(claude_models.values())
     if gemini_models:
         opencode_models["gemini"] = gemini_models
@@ -299,6 +313,7 @@ def configure_shared_state(
     state["base_urls"] = build_shared_base_urls(workspace)
     if want_claude:
         state["claude_models"] = claude_models
+        state["claude_model_options"] = claude_model_options
     if want_gemini:
         state["gemini_models"] = gemini_models
     if want_codex:
@@ -343,6 +358,29 @@ def _configure_shared_workspace_states(
     return states
 
 
+def _prompt_for_selected_models(state: dict, tools: list[str], *, prompt: bool = True) -> dict:
+    """Prompt for and persist per-agent model selections for configured tools."""
+    selected_models_value = state.get("selected_models")
+    selected_models = dict(selected_models_value) if isinstance(selected_models_value, dict) else {}
+
+    for tool in tools:
+        options = available_models_for_tool(tool, state)
+        if not options:
+            continue
+        default = default_model_for_tool(tool, state)
+        if default not in options:
+            default = options[0]
+        if len(options) == 1 or not prompt:
+            selected = default
+        else:
+            selected = prompt_for_model(TOOL_SPECS[tool]["display"], options, default)
+        selected_models[tool] = selected
+
+    if selected_models:
+        state["selected_models"] = selected_models
+    return state
+
+
 def configure_workspace_command(
     tool: str | None = None,
     selected_tools: list[str] | None = None,
@@ -351,6 +389,7 @@ def configure_workspace_command(
     prompt_optional_updates: bool = True,
     use_pat: bool = False,
     skip_validate: bool = False,
+    prompt_models: bool = True,
 ) -> int:
     if tool is not None and selected_tools is not None:
         raise RuntimeError("Use either --agent or --agents, not both.")
@@ -365,6 +404,7 @@ def configure_workspace_command(
             use_pat=use_pat,
         )
         state = states[0]
+        state = _prompt_for_selected_models(state, [tool], prompt=prompt_models)
         state = configure_single_tool(tool, state)
         spec = TOOL_SPECS[tool]
         console.print(
@@ -432,6 +472,8 @@ def configure_workspace_command(
         print_note("No coding agents selected — nothing to configure.")
         return 0
 
+    state = _prompt_for_selected_models(state, picked, prompt=prompt_models)
+
     for tool_name in picked:
         install_tool_binary(
             tool_name,
@@ -495,6 +537,10 @@ def status() -> int:
         print_kv("Coding Agent", spec["display"])
         print_kv("Configured", "yes" if configured else "no")
         print_kv("Base URL", base_url)
+        if configured:
+            model = default_model_for_tool(tool, state)
+            if model:
+                print_kv("Model", model)
         if configured and tool in MCP_CLIENTS:
             tool_mcp_servers = [
                 str(server.get("name"))
@@ -730,7 +776,8 @@ def configure(
         str | None,
         typer.Option(
             "--agents",
-            help="Configure a comma-separated list of agents without prompting (e.g. claude,codex).",
+            help="Configure a comma-separated list of agents without the agent picker "
+            "(e.g. claude,codex).",
         ),
     ] = None,
     workspaces: Annotated[
@@ -824,6 +871,13 @@ def configure(
             skip_kwargs["use_pat"] = True
         if skip_validate:
             skip_kwargs["skip_validate"] = True
+        if (
+            use_pat
+            and skip_validate
+            and workspace_entries is not None
+            and (agent is not None or agents is not None)
+        ):
+            skip_kwargs["prompt_models"] = False
         if agent is not None:
             tool = normalize_tool(agent)
             install_tool_binary(