Add AI Tools support for Gemini CLI and Pi

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"))
 

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:

tests/test_agents_init.py

@@ -64,18 +64,35 @@ 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",
             lambda agents, profile: captured.update(agents=agents, profile=profile),
         )
         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.

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):