From efaec900e25c3ad01b2819c95bab30de14eb6638 Mon Sep 17 00:00:00 2001 From: Bryan Qiu Date: Fri, 12 Jun 2026 14:43:23 -0700 Subject: [PATCH] Support non-interactive configure with --profiles, --use-pat, and --skip-validate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an explicit, fully non-interactive setup path for CI / headless environments (e.g. lakebox sandboxes provisioned with a PAT-backed DEFAULT profile): - `configure --profiles ` resolves workspace URLs from existing ~/.databrickscfg profiles instead of prompting. Auth behaves like --workspaces: OAuth login is forced by default. - `--use-pat` (requires --profiles) authenticates with the profile's personal access token instead of OAuth — ucode never picks up a PAT implicitly. The token is validated, exported as DATABRICKS_BEARER for the configure run and launched agents, and persisted as use_pat in state so launches inherit the mode; written agent configs carry a `databricks auth describe --sensitive`-based auth command so bare claude/codex runs refresh from the PAT too. - `--skip-validate` skips the post-configure test message through each agent; configs are still written with freshly discovered models. Co-authored-by: Isaac Signed-off-by: Bryan Qiu --- README.md | 17 ++ src/ucode/agents/claude.py | 7 +- src/ucode/agents/codex.py | 26 +++- src/ucode/cli.py | 152 +++++++++++++++++- src/ucode/databricks.py | 87 ++++++++++- src/ucode/mcp.py | 2 + src/ucode/state.py | 2 +- src/ucode/tracing.py | 3 + src/ucode/usage.py | 2 + tests/test_cli.py | 312 +++++++++++++++++++++++++++++++++++++ tests/test_databricks.py | 114 ++++++++++++++ tests/test_state.py | 12 ++ 12 files changed, 716 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 77f9b33..0420190 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,20 @@ ucode configure --workspaces https://first.databricks.com,https://second.databri When multiple workspaces are provided, `ucode` logs into and saves state for each workspace. Launch commands such as `ucode codex` use the first workspace in the list. +Alternatively, pass existing Databricks CLI profiles (from `~/.databrickscfg`) instead of workspace URLs — each profile's host supplies the workspace URL: + +```bash +ucode configure --profiles DEFAULT --agents claude,codex +``` + +Auth behaves the same as `--workspaces`: an OAuth `databricks auth login` is forced by default. + +For CI or headless environments where the profile holds a personal access token (`auth_type = pat` in `~/.databrickscfg`), add `--use-pat`. It must be combined with `--profiles` — ucode never picks up a PAT implicitly — and runs no interactive login: the profile's token is used for the whole setup (and by launched agents afterwards), with workspace access verified against the AI Gateway. `--skip-validate` additionally skips the post-configure test message sent through each agent, so configure only writes config files with the freshly discovered models. Together these make setup fully non-interactive: + +```bash +ucode configure --profiles DEFAULT --agents claude,codex --use-pat --skip-validate --skip-upgrade +``` + ### MCP servers (optional) ```bash @@ -90,6 +104,9 @@ Discovered external MCP connections are listed directly. MCP auth uses a Databri | `ucode configure --dry-run` | Preview config files without writing them | | `ucode configure --agents claude,codex` | Configure specific agents without the interactive picker | | `ucode configure --workspaces https://first.databricks.com,https://second.databricks.com` | Configure workspaces without the interactive picker | +| `ucode configure --profiles DEFAULT` | Configure using existing Databricks CLI profiles (hosts come from `~/.databrickscfg`) | +| `ucode configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login | +| `ucode configure --skip-validate` | Write configs without sending a test message through each agent | ## Managed Local Files diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 1afe962..15d17fa 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -116,6 +116,7 @@ def render_overlay( claude_models: dict[str, str] | None = None, disable_web_search: bool = False, profile: str | None = None, + use_pat: bool = False, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -147,7 +148,10 @@ def render_overlay( env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = _maybe_add_1m_suffix(claude_models["sonnet"]) if claude_models.get("haiku"): env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = claude_models["haiku"] - overlay: dict = {"apiKeyHelper": build_auth_shell_command(workspace, profile), "env": env} + overlay: dict = { + "apiKeyHelper": build_auth_shell_command(workspace, profile, use_pat=use_pat), + "env": env, + } keys: list[list[str]] = [["apiKeyHelper"]] + [["env", k] for k in env] # Disable Claude Code's built-in WebSearch (it routes through Anthropic's @@ -226,6 +230,7 @@ def write_tool_config(state: dict, model: str) -> dict: state.get("claude_models") or {}, disable_web_search=web_search_model is not None, profile=state.get("profile"), + use_pat=bool(state.get("use_pat")), ) tracing_env_vars = tracing_env(state, "claude") stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index e8d5eb9..4127afa 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -101,8 +101,8 @@ def _use_legacy_layout() -> bool: return parsed < MINIMUM_CODEX_VERSION -def _provider_block(workspace: str, databricks_profile: str | None) -> dict: - auth_command = build_auth_shell_command(workspace, databricks_profile) +def _provider_block(workspace: str, databricks_profile: str | None, use_pat: bool = False) -> dict: + auth_command = build_auth_shell_command(workspace, databricks_profile, use_pat=use_pat) base_url = build_tool_base_url("codex", workspace) return { "name": "Databricks AI Gateway", @@ -121,19 +121,25 @@ def _provider_block(workspace: str, databricks_profile: str | None) -> dict: def render_overlay( - workspace: str, model: str | None = None, databricks_profile: str | None = None + workspace: str, + model: str | None = None, + databricks_profile: str | None = None, + use_pat: bool = False, ) -> dict: overlay: dict = {"model_provider": CODEX_MODEL_PROVIDER_NAME} if model: overlay["model"] = model overlay["model_providers"] = { - CODEX_MODEL_PROVIDER_NAME: _provider_block(workspace, databricks_profile), + CODEX_MODEL_PROVIDER_NAME: _provider_block(workspace, databricks_profile, use_pat), } return overlay def render_legacy_overlay( - workspace: str, model: str | None = None, databricks_profile: str | None = None + workspace: str, + model: str | None = None, + databricks_profile: str | None = None, + use_pat: bool = False, ) -> dict: """Overlay for Codex CLI < 0.134.0, which only reads `~/.codex/config.toml`. @@ -147,7 +153,7 @@ def render_legacy_overlay( "profile": CODEX_PROFILE_NAME, "profiles": {CODEX_PROFILE_NAME: profile_block}, "model_providers": { - CODEX_MODEL_PROVIDER_NAME: _provider_block(workspace, databricks_profile), + CODEX_MODEL_PROVIDER_NAME: _provider_block(workspace, databricks_profile, use_pat), }, } @@ -295,7 +301,9 @@ def write_tool_config(state: dict, model: str | None = None) -> dict: # and skip the per-profile-file cleanup that would normally strip # ucode's entry from the shared file. backup_existing_file(LEGACY_CODEX_CONFIG_PATH, LEGACY_CODEX_BACKUP_PATH) - overlay = render_legacy_overlay(workspace, chosen_model, databricks_profile) + overlay = render_legacy_overlay( + workspace, chosen_model, databricks_profile, use_pat=bool(state.get("use_pat")) + ) doc = read_toml_safe(LEGACY_CODEX_CONFIG_PATH) deep_merge_dict(doc, overlay) write_toml_file(LEGACY_CODEX_CONFIG_PATH, doc) @@ -305,7 +313,9 @@ def write_tool_config(state: dict, model: str | None = None) -> dict: _remove_legacy_ucode_profile() backup_existing_file(CODEX_CONFIG_PATH, CODEX_BACKUP_PATH) - overlay = render_overlay(workspace, chosen_model, databricks_profile) + overlay = render_overlay( + workspace, chosen_model, databricks_profile, use_pat=bool(state.get("use_pat")) + ) doc = read_toml_safe(CODEX_CONFIG_PATH) deep_merge_dict(doc, overlay) write_toml_file(CODEX_CONFIG_PATH, doc) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 15973b0..76aa47a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os from typing import Annotated import typer @@ -29,6 +30,7 @@ from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import restore_file, set_dry_run from ucode.databricks import ( + apply_pat_environment, build_shared_base_urls, discover_claude_models, discover_codex_models, @@ -40,7 +42,9 @@ get_databricks_profiles, get_databricks_token, install_databricks_cli, + list_profile_entries, normalize_workspace_url, + resolve_pat_token, run_databricks_login, uc_enabled, ) @@ -146,6 +150,46 @@ def _parse_workspaces_option(workspaces: str) -> list[tuple[str, str | None]]: return workspace_entries +def _parse_profiles_option(profiles: str) -> list[tuple[str, str | None]]: + """Parse `--profiles` into [(url, profile_name), ...]. + + Each name must be an existing Databricks CLI profile; its host supplies + the workspace URL. Auth behaves the same as `--workspaces`: OAuth login is + forced unless `--use-pat` is also passed.""" + available = {str(p.get("name")): p for p in list_profile_entries() if p.get("name")} + workspace_entries: list[tuple[str, str | None]] = [] + seen: set[str] = set() + for raw_name in profiles.split(","): + name = raw_name.strip() + if not name: + continue + entry = available.get(name) + if entry is None: + known = ", ".join(sorted(available)) or "none" + raise RuntimeError( + f"Databricks CLI profile '{name}' was not found (available: {known}). " + "Check `databricks auth profiles` or add the profile to ~/.databrickscfg." + ) + host = str(entry.get("host") or "").strip() + if not host: + raise RuntimeError( + f"Databricks CLI profile '{name}' has no host configured in ~/.databrickscfg." + ) + try: + workspace = normalize_workspace_url(host) + except ValueError as exc: + raise RuntimeError(str(exc)) from exc + if workspace not in seen: + seen.add(workspace) + workspace_entries.append((workspace, name)) + if not workspace_entries: + raise RuntimeError( + "No profiles provided for --profiles. Use a comma-separated list like " + "`--profiles DEFAULT`." + ) + return workspace_entries + + def configure_shared_state( workspace: str, profile: str | None = None, @@ -153,11 +197,16 @@ def configure_shared_state( force_login: bool = False, enable_uc: bool | None = None, reset_uc: bool = False, + use_pat: bool | None = None, ) -> dict: """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state. If tools is provided, only fetch models for those tools. Otherwise fetch all. If force_login is True, always run databricks auth login (used by explicit configure). + If use_pat is True (explicit `configure --profiles --use-pat`), the + profile's personal access token from ~/.databrickscfg is used instead of + OAuth and no interactive login ever runs. ``None`` means "inherit": a + launch re-run keeps the mode the workspace was configured with. ``profile`` is the Databricks CLI profile name to address — passed via ``--profile`` to every CLI invocation so ambiguous `~/.databrickscfg` entries (e.g. DEFAULT and a named profile both pointing at the same host) @@ -184,8 +233,27 @@ def configure_shared_state( else: target_ws_state = load_full_state().get("workspaces", {}).get(workspace) or {} enable_uc = uc_enabled(default=bool(target_ws_state.get("uc_enabled"))) + if use_pat is None: + use_pat = bool(prior_state.get("use_pat")) and previous_workspace == workspace fetch_all = tools is None - if force_login: + if use_pat: + if not profile: + raise RuntimeError( + "--use-pat requires a Databricks CLI profile. Pass one via `--profiles `." + ) + pat = resolve_pat_token(profile) + if not pat: + raise RuntimeError( + f"--use-pat: profile '{profile}' has no personal access token in " + "~/.databrickscfg (its auth_type must be `pat`). Add a `token = ` " + f"entry under [{profile}], or re-run without --use-pat to use OAuth." + ) + # Export the PAT for this process and launched agent subprocesses so + # every token fetch takes the static-bearer path; a bearer already in + # the environment wins. + os.environ.setdefault("DATABRICKS_BEARER", pat) + ensure_databricks_auth(workspace, profile) + elif force_login: run_databricks_login(workspace, profile) else: ensure_databricks_auth(workspace, profile) @@ -246,6 +314,12 @@ def configure_shared_state( # Persist the resolved flag so subsequent launches stay on the same # discovery path without the env var or CLI flag being re-passed. state["uc_enabled"] = enable_uc + # Persist the auth mode so launches rebuild the same (PAT-based) agent + # auth command; an explicit re-configure without --use-pat clears it. + if use_pat: + state["use_pat"] = True + else: + state.pop("use_pat", None) state["base_urls"] = build_shared_base_urls(workspace) if want_claude: state["claude_models"] = claude_models @@ -277,6 +351,7 @@ def _configure_shared_workspace_states( force_login: bool, enable_uc: bool | None = None, reset_uc: bool = False, + use_pat: bool = False, ) -> list[dict]: if not workspaces: raise RuntimeError("At least one workspace must be provided.") @@ -290,6 +365,7 @@ def _configure_shared_workspace_states( force_login=force_login, enable_uc=enable_uc, reset_uc=reset_uc, + use_pat=use_pat, ) ) return states @@ -303,6 +379,8 @@ def configure_workspace_command( prompt_optional_updates: bool = True, enable_uc: bool | None = None, reset_uc: bool = False, + use_pat: bool = False, + skip_validate: bool = False, ) -> int: if tool is not None and selected_tools is not None: raise RuntimeError("Use either --agent or --agents, not both.") @@ -311,7 +389,12 @@ def configure_workspace_command( if tool is not None: states = _configure_shared_workspace_states( - workspace_entries, [tool], force_login=True, enable_uc=enable_uc, reset_uc=reset_uc + workspace_entries, + [tool], + force_login=True, + enable_uc=enable_uc, + reset_uc=reset_uc, + use_pat=use_pat, ) state = states[0] state = configure_single_tool(tool, state) @@ -325,6 +408,9 @@ def configure_workspace_command( expand=False, ) ) + if skip_validate: + print_note(f"Skipping {spec['display']} validation (--skip-validate).") + return 0 with spinner(f"Validating {spec['display']}..."): ok, err = validate_tool(tool) if ok: @@ -345,6 +431,7 @@ def configure_workspace_command( force_login=True, enable_uc=enable_uc, reset_uc=reset_uc, + use_pat=use_pat, ) state = states[0] save_state(state) @@ -402,6 +489,9 @@ def configure_workspace_command( ) ) + if skip_validate: + print_note("Skipping agent validation (--skip-validate).") + return 0 # Limit validation to just-configured tools so we don't re-validate # previously-configured tools the user didn't touch this run. validate_state = {**state, "available_tools": picked} @@ -581,6 +671,10 @@ def _launch_tool(tool_name: str, ctx: typer.Context) -> None: try: tool = normalize_tool(tool_name) existing = load_state() + # Workspaces configured with --use-pat export the profile's PAT as + # DATABRICKS_BEARER up front so every auth check below (and the + # launched agent itself) uses the static token instead of OAuth. + apply_pat_environment(existing) needs_auto_configure = not existing.get("workspace") or tool not in ( existing.get("available_tools") or [] ) @@ -680,6 +774,35 @@ def configure( help="Configure a comma-separated list of workspaces without prompting.", ), ] = None, + profiles: Annotated[ + str | None, + typer.Option( + "--profiles", + help="Configure a comma-separated list of existing Databricks CLI profiles " + "without the workspace prompt. Each profile's host from ~/.databrickscfg " + "supplies the workspace URL. Auth behaves like --workspaces: OAuth login " + "is forced unless --use-pat is also passed.", + ), + ] = None, + use_pat: Annotated[ + bool, + typer.Option( + "--use-pat", + help="Authenticate with the personal access token stored in " + "~/.databrickscfg for the selected profile(s) instead of OAuth. " + "Requires --profiles; no interactive login is run. Intended for " + "CI / headless environments.", + ), + ] = False, + skip_validate: Annotated[ + bool, + typer.Option( + "--skip-validate", + help="Skip the post-configure validation step that sends a quick test " + "message through each agent. Config files are still written with the " + "freshly discovered models.", + ), + ] = False, tracing: Annotated[ bool, typer.Option( @@ -734,7 +857,23 @@ def configure( install_databricks_cli() if agent is not None and agents is not None: raise RuntimeError("Use either --agent or --agents, not both.") + if workspaces is not None and profiles is not None: + raise RuntimeError("Use either --workspaces or --profiles, not both.") + if use_pat and profiles is None: + raise RuntimeError( + "--use-pat requires --profiles. Pass the PAT-backed Databricks CLI " + "profile(s) explicitly, e.g. `ucode configure --profiles DEFAULT --use-pat`." + ) workspace_entries = _parse_workspaces_option(workspaces) if workspaces is not None else None + if profiles is not None: + workspace_entries = _parse_profiles_option(profiles) + # Only forward the opt-in flags when set so existing call expectations + # (and defaults) stay unchanged for the common interactive path. + skip_kwargs: dict = {} + if use_pat: + skip_kwargs["use_pat"] = True + if skip_validate: + skip_kwargs["skip_validate"] = True if agent is not None: tool = normalize_tool(agent) install_tool_binary( @@ -744,13 +883,16 @@ def configure( prompt_optional_updates=prompt_optional_updates, ) if workspace_entries is None: - configure_workspace_command(tool, enable_uc=flag_enable_uc, reset_uc=True) + configure_workspace_command( + tool, enable_uc=flag_enable_uc, reset_uc=True, **skip_kwargs + ) else: configure_workspace_command( tool, workspaces=workspace_entries, enable_uc=flag_enable_uc, reset_uc=True, + **skip_kwargs, ) elif agents is not None: selected_tools = _parse_agents_option(agents) @@ -760,6 +902,7 @@ def configure( prompt_optional_updates=prompt_optional_updates, enable_uc=flag_enable_uc, reset_uc=True, + **skip_kwargs, ) else: configure_workspace_command( @@ -768,6 +911,7 @@ def configure( prompt_optional_updates=prompt_optional_updates, enable_uc=flag_enable_uc, reset_uc=True, + **skip_kwargs, ) else: # Tool binaries are installed after the user picks which agents @@ -777,6 +921,7 @@ def configure( prompt_optional_updates=prompt_optional_updates, enable_uc=flag_enable_uc, reset_uc=True, + **skip_kwargs, ) else: configure_workspace_command( @@ -784,6 +929,7 @@ def configure( prompt_optional_updates=prompt_optional_updates, enable_uc=flag_enable_uc, reset_uc=True, + **skip_kwargs, ) if tracing: # The workspaces were just configured, so enable tracing for them diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index db26583..b35c07b 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3,6 +3,7 @@ from __future__ import annotations +import configparser import functools import json import logging @@ -591,8 +592,9 @@ def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> boo return False -def get_databricks_profiles() -> list[tuple[str, str]]: - """Return [(host_url, profile_name), ...] from Databricks CLI profiles. +def list_profile_entries() -> list[dict]: + """Return raw profile dicts ({"name", "host", "auth_type", ...}) from + `databricks auth profiles`. Returns ``[]`` on any failure (CLI missing, timeout, non-zero exit, JSON decode error). When ``UCODE_DEBUG=1`` each dropout path logs *why* the @@ -608,16 +610,22 @@ def get_databricks_profiles() -> list[tuple[str, str]]: timeout=20, ) except (OSError, subprocess.TimeoutExpired) as exc: - _debug("get_databricks_profiles", f"subprocess error: {type(exc).__name__}: {exc}") + _debug("list_profile_entries", f"subprocess error: {type(exc).__name__}: {exc}") return [] if result.returncode != 0: - _debug("get_databricks_profiles", _format_subprocess_result(result)) + _debug("list_profile_entries", _format_subprocess_result(result)) return [] try: profiles = json.loads(result.stdout or "{}").get("profiles") or [] except json.JSONDecodeError as exc: - _debug("get_databricks_profiles", f"json decode error: {exc.msg}") + _debug("list_profile_entries", f"json decode error: {exc.msg}") return [] + return [p for p in profiles if isinstance(p, dict)] + + +def get_databricks_profiles() -> list[tuple[str, str]]: + """Return [(host_url, profile_name), ...] from Databricks CLI profiles.""" + profiles = list_profile_entries() # dict dedupes by host (first non-PAT profile wins). out: dict[str, str] = {} @@ -648,6 +656,61 @@ def find_profile_name_for_host(workspace: str) -> str | None: return None +def profile_auth_type(profile: str) -> str | None: + """Return the auth_type of a Databricks CLI profile (e.g. "pat"), or None.""" + for p in list_profile_entries(): + if p.get("name") == profile: + auth_type = p.get("auth_type") + return auth_type if isinstance(auth_type, str) else None + return None + + +def _read_databrickscfg_token(profile: str) -> str | None: + """Read the static ``token`` value for a profile from ``~/.databrickscfg``. + + `databricks auth token` only knows OAuth caches; for PAT profiles the PAT + itself is the credential, stored in the config file. The parser's default + section is pointed at a name that never appears in the file so a token in + ``[DEFAULT]`` does not leak into every named profile.""" + cfg_path = Path(os.environ.get("DATABRICKS_CONFIG_FILE") or "~/.databrickscfg").expanduser() + parser = configparser.ConfigParser(default_section="@ucode-no-defaults@", interpolation=None) + try: + if not parser.read(cfg_path, encoding="utf-8"): + return None + except (configparser.Error, OSError): + return None + if not parser.has_section(profile): + return None + token = (parser.get(profile, "token", fallback="") or "").strip() + return token or None + + +def resolve_pat_token(profile: str | None) -> str | None: + """Return the static PAT of a PAT-type Databricks CLI profile, or None. + + Only consulted when the user explicitly opted in via + ``ucode configure --profiles --use-pat`` — ucode never picks up a + PAT implicitly.""" + if profile and profile_auth_type(profile) == "pat": + return _read_databrickscfg_token(profile) + return None + + +def apply_pat_environment(state: dict) -> None: + """Export the configured profile's PAT as ``DATABRICKS_BEARER`` when the + workspace was configured with ``--use-pat``. + + Every token fetch in this process (and in launched agent subprocesses, + which inherit the environment) then takes the existing static-bearer + short-circuit instead of the OAuth-only `databricks auth token` path. + A bearer already present in the environment is left untouched.""" + if not state.get("use_pat"): + return + pat = resolve_pat_token(state.get("profile")) + if pat: + os.environ.setdefault("DATABRICKS_BEARER", pat) + + def run_databricks_login(workspace: str, profile: str | None = None) -> None: """Run databricks auth login unconditionally. @@ -955,9 +1018,19 @@ def list_databricks_apps(workspace: str, profile: str | None = None) -> list[dic raise RuntimeError("Databricks apps listing returned invalid JSON.") from exc -def build_auth_shell_command(workspace: str, profile: str | None = None) -> str: +def build_auth_shell_command( + workspace: str, profile: str | None = None, *, use_pat: bool = False +) -> str: workspace_arg = shlex.quote(workspace.rstrip("/")) - if profile: + if use_pat and profile: + # --use-pat profiles have no OAuth cache for `auth token` to read, so + # the persisted command reads the profile's static token instead. + profile_arg = shlex.quote(profile) + cli_command = ( + f"databricks auth describe --profile {profile_arg} --sensitive --output json " + "| jq -r '.details.configuration.token.value'" + ) + elif profile: profile_arg = shlex.quote(profile) cli_command = ( f"databricks auth token --host {workspace_arg} " diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index e961ff8..18fe57d 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -27,6 +27,7 @@ from ucode.agents import copilot, gemini, opencode from ucode.config_io import restore_file from ucode.databricks import ( + apply_pat_environment, build_mcp_service_url, ensure_databricks_auth, get_databricks_token, @@ -965,6 +966,7 @@ def configure_mcp_command() -> int: ] profile = state.get("profile") + apply_pat_environment(state) ensure_databricks_auth(workspace, profile) print_section("MCP Servers") diff --git a/src/ucode/state.py b/src/ucode/state.py index 5ff6ff4..471eae0 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -124,7 +124,7 @@ def build_agent_state(state: dict) -> dict[str, dict]: profile = state.get("profile") if isinstance(state.get("profile"), str) else None base_urls_value = state.get("base_urls") base_urls = base_urls_value if isinstance(base_urls_value, dict) else {} - auth_command = build_auth_shell_command(workspace, profile) + auth_command = build_auth_shell_command(workspace, profile, use_pat=bool(state.get("use_pat"))) claude_models_value = state.get("claude_models") claude_models: dict = claude_models_value if isinstance(claude_models_value, dict) else {} codex_models_value = state.get("codex_models") diff --git a/src/ucode/tracing.py b/src/ucode/tracing.py index 55d99e1..84e81bb 100644 --- a/src/ucode/tracing.py +++ b/src/ucode/tracing.py @@ -20,6 +20,7 @@ from __future__ import annotations from ucode.databricks import ( + apply_pat_environment, ensure_databricks_auth, find_uc_backed_experiment, get_databricks_token, @@ -301,6 +302,7 @@ def _enable_tracing_for_state(state: dict) -> dict: workspace = state["workspace"] configured = _configured_tracing_agents(state) profile = state.get("profile") + apply_pat_environment(state) ensure_databricks_auth(workspace, profile) print_section("MLflow Tracing") @@ -364,6 +366,7 @@ def _disable_tracing_command() -> int: workspace = state["workspace"] profile = state.get("profile") + apply_pat_environment(state) ensure_databricks_auth(workspace, profile) print_section("MLflow Tracing") diff --git a/src/ucode/usage.py b/src/ucode/usage.py index ab9a0c4..0e30aa5 100644 --- a/src/ucode/usage.py +++ b/src/ucode/usage.py @@ -11,6 +11,7 @@ from typing import cast from ucode.databricks import ( + apply_pat_environment, discover_sql_warehouse_http_path, ensure_databricks_auth, get_databricks_token, @@ -446,6 +447,7 @@ def usage() -> int: raise RuntimeError("Workspace is not configured. Run `ucode configure` first.") profile = state.get("profile") + apply_pat_environment(state) ensure_databricks_auth(workspace, profile) with spinner("Retrieving Databricks access token..."): token = get_databricks_token(workspace, profile) diff --git a/tests/test_cli.py b/tests/test_cli.py index 7e2a98c..9a4fe5d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -687,6 +687,7 @@ def fake_configure_shared_state( force_login=False, enable_uc=None, reset_uc=False, + use_pat=False, ): configured_shared.append( (workspace, profile, tuple(tools) if tools is not None else None, force_login) @@ -724,6 +725,317 @@ def fake_configure_shared_state( assert configured_tools == [("https://first.com", ["codex"])] +class TestParseProfilesOption: + @staticmethod + def _patch_profiles(monkeypatch, entries): + import ucode.cli as cli_mod + + monkeypatch.setattr(cli_mod, "list_profile_entries", lambda: entries) + return cli_mod + + def test_resolves_profiles_to_workspace_entries(self, monkeypatch): + cli_mod = self._patch_profiles( + monkeypatch, + [ + {"name": "DEFAULT", "host": "https://first.databricks.com/", "auth_type": "pat"}, + { + "name": "second", + "host": "https://second.databricks.com", + "auth_type": "databricks-cli", + }, + ], + ) + assert cli_mod._parse_profiles_option("DEFAULT, second") == [ + ("https://first.databricks.com", "DEFAULT"), + ("https://second.databricks.com", "second"), + ] + + def test_unknown_profile_raises_with_available_names(self, monkeypatch): + cli_mod = self._patch_profiles( + monkeypatch, + [{"name": "DEFAULT", "host": "https://first.databricks.com", "auth_type": "pat"}], + ) + with pytest.raises(RuntimeError, match=r"'missing' was not found.*DEFAULT"): + cli_mod._parse_profiles_option("missing") + + def test_profile_without_host_raises(self, monkeypatch): + cli_mod = self._patch_profiles(monkeypatch, [{"name": "DEFAULT", "auth_type": "pat"}]) + with pytest.raises(RuntimeError, match="no host configured"): + cli_mod._parse_profiles_option("DEFAULT") + + def test_empty_value_raises(self, monkeypatch): + cli_mod = self._patch_profiles(monkeypatch, []) + with pytest.raises(RuntimeError, match="No profiles provided"): + cli_mod._parse_profiles_option(" , ") + + +class TestConfigureProfilesFlag: + PROFILE_ENTRIES = [ + {"name": "DEFAULT", "host": "https://first.databricks.com", "auth_type": "pat"} + ] + + def test_profiles_flag_resolves_workspaces(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.list_profile_entries", return_value=self.PROFILE_ENTRIES), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke(app, ["configure", "--profiles", "DEFAULT"]) + assert result.exit_code == 0, result.output + # Auth behaves like --workspaces: no skip flags are forwarded, so the + # default forced OAuth login applies. + mock_cfg.assert_called_once_with( + workspaces=[("https://first.databricks.com", "DEFAULT")], + prompt_optional_updates=True, + enable_uc=None, + reset_uc=True, + ) + + def test_profiles_flag_with_agents(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.list_profile_entries", return_value=self.PROFILE_ENTRIES), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, ["configure", "--agents", "claude,codex", "--profiles", "DEFAULT"] + ) + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + selected_tools=["claude", "codex"], + workspaces=[("https://first.databricks.com", "DEFAULT")], + prompt_optional_updates=True, + enable_uc=None, + reset_uc=True, + ) + + def test_profiles_flag_with_agent(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.list_profile_entries", return_value=self.PROFILE_ENTRIES), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke(app, ["configure", "--agent", "claude", "--profiles", "DEFAULT"]) + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + "claude", + workspaces=[("https://first.databricks.com", "DEFAULT")], + enable_uc=None, + reset_uc=True, + ) + + def test_use_pat_and_skip_validate_are_forwarded(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.list_profile_entries", return_value=self.PROFILE_ENTRIES), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, + [ + "configure", + "--agents", + "claude,codex", + "--profiles", + "DEFAULT", + "--use-pat", + "--skip-validate", + ], + ) + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + selected_tools=["claude", "codex"], + workspaces=[("https://first.databricks.com", "DEFAULT")], + prompt_optional_updates=True, + enable_uc=None, + reset_uc=True, + use_pat=True, + skip_validate=True, + ) + + def test_use_pat_requires_profiles(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, + ["configure", "--workspaces", "https://first.databricks.com", "--use-pat"], + ) + assert result.exit_code == 1 + assert "--use-pat requires --profiles" in _strip_ansi(result.output) + mock_cfg.assert_not_called() + + def test_profiles_and_workspaces_are_mutually_exclusive(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, + [ + "configure", + "--profiles", + "DEFAULT", + "--workspaces", + "https://first.databricks.com", + ], + ) + assert result.exit_code == 1 + assert "not both" in _strip_ansi(result.output) + mock_cfg.assert_not_called() + + +class TestConfigureSharedStateUsePat: + """--use-pat reads the profile's PAT from ~/.databrickscfg, exports it as + DATABRICKS_BEARER, persists the mode, and never opens a browser.""" + + WS = "https://example.databricks.com" + + @pytest.fixture(autouse=True) + def _isolated_bearer(self): + # configure_shared_state writes DATABRICKS_BEARER directly; restore it + # since monkeypatch can't track writes made by code under test. + import os as os_mod + + original = os_mod.environ.pop("DATABRICKS_BEARER", None) + yield + if original is None: + os_mod.environ.pop("DATABRICKS_BEARER", None) + else: + os_mod.environ["DATABRICKS_BEARER"] = original + + @staticmethod + def _stub_deps(monkeypatch, *, pat_token, existing_state=None): + import ucode.cli as cli_mod + + logins: list[tuple] = [] + ensures: list[tuple] = [] + saved: list[dict] = [] + monkeypatch.setattr(cli_mod, "load_state", lambda: dict(existing_state or {})) + monkeypatch.setattr(cli_mod, "save_state", lambda s: saved.append(dict(s))) + monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: logins.append((w, p))) + monkeypatch.setattr( + cli_mod, "ensure_databricks_auth", lambda w, p=None: ensures.append((w, p)) + ) + monkeypatch.setattr(cli_mod, "resolve_pat_token", lambda p: pat_token) + monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) + monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) + monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) + return cli_mod, logins, ensures, saved + + def test_use_pat_exports_bearer_and_skips_login(self, monkeypatch): + import os as os_mod + + cli_mod, logins, ensures, saved = self._stub_deps(monkeypatch, pat_token="dapi-pat") + + state = cli_mod.configure_shared_state( + self.WS, profile="DEFAULT", force_login=True, use_pat=True + ) + + assert logins == [] + assert ensures == [(self.WS, "DEFAULT")] + assert os_mod.environ["DATABRICKS_BEARER"] == "dapi-pat" + assert state["use_pat"] is True + assert saved and saved[-1]["use_pat"] is True + + def test_use_pat_without_pat_profile_raises(self, monkeypatch): + cli_mod, logins, _, _ = self._stub_deps(monkeypatch, pat_token=None) + + with pytest.raises(RuntimeError, match="no personal access token"): + cli_mod.configure_shared_state( + self.WS, profile="oauth-profile", force_login=True, use_pat=True + ) + assert logins == [] + + def test_use_pat_without_profile_raises(self, monkeypatch): + cli_mod, _, _, _ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + + with pytest.raises(RuntimeError, match="requires a Databricks CLI profile"): + cli_mod.configure_shared_state(self.WS, force_login=True, use_pat=True) + + def test_launch_inherits_persisted_use_pat(self, monkeypatch): + # A launch re-run passes use_pat=None; the persisted mode for the same + # workspace must apply so no OAuth login is forced. + cli_mod, logins, ensures, _ = self._stub_deps( + monkeypatch, + pat_token="dapi-pat", + existing_state={"workspace": self.WS, "profile": "DEFAULT", "use_pat": True}, + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT", force_login=False) + + assert logins == [] + assert state["use_pat"] is True + + def test_reconfigure_without_flag_clears_use_pat(self, monkeypatch): + cli_mod, logins, _, _ = self._stub_deps( + monkeypatch, + pat_token="dapi-pat", + existing_state={"workspace": self.WS, "profile": "DEFAULT", "use_pat": True}, + ) + + state = cli_mod.configure_shared_state( + self.WS, profile="DEFAULT", force_login=True, use_pat=False + ) + + assert logins == [(self.WS, "DEFAULT")] + assert "use_pat" not in state + + +class TestConfigureSkipValidate: + def test_skip_validate_skips_agent_validation(self, monkeypatch): + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "workspace": "https://first.com"} + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "save_state", lambda s: None) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda s, tools: {**s, "available_tools": tools}, + ) + validated: list = [] + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda s: validated.append(s)) + + result = cli_mod.configure_workspace_command( + selected_tools=["codex"], + workspaces=[("https://first.com", None)], + skip_validate=True, + ) + + assert result == 0 + assert validated == [] + + def test_skip_validate_skips_single_tool_validation(self, monkeypatch): + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "workspace": "https://first.com"} + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "configure_single_tool", lambda t, s: s) + validated: list = [] + monkeypatch.setattr(cli_mod, "validate_tool", lambda t: validated.append(t) or (True, "")) + + result = cli_mod.configure_workspace_command( + "claude", + workspaces=[("https://first.com", None)], + skip_validate=True, + ) + + assert result == 0 + assert validated == [] + + class TestConfigureSharedStateMcpCleanup: """A workspace switch should scrub the previous workspace's MCP entries from installed client configs. Switching to the same workspace must not.""" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 3a850e7..5be51f0 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -488,6 +488,83 @@ def test_codex_discovery_keeps_alphabetical_order(self, monkeypatch): assert models == ["databricks-gpt-4-1", "databricks-gpt-5-2-codex"] +class TestResolvePatToken: + def test_reads_pat_profile_token_from_cfg(self, monkeypatch, tmp_path): + cfg = tmp_path / "databrickscfg" + cfg.write_text(f"[lakebox]\nhost = {WS}\ntoken = dapi-from-cfg\n") + monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg)) + monkeypatch.setattr( + db_mod, + "list_profile_entries", + lambda: [{"name": "lakebox", "host": WS, "auth_type": "pat"}], + ) + assert db_mod.resolve_pat_token("lakebox") == "dapi-from-cfg" + + def test_default_section_token_does_not_leak_into_named_profiles(self, monkeypatch, tmp_path): + cfg = tmp_path / "databrickscfg" + cfg.write_text( + f"[DEFAULT]\nhost = {WS}\ntoken = dapi-default\n" + "[other]\nhost = https://other.databricks.com\n" + ) + monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg)) + monkeypatch.setattr( + db_mod, + "list_profile_entries", + lambda: [ + {"name": "DEFAULT", "host": WS, "auth_type": "pat"}, + {"name": "other", "host": "https://other.databricks.com", "auth_type": "pat"}, + ], + ) + assert db_mod.resolve_pat_token("DEFAULT") == "dapi-default" + assert db_mod.resolve_pat_token("other") is None + + def test_returns_none_for_oauth_profile(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "list_profile_entries", + lambda: [{"name": "oauth", "host": WS, "auth_type": "databricks-cli"}], + ) + assert db_mod.resolve_pat_token("oauth") is None + + def test_returns_none_without_profile(self): + assert db_mod.resolve_pat_token(None) is None + + +class TestApplyPatEnvironment: + @pytest.fixture(autouse=True) + def _isolated_bearer(self): + # apply_pat_environment writes os.environ directly; restore it even + # though monkeypatch can't track writes made by code under test. + original = os.environ.pop("DATABRICKS_BEARER", None) + yield + if original is None: + os.environ.pop("DATABRICKS_BEARER", None) + else: + os.environ["DATABRICKS_BEARER"] = original + + def test_exports_bearer_for_use_pat_state(self, monkeypatch): + monkeypatch.setattr(db_mod, "resolve_pat_token", lambda p: "dapi-pat") + + db_mod.apply_pat_environment({"use_pat": True, "profile": "DEFAULT"}) + + assert os.environ["DATABRICKS_BEARER"] == "dapi-pat" + + def test_noop_without_use_pat(self, monkeypatch): + monkeypatch.setattr(db_mod, "resolve_pat_token", lambda p: "dapi-pat") + + db_mod.apply_pat_environment({"profile": "DEFAULT"}) + + assert "DATABRICKS_BEARER" not in os.environ + + def test_existing_bearer_wins(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "explicit-bearer") + monkeypatch.setattr(db_mod, "resolve_pat_token", lambda p: "dapi-pat") + + db_mod.apply_pat_environment({"use_pat": True, "profile": "DEFAULT"}) + + assert os.environ["DATABRICKS_BEARER"] == "explicit-bearer" + + class TestBuildAuthShellCommand: def test_contains_workspace(self): cmd = build_auth_shell_command(WS) @@ -552,6 +629,43 @@ def test_quotes_profile_shell_metacharacters(self): assert "rm -rf /" in cmd assert "'weird name; rm -rf /'" in cmd + def test_use_pat_reads_profile_token_via_describe(self): + cmd = build_auth_shell_command(WS, profile="DEFAULT", use_pat=True) + assert "auth describe --profile DEFAULT --sensitive" in cmd + assert ".details.configuration.token.value" in cmd + # No OAuth attempt for PAT profiles — `auth token` cannot serve them. + assert "auth token" not in cmd + # The DATABRICKS_BEARER escape hatch still takes precedence. + assert "DATABRICKS_BEARER" in cmd + + def test_use_pat_command_emits_token(self, tmp_path): + fake = tmp_path / "databricks" + fake.write_text( + "#!/bin/sh\n" + 'case "$*" in\n' + ' *"auth describe"*) echo \'{"details": {"configuration": ' + '{"token": {"value": "dapi-pat"}}}}\' ;;\n' + " *) exit 1 ;;\n" + "esac\n" + ) + fake.chmod(0o755) + cmd = build_auth_shell_command(WS, profile="DEFAULT", use_pat=True) + result = subprocess.run( + ["sh", "-c", cmd], + capture_output=True, + text=True, + env={ + **os.environ, + "PATH": f"{tmp_path}:{os.environ['PATH']}", + "DATABRICKS_BEARER": "", + }, + ) + assert result.stdout.strip() == "dapi-pat" + + def test_use_pat_without_profile_falls_back_to_oauth_command(self): + cmd = build_auth_shell_command(WS, use_pat=True) + assert "auth token" in cmd + class TestFormatSubprocessResult: def test_suppresses_stdout_on_success(self): diff --git a/tests/test_state.py b/tests/test_state.py index 6593a65..95d1440 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -205,6 +205,18 @@ def test_returns_empty_without_workspace(self): result = build_agent_state({"base_urls": FAKE_URLS}) assert result == {} + def test_use_pat_state_builds_pat_auth_command(self): + result = build_agent_state( + { + "workspace": "https://example.databricks.com", + "profile": "DEFAULT", + "use_pat": True, + "base_urls": FAKE_URLS, + } + ) + for agent in ("claude", "codex", "pi"): + assert "auth describe --profile DEFAULT --sensitive" in (result[agent]["auth_command"]) + # --------------------------------------------------------------------------- # mark_tool_managed