From 861f4ca8bb3bb39f0b3addfba5d7ee2ad57cb0ad Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 16 Sep 2026 21:27:18 +0000 Subject: [PATCH 1/6] Delegate custom OAuth to Databricks CLI --- src/ucode/custom_oauth.py | 153 +++++++++++++++++-------------- tests/test_custom_oauth.py | 179 ++++++++++++------------------------- 2 files changed, 141 insertions(+), 191 deletions(-) diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index c6b8ac7e..a5413ca6 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -2,24 +2,26 @@ from __future__ import annotations +import hashlib +import json +import os import platform import shlex import subprocess -from collections.abc import Iterator, Sequence -from contextlib import contextmanager -from pathlib import Path +from collections.abc import Sequence from typing import TypedDict from urllib.parse import urlparse -from databricks.sdk import oauth - +from ucode.config_io import APP_DIR from ucode.constants import LOCALHOST, LOOPBACK_HOST from ucode.databricks import build_auth_token_argv -from ucode.ui import err_console, normalize_workspace_url, print_warning_err +from ucode.ui import normalize_workspace_url DEFAULT_REDIRECT_URL = f"http://{LOCALHOST}:8020" # Custom OAuth may need a human to finish browser consent, not just a token fetch. CUSTOM_OAUTH_TIMEOUT_MS = 180_000 +CUSTOM_OAUTH_CLI_VERSION = (1, 17, 0) +CUSTOM_OAUTH_CONFIG_FILE = APP_DIR / "custom-oauth.databrickscfg" class CustomOAuthConfig(TypedDict): @@ -93,23 +95,44 @@ def build_custom_auth_shell_command(workspace: str, config: CustomOAuthConfig) - return shlex.join(argv) -@contextmanager -def _custom_oauth_lock(cache_dir: Path, redirect_url: str) -> Iterator[None]: - """Serialize helpers sharing a callback port with a POSIX file lock. +def _custom_oauth_profile(workspace: str, client_id: str, scopes: Sequence[str]) -> str: + key = "\0".join((workspace, client_id, *scopes)).encode() + return f"ug-custom-oauth-{hashlib.sha256(key).hexdigest()[:12]}" + - Keep the lock file in place: unlinking it could let waiters lock different - inodes. The OS releases the lock even if the helper is killed on timeout. - """ - import fcntl +def _require_custom_oauth_cli() -> None: + from ucode.databricks import databricks_cli_version - cache_dir.mkdir(parents=True, exist_ok=True) - port = urlparse(redirect_url).port - with (cache_dir / f"ug-oauth-{port}.lock").open("a+b") as lock_file: - fcntl.flock(lock_file, fcntl.LOCK_EX) - try: - yield - finally: - fcntl.flock(lock_file, fcntl.LOCK_UN) + version = databricks_cli_version() + if version is None or version < CUSTOM_OAUTH_CLI_VERSION: + required = ".".join(map(str, CUSTOM_OAUTH_CLI_VERSION)) + raise RuntimeError( + f"Custom-client OAuth requires Databricks CLI v{required} or newer. Upgrade the CLI " + "and retry." + ) + + +def _token_from_cli(workspace: str, profile: str, env: dict[str, str], force: bool) -> str: + command = [ + "databricks", + "auth", + "token", + "--host", + workspace, + "--profile", + profile, + "--output", + "json", + ] + if force: + command.append("--force-refresh") + result = subprocess.run(command, capture_output=True, text=True, env=env, timeout=15) + if result.returncode != 0: + return "" + try: + return json.loads(result.stdout or "{}").get("access_token", "") + except json.JSONDecodeError: + return "" def get_custom_client_token( @@ -120,57 +143,51 @@ def get_custom_client_token( scopes: Sequence[str], force_refresh: bool = False, ) -> str: - """Reuse the SDK's PKCE flow and per-workspace/client token cache.""" + """Delegate custom-client U2M login, refresh, and caching to Databricks CLI.""" config = create_custom_oauth_config(client_id, scopes, redirect_url) workspace = normalize_workspace_url(workspace) + _require_custom_oauth_cli() + profile = _custom_oauth_profile(workspace, config["client_id"], config["scopes"]) + CUSTOM_OAUTH_CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True) + env = os.environ.copy() + env["DATABRICKS_CONFIG_FILE"] = str(CUSTOM_OAUTH_CONFIG_FILE) + try: - endpoints = oauth.get_workspace_endpoints(workspace) - cache = oauth.TokenCache( - host=workspace, - oidc_endpoints=endpoints, - client_id=config["client_id"], - redirect_url=config["redirect_url"], - scopes=config["scopes"], + token = _token_from_cli(workspace, profile, env, force_refresh) + if token: + return token + login = subprocess.run( + [ + "databricks", + "auth", + "login", + "--host", + workspace, + "--profile", + profile, + "--client-id", + config["client_id"], + "--scopes", + ",".join(config["scopes"]), + "--timeout", + "3m", + ], + capture_output=True, + text=True, + check=False, + env=env, + timeout=CUSTOM_OAUTH_TIMEOUT_MS / 1000, ) - with _custom_oauth_lock(Path(cache.filename).parent, config["redirect_url"]): - # Read only after acquiring the lock: another helper may have just - # completed login or rotated the refresh token while we waited. - credentials = cache.load() - if credentials is not None: - try: - if force_refresh: - credentials = oauth.SessionCredentials( - token=credentials.refresh(), - token_endpoint=endpoints.token_endpoint, - client_id=config["client_id"], - redirect_url=config["redirect_url"], - ) - credentials.token() - except Exception: - print_warning_err("Cached OAuth token could not be refreshed. Sign in again.") - credentials = None - if credentials is None: - client = oauth.OAuthClient( - oidc_endpoints=endpoints, - client_id=config["client_id"], - redirect_url=config["redirect_url"], - scopes=config["scopes"], - ) - consent = client.initiate_consent() - err_console.print( - f"Sign in using your browser: {consent.authorization_url}", - markup=False, - soft_wrap=True, - ) - credentials = consent.launch_external_browser() - token = credentials.token().access_token - if not token: - raise ValueError("OAuth returned no access token") - cache.save(credentials) + if login.returncode == 0: + token = _token_from_cli(workspace, profile, env, False) + if token: return token - except Exception as exc: + except (OSError, subprocess.TimeoutExpired) as exc: raise RuntimeError( - "Custom-client OAuth failed. Check the workspace, client ID, and registered " - f"redirect URL ({config['redirect_url']}); ensure its local port is available and " - "the SDK token cache is writable, then retry." + "Databricks CLI custom-client OAuth failed. Check the workspace, client ID, and " + "scopes, then retry." ) from exc + raise RuntimeError( + "Databricks CLI returned no custom-client OAuth token. Check the workspace, client ID, " + "and scopes, then retry." + ) diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index ed0d945a..4a69cffb 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -2,140 +2,70 @@ from __future__ import annotations -import json -from datetime import UTC, datetime, timedelta -from pathlib import Path +import subprocess from unittest.mock import Mock, patch -from urllib.parse import parse_qs import pytest -from databricks.sdk import oauth from typer.testing import CliRunner import ucode.cli as cli_mod +import ucode.custom_oauth as oauth_mod import ucode.databricks as db_mod from ucode.cli import app -from ucode.custom_oauth import _custom_oauth_lock, get_custom_client_token +from ucode.custom_oauth import get_custom_client_token WS = "https://example.databricks.com" TEST_SCOPES = ("offline_access", "catalog.catalogs:read") runner = CliRunner() -class TestCustomOAuthLock: - def test_releases_lock_when_login_fails(self, tmp_path): - with pytest.raises(ValueError, match="login failed"): - with _custom_oauth_lock(tmp_path, "http://localhost:8020/callback"): - raise ValueError("login failed") - with _custom_oauth_lock(tmp_path, "http://127.0.0.1:8020/other-callback"): - assert len(list(tmp_path.glob("*.lock"))) == 1 - - class TestCustomClientToken: @pytest.fixture(autouse=True) def _setup(self, tmp_path, monkeypatch): - monkeypatch.setattr(oauth.TokenCache, "BASE_PATH", str(tmp_path / "oauth")) - monkeypatch.setenv("DATABRICKS_BEARER", "unrelated-bearer") - monkeypatch.setattr(db_mod, "run", Mock(side_effect=AssertionError("CLI not expected"))) - monkeypatch.setattr( - db_mod, "find_profile_name_for_host", Mock(side_effect=AssertionError("No profile")) - ) - self.endpoints = oauth.OidcEndpoints( - authorization_endpoint=f"{WS}/oidc/v1/authorize", - token_endpoint=f"{WS}/oidc/v1/token", - ) - self.discovery = Mock(return_value=self.endpoints) - monkeypatch.setattr(oauth, "get_workspace_endpoints", self.discovery) - self.browser = Mock(return_value=self._credentials("browser-token", "browser-refresh")) - monkeypatch.setattr(oauth.Consent, "launch_external_browser", self.browser) - self.refresh = Mock(return_value=self._credentials("refreshed", "rotated-refresh").token()) - monkeypatch.setattr(oauth, "retrieve_token", self.refresh) - - def _credentials(self, access_token, refresh_token): - return oauth.SessionCredentials( - token=oauth.Token( - access_token=access_token, - token_type="Bearer", - refresh_token=refresh_token, - expiry=datetime.now(UTC) + timedelta(hours=1), - ), - token_endpoint=self.endpoints.token_endpoint, - client_id="custom-client", - redirect_url="http://localhost:8020", + monkeypatch.setattr(db_mod, "databricks_cli_version", lambda: (1, 17, 0)) + monkeypatch.setattr(oauth_mod, "CUSTOM_OAUTH_CONFIG_FILE", tmp_path / "oauth.cfg") + + @staticmethod + def _result(returncode=0, stdout=""): + return subprocess.CompletedProcess([], returncode, stdout, "") + + def test_reuses_cli_token_without_login(self, monkeypatch): + run = Mock(return_value=self._result(stdout='{"access_token":"cached"}')) + monkeypatch.setattr(oauth_mod.subprocess, "run", run) + + assert get_custom_client_token(WS + "/", "custom-client", scopes=TEST_SCOPES) == "cached" + command = run.call_args.args[0] + assert command[:3] == ["databricks", "auth", "token"] + assert run.call_args.kwargs["env"]["DATABRICKS_CONFIG_FILE"].endswith("oauth.cfg") + + def test_login_then_retries_token(self, monkeypatch): + run = Mock( + side_effect=[ + self._result(1), + self._result(), + self._result(stdout='{"access_token":"new-token"}'), + ] ) - - def _cache(self, workspace=WS, client_id="custom-client"): - return oauth.TokenCache( - host=workspace, - oidc_endpoints=self.endpoints, - client_id=client_id, - redirect_url="http://localhost:8020", - scopes=list(TEST_SCOPES), - ) - - def test_browser_login_uses_custom_client_and_redirect(self, capsys): - redirect_url = "http://localhost:41735/ai-devtools-workspace-oauth" - token = get_custom_client_token( - WS, client_id="custom-client", redirect_url=redirect_url, scopes=TEST_SCOPES - ) - assert token == "browser-token" - cached = self._cache().load().token() - assert cached.refresh_token == "browser-refresh" - output = capsys.readouterr() - assert output.out == "" - query = parse_qs(output.err.split("?", 1)[1].strip()) - assert query["client_id"] == ["custom-client"] - assert query["redirect_uri"] == [redirect_url] - assert query["scope"][0].split() == list(TEST_SCOPES) - - def test_reuses_cached_token_without_refresh_or_login(self): - self._cache().save(self._credentials("cached", "refresh")) - assert ( - get_custom_client_token(WS + "/", client_id="custom-client", scopes=TEST_SCOPES) - == "cached" - ) - self.browser.assert_not_called() - self.refresh.assert_not_called() - - def test_expired_token_refreshes_with_custom_client_and_saves_rotation(self): - cached = self._credentials("expired", "old-refresh") - self._cache().save(cached) - cache_path = Path(self._cache().filename) - payload = json.loads(cache_path.read_text()) - payload["token"]["expiry"] = (datetime.now(UTC) - timedelta(hours=1)).isoformat() - cache_path.write_text(json.dumps(payload)) - assert ( - get_custom_client_token(WS, client_id="custom-client", scopes=TEST_SCOPES) - == "refreshed" - ) - self.refresh.assert_called_once() - self.browser.assert_not_called() - assert self._cache().load().token().refresh_token == "rotated-refresh" - - def test_force_refresh_bypasses_fresh_access_token(self): - self._cache().save(self._credentials("cached", "refresh")) - assert ( - get_custom_client_token( - WS, client_id="custom-client", scopes=TEST_SCOPES, force_refresh=True - ) - == "refreshed" - ) - self.refresh.assert_called_once() - self.browser.assert_not_called() - - def test_refresh_failure_falls_back_to_browser(self, capsys): - self._cache().save(self._credentials("cached", "revoked-refresh")) - self.refresh.side_effect = ValueError("sensitive server response") - assert ( - get_custom_client_token( - WS, client_id="custom-client", scopes=TEST_SCOPES, force_refresh=True - ) - == "browser-token" - ) - self.browser.assert_called_once() - output = capsys.readouterr() - assert "Sign in again" in output.err - assert "sensitive server response" not in output.err + monkeypatch.setattr(oauth_mod.subprocess, "run", run) + + assert get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) == "new-token" + first, login, retry = [call.args[0] for call in run.call_args_list] + assert first[first.index("--profile") + 1] == retry[retry.index("--profile") + 1] + assert login[:3] == ["databricks", "auth", "login"] + assert login[login.index("--client-id") + 1] == "custom-client" + assert login[login.index("--scopes") + 1] == ",".join(TEST_SCOPES) + assert login[login.index("--timeout") + 1] == "3m" + assert login[login.index("--profile") + 1] == first[first.index("--profile") + 1] + assert login[login.index("--profile") + 1].startswith("ug-custom-oauth-") + assert run.call_args_list[1].kwargs["capture_output"] is True + + def test_force_refresh_is_forwarded(self, monkeypatch): + run = Mock(return_value=self._result(stdout='{"access_token":"fresh"}')) + monkeypatch.setattr(oauth_mod.subprocess, "run", run) + assert get_custom_client_token( + WS, "custom-client", scopes=TEST_SCOPES, force_refresh=True + ) == "fresh" + assert "--force-refresh" in run.call_args.args[0] def test_invalid_redirect_fails_before_network(self): with pytest.raises(RuntimeError, match="--redirect-url must be"): @@ -145,19 +75,22 @@ def test_invalid_redirect_fails_before_network(self): redirect_url="https://example.com/callback", scopes=TEST_SCOPES, ) - self.discovery.assert_not_called() - - def test_login_failure_is_actionable_and_does_not_expose_response(self): - self.browser.side_effect = ValueError("sensitive server response") - with pytest.raises(RuntimeError, match="registered redirect URL") as error: - get_custom_client_token(WS, client_id="custom-client", scopes=TEST_SCOPES) + def test_old_cli_is_actionable(self, monkeypatch): + monkeypatch.setattr(db_mod, "databricks_cli_version", lambda: (1, 16, 1)) + with pytest.raises(RuntimeError, match="v1.17.0 or newer"): + get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) + + def test_login_failure_is_actionable_and_sanitized(self, monkeypatch): + run = Mock(side_effect=[self._result(1), self._result(1, "sensitive response")]) + monkeypatch.setattr(oauth_mod.subprocess, "run", run) + with pytest.raises(RuntimeError, match="returned no custom-client OAuth token") as error: + get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) assert "sensitive server response" not in str(error.value) @pytest.mark.parametrize("scopes", [["offline_access"], ["catalog.catalogs:read"]]) def test_api_scopes_are_required(self, scopes): with pytest.raises(RuntimeError, match="OAuth scopes|API OAuth scope"): get_custom_client_token(WS, client_id="custom-client", scopes=scopes) - self.discovery.assert_not_called() class TestCustomClientCommand: From d4212fd6031f85296f2bb39cec8cd32e41c7ccb3 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 16 Sep 2026 21:32:40 +0000 Subject: [PATCH 2/6] Require Databricks CLI 1.17 --- src/ucode/custom_oauth.py | 7 +++---- src/ucode/databricks.py | 4 ++-- tests/integration/README.md | 2 +- tests/test_databricks.py | 4 ++-- 4 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index a5413ca6..6dcb0e7d 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -14,13 +14,12 @@ from ucode.config_io import APP_DIR from ucode.constants import LOCALHOST, LOOPBACK_HOST -from ucode.databricks import build_auth_token_argv +from ucode.databricks import MIN_DATABRICKS_CLI_VERSION, build_auth_token_argv from ucode.ui import normalize_workspace_url DEFAULT_REDIRECT_URL = f"http://{LOCALHOST}:8020" # Custom OAuth may need a human to finish browser consent, not just a token fetch. CUSTOM_OAUTH_TIMEOUT_MS = 180_000 -CUSTOM_OAUTH_CLI_VERSION = (1, 17, 0) CUSTOM_OAUTH_CONFIG_FILE = APP_DIR / "custom-oauth.databrickscfg" @@ -104,8 +103,8 @@ def _require_custom_oauth_cli() -> None: from ucode.databricks import databricks_cli_version version = databricks_cli_version() - if version is None or version < CUSTOM_OAUTH_CLI_VERSION: - required = ".".join(map(str, CUSTOM_OAUTH_CLI_VERSION)) + if version is None or version < MIN_DATABRICKS_CLI_VERSION: + required = ".".join(map(str, MIN_DATABRICKS_CLI_VERSION)) raise RuntimeError( f"Custom-client OAuth requires Databricks CLI v{required} or newer. Upgrade the CLI " "and retry." diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a7f09957..1416e0d5 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -58,8 +58,8 @@ ) AI_GATEWAY_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" ANTHROPIC_MODELS_PATH = "/ai-gateway/anthropic/v1/models" -# v1.0.0 is the release that ships `databricks aitools`. -MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) +# v1.17.0 adds custom OAuth client IDs to `databricks auth login` and token refresh. +MIN_DATABRICKS_CLI_VERSION = (1, 17, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 # 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 diff --git a/tests/integration/README.md b/tests/integration/README.md index d9056b2c..13bf0890 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -11,7 +11,7 @@ It is not collected by the default `uv run pytest` command. ## Run a specific combination -Prerequisites: Python 3.12+, uv, Node/npm, and Databricks CLI >=1.0.0. The runner +Prerequisites: Python 3.12+, uv, Node/npm, and Databricks CLI >=1.17.0. The runner installs the requested agents into a new npm prefix and ug into a new virtualenv. Pytest and the PTY/screen libraries (pexpect and pyte) live in a different virtualenv, so they cannot accidentally supply a missing application dependency. No packages are installed into your existing diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 6bab2b57..f498c8d5 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2536,12 +2536,12 @@ 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 v1.0.0") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.17.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 v1.8.0") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.18.0") monkeypatch.setattr("os.environ", env) ensure_databricks_cli_version() From c17596fc0f85ee6e8125ea27dc26f78ff50be2b4 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 16 Sep 2026 21:46:40 +0000 Subject: [PATCH 3/6] fix --- tests/test_custom_oauth.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index 4a69cffb..acd8e91d 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -62,9 +62,10 @@ def test_login_then_retries_token(self, monkeypatch): def test_force_refresh_is_forwarded(self, monkeypatch): run = Mock(return_value=self._result(stdout='{"access_token":"fresh"}')) monkeypatch.setattr(oauth_mod.subprocess, "run", run) - assert get_custom_client_token( - WS, "custom-client", scopes=TEST_SCOPES, force_refresh=True - ) == "fresh" + assert ( + get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES, force_refresh=True) + == "fresh" + ) assert "--force-refresh" in run.call_args.args[0] def test_invalid_redirect_fails_before_network(self): @@ -75,6 +76,7 @@ def test_invalid_redirect_fails_before_network(self): redirect_url="https://example.com/callback", scopes=TEST_SCOPES, ) + def test_old_cli_is_actionable(self, monkeypatch): monkeypatch.setattr(db_mod, "databricks_cli_version", lambda: (1, 16, 1)) with pytest.raises(RuntimeError, match="v1.17.0 or newer"): From fde579e7edb4788c36519c80c59c7fb8c62e3c57 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 16 Sep 2026 22:33:51 +0000 Subject: [PATCH 4/6] Gate CLI-managed custom OAuth profiles --- src/ucode/custom_oauth.py | 114 +++++++++++++++++++++++++++++++----- src/ucode/databricks.py | 4 +- tests/integration/README.md | 2 +- tests/test_custom_oauth.py | 20 +++++++ tests/test_databricks.py | 4 +- 5 files changed, 125 insertions(+), 19 deletions(-) diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index 6dcb0e7d..e8e2484d 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -8,19 +8,25 @@ import platform import shlex import subprocess -from collections.abc import Sequence +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from pathlib import Path from typing import TypedDict from urllib.parse import urlparse +from databricks.sdk import oauth + from ucode.config_io import APP_DIR from ucode.constants import LOCALHOST, LOOPBACK_HOST -from ucode.databricks import MIN_DATABRICKS_CLI_VERSION, build_auth_token_argv -from ucode.ui import normalize_workspace_url +from ucode.databricks import build_auth_token_argv +from ucode.ui import err_console, normalize_workspace_url, print_warning_err DEFAULT_REDIRECT_URL = f"http://{LOCALHOST}:8020" # Custom OAuth may need a human to finish browser consent, not just a token fetch. CUSTOM_OAUTH_TIMEOUT_MS = 180_000 -CUSTOM_OAUTH_CONFIG_FILE = APP_DIR / "custom-oauth.databrickscfg" +CUSTOM_OAUTH_CLI_VERSION = (1, 17, 0) +CUSTOM_OAUTH_CONFIG_FILE = APP_DIR / "ug.databrickscfg" +ENABLE_CUSTOM_OAUTH_PROFILE = "ENABLE_CUSTOM_OAUTH_PROFILE" class CustomOAuthConfig(TypedDict): @@ -99,12 +105,26 @@ def _custom_oauth_profile(workspace: str, client_id: str, scopes: Sequence[str]) return f"ug-custom-oauth-{hashlib.sha256(key).hexdigest()[:12]}" +@contextmanager +def _custom_oauth_lock(cache_dir: Path, redirect_url: str) -> Iterator[None]: + import fcntl + + cache_dir.mkdir(parents=True, exist_ok=True) + port = urlparse(redirect_url).port + with (cache_dir / f"ug-oauth-{port}.lock").open("a+b") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(lock_file, fcntl.LOCK_UN) + + def _require_custom_oauth_cli() -> None: from ucode.databricks import databricks_cli_version version = databricks_cli_version() - if version is None or version < MIN_DATABRICKS_CLI_VERSION: - required = ".".join(map(str, MIN_DATABRICKS_CLI_VERSION)) + if version is None or version < CUSTOM_OAUTH_CLI_VERSION: + required = ".".join(map(str, CUSTOM_OAUTH_CLI_VERSION)) raise RuntimeError( f"Custom-client OAuth requires Databricks CLI v{required} or newer. Upgrade the CLI " "and retry." @@ -134,17 +154,12 @@ def _token_from_cli(workspace: str, profile: str, env: dict[str, str], force: bo return "" -def get_custom_client_token( +def _get_custom_client_token_from_cli( workspace: str, - client_id: str, - redirect_url: str = DEFAULT_REDIRECT_URL, - *, - scopes: Sequence[str], - force_refresh: bool = False, + config: CustomOAuthConfig, + force_refresh: bool, ) -> str: """Delegate custom-client U2M login, refresh, and caching to Databricks CLI.""" - config = create_custom_oauth_config(client_id, scopes, redirect_url) - workspace = normalize_workspace_url(workspace) _require_custom_oauth_cli() profile = _custom_oauth_profile(workspace, config["client_id"], config["scopes"]) CUSTOM_OAUTH_CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True) @@ -190,3 +205,74 @@ def get_custom_client_token( "Databricks CLI returned no custom-client OAuth token. Check the workspace, client ID, " "and scopes, then retry." ) + + +def _get_custom_client_token_from_sdk( + workspace: str, + config: CustomOAuthConfig, + force_refresh: bool, +) -> str: + try: + endpoints = oauth.get_workspace_endpoints(workspace) + cache = oauth.TokenCache( + host=workspace, + oidc_endpoints=endpoints, + client_id=config["client_id"], + redirect_url=config["redirect_url"], + scopes=config["scopes"], + ) + with _custom_oauth_lock(Path(cache.filename).parent, config["redirect_url"]): + credentials = cache.load() + if credentials is not None: + try: + if force_refresh: + credentials = oauth.SessionCredentials( + token=credentials.refresh(), + token_endpoint=endpoints.token_endpoint, + client_id=config["client_id"], + redirect_url=config["redirect_url"], + ) + credentials.token() + except Exception: + print_warning_err("Cached OAuth token could not be refreshed. Sign in again.") + credentials = None + if credentials is None: + client = oauth.OAuthClient( + oidc_endpoints=endpoints, + client_id=config["client_id"], + redirect_url=config["redirect_url"], + scopes=config["scopes"], + ) + consent = client.initiate_consent() + err_console.print( + f"Sign in using your browser: {consent.authorization_url}", + markup=False, + soft_wrap=True, + ) + credentials = consent.launch_external_browser() + token = credentials.token().access_token + if not token: + raise ValueError("OAuth returned no access token") + cache.save(credentials) + return token + except Exception as exc: + raise RuntimeError( + "Custom-client OAuth failed. Check the workspace, client ID, and registered " + f"redirect URL ({config['redirect_url']}); ensure its local port is available and " + "the SDK token cache is writable, then retry." + ) from exc + + +def get_custom_client_token( + workspace: str, + client_id: str, + redirect_url: str = DEFAULT_REDIRECT_URL, + *, + scopes: Sequence[str], + force_refresh: bool = False, +) -> str: + config = create_custom_oauth_config(client_id, scopes, redirect_url) + workspace = normalize_workspace_url(workspace) + if os.environ.get(ENABLE_CUSTOM_OAUTH_PROFILE) == "1": + return _get_custom_client_token_from_cli(workspace, config, force_refresh) + return _get_custom_client_token_from_sdk(workspace, config, force_refresh) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 1416e0d5..a7f09957 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -58,8 +58,8 @@ ) AI_GATEWAY_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" ANTHROPIC_MODELS_PATH = "/ai-gateway/anthropic/v1/models" -# v1.17.0 adds custom OAuth client IDs to `databricks auth login` and token refresh. -MIN_DATABRICKS_CLI_VERSION = (1, 17, 0) +# v1.0.0 is the release that ships `databricks aitools`. +MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 # 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 diff --git a/tests/integration/README.md b/tests/integration/README.md index 13bf0890..d9056b2c 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -11,7 +11,7 @@ It is not collected by the default `uv run pytest` command. ## Run a specific combination -Prerequisites: Python 3.12+, uv, Node/npm, and Databricks CLI >=1.17.0. The runner +Prerequisites: Python 3.12+, uv, Node/npm, and Databricks CLI >=1.0.0. The runner installs the requested agents into a new npm prefix and ug into a new virtualenv. Pytest and the PTY/screen libraries (pexpect and pyte) live in a different virtualenv, so they cannot accidentally supply a missing application dependency. No packages are installed into your existing diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index acd8e91d..d7291572 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -22,6 +22,7 @@ class TestCustomClientToken: @pytest.fixture(autouse=True) def _setup(self, tmp_path, monkeypatch): + monkeypatch.setenv(oauth_mod.ENABLE_CUSTOM_OAUTH_PROFILE, "1") monkeypatch.setattr(db_mod, "databricks_cli_version", lambda: (1, 17, 0)) monkeypatch.setattr(oauth_mod, "CUSTOM_OAUTH_CONFIG_FILE", tmp_path / "oauth.cfg") @@ -95,6 +96,25 @@ def test_api_scopes_are_required(self, scopes): get_custom_client_token(WS, client_id="custom-client", scopes=scopes) +@pytest.mark.parametrize("value", [None, "0"]) +def test_cli_profile_is_opt_in(monkeypatch, value): + if value is None: + monkeypatch.delenv(oauth_mod.ENABLE_CUSTOM_OAUTH_PROFILE, raising=False) + else: + monkeypatch.setenv(oauth_mod.ENABLE_CUSTOM_OAUTH_PROFILE, value) + legacy = Mock(return_value="legacy-token") + monkeypatch.setattr(oauth_mod, "_get_custom_client_token_from_sdk", legacy) + cli = Mock(side_effect=AssertionError("CLI profile path must be disabled")) + monkeypatch.setattr(oauth_mod, "_get_custom_client_token_from_cli", cli) + + assert get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) == "legacy-token" + legacy.assert_called_once() + + +def test_custom_oauth_config_filename(): + assert oauth_mod.CUSTOM_OAUTH_CONFIG_FILE.name == "ug.databrickscfg" + + class TestCustomClientCommand: @pytest.fixture(autouse=True) def _no_production_auth(self, monkeypatch): diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f498c8d5..6bab2b57 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2536,12 +2536,12 @@ 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 v1.17.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 v1.18.0") + env = self._fake_databricks(tmp_path, "Databricks CLI v1.8.0") monkeypatch.setattr("os.environ", env) ensure_databricks_cli_version() From 5d853c8532801398c34faa6e98e03033964a2ea2 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 16 Sep 2026 22:36:31 +0000 Subject: [PATCH 5/6] Restore OAuth lock documentation --- src/ucode/custom_oauth.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index e8e2484d..e3465af4 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -107,6 +107,11 @@ def _custom_oauth_profile(workspace: str, client_id: str, scopes: Sequence[str]) @contextmanager def _custom_oauth_lock(cache_dir: Path, redirect_url: str) -> Iterator[None]: + """Serialize helpers sharing a callback port with a POSIX file lock. + + Keep the lock file in place: unlinking it could let waiters lock different + inodes. The OS releases the lock even if the helper is killed on timeout. + """ import fcntl cache_dir.mkdir(parents=True, exist_ok=True) From 30ebf8df840395c2cd5c8b49814356ff8cdcc341 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 16 Sep 2026 22:42:51 +0000 Subject: [PATCH 6/6] hi --- src/ucode/custom_oauth.py | 5 +++-- tests/test_custom_oauth.py | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index e3465af4..1d7d2a7d 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -26,7 +26,6 @@ CUSTOM_OAUTH_TIMEOUT_MS = 180_000 CUSTOM_OAUTH_CLI_VERSION = (1, 17, 0) CUSTOM_OAUTH_CONFIG_FILE = APP_DIR / "ug.databrickscfg" -ENABLE_CUSTOM_OAUTH_PROFILE = "ENABLE_CUSTOM_OAUTH_PROFILE" class CustomOAuthConfig(TypedDict): @@ -217,6 +216,7 @@ def _get_custom_client_token_from_sdk( config: CustomOAuthConfig, force_refresh: bool, ) -> str: + """Reuse the SDK's PKCE flow and per-workspace/client token cache.""" try: endpoints = oauth.get_workspace_endpoints(workspace) cache = oauth.TokenCache( @@ -276,8 +276,9 @@ def get_custom_client_token( scopes: Sequence[str], force_refresh: bool = False, ) -> str: + """Use CLI-managed OAuth profiles when enabled; otherwise retain the SDK flow.""" config = create_custom_oauth_config(client_id, scopes, redirect_url) workspace = normalize_workspace_url(workspace) - if os.environ.get(ENABLE_CUSTOM_OAUTH_PROFILE) == "1": + if os.environ.get("CUSTOM_OAUTH_CONFIG_FILE") == "1": return _get_custom_client_token_from_cli(workspace, config, force_refresh) return _get_custom_client_token_from_sdk(workspace, config, force_refresh) diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index d7291572..ca29fe14 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -22,7 +22,7 @@ class TestCustomClientToken: @pytest.fixture(autouse=True) def _setup(self, tmp_path, monkeypatch): - monkeypatch.setenv(oauth_mod.ENABLE_CUSTOM_OAUTH_PROFILE, "1") + monkeypatch.setenv("CUSTOM_OAUTH_CONFIG_FILE", "1") monkeypatch.setattr(db_mod, "databricks_cli_version", lambda: (1, 17, 0)) monkeypatch.setattr(oauth_mod, "CUSTOM_OAUTH_CONFIG_FILE", tmp_path / "oauth.cfg") @@ -99,9 +99,9 @@ def test_api_scopes_are_required(self, scopes): @pytest.mark.parametrize("value", [None, "0"]) def test_cli_profile_is_opt_in(monkeypatch, value): if value is None: - monkeypatch.delenv(oauth_mod.ENABLE_CUSTOM_OAUTH_PROFILE, raising=False) + monkeypatch.delenv("CUSTOM_OAUTH_CONFIG_FILE", raising=False) else: - monkeypatch.setenv(oauth_mod.ENABLE_CUSTOM_OAUTH_PROFILE, value) + monkeypatch.setenv("CUSTOM_OAUTH_CONFIG_FILE", value) legacy = Mock(return_value="legacy-token") monkeypatch.setattr(oauth_mod, "_get_custom_client_token_from_sdk", legacy) cli = Mock(side_effect=AssertionError("CLI profile path must be disabled"))