Install Databricks AI Tools as part of the configuration step (#230)

* Install Databricks AI Tools during ucode configure

`ucode configure` now runs `databricks aitools install` for each coding
agent it sets up (claude-code, codex, opencode, copilot), so agents work
with Databricks out of the box. Best-effort: a failure only warns with the
CLI's own error, since AI Tools aren't required to launch an agent.

Co-authored-by: Isaac

* Require Databricks CLI v1.0.0 for AI Tools; simplify install

v1.0.0 is the release that ships `databricks aitools install --agents
... --scope`, so bump MIN_DATABRICKS_CLI_VERSION to (1, 0, 0) and drop the
decorative AITOOLS_MIN_CLI_VERSION. `ensure_databricks_cli_version` already
runs (and auto-upgrades) before every path that reaches install_ai_tools,
so the install no longer guesses a version cause on failure: it just
surfaces the CLI's own error.

Co-authored-by: Isaac

* Decode bytes stderr in AI Tools install warning

subprocess.TimeoutExpired.stderr is bytes even when the call used text=True
(a CPython quirk), so a timed-out `databricks aitools install` rendered its
warning as `b'...'`. Decode bytes before building the reason string, and
cover the timeout-with-stderr path with a test.

Co-authored-by: Isaac

* Make Databricks AI Tools install opt-out via configure

Adds `configure --enable-databricks-ai-tools/--disable-databricks-ai-tools`
(installed by default) plus an interactive prompt that defaults yes but
honors a prior opt-out. The choice persists per workspace as
`databricks_ai_tools_enabled` and gates install_ai_tools_for_agents. Interim
client-side control ahead of admin-managed configuration.

Co-authored-by: Isaac

* Ask the AI Tools opt-out last in interactive configure

Move the "Install Databricks AI Tools?" prompt out of the top-level configure
body into configure_workspace_command, so it runs after the workspace and
agent picks instead of before them. Reword it to explain it adds Databricks
skills and plugins. An explicit --enable/--disable flag still answers upfront
and skips the prompt.

Co-authored-by: Isaac

---------

Co-authored-by: Rohit Agrawal <10210921+rohita5l@users.noreply.github.com>

Author: lennartkats-db

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,23 @@
 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
+    (gemini/pi have no ``aitools`` support and are dropped)."""
+    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]
+    install_ai_tools(agents, state.get("profile"))
+
 
 def normalize_tool(tool: str) -> str:
     normalized = TOOL_ALIASES.get(tool.strip().lower())
@@ -380,6 +398,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 +429,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

@@ -80,6 +80,7 @@
     prompt_for_selection,
     prompt_for_tools,
     prompt_for_workspace,
+    prompt_yes_no_default,
     set_verbosity,
     spinner,
     status_badge,
@@ -218,6 +219,7 @@ def configure_shared_state(
     skip_model_discovery: bool = False,
     skip_preflight: bool = False,
     fable_enabled: bool | None = None,
+    databricks_ai_tools_enabled: bool | None = None,
 ) -> dict:
     """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state.
 
@@ -249,6 +251,14 @@ def configure_shared_state(
         use_pat = bool(prior_state.get("use_pat")) and previous_workspace == workspace
     if fable_enabled is None:
         fable_enabled = bool(prior_state.get("fable_enabled")) and previous_workspace == workspace
+    if databricks_ai_tools_enabled is None:
+        # Opt-out: on by default. With no flag, keep this workspace's prior
+        # choice but don't inherit another workspace's opt-out.
+        disabled = (
+            prior_state.get("databricks_ai_tools_enabled") is False
+            and previous_workspace == workspace
+        )
+        databricks_ai_tools_enabled = not disabled
     fetch_all = tools is None
 
     # Assemble the shared workspace state that doesn't depend on model discovery:
@@ -278,6 +288,7 @@ def configure_shared_state(
         state["fable_enabled"] = True
     else:
         state.pop("fable_enabled", None)
+    state["databricks_ai_tools_enabled"] = databricks_ai_tools_enabled
     state["base_urls"] = build_shared_base_urls(workspace)
 
     if skip_preflight:
@@ -436,6 +447,7 @@ def _configure_shared_workspace_states(
     force_login: bool,
     use_pat: bool = False,
     fable_enabled: bool | None = None,
+    databricks_ai_tools_enabled: bool | None = None,
 ) -> list[dict]:
     if not workspaces:
         raise RuntimeError("At least one workspace must be provided.")
@@ -449,6 +461,7 @@ def _configure_shared_workspace_states(
                 force_login=force_login,
                 use_pat=use_pat,
                 fable_enabled=fable_enabled,
+                databricks_ai_tools_enabled=databricks_ai_tools_enabled,
             )
         )
     return states
@@ -529,6 +542,7 @@ def configure_workspace_command(
     use_pat: bool = False,
     skip_validate: bool = False,
     fable_enabled: bool | None = None,
+    databricks_ai_tools_enabled: bool | None = None,
 ) -> int:
     if tool is not None and selected_tools is not None:
         raise RuntimeError("Use either --agent or --agents, not both.")
@@ -547,6 +561,7 @@ def configure_workspace_command(
             force_login=True,
             use_pat=use_pat,
             fable_enabled=fable_enabled,
+            databricks_ai_tools_enabled=databricks_ai_tools_enabled,
         )
         state = states[0]
         state = configure_single_tool(tool, state)
@@ -584,6 +599,7 @@ def configure_workspace_command(
         force_login=True,
         use_pat=use_pat,
         fable_enabled=fable_enabled,
+        databricks_ai_tools_enabled=databricks_ai_tools_enabled,
     )
     state = states[0]
     save_state(state)
@@ -632,6 +648,16 @@ def configure_workspace_command(
         for tool_name in picked:
             state = _maybe_select_provider_service(tool_name, state)
 
+    # Last question in the interactive flow: opt out of AI Tools. When a flag
+    # already decided it, configure_shared_state persisted that; skip the prompt.
+    # The default is the resolved prior choice, so Enter won't undo a past opt-out.
+    if databricks_ai_tools_enabled is None and offer_provider:
+        state["databricks_ai_tools_enabled"] = prompt_yes_no_default(
+            "Install Databricks AI Tools for your coding agents? "
+            "This adds Databricks skills and plugins.",
+            default=state.get("databricks_ai_tools_enabled", True),
+        )
+
     state = configure_selected_tools(state, picked)
 
     summary_lines = [f"[bold]Workspace:[/bold] [cyan]{state['workspace']}[/cyan]"]
@@ -1102,6 +1128,15 @@ def configure(
             "it configures Claude Code directly since Fable is Claude-only.",
         ),
     ] = None,
+    enable_databricks_ai_tools: Annotated[
+        bool | None,
+        typer.Option(
+            "--enable-databricks-ai-tools/--disable-databricks-ai-tools",
+            help="Install Databricks AI Tools (skills + plugins that teach agents to use "
+            "Databricks) for the configured agents. Installed by default; pass "
+            "--disable-databricks-ai-tools to opt out.",
+        ),
+    ] = None,
     tracing: Annotated[
         bool,
         typer.Option(
@@ -1167,6 +1202,8 @@ def configure(
         # target claude instead of dropping into the interactive agent picker.
         if enable_fable is not None and agent is None and agents is None:
             agent = "claude"
+        if enable_databricks_ai_tools is not None:
+            skip_kwargs["databricks_ai_tools_enabled"] = enable_databricks_ai_tools
         if agent is not None:
             tool = normalize_tool(agent)
             install_tool_binary(

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,41 @@ def install_databricks_cli() -> None:
     ensure_databricks_cli_version()
 
 
+def install_ai_tools(agent_tokens: list[str], profile: str | None = None) -> None:
+    """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."""
+    if not agent_tokens:
+        return
+
+    agents_arg = ",".join(agent_tokens)
+    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}")
+    else:
+        print_success("Databricks AI Tools installed")
+
+
 def _profile_args(profile: str | None) -> list[str]:
     """Return ``["--profile", profile]`` when set, otherwise an empty list.
 

src/ucode/ui.py

@@ -302,6 +302,23 @@ def prompt_yes_no(prompt: str) -> bool:
         print_err("Please answer yes or no.")
 
 
+def prompt_yes_no_default(prompt: str, *, default: bool) -> bool:
+    """Empty answer or closed stdin (EOF) takes ``default`` (no abort on piped runs)."""
+    hint = "(Y/n)" if default else "(y/N)"
+    while True:
+        try:
+            response = console.input(f"{label(prompt)} {muted(hint)} {muted('›')} ").strip().lower()
+        except EOFError:
+            return default
+        if not response:
+            return default
+        if response in {"y", "yes"}:
+            return True
+        if response in {"n", "no"}:
+            return False
+        print_err("Please answer yes or no.")
+
+
 def prompt_for_choice(prompt: str, options: list[tuple[str, str]]) -> str:
     console.print()
     for index, (_, option_label) in enumerate(options, start=1):

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,70 @@ def test_each_agent_exposes_update_check(self):
             assert callable(module.is_update_available), f"{tool} missing is_update_available"
 
 
+class TestInstallAiToolsForAgents:
+    def _capture(self, monkeypatch):
+        captured = {}
+        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):
+        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"}
+
+    def test_installed_by_default(self, monkeypatch):
+        # Opt-out: absent flag means install.
+        captured = self._capture(monkeypatch)
+        install_ai_tools_for_agents(["claude"], {"profile": "p"})
+        assert captured == {"agents": ["claude-code"], "profile": "p"}
+
+    def test_skipped_when_disabled(self, monkeypatch):
+        # `configure --disable-databricks-ai-tools` persists this False.
+        captured = self._capture(monkeypatch)
+        install_ai_tools_for_agents(
+            ["claude"], {"profile": "p", "databricks_ai_tools_enabled": False}
+        )
+        assert captured == {}  # install_ai_tools never called
+
+
+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 agents, profile: captured.update(agents=agents, profile=profile),
+        )
+        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 == {"agents": ["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 == {"agents": ["codex"], "profile": "myprof"}
+
+    def test_configure_single_tool_respects_disable(self, monkeypatch):
+        captured = self._stub_configure(monkeypatch)
+        agents_mod.configure_single_tool(
+            "codex",
+            {"codex_models": ["m"], "profile": "myprof", "databricks_ai_tools_enabled": False},
+        )
+        assert captured == {}
+
+
 class TestNormalizeTool:
     @pytest.mark.parametrize(
         "alias,expected",