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..897f4ac 100644 --- a/src/ucode/databricks.py +++ b/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. 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..38649f2 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, @@ -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,72 @@ 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_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]