Install AI Tools once per agent for users who previously configured ucode

src/ucode/agents/__init__.py

@@ -20,6 +20,7 @@
 from ucode.databricks import (
     BEDROCK_PROVIDER_TYPES,
     get_databricks_token,
+    install_ai_tools,
     install_databricks_cli,
     map_bedrock_claude_models,
     resolve_provider_service,
@@ -65,6 +66,30 @@
 DEFAULT_TOOL = "codex"
 BUNDLE_VERSION = 1
 
+# ucode tool -> `databricks aitools` agent id. gemini/pi aren't supported.
+AITOOLS_AGENT_TOKENS = {
+    "claude": "claude-code",
+    "codex": "codex",
+    "opencode": "opencode",
+    "copilot": "copilot",
+}
+
+
+def install_ai_tools_for_agents(tools: list[str], state: dict) -> None:
+    """Install Databricks AI Tools for the coding agents that support them.
+
+    Persists an ``ai_tools_installed`` marker so an already-installed agent is
+    skipped without shelling out, making repeat calls (e.g. on every launch) cheap."""
+    installed = set(state.get("ai_tools_installed") or [])
+    mapped = (AITOOLS_AGENT_TOKENS.get(tool) for tool in tools)
+    tokens = [tok for tok in mapped if tok and tok not in installed]
+    if not tokens:
+        return
+    newly_installed = install_ai_tools(tokens, state.get("profile"))
+    if newly_installed:
+        state["ai_tools_installed"] = sorted(installed | set(newly_installed))
+        save_state(state)
+
 
 def normalize_tool(tool: str) -> str:
     normalized = TOOL_ALIASES.get(tool.strip().lower())
@@ -380,6 +405,7 @@ def configure_single_tool(tool: str, state: dict) -> dict:
     available_tools = list(set((state.get("available_tools") or []) + [tool]))
     state["available_tools"] = available_tools
     save_state(state)
+    install_ai_tools_for_agents([tool], state)
     return state
 
 
@@ -410,6 +436,7 @@ def configure_selected_tools(state: dict, tools: list[str]) -> dict:
     existing = state.get("available_tools") or []
     state["available_tools"] = sorted(set(existing) | set(tools))
     save_state(state)
+    install_ai_tools_for_agents(tools, state)
     return state
 
 

src/ucode/cli.py

@@ -16,6 +16,7 @@
     configure_tool,
     ensure_bootstrap_dependencies,
     ensure_provider_state,
+    install_ai_tools_for_agents,
     install_tool_binary,
     normalize_tool,
     provider_permission_error,
@@ -927,6 +928,7 @@ def _launch_tool(
             skip_model_discovery=bool(provider),
             skip_preflight=skip_preflight,
         )
+        install_ai_tools_for_agents([tool], state)
         if provider:
             # Routing through a Model Provider Service pins no Databricks model;
             # the agent uses its own canonical model names (header selects the

src/ucode/databricks.py

@@ -49,7 +49,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"
-MIN_DATABRICKS_CLI_VERSION = (0, 298, 0)
+# v1.0.0 is the release that ships `databricks aitools`.
+MIN_DATABRICKS_CLI_VERSION = (1, 0, 0)
 TOKEN_REFRESH_INTERVAL_SECONDS = 1800
 
 
@@ -548,6 +549,46 @@ def install_databricks_cli() -> None:
     ensure_databricks_cli_version()
 
 
+def install_ai_tools(agents: list[str], profile: str | None = None) -> list[str]:
+    """Install Databricks AI Tools for the given agents (e.g. ``claude-code``).
+
+    Databricks AI Tools is the set of skills and plugins that teach coding
+    agents how to work with Databricks (installed via ``databricks aitools``).
+    Idempotent and best-effort: any failure only warns (surfacing the CLI's
+    own error), since AI Tools aren't required to launch an agent.
+
+    Returns the agents that were installed (empty on failure or when ``agents``
+    is empty), so callers can record a one-time marker."""
+    if not agents:
+        return []
+
+    agents_arg = ",".join(agents)
+    try:
+        with spinner(f"Installing Databricks AI Tools for {agents_arg}..."):
+            run(
+                ["databricks", "aitools", "install", "--agents", agents_arg, "--scope", "global"]
+                + _profile_args(profile),
+                check=True,
+                capture_output=True,
+                text=True,
+                timeout=300,
+            )
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError) as exc:
+        # The CLI version is already guaranteed by ensure_databricks_cli_version,
+        # so any failure here is something else (e.g. an agent binary missing
+        # from PATH). Surface the CLI's own error rather than guessing a cause.
+        detail = getattr(exc, "stderr", None) or ""
+        if isinstance(detail, bytes):  # TimeoutExpired.stderr is bytes even with text=True
+            detail = detail.decode(errors="replace")
+        detail = detail.strip()
+        reason = detail.splitlines()[-1] if detail else str(exc)
+        print_warning(f"Could not install Databricks AI Tools: {reason}")
+        return []
+    else:
+        print_success("Databricks AI Tools installed")
+        return list(agents)
+
+
 def _profile_args(profile: str | None) -> list[str]:
     """Return ``["--profile", profile]`` when set, otherwise an empty list.
 

tests/test_agents_init.py

@@ -14,6 +14,7 @@
     configure_selected_tools,
     default_model_for_tool,
     ensure_tool_binary_available,
+    install_ai_tools_for_agents,
     install_tool_binary,
     normalize_tool,
     provider_permission_error,
@@ -60,6 +61,80 @@ def test_each_agent_exposes_update_check(self):
             assert callable(module.is_update_available), f"{tool} missing is_update_available"
 
 
+class TestInstallAiToolsForAgents:
+    def _stub_install(self, monkeypatch):
+        """Stub install_ai_tools to record calls and echo tokens back as installed."""
+        calls = []
+
+        def fake(tokens, profile):
+            calls.append({"tokens": tokens, "profile": profile})
+            return list(tokens)  # pretend every token installed successfully
+
+        monkeypatch.setattr(agents_mod, "install_ai_tools", fake)
+        monkeypatch.setattr(agents_mod, "save_state", lambda state: None)
+        return calls
+
+    def test_maps_supported_tools_and_drops_others(self, monkeypatch):
+        calls = self._stub_install(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 calls == [{"tokens": ["claude-code", "codex"], "profile": "prof"}]
+
+    def test_records_marker_on_success(self, monkeypatch):
+        self._stub_install(monkeypatch)
+        state = {"profile": "prof"}
+        install_ai_tools_for_agents(["claude", "codex"], state)
+        assert state["ai_tools_installed"] == ["claude-code", "codex"]
+
+    def test_skips_already_marked_tokens(self, monkeypatch):
+        calls = self._stub_install(monkeypatch)
+        state = {"profile": "prof", "ai_tools_installed": ["claude-code"]}
+        install_ai_tools_for_agents(["claude", "codex"], state)
+        # claude-code already marked -> only codex is installed.
+        assert calls == [{"tokens": ["codex"], "profile": "prof"}]
+        assert state["ai_tools_installed"] == ["claude-code", "codex"]
+
+    def test_no_install_when_all_marked(self, monkeypatch):
+        calls = self._stub_install(monkeypatch)
+        state = {"profile": "prof", "ai_tools_installed": ["claude-code"]}
+        install_ai_tools_for_agents(["claude"], state)
+        assert calls == []  # membership check only, no subprocess
+
+    def test_failed_install_not_recorded(self, monkeypatch):
+        # install_ai_tools returns [] on failure -> marker must not be written,
+        # so the next run retries.
+        monkeypatch.setattr(agents_mod, "install_ai_tools", lambda tokens, profile: [])
+        monkeypatch.setattr(agents_mod, "save_state", lambda state: None)
+        state = {"profile": "prof"}
+        install_ai_tools_for_agents(["claude"], state)
+        assert "ai_tools_installed" not in state
+
+
+class TestConfigureWiresAiToolsInstall:
+    """Both configure chokepoints must trigger AI Tools install."""
+
+    def _stub_configure(self, monkeypatch):
+        captured = {}
+        monkeypatch.setattr(agents_mod, "configure_tool", lambda tool, state, model=None: state)
+        monkeypatch.setattr(agents_mod, "save_state", lambda state: None)
+        monkeypatch.setattr(
+            agents_mod,
+            "install_ai_tools",
+            lambda tokens, profile: captured.update(tokens=tokens, profile=profile) or list(tokens),
+        )
+        return captured
+
+    def test_configure_single_tool_triggers_install(self, monkeypatch):
+        captured = self._stub_configure(monkeypatch)
+        agents_mod.configure_single_tool("codex", {"codex_models": ["m"], "profile": "myprof"})
+        assert captured == {"tokens": ["codex"], "profile": "myprof"}
+
+    def test_configure_selected_tools_triggers_install(self, monkeypatch):
+        captured = self._stub_configure(monkeypatch)
+        agents_mod.configure_selected_tools({"profile": "myprof"}, ["codex"])
+        assert captured == {"tokens": ["codex"], "profile": "myprof"}
+
+
 class TestNormalizeTool:
     @pytest.mark.parametrize(
         "alias,expected",

tests/test_databricks.py

@@ -25,6 +25,7 @@
     ensure_databricks_cli_version,
     ensure_pat_bearer,
     get_databricks_token,
+    install_ai_tools,
     list_databricks_apps,
     list_databricks_connections,
     list_genie_spaces,
@@ -1538,19 +1539,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 v0.298.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 v0.299.2")
+        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):
         import ucode.databricks as db_mod
 
-        env = self._fake_databricks(tmp_path, "Databricks CLI v0.297.0")
+        env = self._fake_databricks(tmp_path, "Databricks CLI v0.299.2")
         monkeypatch.setattr("os.environ", env)
         upgraded = []
         monkeypatch.setattr(
@@ -1699,3 +1700,73 @@ 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 TestInstallAiTools:
+    def _capture_run(self, monkeypatch, *, raises=None):
+        calls = []
+
+        def fake_run(args, **kwargs):
+            calls.append(args)
+            if raises is not None:
+                raise raises
+            return subprocess.CompletedProcess(args, 0, "Installed 1 skill.", "")
+
+        monkeypatch.setattr(db_mod, "run", fake_run)
+        return calls
+
+    def test_no_tokens_skips_entirely(self, monkeypatch):
+        calls = self._capture_run(monkeypatch)
+        assert install_ai_tools([]) == []
+        assert calls == []
+
+    def test_invokes_aitools_install(self, monkeypatch):
+        calls = self._capture_run(monkeypatch)
+        # Returns the tokens it installed, so callers can record a marker.
+        assert install_ai_tools(["claude-code", "codex"]) == ["claude-code", "codex"]
+        assert len(calls) == 1
+        cmd = calls[0]
+        assert cmd[:3] == ["databricks", "aitools", "install"]
+        assert "--agents" in cmd and cmd[cmd.index("--agents") + 1] == "claude-code,codex"
+        assert "--scope" in cmd and cmd[cmd.index("--scope") + 1] == "global"
+        assert "--profile" not in cmd
+
+    def test_passes_profile_when_set(self, monkeypatch):
+        calls = self._capture_run(monkeypatch)
+        install_ai_tools(["claude-code"], profile="myprofile")
+        cmd = calls[0]
+        assert "--profile" in cmd and cmd[cmd.index("--profile") + 1] == "myprofile"
+
+    def test_install_failure_is_non_fatal(self, monkeypatch):
+        self._capture_run(monkeypatch, raises=subprocess.CalledProcessError(1, "databricks"))
+        # Must not raise, and returns [] so no marker is recorded — best-effort.
+        assert install_ai_tools(["claude-code"]) == []
+
+    def test_timeout_is_non_fatal(self, monkeypatch):
+        self._capture_run(monkeypatch, raises=subprocess.TimeoutExpired("databricks", 300))
+        install_ai_tools(["claude-code"])
+
+    def test_timeout_stderr_bytes_decoded_in_warning(self, monkeypatch):
+        # TimeoutExpired.stderr is bytes even with text=True; the warning must
+        # decode it, not render a `b'...'` repr.
+        err = subprocess.TimeoutExpired("databricks", 300)
+        err.stderr = b"resolving agents...\ninstall timed out"
+        self._capture_run(monkeypatch, raises=err)
+        warnings = []
+        monkeypatch.setattr(db_mod, "print_warning", warnings.append)
+        install_ai_tools(["claude-code"])
+        assert len(warnings) == 1
+        assert "install timed out" in warnings[0]
+        assert "b'" not in warnings[0]
+
+    def test_failure_surfaces_cli_stderr(self, monkeypatch):
+        # A modern CLI can still fail (e.g. an agent binary missing from PATH);
+        # the warning must show the CLI's real error, not blame the version.
+        err = subprocess.CalledProcessError(1, "databricks")
+        err.stderr = "resolving agents...\ncopilot: cli-not-on-path: could not resolve copilot"
+        self._capture_run(monkeypatch, raises=err)
+        warnings = []
+        monkeypatch.setattr(db_mod, "print_warning", warnings.append)
+        install_ai_tools(["copilot"])
+        assert len(warnings) == 1
+        assert "copilot: cli-not-on-path: could not resolve copilot" in warnings[0]