Add ucode run: recommend-model-driven session startup

src/ucode/agents/__init__.py

@@ -21,6 +21,7 @@
     BEDROCK_PROVIDER_TYPES,
     get_databricks_token,
     install_databricks_cli,
+    is_oss_model,
     map_bedrock_claude_models,
     resolve_provider_service,
 )
@@ -75,6 +76,34 @@ def normalize_tool(tool: str) -> str:
     return normalized
 
 
+def harness_for_model(model_id: str) -> str | None:
+    """Map a model id to the coding-agent harness that can serve it.
+
+    Used by `ucode run` to pick a harness from the model the user selected:
+
+    - claude models  -> claude (Claude Code)
+    - gpt models     -> codex
+    - OSS (kimi/glm)  -> codex
+    - gemini models  -> gemini
+
+    Returns None when no harness maps to the model so the caller can surface a
+    clear error instead of launching the wrong tool. OSS is checked before the
+    generic families via the shared `is_oss_model` substrings so the map stays in
+    sync with discovery.
+    """
+    if not model_id:
+        return None
+    if is_oss_model(model_id):
+        return "codex"
+    if "claude" in model_id:
+        return "claude"
+    if "gpt-" in model_id:
+        return "codex"
+    if "gemini-" in model_id:
+        return "gemini"
+    return None
+
+
 def _update_installed_tool_binary(tool: str, version: str | None = None) -> bool:
     spec = TOOL_SPECS[tool]
     binary = spec["binary"]

src/ucode/cli.py

@@ -16,6 +16,7 @@
     configure_tool,
     ensure_bootstrap_dependencies,
     ensure_provider_state,
+    harness_for_model,
     install_tool_binary,
     normalize_tool,
     provider_permission_error,
@@ -31,6 +32,7 @@
 from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH
 from ucode.config_io import restore_file, set_dry_run
 from ucode.databricks import (
+    _debug,
     apply_pat_environment,
     build_shared_base_urls,
     discover_claude_models,
@@ -48,6 +50,7 @@
     list_profile_entries,
     list_tool_provider_services,
     normalize_workspace_url,
+    recommend_coding_agent_models,
     resolve_pat_token,
     run_databricks_login,
 )
@@ -1034,6 +1037,158 @@ def pi_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> N
     _launch_tool("pi", ctx, skip_preflight=skip_preflight)
 
 
+def _available_models(state: dict) -> list[str]:
+    """Collect every Databricks model id discovered for this workspace.
+
+    Flattens the per-family state buckets (claude by family alias, plus the flat
+    codex/gemini/oss lists) into a single de-duplicated, order-preserving list to
+    send to the recommendModel endpoint.
+    """
+    models: list[str] = []
+    models.extend((state.get("claude_models") or {}).values())
+    for key in ("codex_models", "gemini_models", "oss_models"):
+        models.extend(state.get(key) or [])
+    seen: set[str] = set()
+    unique: list[str] = []
+    for model in models:
+        if model and model not in seen:
+            seen.add(model)
+            unique.append(model)
+    return unique
+
+
+def _launch_args(tool: str, resolved_model: str | None, tool_args: list[str]) -> list[str]:
+    """Pin the model the user picked in `ucode run` onto the launch argv.
+
+    The claude harness deliberately doesn't pin a model in its settings (doing so
+    dupes rows in Claude Code's /model picker — see claude.render_overlay), so
+    Claude Code would otherwise boot on its own default instead of the model the
+    user selected. Prepend `--model <resolved_model>` for claude so the session
+    starts on the selected model.
+
+    `ucode run` owns model selection via the recommendation flow, so a
+    caller-supplied `--model` is rejected up front (see `_reject_model_override`)
+    rather than reaching here.
+    """
+    if tool != "claude" or not resolved_model:
+        return tool_args
+    return ["--model", resolved_model, *tool_args]
+
+
+def _reject_model_override(tool_args: list[str]) -> None:
+    """`ucode run` selects the model itself, so `--model` on the argv is invalid.
+
+    Fail fast with actionable guidance rather than silently letting a caller
+    `--model` fight the recommendation flow.
+    """
+    if any(arg == "--model" or arg.startswith("--model=") for arg in tool_args):
+        raise RuntimeError(
+            "`ucode run` selects the model for you — passing `--model` is not "
+            "supported. Run `ucode run` and pick from the recommended models, or "
+            "use `ucode <harness>` (e.g. `ucode claude`) to launch a specific "
+            "harness directly."
+        )
+
+
+def _run_session(ctx: typer.Context, skip_preflight: bool = False) -> None:
+    """Recommend a model for the user's tier, then launch the matching harness.
+
+    Discovers the workspace's models, asks the AI Gateway's recommendModel
+    endpoint which ones the caller's tier allows, lets the user pick one, maps it
+    to a harness (claude/codex/gemini) and launches that tool pinned to the model.
+    """
+    try:
+        _reject_model_override(ctx.args)
+        existing = load_state()
+        apply_pat_environment(existing)
+        # Ensure a workspace is configured. Without one, fall back to the same
+        # first-run configuration prompt the per-tool launchers use, defaulting to
+        # the codex harness for the bootstrap install.
+        if not existing.get("workspace"):
+            ensure_bootstrap_dependencies("codex", update_existing=True)
+            _auto_configure_tool("codex")
+        # Refresh model discovery for every family so recommendModel sees the full
+        # set the workspace exposes (tools=None => fetch_all).
+        state = configure_shared_state(
+            load_state()["workspace"],
+            profile=load_state().get("profile"),
+            skip_preflight=skip_preflight,
+        )
+        available = _available_models(state)
+        if not available:
+            raise RuntimeError(
+                "No models available for this workspace. Run `ucode configure` to set it up."
+            )
+        with spinner("Requesting recommended models..."):
+            token = get_databricks_token(state["workspace"], state.get("profile"))
+            recommended, reason = recommend_coding_agent_models(
+                state["workspace"], token, available
+            )
+        if reason is not None:
+            raise RuntimeError(f"Could not fetch recommended models: {reason}")
+        if not recommended:
+            print_warning(
+                "No recommended models — use `ucode <harness>` "
+                "(e.g. `ucode claude` or `ucode codex`) to boot up your preferred harness."
+            )
+            raise typer.Exit(0)
+
+        print_section("ucode")
+        if len(recommended) == 1:
+            # Only one model recommended — no point prompting; launch it directly.
+            chosen = recommended[0]
+            _debug("recommend", f"single recommended model, auto-launching {chosen}")
+        else:
+            chosen = prompt_for_selection(
+                "Select a model", [(model, model) for model in recommended]
+            )
+            if not chosen:
+                print_err("No model selected.")
+                raise typer.Exit(130)
+
+        tool = harness_for_model(chosen)
+        if not tool:
+            raise RuntimeError(
+                f"No coding-agent harness maps to model '{chosen}'. "
+                "Supported families: claude, gpt, kimi/glm, gemini."
+            )
+
+        # A newly-picked harness may not have been configured/installed yet.
+        ensure_bootstrap_dependencies(tool, update_existing=True)
+        if tool not in (state.get("available_tools") or []):
+            _auto_configure_tool(tool)
+            state = load_state()
+
+        state, resolved_model = resolve_launch_model(tool, state, chosen)
+        state = configure_tool(tool, state, resolved_model)
+        print_kv("Model", resolved_model or chosen)
+        print_kv("Harness", TOOL_SPECS[tool]["display"])
+        if tool in ("gemini", "opencode", "copilot", "pi"):
+            print_note(
+                f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically "
+                f"every 30 minutes while the session is running."
+            )
+        print_success(f"Starting {TOOL_SPECS[tool]['display']}")
+        launch_agent(tool, state, _launch_args(tool, resolved_model, ctx.args))
+    except typer.Exit:
+        # Intentional exits (empty recommendation, cancelled selection) subclass
+        # RuntimeError, so re-raise them before the error handler below rewrites
+        # them to a generic exit code 1.
+        raise
+    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
+
+
+@app.command("run", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
+def run_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None:
+    """Recommend a model for your usage tier, then launch the matching harness."""
+    _run_session(ctx, skip_preflight=skip_preflight)
+
+
 @configure_app.callback(invoke_without_command=True)
 def configure(
     ctx: typer.Context,

src/ucode/databricks.py

@@ -1131,6 +1131,17 @@ def model_token_limits(model_id: str) -> dict[str, int] | None:
     return None
 
 
+def is_oss_model(model_id: str | None) -> bool:
+    """Return True when ``model_id`` belongs to a supported OSS chat family.
+
+    Matches the same ``kimi-``/``glm-`` substrings used by discovery so callers
+    (e.g. codex OSS routing, the harness map) stay in sync with what
+    ``discover_model_services`` buckets as OSS."""
+    if not model_id:
+        return False
+    return any(family in model_id for family in _OSS_MODEL_FAMILIES)
+
+
 def _model_service_id(service: dict) -> str | None:
     """Extract the `system.ai.<model-name>` id from one model-service entry.
 
@@ -1276,6 +1287,56 @@ def discover_model_services(
     return claude_models, codex_models, gemini_models, oss_models, None
 
 
+_RECOMMEND_MODEL_API_PATH = "/api/ai-gateway/v2/coding-agent-configs:recommendModel"
+
+
+def recommend_coding_agent_models(
+    workspace: str, token: str, available_models: list[str]
+) -> tuple[list[str], str | None]:
+    """Ask the AI Gateway which of ``available_models`` this user may use.
+
+    A workspace admin defines usage-based tiers via a "coding agent config"; the
+    ``recommendModel`` endpoint filters/orders the caller's ``available_models``
+    down to what their tier allows. Returns ``(models, reason)``:
+
+    - ``(models, None)`` on success — ``models`` may be empty, meaning the tier
+      recommends nothing (the caller should guide the user to pick a harness
+      manually rather than fall back to the full list).
+    - ``([], reason)`` on transport/parse failure — distinct from an intentional
+      empty recommendation so the caller can surface the error.
+    """
+    url = f"https://{workspace_hostname(workspace)}{_RECOMMEND_MODEL_API_PATH}"
+    payload, reason = _http_post_json(url, token, {"available_models": available_models})
+    if reason is not None:
+        return [], reason
+    models = _extract_recommended_models(payload)
+    if models is None:
+        return [], "recommendModel response did not contain a model list"
+    return models, None
+
+
+def _extract_recommended_models(payload: dict | list | None) -> list[str] | None:
+    """Pull the recommended-model list out of a recommendModel response.
+
+    Tolerates either a bare JSON array or an object wrapping the list under a
+    common key. Returns None when the shape is unrecognized so the caller can
+    distinguish a malformed response from an empty recommendation."""
+    if isinstance(payload, list):
+        candidate: list | None = payload
+    elif isinstance(payload, dict):
+        candidate = None
+        for key in ("models", "recommended_models", "available_models"):
+            value = payload.get(key)
+            if isinstance(value, list):
+                candidate = value
+                break
+    else:
+        candidate = None
+    if candidate is None:
+        return None
+    return [m for m in candidate if isinstance(m, str) and m]
+
+
 # --- MCP services (parallel to model services) -----------------------------
 
 

tests/test_agents_init.py

@@ -14,13 +14,37 @@
     configure_selected_tools,
     default_model_for_tool,
     ensure_tool_binary_available,
+    harness_for_model,
     install_tool_binary,
     normalize_tool,
     provider_permission_error,
     resolve_launch_model,
 )
 
 
+class TestHarnessForModel:
+    def test_claude_maps_to_claude(self):
+        assert harness_for_model("system.ai.claude-opus-4-8") == "claude"
+
+    def test_gpt_maps_to_codex(self):
+        assert harness_for_model("system.ai.gpt-5-5") == "codex"
+
+    def test_kimi_maps_to_codex(self):
+        assert harness_for_model("system.ai.kimi-k2") == "codex"
+
+    def test_glm_maps_to_codex(self):
+        assert harness_for_model("system.ai.glm-4-6") == "codex"
+
+    def test_gemini_maps_to_gemini(self):
+        assert harness_for_model("system.ai.gemini-2-5-pro") == "gemini"
+
+    def test_unknown_returns_none(self):
+        assert harness_for_model("system.ai.mystery-model") is None
+
+    def test_empty_returns_none(self):
+        assert harness_for_model("") is None
+
+
 class TestProviderPermissionError:
     _CONN_ERR = (
         "User does not have USE CONNECTION on SCHEMA_CONNECTION "