diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 867d980..b6069c5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,11 +11,7 @@ permissions: jobs: test: - runs-on: - group: databricks-protected-runner-group - labels: linux-ubuntu-latest - env: - UV_INDEX_URL: https://pypi.proxy.cloud.databricks.com/simple + runs-on: ubuntu-latest steps: - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 @@ -23,19 +19,38 @@ jobs: e2e: if: vars.E2E_ENABLED == 'true' - runs-on: - group: databricks-protected-runner-group - labels: linux-ubuntu-latest + runs-on: ubuntu-latest env: - UV_INDEX_URL: https://pypi.proxy.cloud.databricks.com/simple UCODE_TEST_WORKSPACE: ${{ secrets.UCODE_TEST_WORKSPACE }} DATABRICKS_HOST: ${{ secrets.UCODE_TEST_WORKSPACE }} - DATABRICKS_CLIENT_ID: ${{ secrets.DATABRICKS_CLIENT_ID }} - DATABRICKS_CLIENT_SECRET: ${{ secrets.DATABRICKS_CLIENT_SECRET }} + # DATABRICKS_BEARER is the CI escape hatch: `databricks auth token` + # only retrieves cached user-OAuth tokens, so on a hosted runner + # (no databrickscfg, no cached login) it can never produce a bearer. + # Pre-fetch one (e.g. via M2M OAuth client_credentials against + # /oidc/v1/token) and store it as a repo secret. Both + # has_valid_databricks_auth + get_databricks_token + the agents' + # apiKeyHelper short-circuit to this value when set. Tokens are + # short-lived (~1h); rotate when CI starts failing with 401s. + DATABRICKS_BEARER: ${{ secrets.DATABRICKS_BEARER }} steps: - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4 - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 - uses: databricks/setup-cli@bdb89f81c11a5bd647fd55b585b7c396ec68a25a # v1.0.0 + # The agent launch tests `_require_binary("codex")` etc. and skip when + # the CLI isn't on PATH. Install all six so each TestXxxLaunch test + # actually runs instead of skipping. + - name: Install agent CLIs + run: npm install -g + @anthropic-ai/claude-code + @openai/codex + @google/gemini-cli + opencode-ai + @github/copilot + @earendil-works/pi-coding-agent - run: uv tool install . - - run: uv run pytest tests/test_e2e.py -v + # Redirect stdin so any interactive `databricks auth login --no-browser` + # fallback EOFs instead of hanging the runner. With DATABRICKS_BEARER + # set, the auth code path doesn't shell out at all — this is a safety + # net for any code path we may have missed. + - run: uv run pytest tests/test_e2e.py -v < /dev/null diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index adc1824..35c7f31 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -14,7 +14,7 @@ import shutil import subprocess from pathlib import Path -from typing import cast +from typing import Literal, cast, overload from urllib import error as urllib_error from urllib import request as urllib_request from urllib.parse import urlparse @@ -90,9 +90,7 @@ def _debug(label: str, detail: str) -> None: logger.debug("%s: %s", label, detail) -_SECRET_KEY_PATTERN = re.compile( - r"(token|secret|password|bearer|api_key|apikey)", re.IGNORECASE -) +_SECRET_KEY_PATTERN = re.compile(r"(token|secret|password|bearer|api_key|apikey)", re.IGNORECASE) def _format_subprocess_result( @@ -127,7 +125,11 @@ def _scrub_databrickscfg(text: str) -> str: def _scrub_json(value: object) -> object: if isinstance(value, dict): return { - k: ("" if _SECRET_KEY_PATTERN.search(k) else _scrub_json(v)) + k: ( + "" + if isinstance(k, str) and _SECRET_KEY_PATTERN.search(k) + else _scrub_json(v) + ) for k, v in value.items() } if isinstance(value, list): @@ -225,6 +227,30 @@ def _http_get_json( return None, f"network error: {exc.reason}" +@overload +def run( + args: list[str], + *, + check: bool = True, + capture_output: bool = False, + text: Literal[True], + env: dict[str, str] | None = None, + timeout: int | None = None, +) -> subprocess.CompletedProcess[str]: ... + + +@overload +def run( + args: list[str], + *, + check: bool = True, + capture_output: bool = False, + text: Literal[False] = False, + env: dict[str, str] | None = None, + timeout: int | None = None, +) -> subprocess.CompletedProcess[bytes]: ... + + def run( args: list[str], *, @@ -331,6 +357,11 @@ def install_databricks_cli() -> None: def has_valid_databricks_auth(workspace: str) -> bool: + # Honor the CI short-circuit (see ``get_databricks_token``): if a + # pre-fetched bearer is available, treat auth as valid and skip the + # `databricks auth token` shell-out (which only knows user-OAuth). + if os.environ.get("DATABRICKS_BEARER", "").strip(): + return True _log_auth_diagnostics() try: env = build_databricks_cli_env(workspace) @@ -419,6 +450,17 @@ def ensure_databricks_auth(workspace: str) -> None: def get_databricks_token(workspace: str, *, force_refresh: bool = False) -> str: + # ``DATABRICKS_BEARER`` is the CI escape hatch: when set, skip the + # `databricks auth token` subprocess entirely and return the pre-fetched + # bearer directly. Used by the e2e job, where the protected runner has + # no `databricks auth login` cache and `databricks auth token` only knows + # how to read user-OAuth caches (not M2M client_credentials). Mirrors the + # same short-circuit baked into ``build_auth_shell_command``. + bearer = os.environ.get("DATABRICKS_BEARER", "").strip() + if bearer: + _debug("get_databricks_token", "using DATABRICKS_BEARER env var") + return bearer + _log_auth_diagnostics() env = build_databricks_cli_env(workspace) cmd = ["databricks", "auth", "token", "--host", workspace, "--output", "json"] @@ -429,9 +471,7 @@ def get_databricks_token(workspace: str, *, force_refresh: bool = False) -> str: "get_databricks_token.env", "set=" + ",".join( - sorted( - k for k in env if k.startswith("DATABRICKS_") or k in {"BUNDLE_PROFILE"} - ) + sorted(k for k in env if k.startswith("DATABRICKS_") or k in {"BUNDLE_PROFILE"}) ), ) diff --git a/src/ucode/mcp_web_search.py b/src/ucode/mcp_web_search.py index 5173c21..7d08554 100644 --- a/src/ucode/mcp_web_search.py +++ b/src/ucode/mcp_web_search.py @@ -119,7 +119,7 @@ def _call_responses_api(query: str) -> dict[str, Any]: }, ) try: - with urllib_request.urlopen(request, timeout=60) as response: + with urllib_request.urlopen(request, timeout=180) as response: raw = response.read().decode("utf-8") except urllib_error.HTTPError as exc: detail = "" diff --git a/tests/test_cli.py b/tests/test_cli.py index 01cff57..0664921 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re from unittest.mock import patch import pytest @@ -9,6 +10,16 @@ from ucode.cli import app +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") + + +def _strip_ansi(text: str) -> str: + """Drop SGR escape sequences so substring assertions match regardless of + whether the runner forces color rendering (e.g. CI sets FORCE_COLOR=1, + which makes rich split styled tokens like ``--agents`` with ANSI codes).""" + return _ANSI_RE.sub("", text) + + runner = CliRunner() TOOLS = ["codex", "claude", "gemini", "opencode"] @@ -68,9 +79,10 @@ def test_subcommand_help(self, tool): def test_configure_help_lists_agents_flag(self): result = runner.invoke(app, ["configure", "--help"]) assert result.exit_code == 0 - assert "--agents" in result.output - assert "comma-separated list of agents" in result.output - assert "--workspaces" in result.output + output = _strip_ansi(result.output) + assert "--agents" in output + assert "comma-separated list of agents" in output + assert "--workspaces" in output def _patch_launch(tool: str): diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 6e49be5..53ca465 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -213,7 +213,7 @@ def test_only_picks_codex_writes_only_codex_config(self, tmp_path, monkeypatch, monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") # Don't actually run `databricks auth login`; the developer running # this suite is already authenticated. - monkeypatch.setattr("ucode.databricks.run_databricks_login", lambda ws: None) + monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws: None) # Skip the workspace prompt and the multi-select picker. monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: ["codex"]) @@ -248,7 +248,7 @@ def test_rerun_with_different_pick_preserves_previous( self._redirect_config_paths(monkeypatch, tmp_path) monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") - monkeypatch.setattr("ucode.databricks.run_databricks_login", lambda ws: None) + monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws: None) monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) monkeypatch.setattr( cli_mod, "install_tool_binary", lambda tool, strict=False, update_existing=False: True @@ -281,7 +281,7 @@ def test_empty_pick_returns_zero_and_writes_nothing(self, tmp_path, monkeypatch, codex_path = self._redirect_config_paths(monkeypatch, tmp_path) monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") - monkeypatch.setattr("ucode.databricks.run_databricks_login", lambda ws: None) + monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws: None) monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: []) install_calls: list[str] = []