diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a0accb9be..db8cf5d03 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3792,9 +3792,27 @@ def _raise_model_service_permission_failure(workspace: str, model_service_reason ) +# Probes that returned an accessible model service, keyed by (workspace, token) so a re-login or a +# different credential probes again. `ucode codex` configures once for auto-setup and again on the +# launch path; re-probing a gateway that just answered only gives the flaky listing another chance +# to fail. Failures, empty listings, and inconclusive results are never cached. +_GATEWAY_PROBE_CACHE: dict[tuple[str, str], GatewayProbe] = {} + + +def clear_gateway_probe_cache() -> None: + """Forget successful gateway probes (used by tests).""" + _GATEWAY_PROBE_CACHE.clear() + + def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe: """Return the model service probe, raising if model service access can't be verified.""" + cache_key = (workspace, token) + cached = _GATEWAY_PROBE_CACHE.get(cache_key) + if cached is not None: + return cached model_service_probe = _probe_model_services(workspace, token) + if model_service_probe.resource_available: + _GATEWAY_PROBE_CACHE[cache_key] = model_service_probe if model_service_probe.reachable: # resource_available, reachable-but-empty, and inconclusive are all non-fatal: the # caller surfaces the detail as a warning when no accessible model service came back. @@ -3815,9 +3833,14 @@ def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe "--skip-preflight to launch without this check." ) + if _is_not_found_reason(reason): + raise RuntimeError( + "Databricks Unity Gateway is not enabled on this workspace: model services " + f"({reason}) are not available. See {AI_GATEWAY_DOCS_URL}" + ) raise RuntimeError( - "Databricks Unity Gateway is not enabled on this workspace: model services " - f"({reason}) are not available. See {AI_GATEWAY_DOCS_URL}" + f"Could not verify Databricks Unity Gateway model services on {workspace} " + f"({reason}). See {AI_GATEWAY_DOCS_URL}" ) diff --git a/src/ucode/managed_files.py b/src/ucode/managed_files.py index 0eaa760c4..93909f423 100644 --- a/src/ucode/managed_files.py +++ b/src/ucode/managed_files.py @@ -273,6 +273,33 @@ def _sudo_may_prompt() -> bool: return current_os() is not OS.WINDOWS and sys.stdin.isatty() +_sudo_needs_password_cache: bool | None = None + + +def _sudo_needs_password() -> bool: + """False when sudo is passwordless, so the "Enter password" note would be misleading. + + Only informs that note; whether and how writes happen is unchanged. Skips the probe (and + assumes a prompt) on Windows and in dry-run mode, where no sudo is run.""" + global _sudo_needs_password_cache + if current_os() is OS.WINDOWS or is_dry_run(): + return True + if _sudo_needs_password_cache is None: + try: + result = subprocess_cross_os.run( + [_SUDO, "-n", "true"], capture_output=True, timeout=10, check=False + ) + _sudo_needs_password_cache = result.returncode != 0 + except (OSError, subprocess.SubprocessError): + _sudo_needs_password_cache = True + return _sudo_needs_password_cache + + +def _print_password_note(message: str) -> None: + if _sudo_needs_password(): + print_note(message) + + def request_explicit_managed_writes(command: str | None = "ug configure") -> None: """Mark this run as the explicit ``command`` (``ug configure`` or ``ug revert``); ``None`` clears it. Error messages name the command to rerun. @@ -389,17 +416,17 @@ def _print_managed_write_permission(display: str) -> None: if _managed_write_worker is not None and _managed_write_worker.usable: _managed_write_notice_shown = True return - print_note("Enter password once to configure machine-wide coding agent settings.") + _print_password_note("Enter password once to configure machine-wide coding agent settings.") _managed_write_notice_shown = True return if not _managed_write_batch: - print_note(f"Enter password to configure settings for {display}.") + _print_password_note(f"Enter password to configure settings for {display}.") return if _managed_write_notice_shown: return displays = " and ".join(_managed_write_batch) - print_note(f"Enter password to configure settings for {displays}.") + _print_password_note(f"Enter password to configure settings for {displays}.") _managed_write_notice_shown = True diff --git a/tests/conftest.py b/tests/conftest.py index 1f3470782..bfdc36891 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -123,6 +123,9 @@ def reject_privileged_write(path, _desired_text): ) monkeypatch.setattr(managed_files_mod, "_sudo_replace", reject_privileged_write) + # Never shell out to the host's sudo from the password-note check; assume it would prompt. + monkeypatch.setattr(managed_files_mod, "_sudo_needs_password", lambda: True) + databricks_mod.clear_gateway_probe_cache() # Windows CI runners are elevated; never let an unmocked test write C:\Program Files. monkeypatch.setattr(managed_files_mod.system_cross_os, "is_admin", lambda: False) monkeypatch.setattr(managed_files_mod, "_windows_replace", reject_privileged_write) diff --git a/tests/test_cli.py b/tests/test_cli.py index f460fc67b..86404d242 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -6377,6 +6377,24 @@ def test_local_gateway_probe_failures_do_not_print_success( assert "v2" not in message.lower() assert "v3" not in message.lower() + def test_second_configure_in_process_does_not_reprobe_the_gateway(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + calls: list[str] = [] + + def fake_get(url, token, **kwargs): + calls.append(url) + return {"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", db_mod.probe_unity_gateway_capabilities + ) + + cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert len(calls) == 1 + def test_use_pat_without_pat_profile_raises(self, monkeypatch): cli_mod, logins, _, _ = self._stub_deps(monkeypatch, pat_token=None) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index bcf26da70..205bc8561 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -3689,6 +3689,70 @@ def fake_get(url, token, **kwargs): assert "legacy" not in message.lower() assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50"] + def test_non_404_non_transient_failure_does_not_claim_not_enabled(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, **kwargs: (None, "HTTP 400 Bad Request: malformed"), + ) + + with pytest.raises(RuntimeError, match="Could not verify") as excinfo: + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + message = str(excinfo.value) + assert "HTTP 400 Bad Request: malformed" in message + assert "not enabled" not in message + assert "transient" not in message + + def test_successful_probe_is_memoized_per_workspace_and_failures_are_not(self, monkeypatch): + calls: list[str] = [] + outcomes = iter( + [ + (None, "HTTP 400 Bad Request: malformed"), + ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), + ] + ) + + def fake_get(url, token, **kwargs): + calls.append(url) + return next(outcomes) + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + with pytest.raises(RuntimeError): + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + first = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + second = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + assert first.resource_available + assert second == first + assert len(calls) == 2 + + def test_probe_cache_is_keyed_by_token_and_skips_empty_listings(self, monkeypatch): + calls: list[str] = [] + outcomes = iter( + [ + ({"model_services": []}, None), + ({"model_services": []}, None), + ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), + ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), + ] + ) + + def fake_get(url, token, **kwargs): + calls.append(token) + return next(outcomes) + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + assert not db_mod.probe_unity_gateway_capabilities(WS, "token-a").resource_available + assert not db_mod.probe_unity_gateway_capabilities(WS, "token-a").resource_available + assert db_mod.probe_unity_gateway_capabilities(WS, "token-a").resource_available + assert db_mod.probe_unity_gateway_capabilities(WS, "token-a").resource_available + assert db_mod.probe_unity_gateway_capabilities(WS, "token-b").resource_available + + assert calls == ["token-a", "token-a", "token-a", "token-b"] + def test_model_service_auth_failure_raises(self, monkeypatch): calls: list[str] = [] diff --git a/tests/test_managed_files.py b/tests/test_managed_files.py index 11b87d23c..766f501a5 100644 --- a/tests/test_managed_files.py +++ b/tests/test_managed_files.py @@ -24,6 +24,7 @@ _REAL_MANAGED_FILES_SUPPORTED = managed_files.managed_files_supported _REAL_WINDOWS_REPLACE = managed_files._windows_replace _REAL_SUBPROCESS_RUN = managed_files.subprocess_cross_os.run +_REAL_SUDO_NEEDS_PASSWORD = managed_files._sudo_needs_password @pytest.fixture(autouse=True) @@ -2028,3 +2029,65 @@ def test_sudo_command_derives_prompt_policy(monkeypatch, interactive): expected = ["/usr/bin/sudo", *([] if interactive else ["-n"]), "cp", "a", "b"] assert managed_files._sudo_command("cp", "a", "b") == expected + + +class TestPasswordNote: + @pytest.fixture(autouse=True) + def _real_check(self, monkeypatch): + monkeypatch.setattr(managed_files, "_sudo_needs_password", _REAL_SUDO_NEEDS_PASSWORD) + monkeypatch.setattr(managed_files, "_sudo_needs_password_cache", None) + monkeypatch.setattr(managed_files, "current_os", lambda: managed_files.OS.LINUX) + + @staticmethod + def _fake_sudo(monkeypatch, returncode): + commands: list[list[str]] = [] + + def run(command, **kwargs): + commands.append(command) + return subprocess.CompletedProcess(command, returncode) + + monkeypatch.setattr(managed_files.subprocess_cross_os, "run", run) + return commands + + def test_passwordless_sudo_suppresses_the_note_and_is_checked_once(self, monkeypatch): + notes: list[str] = [] + commands = self._fake_sudo(monkeypatch, 0) + monkeypatch.setattr(managed_files, "print_note", notes.append) + + managed_files._print_managed_write_permission("Claude Code") + managed_files._print_managed_write_permission("Codex") + + assert notes == [] + assert commands == [["/usr/bin/sudo", "-n", "true"]] + + def test_sudo_needing_a_password_still_prints_the_note(self, monkeypatch): + notes: list[str] = [] + self._fake_sudo(monkeypatch, 1) + monkeypatch.setattr(managed_files, "print_note", notes.append) + + managed_files._print_managed_write_permission("Codex") + + assert notes == ["Enter password to configure settings for Codex."] + + def test_passwordless_batch_still_reports_success(self, monkeypatch): + successes: list[str] = [] + self._fake_sudo(monkeypatch, 0) + monkeypatch.setattr(managed_files, "print_note", lambda _message: None) + monkeypatch.setattr(managed_files, "print_success", successes.append) + + with managed_files.managed_write_batch(["Codex"]): + managed_files._print_managed_write_permission("Codex") + + assert successes == ["Settings configured for Codex"] + + def test_no_sudo_call_on_windows_or_in_dry_run(self, monkeypatch): + commands = self._fake_sudo(monkeypatch, 0) + + monkeypatch.setattr(managed_files, "current_os", lambda: managed_files.OS.WINDOWS) + assert managed_files._sudo_needs_password() is True + + monkeypatch.setattr(managed_files, "current_os", lambda: managed_files.OS.LINUX) + config_io.set_dry_run(True) + assert managed_files._sudo_needs_password() is True + + assert commands == []