Fix configure/launch when ~/.databrickscfg has duplicate hosts

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])
 
 

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])
 
 

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)

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)

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")))

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")))