Add ucode run: recommend-model-driven session startup

New top-level `ucode run` command that queries the AI Gateway's
recommendModel endpoint for the caller's usage tier, then launches the
matching harness pinned to the chosen model:

- Discover all workspace models, POST them to
  /api/ai-gateway/v2/coding-agent-configs:recommendModel.
- 1 model -> auto-launch (no prompt); 2+ -> interactive picker;
  empty -> guidance to use `ucode <harness>`; error -> exit non-zero.
- Map model family -> harness: claude->claude, gpt->codex,
  kimi/glm->codex, gemini->gemini.
- Extend codex to route OSS models (kimi/glm) via the MLflow
  chat-completions gateway (wire_api="chat") instead of the Responses
  route; GPT models unchanged.

Adds recommend_coding_agent_models, is_oss_model, build_oss_base_url
(databricks.py) and harness_for_model (agents/__init__.py), with unit
tests for the endpoint, the harness map, codex OSS routing, and CLI
routing.

Co-authored-by: Isaac

Author: lilly-luo

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/agents/codex.py

@@ -17,8 +17,10 @@
 )
 from ucode.databricks import (
     build_auth_token_argv,
+    build_oss_base_url,
     build_tool_base_url,
     get_databricks_token,
+    is_oss_model,
 )
 from ucode.launcher import exec_or_spawn
 from ucode.state import mark_tool_managed, save_state
@@ -107,9 +109,14 @@ def _provider_block(
     databricks_profile: str | None,
     use_pat: bool = False,
     provider: str | None = None,
+    model: str | None = None,
 ) -> dict:
     auth_argv = build_auth_token_argv(workspace, databricks_profile, use_pat=use_pat)
-    base_url = build_tool_base_url("codex", workspace)
+    # OSS chat models (kimi/glm) don't speak the Responses API — they're served
+    # by the OpenAI-compatible MLflow chat-completions gateway. Route them there
+    # with `wire_api = "chat"`; everything else uses the codex Responses route.
+    oss = is_oss_model(model)
+    base_url = build_oss_base_url(workspace) if oss else build_tool_base_url("codex", workspace)
     http_headers = {
         "User-Agent": f"ucode/{ucode_version()} codex/{agent_version('codex')}",
     }
@@ -120,7 +127,7 @@ def _provider_block(
     return {
         "name": "Databricks AI Gateway",
         "base_url": base_url,
-        "wire_api": "responses",
+        "wire_api": "chat" if oss else "responses",
         "http_headers": http_headers,
         # Run the `ucode auth-token` executable directly (not via `sh -c`) so the
         # helper works on Windows, where there is no POSIX shell (issue #116).
@@ -145,7 +152,7 @@ def render_overlay(
         overlay["model"] = model
     overlay["model_providers"] = {
         CODEX_MODEL_PROVIDER_NAME: _provider_block(
-            workspace, databricks_profile, use_pat, provider
+            workspace, databricks_profile, use_pat, provider, model
         ),
     }
     return overlay
@@ -171,7 +178,7 @@ def render_legacy_overlay(
         "profiles": {CODEX_PROFILE_NAME: profile_block},
         "model_providers": {
             CODEX_MODEL_PROVIDER_NAME: _provider_block(
-                workspace, databricks_profile, use_pat, provider
+                workspace, databricks_profile, use_pat, provider, model
             ),
         },
     }

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,
@@ -48,6 +49,7 @@
     list_profile_entries,
     list_tool_provider_services,
     normalize_workspace_url,
+    recommend_coding_agent_models,
     resolve_pat_token,
     run_databricks_login,
 )
@@ -1034,6 +1036,124 @@ 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 _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:
+        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]
+            print_note(f"Only one recommended model — 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, 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
 
 
+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)}"
+        "/api/ai-gateway/v2/coding-agent-configs:recommendModel"
+    )
+    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) -----------------------------
 
 
@@ -2117,11 +2178,18 @@ def build_tool_base_url(tool: str, workspace: str) -> str:
     raise RuntimeError(f"Unsupported tool '{tool}'.")
 
 
+def build_oss_base_url(workspace: str) -> str:
+    # OSS chat models (kimi/glm) are served by the OpenAI-compatible MLflow
+    # chat-completions gateway, not the family-specific routes. Shared by every
+    # agent that speaks chat-completions to OSS models.
+    return f"{workspace}/ai-gateway/mlflow/v1"
+
+
 def build_opencode_base_urls(workspace: str) -> dict[str, str]:
     return {
         "anthropic": build_tool_base_url("claude", workspace) + "/v1",
         "gemini": build_tool_base_url("gemini", workspace) + "/v1beta",
-        "oss": f"{workspace}/ai-gateway/mlflow/v1",
+        "oss": build_oss_base_url(workspace),
     }
 
 

tests/test_agent_codex.py

@@ -45,6 +45,28 @@ def test_provider_wire_api(self):
         provider = overlay["model_providers"]["ucode-databricks"]
         assert provider["wire_api"] == "responses"
 
+    def test_oss_model_routes_to_mlflow_chat(self):
+        overlay = codex.render_overlay(WS, "system.ai.kimi-k2")
+        provider = overlay["model_providers"]["ucode-databricks"]
+        assert provider["base_url"] == f"{WS}/ai-gateway/mlflow/v1"
+        assert provider["wire_api"] == "chat"
+
+    def test_oss_glm_model_routes_to_mlflow_chat(self):
+        overlay = codex.render_overlay(WS, "system.ai.glm-4-6")
+        provider = overlay["model_providers"]["ucode-databricks"]
+        assert provider["base_url"] == f"{WS}/ai-gateway/mlflow/v1"
+        assert provider["wire_api"] == "chat"
+
+    def test_oss_model_pinned_verbatim(self):
+        overlay = codex.render_overlay(WS, "system.ai.kimi-k2")
+        assert overlay["model"] == "system.ai.kimi-k2"
+
+    def test_gpt_model_still_uses_responses_route(self):
+        overlay = codex.render_overlay(WS, "system.ai.gpt-5")
+        provider = overlay["model_providers"]["ucode-databricks"]
+        assert provider["base_url"] == f"{WS}/ai-gateway/codex/v1"
+        assert provider["wire_api"] == "responses"
+
     def test_auth_runs_ucode_auth_token(self):
         # The auth command runs the `ucode auth-token` executable directly
         # (not `sh -c`), so it works on Windows where there is no POSIX shell.

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 "