avoid databricks auth failure when databricks token expires

src/ucode/cli.py

@@ -44,6 +44,7 @@
 from ucode.databricks import (
     apply_pat_environment,
     build_shared_base_urls,
+    create_databricks_user_token,
     discover_claude_models,
     discover_codex_models,
     discover_gemini_models,
@@ -1507,7 +1508,8 @@ def auth_token_cmd(
         print_err("No workspace configured. Run `ug configure` first.")
         raise typer.Exit(1)
     profile = profile or state.get("profile")
-    if use_pat or state.get("use_pat"):
+    use_static_token = bool(use_pat or state.get("use_pat"))
+    if use_static_token:
         # --use-pat explicitly means "serve the profile's static PAT". Fail
         # closed if it can't be read rather than falling through to OAuth —
         # `auth token` cannot serve a PAT-only profile, so that path would
@@ -1521,7 +1523,15 @@ def auth_token_cmd(
             )
             raise typer.Exit(1)
     try:
-        token = get_databricks_token(workspace, profile, force_refresh=force_refresh)
+        if use_static_token:
+            token = get_databricks_token(workspace, profile)
+        else:
+            # The OAuth token is used to mint a short-lived user token, so it
+            # must be fresh even when the caller did not explicitly request a
+            # refresh.  ``--force-refresh`` remains accepted for callers (such
+            # as OpenCode's 401 retry path) that already pass it.
+            token = get_databricks_token(workspace, profile, force_refresh=True)
+            token = create_databricks_user_token(workspace, token, lifetime_seconds=10)
     except RuntimeError as exc:
         print_err(str(exc))
         raise typer.Exit(1) from None

src/ucode/databricks.py

@@ -1250,6 +1250,25 @@ def _fetch_with_lock_retry() -> str:
     return token
 
 
+def create_databricks_user_token(
+    workspace: str,
+    auth_token: str,
+    *,
+    lifetime_seconds: int,
+) -> str:
+    payload, reason = _http_post_json(
+        f"{workspace.rstrip('/')}/api/2.0/token/create",
+        auth_token,
+        {"lifetime_seconds": lifetime_seconds},
+    )
+    if reason:
+        raise RuntimeError(f"Failed to create a short-lived Databricks token: {reason}")
+    token = payload.get("token_value") if isinstance(payload, dict) else None
+    if not isinstance(token, str) or not token:
+        raise RuntimeError("Databricks returned no token_value for the short-lived token.")
+    return token
+
+
 def _extract_connection_page(payload: object) -> tuple[list[dict], str | None]:
     if isinstance(payload, list):
         return [item for item in payload if isinstance(item, dict)], None

tests/test_agent_claude.py

@@ -178,6 +178,10 @@ def test_sets_custom_headers(self):
         overlay, _ = claude.render_overlay(WS, "s4")
         assert "x-databricks-use-coding-agent-mode" in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
 
+    def test_auth_token_cache_ttl(self):
+        overlay, _ = claude.render_overlay(WS, "s4")
+        assert overlay["env"]["CLAUDE_CODE_API_KEY_HELPER_TTL_MS"] == "900000"
+
     def test_does_not_disable_experimental_betas(self):
         # Would suppress the beta header 1h prompt caching needs.
         overlay, _ = claude.render_overlay(WS, "s4")

tests/test_cli.py

@@ -966,33 +966,38 @@ def test_prints_only_the_token_to_stdout(self):
         with (
             patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}),
             patch("ucode.cli.get_databricks_token", return_value="tok-123") as fetch,
+            patch("ucode.cli.create_databricks_user_token", return_value="short-tok") as create,
         ):
             result = runner.invoke(app, ["auth-token"])
         assert result.exit_code == 0
         # 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=False)
+        assert result.stdout == "short-tok\n"
+        fetch.assert_called_once_with("https://ws", None, force_refresh=True)
+        create.assert_called_once_with("https://ws", "tok-123", lifetime_seconds=10)
 
     def test_host_and_profile_override_state(self):
         with (
             patch("ucode.cli.load_state", return_value={"workspace": "https://saved"}),
             patch("ucode.cli.get_databricks_token", return_value="tok") as fetch,
+            patch("ucode.cli.create_databricks_user_token", return_value="short-tok"),
         ):
             result = runner.invoke(
                 app, ["auth-token", "--host", "https://override", "--profile", "prod"]
             )
         assert result.exit_code == 0
-        fetch.assert_called_once_with("https://override", "prod", force_refresh=False)
+        fetch.assert_called_once_with("https://override", "prod", force_refresh=True)
 
     def test_force_refresh_is_forwarded(self):
         with (
             patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}),
             patch("ucode.cli.get_databricks_token", return_value="tok") as fetch,
+            patch("ucode.cli.create_databricks_user_token", return_value="short-tok") as create,
         ):
             result = runner.invoke(app, ["auth-token", "--force-refresh"])
         assert result.exit_code == 0
         fetch.assert_called_once_with("https://ws", None, force_refresh=True)
+        create.assert_called_once_with("https://ws", "tok", lifetime_seconds=10)
 
     def test_errors_without_workspace(self):
         with patch("ucode.cli.load_state", return_value={}):

tests/test_databricks.py

@@ -27,6 +27,7 @@
     build_skills_mcp_url,
     build_tool_base_url,
     classify_model_family,
+    create_databricks_user_token,
     databricks_cli_version,
     discover_sql_warehouses,
     ensure_databricks_cli_version,
@@ -63,6 +64,22 @@ def read(self):
         return self._body
 
 
+class TestCreateDatabricksUserToken:
+    def test_creates_token_with_requested_lifetime(self, monkeypatch):
+        calls = []
+
+        def fake_post(url, token, payload, timeout=10):
+            calls.append((url, token, payload, timeout))
+            return {"token_value": "short-lived"}, None
+
+        monkeypatch.setattr(db_mod, "_http_post_json", fake_post)
+
+        assert create_databricks_user_token(WS, "oauth-token", lifetime_seconds=10) == "short-lived"
+        assert calls == [
+            (f"{WS}/api/2.0/token/create", "oauth-token", {"lifetime_seconds": 10}, 10)
+        ]
+
+
 class TestWorkspaceHostname:
     def test_extracts_hostname(self):
         assert workspace_hostname(WS) == "example.databricks.com"

tests/test_state.py

@@ -218,6 +218,8 @@ def test_populates_agent_state_when_workspace_present(self):
 
         assert result["agents"]["claude"]["model"] == "claude-opus"
         assert result["agents"]["claude"]["base_url"] == FAKE_URLS["claude"]
+        assert result["agents"]["claude"]["auth_refresh_interval_ms"] == 900_000
+        assert result["agents"]["claude"]["env"]["CLAUDE_CODE_API_KEY_HELPER_TTL_MS"] == "900000"
         # Cross-platform helper, not the old POSIX `if [ -n ... ]` pipeline (#116).
         assert "auth-token" in result["agents"]["claude"]["auth_command"]
         assert "if [ -n" not in result["agents"]["claude"]["auth_command"]
@@ -227,8 +229,10 @@ def test_populates_agent_state_when_workspace_present(self):
         codex_auth = result["agents"]["codex"]["auth"]
         assert codex_auth["command"] != "sh"
         assert codex_auth["args"][0] == "auth-token"
+        assert codex_auth["refresh_interval_ms"] == 900_000
         assert result["agents"]["pi"]["model"] == "claude-opus"
         assert result["agents"]["pi"]["base_urls"] == FAKE_URLS["pi"]
+        assert result["agents"]["pi"]["auth_refresh_interval_ms"] == 900_000
 
     def test_normalizes_managed_configs_dict_entry(self):
         state = {"managed_configs": {"claude": {"keys": [["env", "X"]]}}}