From 4c4074a6842f138c0850f110d700a0cc320db36a Mon Sep 17 00:00:00 2001 From: Lennart Kats Date: Tue, 21 Jul 2026 22:40:23 +0200 Subject: [PATCH 1/3] 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 --- src/ucode/agents/__init__.py | 20 +++++++++++++ src/ucode/databricks.py | 40 +++++++++++++++++++++++++ tests/test_agents_init.py | 39 ++++++++++++++++++++++++ tests/test_databricks.py | 57 ++++++++++++++++++++++++++++++++++++ 4 files changed, 156 insertions(+) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 127cc7a..92ad2c1 100644 --- a/src/ucode/agents/__init__.py +++ b/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], profile: str | None) -> None: + """Install Databricks AI Tools for the coding agents that support them. + + ``tools`` are ucode tool keys (e.g. ``"claude"``); we map each to its + ``databricks aitools`` agent id and drop the unsupported ones (gemini, pi).""" + tokens = [AITOOLS_AGENT_TOKENS[tool] for tool in tools if tool in AITOOLS_AGENT_TOKENS] + install_ai_tools(tokens, 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.get("profile")) 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.get("profile")) return state diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b38bca9..75d7e5b 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -50,6 +50,8 @@ ) AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" MIN_DATABRICKS_CLI_VERSION = (0, 298, 0) +# Minimum CLI version that ships `databricks aitools` (Databricks AI Tools). +AITOOLS_MIN_CLI_VERSION = (1, 1, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 @@ -548,6 +550,44 @@ 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: + # Surface the CLI's own error rather than assuming a version cause: + # `aitools` also fails when an agent binary isn't on PATH, etc. An old + # CLI (before `aitools` shipped) is only one possible reason, offered + # as a hint rather than asserted as the cause. + required = ".".join(str(n) for n in AITOOLS_MIN_CLI_VERSION) + detail = (getattr(exc, "stderr", None) or "").strip() + reason = detail.splitlines()[-1] if detail else str(exc) + print_warning( + f"Could not install Databricks AI Tools: {reason}. " + f"This needs Databricks CLI v{required}+ (`databricks aitools`); " + "upgrade the CLI if needed, then re-run `ucode configure` to add them." + ) + else: + print_success("Databricks AI Tools installed") + + def _profile_args(profile: str | None) -> list[str]: """Return ``["--profile", profile]`` when set, otherwise an empty list. diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 1a90ff8..c98ae35 100644 --- a/tests/test_agents_init.py +++ b/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,44 @@ def test_each_agent_exposes_update_check(self): assert callable(module.is_update_available), f"{tool} missing is_update_available" +class TestInstallAiToolsForAgents: + def test_maps_supported_tools_and_drops_others(self, monkeypatch): + captured = {} + monkeypatch.setattr( + agents_mod, + "install_ai_tools", + lambda tokens, profile: captured.update(tokens=tokens, profile=profile), + ) + # gemini and pi aren't supported by `databricks aitools`, so they drop. + install_ai_tools_for_agents(["claude", "codex", "gemini", "pi"], "prof") + assert captured == {"tokens": ["claude-code", "codex"], "profile": "prof"} + + +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), + ) + 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", diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 40af048..95b8c03 100644 --- a/tests/test_databricks.py +++ b/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, @@ -1699,3 +1700,59 @@ 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) + install_ai_tools([]) + assert calls == [] + + def test_invokes_aitools_install(self, monkeypatch): + calls = self._capture_run(monkeypatch) + install_ai_tools(["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 — AI Tools are best-effort. + 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_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] From 635e3c799ea71cfecbee0219fb72efb6cf7eb5ce Mon Sep 17 00:00:00 2001 From: Lennart Kats Date: Tue, 21 Jul 2026 22:59:08 +0200 Subject: [PATCH 2/3] 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 --- src/ucode/databricks.py | 19 ++++++------------- tests/test_databricks.py | 6 +++--- 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 75d7e5b..d84760d 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -49,9 +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) -# Minimum CLI version that ships `databricks aitools` (Databricks AI Tools). -AITOOLS_MIN_CLI_VERSION = (1, 1, 0) +# v1.0.0 is the release that ships `databricks aitools`. +MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 @@ -572,18 +571,12 @@ def install_ai_tools(agent_tokens: list[str], profile: str | None = None) -> Non timeout=300, ) except (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError) as exc: - # Surface the CLI's own error rather than assuming a version cause: - # `aitools` also fails when an agent binary isn't on PATH, etc. An old - # CLI (before `aitools` shipped) is only one possible reason, offered - # as a hint rather than asserted as the cause. - required = ".".join(str(n) for n in AITOOLS_MIN_CLI_VERSION) + # 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 "").strip() reason = detail.splitlines()[-1] if detail else str(exc) - print_warning( - f"Could not install Databricks AI Tools: {reason}. " - f"This needs Databricks CLI v{required}+ (`databricks aitools`); " - "upgrade the CLI if needed, then re-run `ucode configure` to add them." - ) + print_warning(f"Could not install Databricks AI Tools: {reason}") else: print_success("Databricks AI Tools installed") diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 95b8c03..d730dee 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1539,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( From b72896d4a869352dffd8499ed92f94aaf8d3ca7d Mon Sep 17 00:00:00 2001 From: Lennart Kats Date: Wed, 22 Jul 2026 11:33:12 +0200 Subject: [PATCH 3/3] 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 --- src/ucode/databricks.py | 5 ++++- tests/test_databricks.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d84760d..897f4ac 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -574,7 +574,10 @@ def install_ai_tools(agent_tokens: list[str], profile: str | None = None) -> Non # 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 "").strip() + 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: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index d730dee..38649f2 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1745,6 +1745,19 @@ 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.