diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1f217b0c..8ce28a26 100644 --- a/src/ucode/cli.py +++ b/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 diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index ee54d153..8fe004f4 100644 --- a/src/ucode/databricks.py +++ b/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 diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 61f07438..0d646bef 100644 --- a/tests/test_agent_claude.py +++ b/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") diff --git a/tests/test_cli.py b/tests/test_cli.py index 9e3e05df..6fe5fd44 100644 --- a/tests/test_cli.py +++ b/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={}): diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7debda88..a9b69c52 100644 --- a/tests/test_databricks.py +++ b/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" diff --git a/tests/test_state.py b/tests/test_state.py index 36c8ce4f..390cad0f 100644 --- a/tests/test_state.py +++ b/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"]]}}}