From a47b1de2690ec07fc883984e040088bfa5f6d7f2 Mon Sep 17 00:00:00 2001 From: Lennart Kats Date: Fri, 14 Aug 2026 14:45:21 +0200 Subject: [PATCH 1/2] Install AI Tools for Gemini CLI and Pi --- src/ucode/agents/__init__.py | 7 ++++--- src/ucode/databricks.py | 4 ++-- tests/test_agents_init.py | 13 +++++++++---- tests/test_databricks.py | 6 +++--- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index ecb57f7d..ab28ce8d 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -70,18 +70,19 @@ DEFAULT_TOOL = "codex" BUNDLE_VERSION = 1 -# 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", } 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 agents = [AITOOLS_AGENT_TOKENS[tool] for tool in tools if tool in AITOOLS_AGENT_TOKENS] diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index e017be29..b9a2023d 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -51,8 +51,8 @@ "https://raw.githubusercontent.com/databricks/setup-cli/main/install.ps1" ) AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" -# v1.0.0 is the release that ships `databricks aitools`. -MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) +# v1.12.0 adds AI Tools support for every coding agent supported by ucode. +MIN_DATABRICKS_CLI_VERSION = (1, 12, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 993c967d..8e8e2ab1 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -71,11 +71,16 @@ 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) - # 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( + ["claude", "codex", "gemini", "opencode", "copilot", "pi"], + {"profile": "prof"}, + ) + assert captured == { + "agents": ["claude-code", "codex", "gemini-cli", "opencode", "copilot", "pi"], + "profile": "prof", + } 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 d7b3dc35..c433a6f0 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1928,19 +1928,19 @@ def _fake_databricks(self, tmp_path, version_output: str) -> dict: return {**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}"} def test_passes_when_version_meets_minimum(self, tmp_path, monkeypatch): - env = self._fake_databricks(tmp_path, "Databricks CLI v1.0.0") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.12.0") monkeypatch.setattr("os.environ", env) ensure_databricks_cli_version() # should not raise def test_passes_when_version_exceeds_minimum(self, tmp_path, monkeypatch): - env = self._fake_databricks(tmp_path, "Databricks CLI v1.8.0") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.12.1") monkeypatch.setattr("os.environ", env) ensure_databricks_cli_version() def test_auto_upgrades_when_version_too_old(self, tmp_path, monkeypatch): import ucode.databricks as db_mod - env = self._fake_databricks(tmp_path, "Databricks CLI v0.299.2") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.11.0") monkeypatch.setattr("os.environ", env) upgraded = [] monkeypatch.setattr( From 66cccf43e4c541b7fa614ca04fc7c74e6036ddba Mon Sep 17 00:00:00 2001 From: Lennart Kats Date: Sun, 16 Aug 2026 15:25:46 +0200 Subject: [PATCH 2/2] Require CLI v1.12 only for Gemini CLI and Pi --- src/ucode/agents/__init__.py | 13 +++++++++++++ src/ucode/databricks.py | 14 ++++++++------ tests/test_agents_init.py | 12 ++++++++++++ tests/test_databricks.py | 24 ++++++++++++++++++------ 4 files changed, 51 insertions(+), 12 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index ab28ce8d..6bae599a 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,11 +81,23 @@ "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.""" 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 b9a2023d..37deab5f 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -51,8 +51,8 @@ "https://raw.githubusercontent.com/databricks/setup-cli/main/install.ps1" ) AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" -# v1.12.0 adds AI Tools support for every coding agent supported by ucode. -MIN_DATABRICKS_CLI_VERSION = (1, 12, 0) +# v1.0.0 is the first release with `databricks aitools`. +MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 @@ -679,7 +679,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"], @@ -698,14 +700,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 8e8e2ab1..be1b70b2 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", @@ -79,9 +84,16 @@ def test_maps_supported_tools(self, monkeypatch): ) 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) + 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. captured = self._capture(monkeypatch) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index c433a6f0..1dee149e 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1928,19 +1928,28 @@ def _fake_databricks(self, tmp_path, version_output: str) -> dict: return {**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}"} def test_passes_when_version_meets_minimum(self, tmp_path, monkeypatch): - env = self._fake_databricks(tmp_path, "Databricks CLI v1.12.0") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.0.0") monkeypatch.setattr("os.environ", env) ensure_databricks_cli_version() # should not raise def test_passes_when_version_exceeds_minimum(self, tmp_path, monkeypatch): - env = self._fake_databricks(tmp_path, "Databricks CLI v1.12.1") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.8.0") 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 v1.11.0") + env = self._fake_databricks(tmp_path, f"Databricks CLI v{installed}") monkeypatch.setattr("os.environ", env) upgraded = [] monkeypatch.setattr( @@ -1955,10 +1964,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):