Force token refresh in the auth-token apiKeyHelper

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

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"])

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)
@@ -1141,6 +1146,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,