Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 <catalog>.<schema>
```

Add the Unity Catalog Skills MCP server for a `<catalog>.<schema>` 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
Expand Down
29 changes: 28 additions & 1 deletion src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
from ucode.mcp import (
MCP_CLIENTS,
configure_mcp_command,
configure_skills_command,
purge_cross_workspace_mcp_residue,
revert_mcp_configs,
)
Expand Down Expand Up @@ -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")
]
Expand Down Expand Up @@ -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 <catalog>.<schema>` 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
Expand Down Expand Up @@ -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 `<catalog>.<schema>` 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[
Expand Down
4 changes: 4 additions & 0 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
133 changes: 108 additions & 25 deletions src/ucode/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 `<catalog>.<schema>`, got `{location}`.")
_split_catalog_schema(location, "--location")

token = get_databricks_token(workspace, profile)
with spinner(f"Discovering MCP services in {location}..."):
Expand Down Expand Up @@ -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
# `<catalog>.<schema>` 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.<name>`): {', '.join(bare)}"
)
if len(schemas) != 1:
raise RuntimeError(
"--services without --location must all share one `<catalog>.<schema>` "
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 ``<catalog>.<schema>`` 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 `<catalog>.<schema>`, 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.")
Expand Down Expand Up @@ -1223,14 +1228,37 @@ 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:
print_warning(
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
# `<catalog>.<schema>` 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.<name>`): {', '.join(bare)}"
)
if len(schemas) != 1:
raise RuntimeError(
"--services without --location must all share one `<catalog>.<schema>` "
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:
Expand Down Expand Up @@ -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 ``<catalog>.<schema>``)
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
24 changes: 24 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)."""

Expand Down
14 changes: 14 additions & 0 deletions tests/test_databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down
Loading