From 51cf4527236a56a0dc74f813757985d016f613d6 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 21 Jul 2026 15:11:42 +0000 Subject: [PATCH 1/6] 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 `; 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 --- src/ucode/agents/__init__.py | 29 +++++++ src/ucode/agents/codex.py | 15 +++- src/ucode/cli.py | 120 ++++++++++++++++++++++++++++ src/ucode/databricks.py | 70 +++++++++++++++- tests/test_agent_codex.py | 22 +++++ tests/test_agents_init.py | 24 ++++++ tests/test_cli.py | 151 +++++++++++++++++++++++++++++++++++ tests/test_databricks.py | 83 +++++++++++++++++++ 8 files changed, 509 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 127cc7a..c56e5cf 100644 --- a/src/ucode/agents/__init__.py +++ b/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"] diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 2a5c657..8b098a9 100644 --- a/src/ucode/agents/codex.py +++ b/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 ), }, } diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a0ef7c4..2bc17e8 100644 --- a/src/ucode/cli.py +++ b/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 ` " + "(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, diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b38bca9..a6c67e8 100644 --- a/src/ucode/databricks.py +++ b/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.` 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), } diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 9a68e03..48e4072 100644 --- a/tests/test_agent_codex.py +++ b/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. diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 1a90ff8..43a6df9 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -14,6 +14,7 @@ configure_selected_tools, default_model_for_tool, ensure_tool_binary_available, + harness_for_model, install_tool_binary, normalize_tool, provider_permission_error, @@ -21,6 +22,29 @@ ) +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 " diff --git a/tests/test_cli.py b/tests/test_cli.py index 76a2ec5..1a64ca8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -145,6 +145,157 @@ def test_no_agent_flag(self): assert result.exit_code != 0 +class TestRunCommand: + """`ucode run` recommends a model, then launches the matching harness.""" + + def _patch_run(self, *, recommended, reason=None, chosen=None): + return [ + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE), + patch("ucode.cli.get_databricks_token", return_value="tok"), + patch( + "ucode.cli.recommend_coding_agent_models", + return_value=(recommended, reason), + ), + patch("ucode.cli.prompt_for_selection", return_value=chosen), + patch( + "ucode.cli.resolve_launch_model", + side_effect=lambda tool, state, model: (state, model), + ), + patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE), + patch("ucode.cli.launch_agent"), + ] + + def test_recommends_and_launches_claude(self): + patches = self._patch_run( + recommended=["system.ai.claude-opus-4-8", "system.ai.gpt-5"], + chosen="system.ai.claude-opus-4-8", + ) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 0, result.output + mock_launch.assert_called_once() + assert mock_launch.call_args[0][0] == "claude" + + def test_gpt_launches_codex(self): + patches = self._patch_run(recommended=["system.ai.gpt-5"], chosen="system.ai.gpt-5") + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[0][0] == "codex" + + def test_oss_launches_codex(self): + patches = self._patch_run(recommended=["system.ai.kimi-k2"], chosen="system.ai.kimi-k2") + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[0][0] == "codex" + + def test_single_recommendation_auto_launches_without_prompt(self): + patches = self._patch_run(recommended=["system.ai.claude-opus-4-8"], chosen=None) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5] as mock_prompt, + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 0, result.output + mock_prompt.assert_not_called() + assert mock_launch.call_args[0][0] == "claude" + + def test_empty_recommendation_exits_zero_with_guidance(self): + patches = self._patch_run(recommended=[]) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 0, result.output + assert "No recommended models" in _strip_ansi(result.output) + mock_launch.assert_not_called() + + def test_endpoint_failure_exits_nonzero(self): + patches = self._patch_run(recommended=[], reason="HTTP 500 Server Error") + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 1 + mock_launch.assert_not_called() + + def test_cancelled_selection_exits(self): + # Two models => the picker is shown; a None result means the user cancelled. + patches = self._patch_run( + recommended=["system.ai.gpt-5", "system.ai.claude-opus-4-8"], chosen=None + ) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 130 + mock_launch.assert_not_called() + + class TestMcpSubcommands: def test_web_search_subcommand_help(self): result = runner.invoke(app, ["mcp", "web-search", "--help"]) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 40af048..b4b850c 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -20,14 +20,17 @@ build_auth_token_argv, build_databricks_cli_env, build_opencode_base_urls, + build_oss_base_url, build_shared_base_urls, build_tool_base_url, ensure_databricks_cli_version, ensure_pat_bearer, get_databricks_token, + is_oss_model, list_databricks_apps, list_databricks_connections, list_genie_spaces, + recommend_coding_agent_models, workspace_hostname, ) @@ -1699,3 +1702,83 @@ def test_falls_through_for_unrelated_permission_error(self, monkeypatch): db_mod.run_usage_query(WS, "/sql/1.0/warehouses/abc", "tok", "SELECT 1") assert "Ask your workspace admin" not in str(exc_info.value) assert str(exc_info.value).startswith("Usage query failed:") + + +class TestIsOSSModel: + def test_kimi(self): + assert is_oss_model("system.ai.kimi-k2") is True + + def test_glm(self): + assert is_oss_model("system.ai.glm-4-6") is True + + def test_non_oss(self): + assert is_oss_model("system.ai.gpt-5") is False + + def test_none(self): + assert is_oss_model(None) is False + + +class TestBuildOSSBaseUrl: + def test_mlflow_route(self): + assert build_oss_base_url(WS) == f"{WS}/ai-gateway/mlflow/v1" + + +class TestRecommendCodingAgentModels: + URL = "https://example.databricks.com/api/ai-gateway/v2/coding-agent-configs:recommendModel" + + def test_bare_list_response(self, monkeypatch): + captured = {} + + def fake_post(url, token, payload, timeout=10): + captured["url"] = url + captured["payload"] = payload + return ["system.ai.claude-opus-4-8", "system.ai.gpt-5"], None + + monkeypatch.setattr(db_mod, "_http_post_json", fake_post) + models, reason = recommend_coding_agent_models( + WS, "tok", ["system.ai.claude-opus-4-8", "system.ai.gpt-5"] + ) + assert models == ["system.ai.claude-opus-4-8", "system.ai.gpt-5"] + assert reason is None + assert captured["url"] == self.URL + assert captured["payload"] == { + "available_models": ["system.ai.claude-opus-4-8", "system.ai.gpt-5"] + } + + def test_wrapped_models_key(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, payload, timeout=10: ({"models": ["system.ai.gpt-5"]}, None), + ) + models, reason = recommend_coding_agent_models(WS, "tok", ["system.ai.gpt-5"]) + assert models == ["system.ai.gpt-5"] + assert reason is None + + def test_empty_recommendation_is_success(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_post_json", lambda url, token, payload, timeout=10: ([], None) + ) + models, reason = recommend_coding_agent_models(WS, "tok", ["system.ai.gpt-5"]) + assert models == [] + assert reason is None + + def test_transport_error(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, payload, timeout=10: (None, "HTTP 500 Server Error"), + ) + models, reason = recommend_coding_agent_models(WS, "tok", ["system.ai.gpt-5"]) + assert models == [] + assert reason == "HTTP 500 Server Error" + + def test_unrecognized_shape_is_error(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, payload, timeout=10: ({"unexpected": True}, None), + ) + models, reason = recommend_coding_agent_models(WS, "tok", ["system.ai.gpt-5"]) + assert models == [] + assert reason is not None From 8e1d06a3e5ed7f935efd27ee11a196a2811f0af0 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 21 Jul 2026 16:02:43 +0000 Subject: [PATCH 2/6] ucode run: don't print the single-model auto-launch note MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When only one model is recommended, launch it silently instead of printing a note — the Model/Harness banner already shows what launched. Record the auto-launch in the debug log (UCODE_DEBUG=1) instead. Co-authored-by: Isaac --- src/ucode/cli.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 2bc17e8..e490318 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -32,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, @@ -1101,8 +1102,10 @@ def _run_session(ctx: typer.Context, skip_preflight: bool = False) -> None: print_section("ucode") if len(recommended) == 1: # Only one model recommended — no point prompting; launch it directly. + # Keep the console quiet (the Model/Harness banner below still shows + # what launched); note it only in the debug log. chosen = recommended[0] - print_note(f"Only one recommended model — launching {chosen}.") + _debug("recommend", f"single recommended model, auto-launching {chosen}") else: chosen = prompt_for_selection( "Select a model", [(model, model) for model in recommended] From 9e722d5ffbd0b5d06c8c7a3b220cf4ebecd08944 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 21 Jul 2026 16:02:56 +0000 Subject: [PATCH 3/6] ucode run: trim verbose comment on single-model auto-launch Co-authored-by: Isaac --- src/ucode/cli.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index e490318..f29932a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1102,8 +1102,6 @@ def _run_session(ctx: typer.Context, skip_preflight: bool = False) -> None: print_section("ucode") if len(recommended) == 1: # Only one model recommended — no point prompting; launch it directly. - # Keep the console quiet (the Model/Harness banner below still shows - # what launched); note it only in the debug log. chosen = recommended[0] _debug("recommend", f"single recommended model, auto-launching {chosen}") else: From ce764db5b0ac64e68fa8c18308a70e90f1ce7b68 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 21 Jul 2026 16:43:50 +0000 Subject: [PATCH 4/6] codex: route OSS models through responses, not chat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex CLI dropped support for `wire_api = "chat"` (openai/codex discussion 7782), so the OSS chat-completions route added for kimi/glm now fails to load with "`wire_api = "chat"` is no longer supported". Route every codex-harness model — including OSS families — through the Responses gateway, matching what all other codex models already use. Co-authored-by: Isaac --- src/ucode/agents/codex.py | 15 ++++----------- tests/test_agent_codex.py | 22 ---------------------- 2 files changed, 4 insertions(+), 33 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 8b098a9..2a5c657 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -17,10 +17,8 @@ ) 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 @@ -109,14 +107,9 @@ 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) - # 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) + base_url = build_tool_base_url("codex", workspace) http_headers = { "User-Agent": f"ucode/{ucode_version()} codex/{agent_version('codex')}", } @@ -127,7 +120,7 @@ def _provider_block( return { "name": "Databricks AI Gateway", "base_url": base_url, - "wire_api": "chat" if oss else "responses", + "wire_api": "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). @@ -152,7 +145,7 @@ def render_overlay( overlay["model"] = model overlay["model_providers"] = { CODEX_MODEL_PROVIDER_NAME: _provider_block( - workspace, databricks_profile, use_pat, provider, model + workspace, databricks_profile, use_pat, provider ), } return overlay @@ -178,7 +171,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, model + workspace, databricks_profile, use_pat, provider ), }, } diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 48e4072..9a68e03 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -45,28 +45,6 @@ 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. From 1edbbf5baebe200dbd07b3c5bb8b3d08d3c04590 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 21 Jul 2026 18:57:18 +0000 Subject: [PATCH 5/6] ucode run: reject --model override and pin selected model on launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode run` owns model selection via the recommendation flow, so a caller-supplied `--model` is rejected up front with actionable guidance. For the claude harness — which deliberately doesn't pin a model in its settings — prepend `--model ` to the launch argv so the session starts on the model the user picked instead of Claude Code's default. Co-authored-by: Isaac --- src/ucode/cli.py | 36 +++++++++++++++++++- tests/test_cli.py | 86 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index f29932a..8bedf19 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1057,6 +1057,39 @@ def _available_models(state: dict) -> list[str]: 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 ` 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 ` (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. @@ -1065,6 +1098,7 @@ def _run_session(ctx: typer.Context, skip_preflight: bool = False) -> None: 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 @@ -1135,7 +1169,7 @@ def _run_session(ctx: typer.Context, skip_preflight: bool = False) -> None: f"every 30 minutes while the session is running." ) print_success(f"Starting {TOOL_SPECS[tool]['display']}") - launch_agent(tool, state, ctx.args) + 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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 1a64ca8..e7a82c3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -188,6 +188,92 @@ def test_recommends_and_launches_claude(self): mock_launch.assert_called_once() assert mock_launch.call_args[0][0] == "claude" + def test_claude_launch_pins_selected_model(self): + # The claude harness doesn't pin a model in settings, so `ucode run` + # must pass `--model ` on the argv or Claude Code boots on its + # own default instead of the model the user picked. + patches = self._patch_run( + recommended=["system.ai.claude-sonnet-5", "system.ai.gpt-5"], + chosen="system.ai.claude-sonnet-5", + ) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[0][0] == "claude" + assert mock_launch.call_args[0][2] == ["--model", "system.ai.claude-sonnet-5"] + + def test_caller_model_override_is_rejected(self): + # `ucode run` selects the model itself, so a caller `--model` is invalid. + patches = self._patch_run( + recommended=["system.ai.claude-sonnet-5"], + chosen=None, + ) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run", "--", "--model", "system.ai.claude-opus-4-8"]) + assert result.exit_code == 1 + assert "--model" in _strip_ansi(result.output) + mock_launch.assert_not_called() + + def test_caller_model_override_equals_form_is_rejected(self): + patches = self._patch_run( + recommended=["system.ai.claude-sonnet-5"], + chosen=None, + ) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run", "--", "--model=system.ai.claude-opus-4-8"]) + assert result.exit_code == 1 + mock_launch.assert_not_called() + + def test_codex_launch_does_not_pin_model(self): + # Non-claude harnesses already pin the model in their own config, so the + # argv stays untouched. + patches = self._patch_run(recommended=["system.ai.gpt-5"], chosen="system.ai.gpt-5") + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patches[7], + patches[8] as mock_launch, + ): + result = runner.invoke(app, ["run"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[0][0] == "codex" + assert mock_launch.call_args[0][2] == [] + def test_gpt_launches_codex(self): patches = self._patch_run(recommended=["system.ai.gpt-5"], chosen="system.ai.gpt-5") with ( From bdf4c55ca33327dd33ddb4c32162c49633a104cf Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 21 Jul 2026 18:57:24 +0000 Subject: [PATCH 6/6] databricks: drop dead build_oss_base_url, hoist recommendModel path to constant `build_oss_base_url` lost its only shared caller when codex OSS routing was reverted, leaving a single-use indirection; inline it back into build_opencode_base_urls. Extract the recommendModel API path into a module constant. Co-authored-by: Isaac --- src/ucode/databricks.py | 17 +++++------------ tests/test_databricks.py | 6 ------ 2 files changed, 5 insertions(+), 18 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a6c67e8..acec3cf 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1287,6 +1287,9 @@ 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]: @@ -1302,10 +1305,7 @@ def recommend_coding_agent_models( - ``([], 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" - ) + 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 @@ -2178,18 +2178,11 @@ 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": build_oss_base_url(workspace), + "oss": f"{workspace}/ai-gateway/mlflow/v1", } diff --git a/tests/test_databricks.py b/tests/test_databricks.py index b4b850c..c5d143d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -20,7 +20,6 @@ build_auth_token_argv, build_databricks_cli_env, build_opencode_base_urls, - build_oss_base_url, build_shared_base_urls, build_tool_base_url, ensure_databricks_cli_version, @@ -1718,11 +1717,6 @@ def test_none(self): assert is_oss_model(None) is False -class TestBuildOSSBaseUrl: - def test_mlflow_route(self): - assert build_oss_base_url(WS) == f"{WS}/ai-gateway/mlflow/v1" - - class TestRecommendCodingAgentModels: URL = "https://example.databricks.com/api/ai-gateway/v2/coding-agent-configs:recommendModel"