diff --git a/README.md b/README.md index 0420190..8fa837e 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,18 @@ 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. +### Skills (optional) + +```bash +ucode configure skills --location . +``` + +Add the Unity Catalog Skills MCP server for a `.` to every configured +MCP-capable tool. Claude Code additionally loads it eagerly (`alwaysLoad`) so the skills in that +schema are discoverable; other agents get the same endpoint without eager loading. Uses the same +Databricks token as `configure mcp`. Skills and other MCP servers coexist; re-running with a new +location repoints skills at that schema. + --- ## Other Commands diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 9ebdae9..6a04622 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -54,6 +54,7 @@ from ucode.mcp import ( MCP_CLIENTS, configure_mcp_command, + configure_skills_command, purge_cross_workspace_mcp_residue, revert_mcp_configs, ) @@ -671,7 +672,7 @@ def status() -> int: print_kv("Base URL", base_url) if configured and tool in MCP_CLIENTS: tool_mcp_servers = [ - str(server.get("name")) + str(server.get("name")) + (" (skills)" if server.get("kind") == "skills" else "") for server in mcp_servers if tool in (server.get("clients") or []) and server.get("name") ] @@ -707,6 +708,9 @@ def status() -> int: print_note( "Use `ucode configure mcp` to add Databricks MCP servers to configured coding tools." ) + print_note( + "Use `ucode configure skills --location .` to add the UC Skills MCP." + ) print_note("Use `ucode configure tracing` to log coding sessions to an MLflow experiment.") print_note("Use `ucode revert` to clear managed configs and restore prior files.") return 0 @@ -1222,6 +1226,29 @@ def configure_mcp( raise typer.Exit(130) from None +@configure_app.command("skills") +def configure_skills( + location: Annotated[ + str, + typer.Option( + "--location", + help="Unity Catalog `.` whose UC Skills to expose (e.g. " + "`main.default`). The skills MCP is added to every configured agent; Claude " + "additionally loads it eagerly so skills are discoverable.", + ), + ], +) -> None: + """Add the Unity Catalog Skills MCP server to installed coding tools.""" + try: + configure_skills_command(location=location) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + except KeyboardInterrupt: + print_err("Interrupted.") + raise typer.Exit(130) from None + + @configure_app.command("tracing") def configure_tracing( disable: Annotated[ diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9f89c6b..805c19c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1322,6 +1322,10 @@ def build_mcp_service_url(workspace: str, full_name: str) -> str: return f"{workspace}/ai-gateway/mcp-services/{full_name}" +def build_skills_mcp_url(workspace: str, catalog: str, schema: str) -> str: + return f"{workspace}/ai-gateway/skills/{catalog}/{schema}" + + # Maps the gateway routing dialect a coding tool speaks to the Model Provider # Service `provider_type`s it can be backed by. claude speaks Anthropic's API, # which both the `anthropic` and `amazon_bedrock` provider types serve (Bedrock diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 6257458..35fff70 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -29,6 +29,7 @@ from ucode.databricks import ( apply_pat_environment, build_mcp_service_url, + build_skills_mcp_url, ensure_databricks_auth, get_databricks_token, list_databricks_apps, @@ -49,6 +50,10 @@ ) MCP_AUTH_TOKEN_ENV_VAR = "OAUTH_TOKEN" +# `kind` tag on a saved MCP entry marking it as a UC Skills server. Skills +# entries are eagerly loaded in Claude (see `apply_mcp_server_changes`); every +# other entry is untagged and treated as a plain mcp-service. +SKILLS_MCP_KIND = "skills" MCP_USER_SCOPE = "user" MCP_CLEANUP_SCOPES = ("local", "project", MCP_USER_SCOPE) MCP_PICKER_VISIBLE_ROWS = 10 @@ -96,14 +101,20 @@ ) -def build_mcp_http_entry(url: str) -> dict: - return { +def build_mcp_http_entry(url: str, *, always_load: bool = False) -> dict: + entry: dict = { "type": "http", "url": url, "headers": { "Authorization": f"Bearer ${{{MCP_AUTH_TOKEN_ENV_VAR}}}", }, } + # `alwaysLoad` makes Claude Code eagerly load the server's tool names and + # descriptions, which skills rely on to be discoverable. It is Claude-only; + # other agents receive just the URL and never see this field. + if always_load: + entry["alwaysLoad"] = True + return entry def add_claude_mcp_server(name: str, entry: dict, scope: str = MCP_USER_SCOPE) -> None: @@ -1037,7 +1048,7 @@ def apply_mcp_server_changes( url = server.get("url") if not isinstance(url, str) or not url: continue - entry = build_mcp_http_entry(url) + entry = build_mcp_http_entry(url, always_load=server.get("kind") == SKILLS_MCP_KIND) for client in clients: configure_client_mcp_server(client, name, url, entry) changed = True @@ -1125,8 +1136,7 @@ def _resolve_location_mcp_servers( since-removed service still configures the rest. An empty set selects nothing (every previously-registered service in the location is removed). ``None`` keeps the whole schema.""" - if location.count(".") != 1 or not all(part.strip() for part in location.split(".")): - raise RuntimeError(f"--location must be `.`, got `{location}`.") + _split_catalog_schema(location, "--location") token = get_databricks_token(workspace, profile) with spinner(f"Discovering MCP services in {location}..."): @@ -1177,25 +1187,20 @@ def _resolve_location_mcp_servers( return working_servers -def configure_mcp_command(location: str | None = None, services: set[str] | None = None) -> int: - if services is not None and location is None: - # `--services` works standalone with full names (`system.ai.github`): the - # `.` to configure is derived from them. Bare short names - # (`github`) can't be located without `--location`. - schemas = {".".join(s.split(".")[:2]) for s in services if s.count(".") >= 2} - bare = sorted(s for s in services if s.count(".") < 2) - if bare: - raise RuntimeError( - "--services short names need --location (or pass full names like " - f"`system.ai.`): {', '.join(bare)}" - ) - if len(schemas) != 1: - raise RuntimeError( - "--services without --location must all share one `.` " - f"(got: {', '.join(sorted(schemas)) or 'none'}); pass --location instead." - ) - location = next(iter(schemas)) - state = load_state() +def _split_catalog_schema(location: str, flag: str) -> tuple[str, str]: + """Validate and split a ``.`` location. ``flag`` names the + option being parsed so the error points at the right one.""" + if location.count(".") != 1 or not all(part.strip() for part in location.split(".")): + raise RuntimeError(f"{flag} must be `.`, got `{location}`.") + catalog, schema = (part.strip() for part in location.split(".")) + return catalog, schema + + +def _setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[str]]: + """Shared setup for the non-interactive MCP commands: validate the workspace, + drop cross-workspace residue, resolve the configured-and-installed clients, + authenticate, and announce the target agents. Returns + ``(workspace, profile, clients)``.""" workspace = state.get("workspace") if not workspace: raise RuntimeError("Workspace is not configured. Run `ucode configure` first.") @@ -1223,7 +1228,7 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None apply_pat_environment(state) ensure_databricks_auth(workspace, profile) - print_section("MCP Servers") + print_section(section) client_names = ", ".join(str(MCP_CLIENTS[client]["display"]) for client in clients) print_note(f"Configuring for: {client_names}") for client in missing_clients: @@ -1231,6 +1236,29 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None f"{MCP_CLIENTS[client]['display']} is configured in ucode but not installed; " "skipping MCP config." ) + return workspace, profile, clients + + +def configure_mcp_command(location: str | None = None, services: set[str] | None = None) -> int: + if services is not None and location is None: + # `--services` works standalone with full names (`system.ai.github`): the + # `.` to configure is derived from them. Bare short names + # (`github`) can't be located without `--location`. + schemas = {".".join(s.split(".")[:2]) for s in services if s.count(".") >= 2} + bare = sorted(s for s in services if s.count(".") < 2) + if bare: + raise RuntimeError( + "--services short names need --location (or pass full names like " + f"`system.ai.`): {', '.join(bare)}" + ) + if len(schemas) != 1: + raise RuntimeError( + "--services without --location must all share one `.` " + f"(got: {', '.join(sorted(schemas)) or 'none'}); pass --location instead." + ) + location = next(iter(schemas)) + state = load_state() + workspace, profile, clients = _setup_mcp_clients(state, "MCP Servers") original_mcp_servers_for_location: list[dict] = list(state.get("mcp_servers") or []) if location is not None: @@ -1335,3 +1363,58 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None # User submitted the picker without toggling anything --> make it clear nothing was selected print_note("No MCP servers selected. Press space to toggle an item, then enter to save.") return 0 + + +def _resolve_skills_mcp_servers( + workspace: str, + clients: list[str], + location: str, + original_servers: list[dict], +) -> list[dict]: + """Build the desired MCP server list for ``ucode configure skills``. + + Registers exactly one skills entry for ``location`` (a ``.``) + and drops any previously-registered skills entry, so a re-run repoints skills + at the new schema. Non-skills entries (plain mcp-services) are preserved: + skills and mcp-services can coexist.""" + catalog, schema = _split_catalog_schema(location, "--location") + entry_name = _catalog_schema_server_name("databricks-skills", catalog, schema, set()) + + original = _servers_by_name(original_servers).get(entry_name) + original_clients = list((original or {}).get("clients") or []) + merged_clients = original_clients + [c for c in clients if c not in original_clients] + skills_entry = { + "name": entry_name, + "url": build_skills_mcp_url(workspace, catalog, schema), + "auth": f"env:{MCP_AUTH_TOKEN_ENV_VAR}", + "clients": merged_clients, + "kind": SKILLS_MCP_KIND, + } + # `apply_mcp_server_changes` compares by value, so re-emitting an unchanged + # entry is a no-op; only a genuinely repointed skills entry re-registers. + kept = [ + s + for s in original_servers + if s.get("kind") != SKILLS_MCP_KIND and _server_name(s) != entry_name + ] + return [*kept, skills_entry] + + +def configure_skills_command(location: str) -> int: + """Register the UC Skills MCP for ``location`` across configured agents. + + Skills are eagerly loaded in Claude (via ``alwaysLoad``); other agents get + the same endpoint without eager loading, which their CLIs cannot express.""" + state = load_state() + workspace, _profile, clients = _setup_mcp_clients(state, "Skills MCP") + + original_mcp_servers: list[dict] = list(state.get("mcp_servers") or []) + working_mcp_servers = _resolve_skills_mcp_servers( + workspace, clients, location, original_mcp_servers + ) + changed = apply_mcp_server_changes(original_mcp_servers, working_mcp_servers, clients) + if changed or original_mcp_servers != working_mcp_servers: + state["mcp_servers"] = working_mcp_servers + save_state(state) + print_success("Saved") + return 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index 7bdb0e4..e9d930a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -157,6 +157,30 @@ def test_mcp_group_lists_web_search(self): assert "web-search" in result.output +class TestConfigureSkillsCommand: + def test_dispatches_location_to_command(self, monkeypatch): + calls: list[str] = [] + monkeypatch.setattr( + "ucode.cli.configure_skills_command", lambda location: calls.append(location) or 0 + ) + result = runner.invoke(app, ["configure", "skills", "--location", "main.default"]) + assert result.exit_code == 0 + assert calls == ["main.default"] + + def test_requires_location(self): + result = runner.invoke(app, ["configure", "skills"]) + assert result.exit_code != 0 + + def test_runtime_error_exits_nonzero(self, monkeypatch): + def boom(location): + raise RuntimeError("Workspace is not configured.") + + monkeypatch.setattr("ucode.cli.configure_skills_command", boom) + result = runner.invoke(app, ["configure", "skills", "--location", "main.default"]) + assert result.exit_code == 1 + assert "Workspace is not configured." in result.output + + class TestAuthTokenCommand: """`ucode auth-token` is the cross-platform apiKeyHelper (#116).""" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f59aa1d..2106d32 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -512,6 +512,20 @@ def test_false_for_other_errors(self): assert db_mod.is_model_provider_feature_unavailable(None) is False +class TestBuildMcpUrls: + def test_mcp_service_url_keeps_dotted_name(self): + assert ( + db_mod.build_mcp_service_url(WS, "system.ai.github") + == f"{WS}/ai-gateway/mcp-services/system.ai.github" + ) + + def test_skills_url_uses_path_segments(self): + assert ( + db_mod.build_skills_mcp_url(WS, "main", "default") + == f"{WS}/ai-gateway/skills/main/default" + ) + + class TestListMcpServices: def test_accepts_entries_without_connection_status(self, monkeypatch): payload = { diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 81333c8..df1f049 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -25,6 +25,13 @@ def test_uses_oauth_token_header_reference(self): assert "oauth" not in entry assert "headersHelper" not in entry + def test_omits_always_load_by_default(self): + assert "alwaysLoad" not in mcp.build_mcp_http_entry(f"{WS}/api/2.0/mcp/external/github") + + def test_sets_always_load_when_requested(self): + entry = mcp.build_mcp_http_entry(f"{WS}/ai-gateway/skills/main/default", always_load=True) + assert entry["alwaysLoad"] is True + class TestAddClaudeMcpServer: def test_adds_user_scoped_json(self, monkeypatch): @@ -222,6 +229,36 @@ def test_configures_copilot_mcp_server(self, monkeypatch): assert removed_scopes == [] assert calls == [("github", f"{WS}/api/2.0/mcp/external/github")] + def test_claude_receives_entry_with_always_load(self, monkeypatch): + added: dict = {} + url = f"{WS}/ai-gateway/skills/main/default" + entry = mcp.build_mcp_http_entry(url, always_load=True) + monkeypatch.setattr(mcp, "remove_claude_mcp_server", lambda name, scope: False) + monkeypatch.setattr( + mcp, + "add_claude_mcp_server", + lambda name, entry, scope: added.update(name=name, entry=entry), + ) + + mcp.configure_client_mcp_server("claude", "databricks-skills-main-default", url, entry) + + assert added["entry"]["alwaysLoad"] is True + + def test_codex_receives_only_the_url(self, monkeypatch): + added: dict = {} + url = f"{WS}/ai-gateway/skills/main/default" + entry = mcp.build_mcp_http_entry(url, always_load=True) + monkeypatch.setattr(mcp, "remove_codex_mcp_server", lambda name: False) + monkeypatch.setattr( + mcp, "add_codex_mcp_server", lambda name, url: added.update(name=name, url=url) + ) + + mcp.configure_client_mcp_server("codex", "databricks-skills-main-default", url, entry) + + # codex's add signature takes no entry, so alwaysLoad structurally cannot + # reach it -- it is registered from the URL alone. + assert added == {"name": "databricks-skills-main-default", "url": url} + class TestMcpPicker: def test_prompt_uses_scrolling_checkbox_selector(self, monkeypatch): @@ -1642,6 +1679,93 @@ def test_full_names_spanning_multiple_schemas_without_location_raises(self): ) +class TestConfigureSkills: + def test_rejects_malformed_location(self, monkeypatch): + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + for bad in ("main", "cat.sch.extra", ".sch", "cat.", ""): + try: + mcp.configure_skills_command(location=bad) + except RuntimeError as exc: + assert "--location" in str(exc) + else: + raise AssertionError(f"expected RuntimeError for `{bad}`") + + def test_registers_skills_entry_for_every_agent(self, monkeypatch): + saved_states: list[dict] = [] + configured: list[tuple[str, str, str, dict]] = [] + state = {"workspace": WS, "available_tools": ALL_MCP_CLIENTS} + _stub_location_base(monkeypatch, state) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: list(ALL_MCP_CLIENTS)) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_skills_command(location="main.default") == 0 + + assert {c[0] for c in configured} == set(ALL_MCP_CLIENTS) + assert {c[1] for c in configured} == {"databricks-skills-main-default"} + assert {c[2] for c in configured} == {f"{WS}/ai-gateway/skills/main/default"} + # The skills entry carries alwaysLoad; only Claude's CLI consumes it (see + # TestConfigureClientMcpServerAlwaysLoad for the per-agent dispatch). + assert all(c[3].get("alwaysLoad") is True for c in configured) + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "databricks-skills-main-default", + "url": f"{WS}/ai-gateway/skills/main/default", + "auth": "env:OAUTH_TOKEN", + "clients": ALL_MCP_CLIENTS, + "kind": "skills", + } + ] + + def test_repoints_prior_skills_entry_but_keeps_mcp_services(self, monkeypatch): + saved_states: list[dict] = [] + removed: list[tuple[str, str]] = [] + service_entry = { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + stale_skills = { + "name": "databricks-skills-main-old", + "url": f"{WS}/ai-gateway/skills/main/old", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + "kind": "skills", + } + _stub_location_base( + monkeypatch, + {**CLAUDE_STATE, "mcp_servers": [service_entry, stale_skills]}, + ) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda client, name, url, entry: []) + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_skills_command(location="main.default") == 0 + + # The stale skills entry is removed; the mcp-service is untouched. + assert removed == [("claude", "databricks-skills-main-old")] + saved = saved_states[-1]["mcp_servers"] + assert service_entry in saved + assert [s for s in saved if s.get("kind") == "skills"] == [ + { + "name": "databricks-skills-main-default", + "url": f"{WS}/ai-gateway/skills/main/default", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + "kind": "skills", + } + ] + + class TestRevertMcpConfigs: def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkeypatch): removed: list[tuple[str, str]] = []