Merge origin/main into make-mcp-heavy-sources-opt-in

Resolve conflicts:
- cli.py: keep both the new --mcp flag (from main) and the interactive
  "Configure MCP servers now?" step; both call configure_mcp_command.
- mcp.py: keep main's setup_mcp_clients refactor and skills-connection
  handling; feed skills-excluded picker_servers into the new two-step
  search wizard.
- ui.py: union the Callable/Iterator imports.
- tests: keep both new test classes; add the new search-sources and MCP
  prompt mocks to two main-origin tests so they pass under the wizard flow.

Co-authored-by: Isaac

Author: Achille Negrier

AGENTS.md

@@ -20,7 +20,7 @@ Tests live in `tests/`.
 
 - Use Python 3.12+.
 - Keep changes scoped to the requested behavior.
-- Follow the existing module boundaries: CLI orchestration in `cli.py`, agent-specific behavior in `agents/<name>.py`, shared agent dispatch in `agents/__init__.py`, Databricks calls in `databricks.py`, and presentation helpers in `ui.py`.
+- Follow the existing module boundaries: CLI orchestration in `cli.py`, agent-specific behavior in `agents/<name>.py`, shared agent dispatch in `agents/__init__.py`, Databricks calls in `databricks.py`, skill download (UC fetch client + on-disk writer + download orchestration) in `skills_download.py`, MCP-connection state glue in `mcp.py`, and presentation helpers in `ui.py`. Skill download persists no disk state — it writes files to `--path` (or the home dir) and registers only the schema-less skills MCP connection.
 - Prefer existing helpers for config file writes, state persistence, UI messages, and Databricks authentication.
 - Add or update focused tests for behavior changes.
 - Do not modify generated or lock files unless the dependency graph intentionally changes.

README.md

@@ -92,6 +92,38 @@ Options are shown in this order:
 Discovered external MCP connections are listed directly. MCP auth uses a Databricks token that
 `ucode` sets when launching each tool.
 
+To set up an agent and its MCP server(s) in one command, pass `--mcp` with fully-qualified
+service name(s) to `ucode configure`:
+
+```bash
+ucode configure --agents claude --mcp system.ai.slack
+```
+
+`--mcp` also works without `--agents` for MCP-only clients (it configures just the workspace,
+then registers the servers); pass a comma-separated list to register several at once.
+
+### Skills (optional)
+
+Configure Unity Catalog Skills for your coding tools with `ucode configure skills`. It has two
+mutually-exclusive modes, both scoped by `--location <catalog>.<schema>` (comma-separated for
+multiple schemas):
+
+```bash
+# Download mode (default): fetch every skill in the schema to disk.
+ucode configure skills --location main.default --path /abs/project/dir
+
+# MCP mode: expose the schema's skills as MCP tools instead of downloading.
+ucode configure skills --location main.default,ml.prod --mcp
+```
+
+- **Download mode** writes each skill flat as `<leaf>/SKILL.md` (plus its bundled files) into both
+  `.claude/skills/` and `.agents/skills/`. `--path` (an existing absolute directory) is optional;
+  when omitted, skills are written under your home directory. Any pre-existing skill dir prompts
+  before it's overwritten. It then registers a schema-less skills MCP connection (utility tools
+  only), leaving any prior `--mcp` scope untouched.
+- **MCP mode** sets the connection's location set to exactly `<list>` (override-only) and rebuilds
+  its `?schema=` URL; no files are downloaded and `--path` is rejected.
+
 ---
 
 ## Other Commands
@@ -107,6 +139,9 @@ Discovered external MCP connections are listed directly. MCP auth uses a Databri
 | `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 |
+| `ucode configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command |
+| `ucode configure skills --location main.default [--path <dir>]` | Download a schema's skills to disk (under `<dir>`, or your home dir) and register a schema-less skills MCP connection |
+| `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading |
 
 ## Managed Local Files
 

src/ucode/agents/__init__.py

@@ -20,6 +20,7 @@
 from ucode.databricks import (
     BEDROCK_PROVIDER_TYPES,
     get_databricks_token,
+    install_ai_tools,
     install_databricks_cli,
     map_bedrock_claude_models,
     resolve_provider_service,
@@ -65,6 +66,23 @@
 DEFAULT_TOOL = "codex"
 BUNDLE_VERSION = 1
 
+# ucode tool -> `databricks aitools` agent id. gemini/pi aren't supported.
+AITOOLS_AGENT_TOKENS = {
+    "claude": "claude-code",
+    "codex": "codex",
+    "opencode": "opencode",
+    "copilot": "copilot",
+}
+
+
+def install_ai_tools_for_agents(tools: list[str], state: dict) -> None:
+    """Install Databricks AI Tools for the coding agents that support them
+    (gemini/pi have no ``aitools`` support and are dropped)."""
+    if state.get("databricks_ai_tools_enabled", True) is False:
+        return
+    agents = [AITOOLS_AGENT_TOKENS[tool] for tool in tools if tool in AITOOLS_AGENT_TOKENS]
+    install_ai_tools(agents, state.get("profile"))
+
 
 def normalize_tool(tool: str) -> str:
     normalized = TOOL_ALIASES.get(tool.strip().lower())
@@ -380,6 +398,7 @@ def configure_single_tool(tool: str, state: dict) -> dict:
     available_tools = list(set((state.get("available_tools") or []) + [tool]))
     state["available_tools"] = available_tools
     save_state(state)
+    install_ai_tools_for_agents([tool], state)
     return state
 
 
@@ -410,6 +429,7 @@ def configure_selected_tools(state: dict, tools: list[str]) -> dict:
     existing = state.get("available_tools") or []
     state["available_tools"] = sorted(set(existing) | set(tools))
     save_state(state)
+    install_ai_tools_for_agents(tools, state)
     return state
 
 

src/ucode/agents/claude.py

@@ -93,6 +93,9 @@ def _resolve_web_search_model(state: dict) -> str | None:
     "ANTHROPIC_DEFAULT_HAIKU_MODEL",
     "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME",
 )
+# Env keys ucode used to write but no longer does; stripped from the managed
+# settings file on every launch so stale values never linger.
+CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",)
 CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook"
 # Tracing is driven by an `mlflow autolog claude stop-hook` Stop hook, run by
 # the `mlflow` CLI on each session end. Pin to 3.11.x: 3.12 dropped the Unity
@@ -166,8 +169,13 @@ def render_overlay(
     env: dict[str, str] = {
         "ANTHROPIC_BASE_URL": base_url,
         "ANTHROPIC_CUSTOM_HEADERS": custom_headers,
-        "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
         "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": "900000",
+        # 1h prompt caching needs the extended-cache-ttl beta header, which
+        # Claude Code only sends when experimental betas are enabled — so we must
+        # not set CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS (see CLAUDE_REMOVED_ENV_KEYS).
+        "ENABLE_PROMPT_CACHING_1H": "1",
+        "ENABLE_TOOL_SEARCH": "1",
+        "CLAUDE_CODE_USE_GATEWAY": "1",
     }
     # Intentionally NOT setting ANTHROPIC_MODEL. Setting it produces a duplicate
     # catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M context) ✓")
@@ -343,6 +351,11 @@ def write_tool_config(
         for key in CLAUDE_MANAGED_MODEL_ENV_KEYS:
             if key not in overlay_env:
                 merged_env.pop(key, None)
+    # deep_merge_dict keeps keys already in the file, so drop the ones ucode no
+    # longer writes.
+    if isinstance(merged_env, dict):
+        for key in CLAUDE_REMOVED_ENV_KEYS:
+            merged_env.pop(key, None)
     write_json_file(CLAUDE_SETTINGS_PATH, merged)
 
     if web_search_model: