From f6e06e8f8b8d9b1ada942bafae45e7e549a43a0d Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 18:11:01 +0000 Subject: [PATCH 1/3] avoid auth problem --- src/ucode/cli.py | 9 +++++++-- tests/test_agent_claude.py | 4 ++++ tests/test_cli.py | 15 +++++++++++++-- tests/test_state.py | 4 ++++ 4 files changed, 28 insertions(+), 4 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a1662ae..fdf32163 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1382,7 +1382,9 @@ def auth_token_cmd( print_err("No workspace configured. Run `ucode 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")) + preset_bearer = bool(os.environ.get("DATABRICKS_BEARER", "").strip()) + 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 @@ -1396,7 +1398,10 @@ def auth_token_cmd( ) raise typer.Exit(1) try: - token = get_databricks_token(workspace, profile) + if use_static_token or preset_bearer: + token = get_databricks_token(workspace, profile) + else: + 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_agent_claude.py b/tests/test_agent_claude.py index f58fed37..d25b9d00 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -177,6 +177,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 8bb1ea9b..9c2e4dbe 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -715,7 +715,7 @@ 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) + fetch.assert_called_once_with("https://ws", None, force_refresh=True) def test_host_and_profile_override_state(self): with ( @@ -726,7 +726,18 @@ 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_preset_bearer_does_not_force_oauth_refresh(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "preset-token") + with ( + patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), + patch("ucode.cli.get_databricks_token", return_value="preset-token") as fetch, + ): + result = runner.invoke(app, ["auth-token"]) + assert result.exit_code == 0 + assert result.stdout == "preset-token\n" + fetch.assert_called_once_with("https://ws", None) def test_errors_without_workspace(self): with patch("ucode.cli.load_state", return_value={}): 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"]]}}} From 981df2e8304c4f70f349f9172c2cde60e74250ff Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 18:13:38 +0000 Subject: [PATCH 2/3] fix --- src/ucode/cli.py | 3 +-- tests/test_cli.py | 11 ----------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index fdf32163..4650fe86 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1383,7 +1383,6 @@ def auth_token_cmd( raise typer.Exit(1) profile = profile or state.get("profile") use_static_token = bool(use_pat or state.get("use_pat")) - preset_bearer = bool(os.environ.get("DATABRICKS_BEARER", "").strip()) 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 — @@ -1398,7 +1397,7 @@ def auth_token_cmd( ) raise typer.Exit(1) try: - if use_static_token or preset_bearer: + if use_static_token: token = get_databricks_token(workspace, profile) else: token = get_databricks_token(workspace, profile, force_refresh=True) diff --git a/tests/test_cli.py b/tests/test_cli.py index 9c2e4dbe..449f3be0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -728,17 +728,6 @@ def test_host_and_profile_override_state(self): assert result.exit_code == 0 fetch.assert_called_once_with("https://override", "prod", force_refresh=True) - def test_preset_bearer_does_not_force_oauth_refresh(self, monkeypatch): - monkeypatch.setenv("DATABRICKS_BEARER", "preset-token") - with ( - patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), - patch("ucode.cli.get_databricks_token", return_value="preset-token") as fetch, - ): - result = runner.invoke(app, ["auth-token"]) - assert result.exit_code == 0 - assert result.stdout == "preset-token\n" - fetch.assert_called_once_with("https://ws", None) - def test_errors_without_workspace(self): with patch("ucode.cli.load_state", return_value={}): result = runner.invoke(app, ["auth-token"]) From 685d6233957f0935e707ed2ab741ad3d69d77d3a Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 22:06:59 +0000 Subject: [PATCH 3/3] temp --- src/ucode/cli.py | 2 ++ src/ucode/databricks.py | 19 +++++++++++++++++++ tests/test_cli.py | 5 ++++- tests/test_databricks.py | 17 +++++++++++++++++ 4 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 4650fe86..20d7749c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -41,6 +41,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, @@ -1401,6 +1402,7 @@ def auth_token_cmd( token = get_databricks_token(workspace, profile) else: 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 f875a5c7..893d4ed8 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1240,6 +1240,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_cli.py b/tests/test_cli.py index 449f3be0..29a1500c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -709,18 +709,21 @@ 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" + 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"] diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 2d9f61a5..72c06e61 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"