diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 0785342f..4e7cc640 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -264,6 +264,14 @@ def ensure_bootstrap_dependencies( update_existing: bool = False, prompt_optional_updates: bool = True, ) -> None: + if not update_existing: + if not shutil.which("databricks"): + raise RuntimeError( + "Databricks CLI is not installed (`databricks` was not found on PATH). " + "Run `ucode configure` to install it." + ) + ensure_tool_binary_available(tool) + return install_databricks_cli() install_tool_binary( tool, diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index af6faf19..e1705a40 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -50,6 +50,7 @@ CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" +CLAUDE_USER_CONFIG_PATH = Path.home() / ".claude.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" @@ -246,6 +247,41 @@ def _web_search_mcp_entry(workspace: str, search_model: str, profile: str | None } +def _cached_claude_version(state: dict) -> str: + """Return Claude's version without spawning Node when the binary is unchanged.""" + binary = shutil.which(SPEC["binary"]) + try: + binary_mtime_ns = Path(binary).stat().st_mtime_ns if binary else None + except OSError: + binary_mtime_ns = None + if ( + binary_mtime_ns is not None + and state.get("claude_binary_mtime_ns") == binary_mtime_ns + and isinstance(state.get("claude_version"), str) + ): + return state["claude_version"] + if binary_mtime_ns is not None: + try: + settings_are_current = CLAUDE_SETTINGS_PATH.stat().st_mtime_ns >= binary_mtime_ns + except OSError: + settings_are_current = False + if settings_are_current: + env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env") + headers = env.get("ANTHROPIC_CUSTOM_HEADERS") if isinstance(env, dict) else None + match = re.search(r"(?:^|\s)claude/([^\s]+)", headers or "") + if match: + state["claude_version"] = match.group(1) + state["claude_binary_mtime_ns"] = binary_mtime_ns + return match.group(1) + version = agent_version(SPEC["binary"]) + state["claude_version"] = version + if binary_mtime_ns is not None: + state["claude_binary_mtime_ns"] = binary_mtime_ns + else: + state.pop("claude_binary_mtime_ns", None) + return version + + def render_overlay( workspace: str, model: str | None, @@ -260,6 +296,7 @@ def render_overlay( relayed_base_url: str | None = None, route_root_model: str | None = None, custom_model: str | None = None, + agent_version_value: str | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -293,7 +330,8 @@ def render_overlay( # traffic to ucode. header_lines = [ "x-databricks-use-coding-agent-mode: true", - f"User-Agent: ucode/{ucode_version()} claude/{agent_version('claude')}", + f"User-Agent: ucode/{ucode_version()} claude/" + f"{agent_version_value or agent_version('claude')}", ] if provider: header_lines.append(f"Databricks-Model-Provider-Service: {provider}") @@ -420,14 +458,24 @@ def _maybe_add_1m_suffix(model: str) -> str: return f"{model}[1m]" if should_suffix else model -def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool: +def _register_web_search_mcp( + workspace: str, + search_model: str, + profile: str | None = None, +) -> bool: """Register (or replace) the web_search MCP server in Claude Code's user - scope via `claude mcp add-json`. Removes any prior entry first so re-runs - pick up changes to the workspace, model, or ucode binary path. + scope via `claude mcp add-json`. An unchanged, verified user entry is left + alone; changed entries are removed from every scope before being re-added. Returns True if registration succeeded. Failures are non-blocking: we warn and return False so the rest of `ucode claude` setup can complete. """ + entry = _web_search_mcp_entry(workspace, search_model, profile) + user_config = read_json_safe(CLAUDE_USER_CONFIG_PATH) + servers = user_config.get("mcpServers") + if isinstance(servers, dict) and servers.get(WEB_SEARCH_MCP_NAME) == entry: + return True + # Imported lazily to avoid a circular import via ucode.mcp -> ucode.agents. from ucode.mcp import ( MCP_CLEANUP_SCOPES, @@ -441,7 +489,6 @@ def _register_web_search_mcp(workspace: str, search_model: str, profile: str | N except RuntimeError: # Best-effort cleanup of stale entries — keep going. pass - entry = _web_search_mcp_entry(workspace, search_model, profile) try: add_claude_mcp_server(WEB_SEARCH_MCP_NAME, entry) except RuntimeError as exc: @@ -517,6 +564,7 @@ def write_tool_config( relayed_base_url=relayed_base_url, route_root_model=route_root_model, custom_model=custom_model, + agent_version_value=_cached_claude_version(state), ) tracing_env_vars = tracing_env(state, "claude") stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index e98faf95..d4dcd66e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -55,6 +55,8 @@ # v1.0.0 is the release that ships `databricks aitools`. MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 +_TOKEN_REUSE_SECONDS = 30.0 +_TOKEN_CACHE: dict[tuple[str, str | None], tuple[str, float]] = {} # Substrings the Databricks CLI emits when it loses the token-cache write lock # to a concurrent `databricks auth token` (e.g. another ucode helper process or # MLflow tracing refreshing the shared ~/.databricks/token-cache.json at the same @@ -677,6 +679,22 @@ def build_databricks_cli_env(workspace: str, profile: str | None = None) -> dict return env +def clear_token_cache() -> None: + """Clear process-local token reuse. Primarily useful for tests and logout flows.""" + _TOKEN_CACHE.clear() + + +def _cached_token(workspace: str, profile: str | None) -> str | None: + cached = _TOKEN_CACHE.get((workspace, profile)) + if cached is None: + return None + token, cached_at = cached + if time.monotonic() - cached_at >= _TOKEN_REUSE_SECONDS: + _TOKEN_CACHE.pop((workspace, profile), None) + return None + return token + + def workspace_hostname(workspace: str) -> str: parsed = urlparse(normalize_workspace_url(workspace)) if not parsed.hostname: @@ -812,6 +830,8 @@ def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> boo # profiles for the same host, `databricks auth token --host …` refuses # to disambiguate without --profile, so resolve it from the host here. profile = profile or find_profile_name_for_host(workspace) + if _cached_token(workspace, profile): + return True try: env = build_databricks_cli_env(workspace, profile) result = run( @@ -838,7 +858,11 @@ def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> boo if result.returncode != 0: return False data = json.loads(result.stdout or "{}") - return bool(data.get("access_token")) + token = data.get("access_token") + if not isinstance(token, str) or not token: + return False + _TOKEN_CACHE[(workspace, profile)] = (token, time.monotonic()) + return True except (json.JSONDecodeError, OSError, subprocess.TimeoutExpired) as exc: _debug("has_valid_databricks_auth", f"exception: {type(exc).__name__}: {exc}") return False @@ -1053,6 +1077,11 @@ def get_databricks_token( # See has_valid_databricks_auth: resolve the profile from the host when # the caller didn't supply one, so duplicate-host cfgs don't break us. profile = profile or find_profile_name_for_host(workspace) + if not force_refresh: + cached = _cached_token(workspace, profile) + if cached: + _debug("get_databricks_token", "using process-local cached token") + return cached env = build_databricks_cli_env(workspace, profile) cmd = [ "databricks", @@ -1151,6 +1180,7 @@ def _fetch_with_lock_retry() -> str: "Run `databricks auth login` to re-authenticate." f"{stale_profile_hint}" ) + _TOKEN_CACHE[(workspace, profile)] = (token, time.monotonic()) return token diff --git a/tests/conftest.py b/tests/conftest.py index 4ad5f4d1..6852b848 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -48,6 +48,7 @@ def reject_privileged_write(path, _desired_text): # The model-services listing is memoized for the life of the process, so without this a cached # result would leak into the next test and make a stubbed listing look like it was never called. databricks_mod.clear_model_services_cache() + databricks_mod.clear_token_cache() def _workspace() -> str: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 112ea05c..7e0d7319 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -337,6 +337,39 @@ def test_headers_newline_delimited(self, monkeypatch): assert "\n" in self._ua(monkeypatch) +class TestCachedClaudeVersion: + def test_reuses_version_while_binary_is_unchanged(self, monkeypatch, tmp_path): + binary = tmp_path / "claude" + binary.write_text("stub") + calls: list[str] = [] + monkeypatch.setattr(claude.shutil, "which", lambda _name: str(binary)) + monkeypatch.setattr(claude, "agent_version", lambda name: calls.append(name) or "2.1.136") + state: dict = {} + + assert claude._cached_claude_version(state) == "2.1.136" + assert claude._cached_claude_version(state) == "2.1.136" + assert calls == ["claude"] + + def test_reuses_version_from_current_settings_on_first_launch(self, monkeypatch, tmp_path): + binary = tmp_path / "claude" + binary.write_text("stub") + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text( + json.dumps( + {"env": {"ANTHROPIC_CUSTOM_HEADERS": "User-Agent: ucode/1.0 claude/2.1.136"}} + ) + ) + monkeypatch.setattr(claude.shutil, "which", lambda _name: str(binary)) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr( + claude, + "agent_version", + lambda _name: pytest.fail("must not start Claude for an existing version"), + ) + + assert claude._cached_claude_version({}) == "2.1.136" + + class TestRenderOverlayWebSearchDisable: def test_settings_overlay_never_includes_mcp_servers(self): # MCP servers belong in ~/.claude.json, not settings.json. @@ -436,7 +469,9 @@ def _common_patches(self, monkeypatch, calls): monkeypatch.setattr( claude, "_register_web_search_mcp", - lambda ws, model, profile=None: calls.append(("register", ws, model)), + lambda ws, model, profile=None, **_kwargs: ( + calls.append(("register", ws, model)) or True + ), ) def test_registers_mcp_when_codex_model_available(self, monkeypatch): @@ -566,6 +601,30 @@ def test_relayed_skips_managed_write(self, monkeypatch): class TestRegisterWebSearchMcp: + @pytest.fixture(autouse=True) + def _isolated_user_config(self, monkeypatch, tmp_path): + monkeypatch.setattr(claude, "CLAUDE_USER_CONFIG_PATH", tmp_path / ".claude.json") + + def test_current_user_entry_skips_all_claude_subprocesses(self, monkeypatch, tmp_path): + import ucode.mcp as mcp_mod + + entry = claude._web_search_mcp_entry(WS, "databricks-gpt-5", "profile") + config_path = tmp_path / ".claude.json" + config_path.write_text(json.dumps({"mcpServers": {"web_search": entry}})) + monkeypatch.setattr(claude, "CLAUDE_USER_CONFIG_PATH", config_path) + monkeypatch.setattr( + mcp_mod, + "remove_claude_mcp_server", + lambda *a, **k: pytest.fail("must not remove an unchanged entry"), + ) + monkeypatch.setattr( + mcp_mod, + "add_claude_mcp_server", + lambda *a, **k: pytest.fail("must not add an unchanged entry"), + ) + + assert claude._register_web_search_mcp(WS, "databricks-gpt-5", "profile") + def test_clears_existing_then_adds(self, monkeypatch): import ucode.mcp as mcp_mod diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index f346af90..d1188b6e 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -592,6 +592,24 @@ def test_ensure_tool_binary_available_raises_when_missing(self, monkeypatch): ensure_tool_binary_available("opencode") +class TestBootstrapDependencies: + @pytest.mark.parametrize("tool", TOOL_SPECS) + def test_established_launch_only_checks_path(self, monkeypatch, tool): + monkeypatch.setattr(agents_mod.shutil, "which", lambda binary: f"/usr/bin/{binary}") + monkeypatch.setattr( + agents_mod, + "install_databricks_cli", + lambda: pytest.fail("must not probe the Databricks CLI version"), + ) + monkeypatch.setattr( + agents_mod, + "install_tool_binary", + lambda *a, **k: pytest.fail("must not start the agent CLI"), + ) + + agents_mod.ensure_bootstrap_dependencies(tool, update_existing=False) + + class TestConfigureSelectedTools: def test_merges_with_existing_available_tools(self, monkeypatch): """Configuring a new tool should not drop previously-configured tools diff --git a/tests/test_databricks.py b/tests/test_databricks.py index ffd9636b..0047b665 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1482,6 +1482,34 @@ def test_returns_token_on_success(self, tmp_path, monkeypatch): token = get_databricks_token(WS) assert token == "good-token" + def test_reuses_token_within_one_launch(self, monkeypatch): + calls: list[list[str]] = [] + + def fake_run(args, **kwargs): + calls.append(args) + return subprocess.CompletedProcess( + args, 0, stdout='{"access_token": "good-token"}', stderr="" + ) + + monkeypatch.setattr(db_mod, "run", fake_run) + assert get_databricks_token(WS, profile="stablebox") == "good-token" + assert get_databricks_token(WS, profile="stablebox") == "good-token" + assert len(calls) == 1 + + def test_auth_validation_token_is_reused(self, monkeypatch): + calls: list[list[str]] = [] + + def fake_run(args, **kwargs): + calls.append(args) + return subprocess.CompletedProcess( + args, 0, stdout='{"access_token": "good-token"}', stderr="" + ) + + monkeypatch.setattr(db_mod, "run", fake_run) + assert db_mod.has_valid_databricks_auth(WS, profile="stablebox") + assert get_databricks_token(WS, profile="stablebox") == "good-token" + assert len(calls) == 1 + def test_strips_ambient_profile_when_profile_not_provided(self, tmp_path, monkeypatch): profile_log = tmp_path / "profile" env = self._fake_databricks(