diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 127cc7a..e91c9a3 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,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 diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a0ef7c4..3dd9c54 100644 --- a/src/ucode/cli.py +++ b/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 diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b38bca9..48fd8f2 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,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. diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 1a90ff8..78550bb 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,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", diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 40af048..53d17e0 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,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]