From a4b9904d32fe3aae8df739e063e1b059f18be7ed Mon Sep 17 00:00:00 2001 From: Brad Date: Wed, 8 Jul 2026 08:50:59 -0700 Subject: [PATCH 1/4] Force token refresh in the auth-token apiKeyHelper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode auth-token` is the command Claude Code's apiKeyHelper (and Codex's auth command) re-invokes every 15 min via CLAUDE_CODE_API_KEY_HELPER_TTL_MS. It called `databricks auth token` without --force-refresh, so the token was only renewed when the CLI independently judged the cached token "close to expiry" — which can lapse between invocations, killing a long session at the original token's ~1h expiry. Pass force_refresh=True so each refresh cycle hands back a token with a full fresh lifetime, matching the force_refresh=True already used by the gemini/opencode/copilot/pi background refreshers. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/ucode/cli.py | 10 +++++++++- tests/test_cli.py | 18 +++++++++++++----- tests/test_databricks.py | 26 ++++++++++++++++++++++++++ 3 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b2f23be..36c005d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -776,7 +776,15 @@ def auth_token_cmd( ) raise typer.Exit(1) try: - token = get_databricks_token(workspace, profile) + # force_refresh: this helper is re-invoked by Claude Code's apiKeyHelper + # (and Codex's auth command) on the CLAUDE_CODE_API_KEY_HELPER_TTL_MS + # interval (15 min). Forcing a refresh each cycle hands back a token with + # a full fresh lifetime, so a long session never dies at the original + # token's ~1h expiry. Without it, `databricks auth token` only refreshes + # when it independently judges the cached token "close to expiry", which + # can lapse between invocations. Mirrors the force_refresh=True used by + # the gemini/opencode/copilot/pi background refreshers. + token = get_databricks_token(workspace, profile, force_refresh=True) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None diff --git a/tests/test_cli.py b/tests/test_cli.py index 3c5d404..be778b6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -180,7 +180,9 @@ def test_prints_only_the_token_to_stdout(self): # Nothing but the bare token (plus trailing newline) may reach stdout, # or the consuming agent will treat the noise as part of the token. assert result.stdout == "tok-123\n" - fetch.assert_called_once_with("https://ws", None) + # force_refresh=True so each 15-min apiKeyHelper cycle yields a token with + # a full fresh lifetime (long sessions otherwise die at ~1h expiry). + fetch.assert_called_once_with("https://ws", None, force_refresh=True) def test_host_and_profile_override_state(self): with ( @@ -191,7 +193,7 @@ def test_host_and_profile_override_state(self): app, ["auth-token", "--host", "https://override", "--profile", "prod"] ) assert result.exit_code == 0 - fetch.assert_called_once_with("https://override", "prod") + fetch.assert_called_once_with("https://override", "prod", force_refresh=True) def test_errors_without_workspace(self): with patch("ucode.cli.load_state", return_value={}): @@ -213,7 +215,9 @@ def test_use_pat_emits_resolved_pat(self, monkeypatch): patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch( "ucode.cli.get_databricks_token", - side_effect=lambda w, p: os.environ.get("DATABRICKS_BEARER", ""), + side_effect=lambda w, p, force_refresh=False: os.environ.get( + "DATABRICKS_BEARER", "" + ), ), ): result = runner.invoke(app, ["auth-token", "--use-pat", "--profile", "p"]) @@ -229,7 +233,9 @@ def test_use_pat_ignores_empty_bearer_env(self, monkeypatch): patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch( "ucode.cli.get_databricks_token", - side_effect=lambda w, p: os.environ.get("DATABRICKS_BEARER", ""), + side_effect=lambda w, p, force_refresh=False: os.environ.get( + "DATABRICKS_BEARER", "" + ), ), ): result = runner.invoke(app, ["auth-token", "--use-pat", "--profile", "p"]) @@ -259,7 +265,9 @@ def test_use_pat_honors_non_empty_bearer_env(self, monkeypatch): patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch( "ucode.cli.get_databricks_token", - side_effect=lambda w, p: os.environ.get("DATABRICKS_BEARER", ""), + side_effect=lambda w, p, force_refresh=False: os.environ.get( + "DATABRICKS_BEARER", "" + ), ), ): result = runner.invoke(app, ["auth-token", "--use-pat", "--profile", "p"]) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f59aa1d..1e37e2b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1141,6 +1141,32 @@ def test_passes_profile_flag_when_provided(self, tmp_path, monkeypatch): assert "--profile" in argv assert argv[argv.index("--profile") + 1] == "stablebox" + def test_force_refresh_appends_flag(self, tmp_path, monkeypatch): + # The apiKeyHelper (`ucode auth-token`) passes force_refresh=True so each + # refresh cycle bypasses `databricks auth token`'s "close to expiry" + # heuristic and hands back a token with a full fresh lifetime. + argv_log = tmp_path / "argv" + env = self._fake_databricks( + tmp_path, + f'printf "%s\\n" "$@" >> {argv_log}\n' + 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', + ) + monkeypatch.setattr("os.environ", env) + token = get_databricks_token(WS, force_refresh=True) + assert token == "good-token" + assert "--force-refresh" in argv_log.read_text().splitlines() + + def test_omits_force_refresh_flag_by_default(self, tmp_path, monkeypatch): + argv_log = tmp_path / "argv" + env = self._fake_databricks( + tmp_path, + f'printf "%s\\n" "$@" >> {argv_log}\n' + 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', + ) + monkeypatch.setattr("os.environ", env) + get_databricks_token(WS) + assert "--force-refresh" not in argv_log.read_text().splitlines() + def test_error_suggests_logout_when_matching_profile_exists(self, tmp_path, monkeypatch): env = self._fake_databricks( tmp_path, From 08bd1d45203c8e9009fd67021772f95160a0f391 Mon Sep 17 00:00:00 2001 From: Brad Date: Wed, 8 Jul 2026 08:56:45 -0700 Subject: [PATCH 2/4] Make force_refresh keyword-only in auth-token test mocks Mirror the real get_databricks_token signature (force_refresh is keyword-only) in the --use-pat test lambdas, so the test would fail if the production call site ever passed force_refresh positionally. Addresses PR review feedback. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_cli.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index be778b6..7c0f7b9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -215,7 +215,7 @@ def test_use_pat_emits_resolved_pat(self, monkeypatch): patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch( "ucode.cli.get_databricks_token", - side_effect=lambda w, p, force_refresh=False: os.environ.get( + side_effect=lambda w, p, *, force_refresh=False: os.environ.get( "DATABRICKS_BEARER", "" ), ), @@ -233,7 +233,7 @@ def test_use_pat_ignores_empty_bearer_env(self, monkeypatch): patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch( "ucode.cli.get_databricks_token", - side_effect=lambda w, p, force_refresh=False: os.environ.get( + side_effect=lambda w, p, *, force_refresh=False: os.environ.get( "DATABRICKS_BEARER", "" ), ), @@ -265,7 +265,7 @@ def test_use_pat_honors_non_empty_bearer_env(self, monkeypatch): patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch( "ucode.cli.get_databricks_token", - side_effect=lambda w, p, force_refresh=False: os.environ.get( + side_effect=lambda w, p, *, force_refresh=False: os.environ.get( "DATABRICKS_BEARER", "" ), ), From 558abf8398e9e15b1b3add9cfb59116845a76296 Mon Sep 17 00:00:00 2001 From: Brad Rodriguez Date: Wed, 8 Jul 2026 09:00:06 -0700 Subject: [PATCH 3/4] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- tests/test_databricks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 1e37e2b..cefd52f 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1148,7 +1148,7 @@ def test_force_refresh_appends_flag(self, tmp_path, monkeypatch): argv_log = tmp_path / "argv" env = self._fake_databricks( tmp_path, - f'printf "%s\\n" "$@" >> {argv_log}\n' + f'printf "%s\\n" "$@" >> "{argv_log}"\n' 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', ) monkeypatch.setattr("os.environ", env) @@ -1160,7 +1160,7 @@ def test_omits_force_refresh_flag_by_default(self, tmp_path, monkeypatch): argv_log = tmp_path / "argv" env = self._fake_databricks( tmp_path, - f'printf "%s\\n" "$@" >> {argv_log}\n' + f'printf "%s\\n" "$@" >> "{argv_log}"\n' 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', ) monkeypatch.setattr("os.environ", env) From 54b443213adc1dbeaac74e4488200ef3d1856c2c Mon Sep 17 00:00:00 2001 From: Brad Date: Wed, 8 Jul 2026 09:21:01 -0700 Subject: [PATCH 4/4] Harden get_databricks_token tests against ambient env and paths with spaces - Strip DATABRICKS_BEARER in the _fake_databricks helper: get_databricks_token short-circuits on it, so a leaked CI/developer bearer would bypass the fake CLI shim and make the whole TestGetDatabricksToken class's assertions meaningless. - Quote the argv-log redirection target in the fake shim so tests don't break when the temp path contains spaces. Addresses PR review feedback. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_databricks.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index cefd52f..9f1e5a8 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1055,7 +1055,12 @@ def _fake_databricks(self, tmp_path, script: str) -> dict: fake = tmp_path / "databricks" fake.write_text(f"#!/bin/sh\n{script}\n") fake.chmod(0o755) - return {**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}"} + env = {**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}"} + # Drop any ambient DATABRICKS_BEARER: get_databricks_token short-circuits + # on it and would never invoke the fake CLI, so these tests must exercise + # the shim, not a leaked CI/developer bearer. + env.pop("DATABRICKS_BEARER", None) + return env def test_returns_token_on_success(self, tmp_path, monkeypatch): env = self._fake_databricks( @@ -1131,7 +1136,7 @@ def test_passes_profile_flag_when_provided(self, tmp_path, monkeypatch): argv_log = tmp_path / "argv" env = self._fake_databricks( tmp_path, - f'printf "%s\\n" "$@" >> {argv_log}\n' + f'printf "%s\\n" "$@" >> "{argv_log}"\n' 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', ) monkeypatch.setattr("os.environ", env)