diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b54e136..e7d3808 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -19,6 +19,7 @@ from ucode.config_io import ToolSpec from ucode.databricks import ( BEDROCK_PROVIDER_TYPES, + ensure_databricks_cli_version, get_databricks_token, install_ai_tools, install_databricks_cli, @@ -80,20 +81,33 @@ # ask them the machine-wide question. GLOBAL_SETTINGS_AGENTS = frozenset({"claude", "codex"}) -# ucode tool -> `databricks aitools` agent id. gemini/pi aren't supported. +# ucode tool -> `databricks aitools` agent id. AITOOLS_AGENT_TOKENS = { "claude": "claude-code", "codex": "codex", + "gemini": "gemini-cli", "opencode": "opencode", "copilot": "copilot", + "pi": "pi", +} + +AITOOLS_AGENT_MIN_CLI_VERSIONS = { + "gemini": (1, 12, 0), + "pi": (1, 12, 0), } def install_ai_tools_for_agents(tools: list[str], state: dict) -> None: - """Install Databricks AI Tools for the coding agents that support them - (gemini/pi have no ``aitools`` support and are dropped).""" + """Install Databricks AI Tools for the coding agents that support them.""" if state.get("databricks_ai_tools_enabled", True) is False: return + required_versions = [ + AITOOLS_AGENT_MIN_CLI_VERSIONS[tool] + for tool in tools + if tool in AITOOLS_AGENT_MIN_CLI_VERSIONS + ] + if required_versions: + ensure_databricks_cli_version(max(required_versions)) agents = [AITOOLS_AGENT_TOKENS[tool] for tool in tools if tool in AITOOLS_AGENT_TOKENS] install_ai_tools(agents, state.get("profile")) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 88935dd..6219812 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -704,7 +704,9 @@ def _run_databricks_cli_installer(brew_subcommand: str = "install") -> None: raise RuntimeError("Failed to install/upgrade Databricks CLI automatically.") from exc -def ensure_databricks_cli_version() -> None: +def ensure_databricks_cli_version( + minimum_version: tuple[int, int, int] = MIN_DATABRICKS_CLI_VERSION, +) -> None: try: result = run( ["databricks", "--version"], @@ -723,14 +725,14 @@ def ensure_databricks_cli_version() -> None: raise RuntimeError( f"Could not parse Databricks CLI version from `databricks --version` output: {output!r}" ) - if version < MIN_DATABRICKS_CLI_VERSION: + if version < minimum_version: current = ".".join(str(n) for n in version) - required = ".".join(str(n) for n in MIN_DATABRICKS_CLI_VERSION) + required = ".".join(str(n) for n in minimum_version) print_warning( f"Databricks CLI v{current} is too old (need v{required} or newer). Upgrading..." ) _run_databricks_cli_installer(brew_subcommand="upgrade") - ensure_databricks_cli_version() + ensure_databricks_cli_version(minimum_version) def install_databricks_cli() -> None: diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 993c967..be1b70b 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -64,6 +64,11 @@ def test_each_agent_exposes_update_check(self): class TestInstallAiToolsForAgents: def _capture(self, monkeypatch): captured = {} + monkeypatch.setattr( + agents_mod, + "ensure_databricks_cli_version", + lambda minimum: captured.update(minimum_cli_version=minimum), + ) monkeypatch.setattr( agents_mod, "install_ai_tools", @@ -71,11 +76,23 @@ def _capture(self, monkeypatch): ) return captured - def test_maps_supported_tools_and_drops_others(self, monkeypatch): + def test_maps_supported_tools(self, monkeypatch): + captured = self._capture(monkeypatch) + install_ai_tools_for_agents( + ["claude", "codex", "gemini", "opencode", "copilot", "pi"], + {"profile": "prof"}, + ) + assert captured == { + "agents": ["claude-code", "codex", "gemini-cli", "opencode", "copilot", "pi"], + "minimum_cli_version": (1, 12, 0), + "profile": "prof", + } + + @pytest.mark.parametrize("tool", ["gemini", "pi"]) + def test_requires_new_cli_for_extended_agents(self, monkeypatch, tool): captured = self._capture(monkeypatch) - # gemini and pi aren't supported by `databricks aitools`, so they drop. - install_ai_tools_for_agents(["claude", "codex", "gemini", "pi"], {"profile": "prof"}) - assert captured == {"agents": ["claude-code", "codex"], "profile": "prof"} + install_ai_tools_for_agents([tool], {"profile": "prof"}) + assert captured["minimum_cli_version"] == (1, 12, 0) def test_installed_by_default(self, monkeypatch): # Opt-out: absent flag means install. diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 64b4dcc..020d90d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1995,10 +1995,19 @@ def test_passes_when_version_exceeds_minimum(self, tmp_path, monkeypatch): monkeypatch.setattr("os.environ", env) ensure_databricks_cli_version() - def test_auto_upgrades_when_version_too_old(self, tmp_path, monkeypatch): + @pytest.mark.parametrize( + ("installed", "minimum_version"), + [ + ("0.299.2", None), + ("1.11.0", (1, 12, 0)), + ], + ) + def test_auto_upgrades_when_version_too_old( + self, tmp_path, monkeypatch, installed, minimum_version + ): import ucode.databricks as db_mod - env = self._fake_databricks(tmp_path, "Databricks CLI v0.299.2") + env = self._fake_databricks(tmp_path, f"Databricks CLI v{installed}") monkeypatch.setattr("os.environ", env) upgraded = [] monkeypatch.setattr( @@ -2013,10 +2022,13 @@ def test_auto_upgrades_when_version_too_old(self, tmp_path, monkeypatch): def once(*a, **kw): call_count[0] += 1 if call_count[0] == 1: - original() + original(*a, **kw) monkeypatch.setattr(db_mod, "ensure_databricks_cli_version", once) - once() + if minimum_version is None: + once() + else: + once(minimum_version) assert upgraded == ["upgrade"] def test_raises_when_version_unparseable(self, tmp_path, monkeypatch):