From fb1dac0056d69739df140a534b2b10daf9d32fe5 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Fri, 22 May 2026 16:43:58 -0400 Subject: [PATCH 1/4] Pass --profile through all databricks CLI calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When ~/.databrickscfg has multiple profiles pointing at the same host (e.g. DEFAULT and an aliased name), `databricks auth token --host ` errors with "DEFAULT and match ... Use --profile to specify which profile to use" because the CLI refuses to disambiguate from the host alone. Only run_databricks_login passed --profile; has_valid_databricks_auth, get_databricks_token, the inline auth login --no-browser retry, and build_auth_shell_command did not — so a workspace with ambiguous profiles could never complete configure or launch. This threads the profile name from the picker through state ("profile" key) into every downstream CLI invocation, and falls back to find_profile_name_for_host when configure is given a bare URL via --workspaces. The web-search MCP subprocess receives the profile via DATABRICKS_CONFIG_PROFILE so the token refresh inside Claude Code also disambiguates correctly. --- src/ucode/agents/claude.py | 25 +++++---- src/ucode/agents/codex.py | 18 ++++--- src/ucode/agents/copilot.py | 6 ++- src/ucode/agents/gemini.py | 6 ++- src/ucode/agents/opencode.py | 6 ++- src/ucode/agents/pi.py | 6 ++- src/ucode/cli.py | 90 +++++++++++++++++++++---------- src/ucode/databricks.py | 102 +++++++++++++++++++++++++++-------- src/ucode/mcp.py | 2 +- src/ucode/mcp_web_search.py | 3 +- src/ucode/state.py | 3 +- src/ucode/ui.py | 19 ++++--- src/ucode/usage.py | 5 +- tests/test_agent_claude.py | 6 ++- tests/test_agent_codex.py | 4 +- tests/test_agent_gemini.py | 4 +- tests/test_cli.py | 30 ++++++----- tests/test_databricks.py | 30 +++++++++++ tests/test_e2e.py | 14 ++--- tests/test_e2e_user_agent.py | 4 +- tests/test_mcp.py | 16 +++--- tests/test_mcp_web_search.py | 2 +- 22 files changed, 279 insertions(+), 122 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index ed70c80..2f79f7d 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -59,19 +59,22 @@ def _resolve_web_search_model(state: dict) -> str | None: WEB_SEARCH_MCP_NAME = "web_search" -def _web_search_mcp_entry(workspace: str, search_model: str) -> dict: +def _web_search_mcp_entry(workspace: str, search_model: str, profile: str | None = None) -> dict: """Stdio MCP server entry pointing at `ucode mcp web-search`. Resolves the absolute path to the `ucode` binary so launchers without the right PATH (e.g. desktop GUI launchers) still find it.""" ucode_binary = shutil.which("ucode") or "ucode" + env: dict[str, str] = { + "DATABRICKS_HOST": workspace, + "UCODE_WEB_SEARCH_MODEL": search_model, + } + if profile: + env["DATABRICKS_CONFIG_PROFILE"] = profile return { "type": "stdio", "command": ucode_binary, "args": ["mcp", "web-search"], - "env": { - "DATABRICKS_HOST": workspace, - "UCODE_WEB_SEARCH_MODEL": search_model, - }, + "env": env, } @@ -80,6 +83,7 @@ def render_overlay( model: str, claude_models: dict[str, str] | None = None, disable_web_search: bool = False, + profile: str | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -111,7 +115,7 @@ def render_overlay( env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = claude_models["sonnet"] if claude_models.get("haiku"): env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = claude_models["haiku"] - overlay: dict = {"apiKeyHelper": build_auth_shell_command(workspace), "env": env} + overlay: dict = {"apiKeyHelper": build_auth_shell_command(workspace, profile), "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 @@ -124,7 +128,7 @@ def render_overlay( return overlay, keys -def _register_web_search_mcp(workspace: str, search_model: str) -> None: +def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> None: """Register (or replace) the web_search MCP server in Claude Code's user scope via `claude mcp add-json`. Removes any prior entry first so re-runs pick up changes to the workspace, model, or ucode binary path.""" @@ -141,7 +145,7 @@ def _register_web_search_mcp(workspace: str, search_model: str) -> None: except RuntimeError: # Best-effort cleanup of stale entries — keep going. pass - entry = _web_search_mcp_entry(workspace, search_model) + entry = _web_search_mcp_entry(workspace, search_model, profile) add_claude_mcp_server(WEB_SEARCH_MCP_NAME, entry) @@ -164,13 +168,14 @@ def write_tool_config(state: dict, model: str) -> dict: model, state.get("claude_models") or {}, disable_web_search=web_search_model is not None, + profile=state.get("profile"), ) existing = read_json_safe(CLAUDE_SETTINGS_PATH) merged = deep_merge_dict(existing, overlay) write_json_file(CLAUDE_SETTINGS_PATH, merged) if web_search_model: - _register_web_search_mcp(state["workspace"], web_search_model) + _register_web_search_mcp(state["workspace"], web_search_model, state.get("profile")) state = mark_tool_managed(state, "claude", managed_keys) save_state(state) @@ -186,7 +191,7 @@ def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") if workspace: - os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace) + os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) os.execvp(binary, [binary, "--settings", str(CLAUDE_SETTINGS_PATH), *tool_args]) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index c40428c..3a1b2aa 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -47,14 +47,16 @@ def is_update_available() -> tuple[str, str] | None: return available_npm_package_update(SPEC["package"]) -def render_overlay(workspace: str, model: str | None = None) -> dict: - auth_command = build_auth_shell_command(workspace) +def render_overlay( + workspace: str, model: str | None = None, databricks_profile: str | None = None +) -> dict: + auth_command = build_auth_shell_command(workspace, databricks_profile) base_url = build_tool_base_url("codex", workspace) - profile = {"model_provider": CODEX_MODEL_PROVIDER_NAME} + codex_profile_cfg: dict[str, str] = {"model_provider": CODEX_MODEL_PROVIDER_NAME} if model: - profile["model"] = model + codex_profile_cfg["model"] = model return { - "profiles": {CODEX_PROFILE_NAME: profile}, + "profiles": {CODEX_PROFILE_NAME: codex_profile_cfg}, "model_providers": { CODEX_MODEL_PROVIDER_NAME: { "name": "Databricks AI Gateway", @@ -76,7 +78,9 @@ def render_overlay(workspace: str, model: str | None = None) -> dict: def write_tool_config(state: dict, model: str | None = None) -> dict: backup_existing_file(CODEX_CONFIG_PATH, CODEX_BACKUP_PATH) - overlay = render_overlay(state["workspace"], model or default_model(state)) + overlay = render_overlay( + state["workspace"], model or default_model(state), state.get("profile") + ) doc = read_toml_safe(CODEX_CONFIG_PATH) deep_merge_dict(doc, overlay) write_toml_file(CODEX_CONFIG_PATH, doc) @@ -94,7 +98,7 @@ def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") if workspace: - os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace) + os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) os.execvp(binary, [binary, "--profile", CODEX_PROFILE_NAME, *tool_args]) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 18ef291..91192e4 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -143,7 +143,9 @@ def write_tool_config( ) -> tuple[dict, str]: backup_existing_file(COPILOT_ENV_PATH, COPILOT_BACKUP_PATH) if token is None: - token = get_databricks_token(state["workspace"], force_refresh=force_refresh) + token = get_databricks_token( + state["workspace"], state.get("profile"), force_refresh=force_refresh + ) overlay = render_env_overlay(state["workspace"], model, token) existing = parse_dotenv(COPILOT_ENV_PATH) for key in LEGACY_ENV_KEYS: @@ -220,5 +222,5 @@ def validate_env(state: dict) -> dict[str, str]: model = default_model(state) if not model: raise RuntimeError("No Copilot model is available on this workspace.") - token = get_databricks_token(workspace) + token = get_databricks_token(workspace, state.get("profile")) return build_runtime_env(workspace, model, token) diff --git a/src/ucode/agents/gemini.py b/src/ucode/agents/gemini.py index 5649021..3415104 100644 --- a/src/ucode/agents/gemini.py +++ b/src/ucode/agents/gemini.py @@ -84,7 +84,9 @@ def write_tool_config( ) -> tuple[dict, str]: backup_existing_file(GEMINI_ENV_PATH, GEMINI_BACKUP_PATH) if token is None: - token = get_databricks_token(state["workspace"], force_refresh=force_refresh) + token = get_databricks_token( + state["workspace"], state.get("profile"), force_refresh=force_refresh + ) overlay = render_env_overlay(state["workspace"], model, token) existing = parse_dotenv(GEMINI_ENV_PATH) existing.update(overlay) @@ -159,5 +161,5 @@ def validate_env(state: dict) -> dict[str, str]: model = default_model(state) if not model: raise RuntimeError("No Gemini model is configured.") - token = get_databricks_token(workspace) + token = get_databricks_token(workspace, state.get("profile")) return build_runtime_env(workspace, model, token) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 7215388..515fa92 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -132,7 +132,9 @@ def write_tool_config( ) -> tuple[dict, str]: backup_existing_file(OPENCODE_CONFIG_PATH, OPENCODE_BACKUP_PATH) if token is None: - token = get_databricks_token(state["workspace"], force_refresh=force_refresh) + token = get_databricks_token( + state["workspace"], state.get("profile"), force_refresh=force_refresh + ) opencode_base_urls = state.get("base_urls", {}).get("opencode") or build_opencode_base_urls( state["workspace"] ) @@ -255,4 +257,4 @@ def validate_env(state: dict) -> dict[str, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError("No workspace configured.") - return build_runtime_env(get_databricks_token(workspace)) + return build_runtime_env(get_databricks_token(workspace, state.get("profile"))) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index 995c1a5..aab1ec3 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -165,7 +165,9 @@ def write_tool_config( ) -> tuple[dict, str]: backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH) if token is None: - token = get_databricks_token(state["workspace"], force_refresh=force_refresh) + token = get_databricks_token( + state["workspace"], state.get("profile"), force_refresh=force_refresh + ) pi_base_urls = state.get("base_urls", {}).get("pi") or build_pi_base_urls(state["workspace"]) overlay, managed_keys = render_overlay( model, @@ -256,4 +258,4 @@ def validate_env(state: dict) -> dict[str, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError("No workspace configured.") - return build_runtime_env(get_databricks_token(workspace)) + return build_runtime_env(get_databricks_token(workspace, state.get("profile"))) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6bc3f93..4fd20e1 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -33,6 +33,7 @@ discover_gemini_models, ensure_ai_gateway_v2, ensure_databricks_auth, + find_profile_name_for_host, get_databricks_profiles, get_databricks_token, install_databricks_cli, @@ -81,7 +82,7 @@ def _print_discovery_diagnostics(state: dict) -> None: print_note("Re-run with `UCODE_DEBUG=1` to log raw discovery responses to ~/.ucode/debug.log.") -def _prompt_for_configuration(tool: str | None = None) -> str: +def _prompt_for_configuration(tool: str | None = None) -> tuple[str, str | None]: if tool is None: desc = "Configure your Databricks workspace" else: @@ -106,8 +107,14 @@ def _parse_agents_option(agents: str) -> list[str]: return tools -def _parse_workspaces_option(workspaces: str) -> list[str]: - workspace_urls: list[str] = [] +def _parse_workspaces_option(workspaces: str) -> list[tuple[str, str | None]]: + """Parse `--workspaces` into [(url, profile_name | None), ...]. + + `--workspaces` supplies bare URLs; the matching profile (if any) is + resolved later via `find_profile_name_for_host`. + """ + workspace_entries: list[tuple[str, str | None]] = [] + seen: set[str] = set() for raw_workspace in workspaces.split(","): raw_workspace = raw_workspace.strip() if not raw_workspace: @@ -116,32 +123,44 @@ def _parse_workspaces_option(workspaces: str) -> list[str]: workspace = normalize_workspace_url(raw_workspace) except ValueError as exc: raise RuntimeError(str(exc)) from exc - if workspace not in workspace_urls: - workspace_urls.append(workspace) - if not workspace_urls: + if workspace not in seen: + seen.add(workspace) + workspace_entries.append((workspace, None)) + if not workspace_entries: raise RuntimeError( "No workspaces provided for --workspaces. Use a comma-separated list like " "`--workspaces https://workspace.databricks.com`." ) - return workspace_urls + return workspace_entries def configure_shared_state( - workspace: str, tools: list[str] | None = None, force_login: bool = False + workspace: str, + profile: str | None = None, + tools: list[str] | None = None, + force_login: bool = False, ) -> 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). + ``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) + don't error out. If ``None``, we resolve it from the host after login. """ workspace = normalize_workspace_url(workspace) fetch_all = tools is None if force_login: - run_databricks_login(workspace) + run_databricks_login(workspace, profile) else: - ensure_databricks_auth(workspace) + ensure_databricks_auth(workspace, profile) + # After login the profile exists in ~/.databrickscfg, so a host->profile + # lookup is reliable. Persist it so subsequent CLI calls disambiguate. + if profile is None: + profile = find_profile_name_for_host(workspace) with spinner("Verifying Unity AI Gateway..."): - token = get_databricks_token(workspace) + token = get_databricks_token(workspace, profile) ensure_ai_gateway_v2(workspace, token) print_success("Unity AI Gateway detected") @@ -176,6 +195,10 @@ def configure_shared_state( # Merge into existing workspace state so prior tool configs are preserved. state = load_state() state["workspace"] = workspace + if profile: + state["profile"] = profile + else: + state.pop("profile", None) state["base_urls"] = build_shared_base_urls(workspace) if want_claude: state["claude_models"] = claude_models @@ -197,28 +220,33 @@ def configure_shared_state( def _configure_shared_workspace_states( - workspaces: list[str], tools: list[str] | None, *, force_login: bool + workspaces: list[tuple[str, str | None]], + tools: list[str] | None, + *, + force_login: bool, ) -> list[dict]: if not workspaces: raise RuntimeError("At least one workspace must be provided.") states: list[dict] = [] - for workspace in workspaces: - states.append(configure_shared_state(workspace, tools=tools, force_login=force_login)) + for workspace, profile in workspaces: + states.append( + configure_shared_state(workspace, profile=profile, tools=tools, force_login=force_login) + ) return states def configure_workspace_command( tool: str | None = None, selected_tools: list[str] | None = None, - workspaces: list[str] | None = None, + workspaces: list[tuple[str, str | None]] | None = None, ) -> int: if tool is not None and selected_tools is not None: raise RuntimeError("Use either --agent or --agents, not both.") - workspace_urls = workspaces or [_prompt_for_configuration(tool)] + workspace_entries = workspaces or [_prompt_for_configuration(tool)] if tool is not None: - states = _configure_shared_workspace_states(workspace_urls, [tool], force_login=True) + states = _configure_shared_workspace_states(workspace_entries, [tool], force_login=True) state = states[0] state = configure_single_tool(tool, state) spec = TOOL_SPECS[tool] @@ -245,7 +273,7 @@ def configure_workspace_command( raise RuntimeError(f"{spec['display']} validation failed — config reverted.") return 0 - states = _configure_shared_workspace_states(workspace_urls, selected_tools, force_login=True) + states = _configure_shared_workspace_states(workspace_entries, selected_tools, force_login=True) state = states[0] save_state(state) @@ -318,6 +346,9 @@ def status() -> int: print_heading("Provider") print_kv("Workspace URL", workspace or "not configured") + profile = state.get("profile") + if profile: + print_kv("CLI profile", profile) print_heading("Coding Agents") for tool, spec in TOOL_SPECS.items(): @@ -409,9 +440,10 @@ def _auto_configure_tool(tool: str) -> None: """First-time setup for a single tool — mirrors configure_workspace_command.""" existing = load_state() workspace = existing.get("workspace") + profile = existing.get("profile") if not workspace: - workspace = _prompt_for_configuration(tool) - state = configure_shared_state(workspace, tools=[tool]) + workspace, profile = _prompt_for_configuration(tool) + state = configure_shared_state(workspace, profile=profile, tools=[tool]) state = configure_single_tool(tool, state) @@ -455,7 +487,9 @@ def _launch_tool(tool_name: str, ctx: typer.Context) -> None: # endpoints show up without a manual `ucode configure` (and so that # tools like pi which read multiple model bundles never run on # stale state from before a tool added a new bundle). - state = configure_shared_state(state["workspace"], tools=[tool]) + state = configure_shared_state( + state["workspace"], profile=state.get("profile"), tools=[tool] + ) state, resolved_model = resolve_launch_model(tool, state, None) state = configure_tool(tool, state, resolved_model) print_section(f"ucode with {TOOL_SPECS[tool]['display']}") @@ -550,29 +584,29 @@ def configure( install_databricks_cli() if agent is not None and agents is not None: raise RuntimeError("Use either --agent or --agents, not both.") - workspace_urls = _parse_workspaces_option(workspaces) if workspaces is not None else None + workspace_entries = _parse_workspaces_option(workspaces) if workspaces is not None else None if agent is not None: tool = normalize_tool(agent) install_tool_binary(tool, strict=True, update_existing=True) - if workspace_urls is None: + if workspace_entries is None: configure_workspace_command(tool) else: - configure_workspace_command(tool, workspaces=workspace_urls) + configure_workspace_command(tool, workspaces=workspace_entries) elif agents is not None: selected_tools = _parse_agents_option(agents) - if workspace_urls is None: + if workspace_entries is None: configure_workspace_command(selected_tools=selected_tools) else: configure_workspace_command( - selected_tools=selected_tools, workspaces=workspace_urls + selected_tools=selected_tools, workspaces=workspace_entries ) else: # Tool binaries are installed after the user picks which agents # they want, in configure_workspace_command. - if workspace_urls is None: + if workspace_entries is None: configure_workspace_command() else: - configure_workspace_command(workspaces=workspace_urls) + configure_workspace_command(workspaces=workspace_entries) 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 35c7f31..f3f7205 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -356,7 +356,16 @@ def install_databricks_cli() -> None: ensure_databricks_cli_version() -def has_valid_databricks_auth(workspace: str) -> bool: +def _profile_args(profile: str | None) -> list[str]: + """Return ``["--profile", profile]`` when set, otherwise an empty list. + + Centralizing this keeps every `databricks` CLI invocation in this module + consistent when a workspace's `~/.databrickscfg` has more than one profile + pointing at the same host.""" + return ["--profile", profile] if profile else [] + + +def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> bool: # Honor the CI short-circuit (see ``get_databricks_token``): if a # pre-fetched bearer is available, treat auth as valid and skip the # `databricks auth token` shell-out (which only knows user-OAuth). @@ -366,7 +375,16 @@ def has_valid_databricks_auth(workspace: str) -> bool: try: env = build_databricks_cli_env(workspace) result = run( - ["databricks", "auth", "token", "--host", workspace, "--output", "json"], + [ + "databricks", + "auth", + "token", + "--host", + workspace, + *_profile_args(profile), + "--output", + "json", + ], check=False, capture_output=True, text=True, @@ -421,16 +439,25 @@ def find_profile_name_for_host(workspace: str) -> str | None: return None -def run_databricks_login(workspace: str) -> None: - """Run databricks auth login unconditionally.""" +def run_databricks_login(workspace: str, profile: str | None = None) -> None: + """Run databricks auth login unconditionally. + + When ``profile`` is provided, it is passed via ``--profile``. Otherwise we + fall back to looking up an existing profile by host so a stored session is + refreshed in place rather than overwriting another profile's tokens.""" print_section("Databricks Login") print_kv("Workspace", workspace) print_note("A browser may open for `databricks auth login`.") try: - cmd = ["databricks", "auth", "login", "--host", workspace] - profile_name = find_profile_name_for_host(workspace) - if profile_name: - cmd += ["--profile", profile_name] + profile_name = profile or find_profile_name_for_host(workspace) + cmd = [ + "databricks", + "auth", + "login", + "--host", + workspace, + *_profile_args(profile_name), + ] run(cmd, env=build_databricks_cli_env(workspace), timeout=300) except subprocess.CalledProcessError as exc: raise RuntimeError("`databricks auth login` failed.") from exc @@ -439,17 +466,22 @@ def run_databricks_login(workspace: str) -> None: print_success("Databricks authentication complete") -def ensure_databricks_auth(workspace: str) -> None: +def ensure_databricks_auth(workspace: str, profile: str | None = None) -> None: """Check auth and login only if needed (used by launch path).""" with spinner("Checking Databricks auth..."): - auth_is_valid = has_valid_databricks_auth(workspace) + auth_is_valid = has_valid_databricks_auth(workspace, profile) if auth_is_valid: print_success(f"Databricks auth already available for {workspace}") return - run_databricks_login(workspace) + run_databricks_login(workspace, profile) -def get_databricks_token(workspace: str, *, force_refresh: bool = False) -> str: +def get_databricks_token( + workspace: str, + profile: str | None = None, + *, + force_refresh: bool = False, +) -> str: # ``DATABRICKS_BEARER`` is the CI escape hatch: when set, skip the # `databricks auth token` subprocess entirely and return the pre-fetched # bearer directly. Used by the e2e job, where the protected runner has @@ -463,7 +495,16 @@ def get_databricks_token(workspace: str, *, force_refresh: bool = False) -> str: _log_auth_diagnostics() env = build_databricks_cli_env(workspace) - cmd = ["databricks", "auth", "token", "--host", workspace, "--output", "json"] + cmd = [ + "databricks", + "auth", + "token", + "--host", + workspace, + *_profile_args(profile), + "--output", + "json", + ] if force_refresh: cmd.append("--force-refresh") @@ -472,7 +513,8 @@ def get_databricks_token(workspace: str, *, force_refresh: bool = False) -> str: "set=" + ",".join( sorted(k for k in env if k.startswith("DATABRICKS_") or k in {"BUNDLE_PROFILE"}) - ), + ) + + f" profile={profile or ''}", ) def _fetch() -> str: @@ -498,7 +540,15 @@ def _fetch() -> str: _debug("auth token", "empty on first fetch; attempting auth login --no-browser") try: reauth = run( - ["databricks", "auth", "login", "--host", workspace, "--no-browser"], + [ + "databricks", + "auth", + "login", + "--host", + workspace, + *_profile_args(profile), + "--no-browser", + ], capture_output=True, text=True, env=env, @@ -510,7 +560,7 @@ def _fetch() -> str: token = _fetch() if not token: - profile_name = find_profile_name_for_host(workspace) + profile_name = profile or find_profile_name_for_host(workspace) stale_profile_hint = "" if profile_name: stale_profile_hint = ( @@ -692,13 +742,21 @@ def list_databricks_apps(workspace: str) -> list[dict]: raise RuntimeError("Databricks apps listing returned invalid JSON.") from exc -def build_auth_shell_command(workspace: str) -> str: +def build_auth_shell_command(workspace: str, profile: str | None = None) -> str: workspace_arg = shlex.quote(workspace.rstrip("/")) - cli_command = ( - "env -u DATABRICKS_CONFIG_PROFILE " - f"databricks auth token --host {workspace_arg} --force-refresh --output json " - "| jq -r '.access_token'" - ) + if profile: + profile_arg = shlex.quote(profile) + cli_command = ( + f"databricks auth token --host {workspace_arg} " + f"--profile {profile_arg} --force-refresh --output json " + "| jq -r '.access_token'" + ) + else: + cli_command = ( + "env -u DATABRICKS_CONFIG_PROFILE " + f"databricks auth token --host {workspace_arg} --force-refresh --output json " + "| jq -r '.access_token'" + ) return ( 'if [ -n "${DATABRICKS_BEARER:-}" ]; then ' 'printf "%s\\n" "$DATABRICKS_BEARER"; ' diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index a5c1552..cb8a557 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -766,7 +766,7 @@ def configure_mcp_command() -> int: client for client in MCP_CLIENTS if client in configured_tools and client not in clients ] - ensure_databricks_auth(workspace) + ensure_databricks_auth(workspace, state.get("profile")) print_section("MCP Servers") client_names = ", ".join(str(MCP_CLIENTS[client]["display"]) for client in clients) diff --git a/src/ucode/mcp_web_search.py b/src/ucode/mcp_web_search.py index 7d08554..4349f33 100644 --- a/src/ucode/mcp_web_search.py +++ b/src/ucode/mcp_web_search.py @@ -89,13 +89,14 @@ def _call_responses_api(query: str) -> dict[str, Any]: suitable for surfacing as a tool error.""" workspace = os.environ.get("DATABRICKS_HOST", "").strip() model = os.environ.get("UCODE_WEB_SEARCH_MODEL", "").strip() + profile = os.environ.get("DATABRICKS_CONFIG_PROFILE", "").strip() or None if not workspace: raise RuntimeError("DATABRICKS_HOST env var is not set.") if not model: raise RuntimeError("UCODE_WEB_SEARCH_MODEL env var is not set.") try: - token = get_databricks_token(workspace) + token = get_databricks_token(workspace, profile) except RuntimeError as exc: raise RuntimeError(f"Failed to acquire Databricks token: {exc}") from exc diff --git a/src/ucode/state.py b/src/ucode/state.py index 5684eb7..239c859 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -102,9 +102,10 @@ def build_agent_state(state: dict) -> dict[str, dict]: if not isinstance(workspace, str) or not workspace: return {} + 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) + auth_command = build_auth_shell_command(workspace, profile) 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/ui.py b/src/ucode/ui.py index f54352e..249082c 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -177,18 +177,22 @@ def normalize_workspace_url(workspace: str) -> str: def prompt_for_workspace( description: str, profiles: list[tuple[str, str]] | None = None, -) -> str: +) -> tuple[str, str | None]: """Ask the user for a workspace URL, offering profiles as quick-select. `profiles` is a list of (host_url, profile_name) tuples. Caller fetches - them — `ui.py` stays Databricks-agnostic. Returns a normalized URL. + them — `ui.py` stays Databricks-agnostic. Returns ``(url, profile_name)``; + profile_name is ``None`` when the user typed a URL manually. """ console.print() console.print(Panel(description, title="ucode setup", style="bold blue", expand=False)) if profiles: - choices = [questionary.Choice(title=host, value=host) for host, _ in profiles] - choices.append(questionary.Choice(title="Enter a different URL", value="__manual__")) + choices = [ + questionary.Choice(title=host, value=(host, profile_name)) + for host, profile_name in profiles + ] + choices.append(questionary.Choice(title="Enter a different URL", value=None)) style = questionary.Style( [ ("highlighted", "fg:cyan bold"), @@ -199,13 +203,14 @@ def prompt_for_workspace( choice = questionary.select( "Select workspace:", choices=choices, style=style, pointer="›", qmark="" ).ask() - if choice is not None and choice != "__manual__": - return normalize_workspace_url(choice) + if choice is not None: + host, profile_name = choice + return normalize_workspace_url(host), profile_name while True: raw_value = console.input(f" [bold]Workspace URL[/bold] {muted('›')} ").strip() try: - return normalize_workspace_url(raw_value) + return normalize_workspace_url(raw_value), None except ValueError as exc: print_err(str(exc)) diff --git a/src/ucode/usage.py b/src/ucode/usage.py index 9020915..e20c3dd 100644 --- a/src/ucode/usage.py +++ b/src/ucode/usage.py @@ -272,9 +272,10 @@ def usage() -> int: if not workspace: raise RuntimeError("Workspace is not configured. Run `ucode configure` first.") - ensure_databricks_auth(workspace) + profile = state.get("profile") + ensure_databricks_auth(workspace, profile) with spinner("Retrieving Databricks access token..."): - token = get_databricks_token(workspace) + token = get_databricks_token(workspace, profile) with spinner("Discovering SQL warehouse..."): resolved_http_path = discover_sql_warehouse_http_path(workspace, token, quiet=False) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 95f1ec8..07d337b 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -186,7 +186,7 @@ def _common_patches(self, monkeypatch, calls): monkeypatch.setattr( claude, "_register_web_search_mcp", - lambda ws, model: calls.append(("register", ws, model)), + lambda ws, model, profile=None: calls.append(("register", ws, model)), ) def test_registers_mcp_when_codex_model_available(self, monkeypatch): @@ -262,7 +262,9 @@ def fake_execvp(binary: str, args: list[str]) -> None: raise RuntimeError("stop") monkeypatch.delenv("OAUTH_TOKEN", raising=False) - monkeypatch.setattr(claude, "get_databricks_token", lambda workspace: "fresh-token") + monkeypatch.setattr( + claude, "get_databricks_token", lambda workspace, profile=None: "fresh-token" + ) monkeypatch.setattr(os, "execvp", fake_execvp) try: diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 32bcae2..f9cbc3f 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -115,7 +115,9 @@ def fake_execvp(binary: str, args: list[str]) -> None: raise RuntimeError("stop") monkeypatch.delenv("OAUTH_TOKEN", raising=False) - monkeypatch.setattr(codex, "get_databricks_token", lambda workspace: "fresh-token") + monkeypatch.setattr( + codex, "get_databricks_token", lambda workspace, profile=None: "fresh-token" + ) monkeypatch.setattr(os, "execvp", fake_execvp) try: diff --git a/tests/test_agent_gemini.py b/tests/test_agent_gemini.py index 07e6432..6b2d321 100644 --- a/tests/test_agent_gemini.py +++ b/tests/test_agent_gemini.py @@ -121,7 +121,7 @@ def test_writes_ucode_env_file(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr("ucode.agents.gemini.save_state", lambda s: None) monkeypatch.setattr( - "ucode.agents.gemini.get_databricks_token", lambda ws, **kwargs: "fake-token" + "ucode.agents.gemini.get_databricks_token", lambda ws, profile=None, **kwargs: "fake-token" ) gemini.write_tool_config({"workspace": WS}, "some-model") @@ -140,7 +140,7 @@ def test_does_not_write_settings_json(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr("ucode.agents.gemini.save_state", lambda s: None) monkeypatch.setattr( - "ucode.agents.gemini.get_databricks_token", lambda ws, **kwargs: "fake-token" + "ucode.agents.gemini.get_databricks_token", lambda ws, profile=None, **kwargs: "fake-token" ) gemini.write_tool_config({"workspace": WS}, "some-model") diff --git a/tests/test_cli.py b/tests/test_cli.py index 0664921..9dc50a8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -420,8 +420,8 @@ def test_workspaces_flag_calls_configure_with_workspaces(self): assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with( workspaces=[ - "https://first.databricks.com", - "https://second.databricks.com", + ("https://first.databricks.com", None), + ("https://second.databricks.com", None), ] ) @@ -437,7 +437,7 @@ def test_agents_and_workspaces_flags_call_configure_with_both(self): ) assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with( - selected_tools=["claude", "codex"], workspaces=["https://first.com"] + selected_tools=["claude", "codex"], workspaces=[("https://first.com", None)] ) def test_agent_and_workspaces_flags_call_configure_with_both(self): @@ -452,7 +452,7 @@ def test_agent_and_workspaces_flags_call_configure_with_both(self): ) assert result.exit_code == 0, result.output mock_install.assert_called_once_with("claude", strict=True, update_existing=True) - mock_cfg.assert_called_once_with("claude", workspaces=["https://first.com"]) + mock_cfg.assert_called_once_with("claude", workspaces=[("https://first.com", None)]) def test_agent_flag_calls_configure_with_tool(self): with ( @@ -550,7 +550,9 @@ def test_selected_tools_skip_picker(self, monkeypatch): state = {**MINIMAL_STATE, "available_tools": []} monkeypatch.setattr( - cli_mod, "_prompt_for_configuration", lambda tool=None: "https://example.com" + cli_mod, + "_prompt_for_configuration", + lambda tool=None: ("https://example.com", None), ) monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *args, **kwargs: state) monkeypatch.setattr( @@ -584,7 +586,9 @@ def test_unavailable_selected_tool_errors_before_configure(self, monkeypatch): state = {**MINIMAL_STATE, "available_tools": []} monkeypatch.setattr( - cli_mod, "_prompt_for_configuration", lambda tool=None: "https://example.com" + cli_mod, + "_prompt_for_configuration", + lambda tool=None: ("https://example.com", None), ) monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *args, **kwargs: state) monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: tool == "claude") @@ -605,11 +609,13 @@ def test_multiple_workspaces_configure_all_and_use_first(self, monkeypatch): "https://first.com": {**MINIMAL_STATE, "workspace": "https://first.com"}, "https://second.com": {**MINIMAL_STATE, "workspace": "https://second.com"}, } - configured_shared: list[tuple[str, tuple[str, ...] | None, bool]] = [] + configured_shared: list[tuple[str, str | None, tuple[str, ...] | None, bool]] = [] - def fake_configure_shared_state(workspace, tools=None, force_login=False): + def fake_configure_shared_state( + workspace, profile=None, tools=None, force_login=False + ): configured_shared.append( - (workspace, tuple(tools) if tools is not None else None, force_login) + (workspace, profile, tuple(tools) if tools is not None else None, force_login) ) return states[workspace] @@ -632,13 +638,13 @@ def fake_configure_shared_state(workspace, tools=None, force_login=False): assert ( cli_mod.configure_workspace_command( - workspaces=["https://first.com", "https://second.com"] + workspaces=[("https://first.com", None), ("https://second.com", None)] ) == 0 ) assert configured_shared == [ - ("https://first.com", None, True), - ("https://second.com", None, True), + ("https://first.com", None, None, True), + ("https://second.com", None, None, True), ] assert saved == ["https://first.com"] assert configured_tools == [("https://first.com", ["codex"])] diff --git a/tests/test_databricks.py b/tests/test_databricks.py index dfcf17c..16cd749 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -148,6 +148,20 @@ def test_prefers_databricks_bearer(self, tmp_path): ) assert result.stdout.strip() == "bearer-token" + def test_embeds_profile_when_provided(self): + cmd = build_auth_shell_command(WS, profile="stablebox") + assert "--profile stablebox" in cmd + # We do not strip DATABRICKS_CONFIG_PROFILE when we are explicit about + # which profile to use — the --profile flag wins. + assert "env -u DATABRICKS_CONFIG_PROFILE" not in cmd + + def test_quotes_profile_shell_metacharacters(self): + cmd = build_auth_shell_command(WS, profile="weird name; rm -rf /") + # shlex.quote should wrap the value so the rest of the command cannot + # be interpreted as a shell injection. + assert "rm -rf /" in cmd + assert "'weird name; rm -rf /'" in cmd + class TestFormatSubprocessResult: def test_suppresses_stdout_on_success(self): @@ -290,6 +304,22 @@ def test_raises_when_reauth_also_fails(self, tmp_path, monkeypatch): with pytest.raises(RuntimeError, match="no access token"): get_databricks_token(WS) + def test_passes_profile_flag_when_provided(self, tmp_path, monkeypatch): + # Fake CLI that records its argv to a file so we can assert the + # --profile flag is forwarded to `databricks auth token`. + 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, profile="stablebox") + assert token == "good-token" + argv = argv_log.read_text().splitlines() + assert "--profile" in argv + assert argv[argv.index("--profile") + 1] == "stablebox" + def test_error_suggests_logout_when_matching_profile_exists(self, tmp_path, monkeypatch): env = self._fake_databricks( tmp_path, diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 53ca465..f68457c 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -213,7 +213,7 @@ def test_only_picks_codex_writes_only_codex_config(self, tmp_path, monkeypatch, monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") # Don't actually run `databricks auth login`; the developer running # this suite is already authenticated. - monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws: None) + monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws, profile=None: None) # Skip the workspace prompt and the multi-select picker. monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: ["codex"]) @@ -248,7 +248,7 @@ def test_rerun_with_different_pick_preserves_previous( self._redirect_config_paths(monkeypatch, tmp_path) monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") - monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws: None) + monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws, profile=None: None) monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) monkeypatch.setattr( cli_mod, "install_tool_binary", lambda tool, strict=False, update_existing=False: True @@ -281,7 +281,7 @@ def test_empty_pick_returns_zero_and_writes_nothing(self, tmp_path, monkeypatch, codex_path = self._redirect_config_paths(monkeypatch, tmp_path) monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") - monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws: None) + monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws, profile=None: None) monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: []) install_calls: list[str] = [] @@ -447,7 +447,7 @@ def test_launch_gemini_per_model( mp.setattr("ucode.state.save_state", lambda s: None) mp.setattr( "ucode.agents.gemini.get_databricks_token", - lambda ws, **kwargs: e2e_token, + lambda ws, profile=None, **kwargs: e2e_token, ) state = {**e2e_state, "workspace": e2e_workspace} gemini.write_tool_config(state, model, token=e2e_token) @@ -530,7 +530,7 @@ def test_launch_opencode_per_model( mp.setattr("ucode.state.save_state", lambda s: None) mp.setattr( "ucode.agents.opencode.get_databricks_token", - lambda ws, **kwargs: e2e_token, + lambda ws, profile=None, **kwargs: e2e_token, ) opencode.write_tool_config( {**e2e_state, "workspace": e2e_workspace}, @@ -605,7 +605,7 @@ def test_launch_copilot_per_model( mp.setattr("ucode.state.save_state", lambda s: None) mp.setattr( "ucode.agents.copilot.get_databricks_token", - lambda ws: e2e_token, + lambda ws, profile=None, **kwargs: e2e_token, ) copilot.write_tool_config( {**e2e_state, "workspace": e2e_workspace}, model, token=e2e_token @@ -672,7 +672,7 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa mp.setattr("ucode.state.save_state", lambda s: None) mp.setattr( "ucode.agents.pi.get_databricks_token", - lambda ws, **kwargs: e2e_token, + lambda ws, profile=None, **kwargs: e2e_token, ) pi.write_tool_config( {**e2e_state, "workspace": e2e_workspace}, diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index 0b5d33c..1bf359f 100644 --- a/tests/test_e2e_user_agent.py +++ b/tests/test_e2e_user_agent.py @@ -267,7 +267,7 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) mp.setattr( - "ucode.agents.opencode.get_databricks_token", lambda ws, **kwargs: "test-token" + "ucode.agents.opencode.get_databricks_token", lambda ws, profile=None, **kwargs: "test-token" ) opencode.write_tool_config(state, "test-claude-model", token="test-token") @@ -338,7 +338,7 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv } with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr("ucode.agents.pi.get_databricks_token", lambda ws, **kwargs: "test-token") + mp.setattr("ucode.agents.pi.get_databricks_token", lambda ws, profile=None, **kwargs: "test-token") pi.write_tool_config(state, "test-claude-model", token="test-token") env = pi.build_runtime_env("test-token") diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 3075704..445018c 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -404,7 +404,7 @@ def test_skips_existing_server_state_by_name(self, monkeypatch): }, ) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) monkeypatch.setattr(mcp, "discover_external_mcp_connection_names", lambda workspace: []) monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace: []) @@ -428,7 +428,7 @@ def test_registers_discovered_external_server(self, monkeypatch): lambda: {"workspace": WS, "available_tools": ALL_MCP_CLIENTS}, ) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr( mcp, "available_mcp_clients", @@ -484,7 +484,7 @@ def test_registers_discovered_genie_space_server(self, monkeypatch): monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) monkeypatch.setattr(mcp, "discover_external_mcp_connection_names", lambda workspace: []) monkeypatch.setattr( @@ -536,7 +536,7 @@ def test_registers_discovered_app_mcp_server(self, monkeypatch): monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) monkeypatch.setattr(mcp, "discover_external_mcp_connection_names", lambda workspace: []) monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace: []) @@ -588,7 +588,7 @@ def test_continues_when_optional_discovery_fails(self, monkeypatch, capsys): monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) monkeypatch.setattr( mcp, @@ -631,7 +631,7 @@ def test_configures_only_ucode_configured_clients(self, monkeypatch, capsys): lambda: {"workspace": WS, "available_tools": ["claude", "codex"]}, ) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ALL_MCP_CLIENTS) monkeypatch.setattr(mcp, "discover_external_mcp_connection_names", lambda workspace: []) monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace: []) @@ -663,7 +663,7 @@ def test_registers_databricks_sql_server(self, monkeypatch): configured: list[tuple[str, str, str, dict]] = [] monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) monkeypatch.setattr(mcp, "discover_external_mcp_connection_names", lambda workspace: []) monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace: []) @@ -717,7 +717,7 @@ def test_removes_saved_server(self, monkeypatch): monkeypatch.setattr(mcp, "load_state", lambda: state) monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") - monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace: None) + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) monkeypatch.setattr(mcp, "discover_external_mcp_connection_names", lambda workspace: []) monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace: []) diff --git a/tests/test_mcp_web_search.py b/tests/test_mcp_web_search.py index 2f5cb89..e03889d 100644 --- a/tests/test_mcp_web_search.py +++ b/tests/test_mcp_web_search.py @@ -206,7 +206,7 @@ def test_missing_model_env(self, monkeypatch): def test_posts_to_responses_endpoint(self, monkeypatch): monkeypatch.setenv("DATABRICKS_HOST", WS) monkeypatch.setenv("UCODE_WEB_SEARCH_MODEL", "databricks-gpt-5") - monkeypatch.setattr(mcp_web_search, "get_databricks_token", lambda ws: "tok") + monkeypatch.setattr(mcp_web_search, "get_databricks_token", lambda ws, profile=None: "tok") captured: dict[str, Any] = {} From ec1cfa0916fa0aad62cee380260b3139c480ebad Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Tue, 26 May 2026 10:10:20 -0400 Subject: [PATCH 2/4] Apply ruff format and fix e2e mock for tuple workspace return Format-only changes in databricks.py, test_agent_gemini.py, test_cli.py, and test_e2e_user_agent.py. In test_e2e.py, update the `_prompt_for_configuration` mock to return `(workspace, None)` so the new tuple contract from this branch is honored. Co-Authored-By: Claude Opus 4.7 --- src/ucode/databricks.py | 4 +--- tests/test_agent_gemini.py | 6 ++++-- tests/test_cli.py | 4 +--- tests/test_e2e.py | 12 +++++++++--- tests/test_e2e_user_agent.py | 8 ++++++-- 5 files changed, 21 insertions(+), 13 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f3f7205..928e0d8 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -511,9 +511,7 @@ def get_databricks_token( _debug( "get_databricks_token.env", "set=" - + ",".join( - sorted(k for k in env if k.startswith("DATABRICKS_") or k in {"BUNDLE_PROFILE"}) - ) + + ",".join(sorted(k for k in env if k.startswith("DATABRICKS_") or k in {"BUNDLE_PROFILE"})) + f" profile={profile or ''}", ) diff --git a/tests/test_agent_gemini.py b/tests/test_agent_gemini.py index 6b2d321..fef70f5 100644 --- a/tests/test_agent_gemini.py +++ b/tests/test_agent_gemini.py @@ -121,7 +121,8 @@ def test_writes_ucode_env_file(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr("ucode.agents.gemini.save_state", lambda s: None) monkeypatch.setattr( - "ucode.agents.gemini.get_databricks_token", lambda ws, profile=None, **kwargs: "fake-token" + "ucode.agents.gemini.get_databricks_token", + lambda ws, profile=None, **kwargs: "fake-token", ) gemini.write_tool_config({"workspace": WS}, "some-model") @@ -140,7 +141,8 @@ def test_does_not_write_settings_json(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr("ucode.agents.gemini.save_state", lambda s: None) monkeypatch.setattr( - "ucode.agents.gemini.get_databricks_token", lambda ws, profile=None, **kwargs: "fake-token" + "ucode.agents.gemini.get_databricks_token", + lambda ws, profile=None, **kwargs: "fake-token", ) gemini.write_tool_config({"workspace": WS}, "some-model") diff --git a/tests/test_cli.py b/tests/test_cli.py index 9dc50a8..6b335a8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -611,9 +611,7 @@ def test_multiple_workspaces_configure_all_and_use_first(self, monkeypatch): } configured_shared: list[tuple[str, str | None, tuple[str, ...] | None, bool]] = [] - def fake_configure_shared_state( - workspace, profile=None, tools=None, force_login=False - ): + def fake_configure_shared_state(workspace, profile=None, tools=None, force_login=False): configured_shared.append( (workspace, profile, tuple(tools) if tools is not None else None, force_login) ) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index f68457c..8a10182 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -215,7 +215,9 @@ def test_only_picks_codex_writes_only_codex_config(self, tmp_path, monkeypatch, # this suite is already authenticated. monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws, profile=None: None) # Skip the workspace prompt and the multi-select picker. - monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) + monkeypatch.setattr( + cli_mod, "_prompt_for_configuration", lambda tool=None: (e2e_workspace, None) + ) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: ["codex"]) # Skip binary install + post-config validation; we're testing the # selection plumbing, not the agent binaries themselves. @@ -249,7 +251,9 @@ def test_rerun_with_different_pick_preserves_previous( self._redirect_config_paths(monkeypatch, tmp_path) monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws, profile=None: None) - monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) + monkeypatch.setattr( + cli_mod, "_prompt_for_configuration", lambda tool=None: (e2e_workspace, None) + ) monkeypatch.setattr( cli_mod, "install_tool_binary", lambda tool, strict=False, update_existing=False: True ) @@ -282,7 +286,9 @@ def test_empty_pick_returns_zero_and_writes_nothing(self, tmp_path, monkeypatch, codex_path = self._redirect_config_paths(monkeypatch, tmp_path) monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") monkeypatch.setattr("ucode.cli.run_databricks_login", lambda ws, profile=None: None) - monkeypatch.setattr(cli_mod, "_prompt_for_configuration", lambda tool=None: e2e_workspace) + monkeypatch.setattr( + cli_mod, "_prompt_for_configuration", lambda tool=None: (e2e_workspace, None) + ) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: []) install_calls: list[str] = [] monkeypatch.setattr( diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index 1bf359f..a74fb9d 100644 --- a/tests/test_e2e_user_agent.py +++ b/tests/test_e2e_user_agent.py @@ -267,7 +267,8 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) mp.setattr( - "ucode.agents.opencode.get_databricks_token", lambda ws, profile=None, **kwargs: "test-token" + "ucode.agents.opencode.get_databricks_token", + lambda ws, profile=None, **kwargs: "test-token", ) opencode.write_tool_config(state, "test-claude-model", token="test-token") @@ -338,7 +339,10 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv } with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr("ucode.agents.pi.get_databricks_token", lambda ws, profile=None, **kwargs: "test-token") + mp.setattr( + "ucode.agents.pi.get_databricks_token", + lambda ws, profile=None, **kwargs: "test-token", + ) pi.write_tool_config(state, "test-claude-model", token="test-token") env = pi.build_runtime_env("test-token") From 2a60c5194ee716be95c138b637031e585bc84139 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Tue, 26 May 2026 10:30:09 -0400 Subject: [PATCH 3/4] Log why get_databricks_profiles returns empty The configure-time workspace picker silently degrades to a blank URL prompt whenever this returns []. Log each dropout path (subprocess error, non-zero exit, JSON decode, all-PAT) under UCODE_DEBUG so the case can be diagnosed from ~/.ucode/debug.log instead of guessing. Also bump the timeout from 10s to 20s and simplify the dedupe loop. Co-Authored-By: Claude Opus 4.7 --- src/ucode/databricks.py | 52 +++++++++++++++++++++++++++++------------ 1 file changed, 37 insertions(+), 15 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 928e0d8..0d021e0 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -405,29 +405,51 @@ def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> boo def get_databricks_profiles() -> list[tuple[str, str]]: - """Return [(host_url, profile_name), ...] from Databricks CLI profiles.""" + """Return [(host_url, profile_name), ...] from Databricks CLI profiles. + + Returns ``[]`` on any failure (CLI missing, timeout, non-zero exit, JSON + decode error). When ``UCODE_DEBUG=1`` each dropout path logs *why* the + result was empty so a silently-disappearing workspace picker is + diagnosable from ``~/.ucode/debug.log``. + """ try: result = run( ["databricks", "auth", "profiles", "--output", "json"], check=False, capture_output=True, text=True, - timeout=10, + timeout=20, ) - if result.returncode != 0: - return [] - data = json.loads(result.stdout or "{}") - profiles = data.get("profiles") or [] - seen: set[str] = set() - out: list[tuple[str, str]] = [] - for p in profiles: - host = p.get("host", "").rstrip("/") - if host and host not in seen and p.get("auth_type") != "pat": - seen.add(host) - out.append((host, p["name"])) - return out - except (json.JSONDecodeError, OSError, subprocess.TimeoutExpired, KeyError): + except (OSError, subprocess.TimeoutExpired) as exc: + _debug("get_databricks_profiles", f"subprocess error: {type(exc).__name__}: {exc}") + return [] + if result.returncode != 0: + _debug("get_databricks_profiles", _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}") + return [] + + # dict dedupes by host (first non-PAT profile wins). + out: dict[str, str] = {} + pat = 0 + for p in profiles: + host = (p.get("host") or "").rstrip("/") + name = p.get("name") + if not host or not name: + continue + if p.get("auth_type") == "pat": + pat += 1 + continue + out.setdefault(host, name) + + _debug( + "get_databricks_profiles", + f"returned={len(out)} total={len(profiles)} pat={pat}", + ) + return list(out.items()) def find_profile_name_for_host(workspace: str) -> str | None: From a33f17ad77a1eb91aaeb1e7e1af3e9c278fa38ce Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Tue, 26 May 2026 10:47:18 -0400 Subject: [PATCH 4/4] add spinner --- src/ucode/cli.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 4fd20e1..1530609 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -87,7 +87,8 @@ def _prompt_for_configuration(tool: str | None = None) -> tuple[str, str | None] desc = "Configure your Databricks workspace" else: desc = f"Configure {TOOL_SPECS[tool]['display']} to use your Databricks endpoint." - profiles = get_databricks_profiles() + with spinner("Loading Databricks workspaces and profiles..."): + profiles = get_databricks_profiles() return prompt_for_workspace(desc, profiles)