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/cli.py b/src/ucode/cli.py index a0ef7c4..8bedf19 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, @@ -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 ` 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. + + 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 ` " + "(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, diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b38bca9..acec3cf 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 +_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) ----------------------------- 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..e7a82c3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -145,6 +145,243 @@ 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_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 ( + 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..c5d143d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -25,9 +25,11 @@ 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 +1701,78 @@ 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 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