Skip to content
78 changes: 77 additions & 1 deletion src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,9 @@
)
from ucode.mcp import (
MCP_CLIENTS,
SKILLS_MCP_KIND,
configure_mcp_command,
configure_skills_mcp_command,
purge_cross_workspace_mcp_residue,
revert_mcp_configs,
)
Expand Down Expand Up @@ -143,6 +145,27 @@ def _parse_agents_option(agents: str) -> list[str]:
return tools


def _parse_skill_locations(location: str) -> list[str]:
"""Parse a comma-separated `--location` into `<catalog>.<schema>` refs,
dropping duplicates while preserving order."""
locations: list[str] = []
for raw in location.split(","):
raw = raw.strip()
if not raw:
continue
parts = raw.split(".")
if len(parts) != 2 or not all(part.strip() for part in parts):
raise RuntimeError(f"--location entries must be `<catalog>.<schema>`, got `{raw}`.")
if raw not in locations:
locations.append(raw)
if not locations:
raise RuntimeError(
"No schemas provided for --location. Use `<catalog>.<schema>`, "
"comma-separated for multiple."
)
return locations


def _parse_workspaces_option(workspaces: str) -> list[tuple[str, str | None]]:
"""Parse `--workspaces` into [(url, profile_name | None), ...].

Expand Down Expand Up @@ -723,7 +746,9 @@ def status() -> int:
tool_mcp_servers = [
str(server.get("name"))
for server in mcp_servers
if tool in (server.get("clients") or []) and server.get("name")
if tool in (server.get("clients") or [])
and server.get("name")
and server.get("kind") != SKILLS_MCP_KIND
]
print_kv("MCP list command", str(MCP_CLIENTS[tool]["list_command"]))
print_kv(
Expand All @@ -733,6 +758,23 @@ def status() -> int:
print_kv("Config file", str(config_path) if config_path.exists() else "missing")
console.print()

print_heading("Skills")
skill_mcp_entry = next((s for s in mcp_servers if s.get("kind") == SKILLS_MCP_KIND), None)
if not skill_mcp_entry:
print_kv("Skills", "not configured")
else:
locations = skill_mcp_entry.get("skill_locations") or []
print_kv(
"Skill MCP Locations",
", ".join(locations) if locations else "none — utility tools only",
)
configured_agents = [
str(MCP_CLIENTS[client]["display"])
for client in (skill_mcp_entry.get("clients") or [])
if client in MCP_CLIENTS
]
print_kv("Configured", ", ".join(configured_agents) if configured_agents else "none")

print_heading("Tracing")
tracing = state.get("tracing") or {}
if tracing.get("enabled"):
Expand All @@ -757,6 +799,10 @@ 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> --mcp` to connect Unity "
"Catalog Skills."
)
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 @@ -1349,6 +1395,36 @@ def configure_mcp(
raise typer.Exit(130) from None


@configure_app.command("skills")
def configure_skills(
location: Annotated[
str,
typer.Option("--location", help="Comma-separated `<catalog>.<schema>` skill scopes."),
],
mcp: Annotated[
bool,
typer.Option(
"--mcp", help="Manage the skills MCP connection (required until download mode lands)."
),
] = False,
) -> None:
"""Configure Databricks Skills for your coding tools.

``--location`` sets the skills MCP connection's scope to exactly the listed
schemas, replacing any previous set.
"""
try:
if not mcp:
raise RuntimeError("Download mode is not available yet; pass --mcp for now.")
configure_skills_mcp_command(_parse_skill_locations(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
13 changes: 13 additions & 0 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1365,6 +1365,19 @@ 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, locations: list[str]) -> str:
"""Skills route with one ``?schema=`` scope per location. The trailing slash
is required by the Envoy prefix even with no query params.

[] -> ``.../ai-gateway/skills/``
["main.default", "ml.a"] -> ``.../ai-gateway/skills/?schema=main.default&schema=ml.a``
"""
base = f"{workspace}/ai-gateway/skills/"
if not locations:
return base
return base + "?" + urlencode([("schema", loc) for loc in locations])


# 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
152 changes: 122 additions & 30 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 Down Expand Up @@ -79,6 +80,8 @@
"list_command": "copilot mcp list",
},
}
SKILLS_MCP_KIND = "skills"
SKILLS_MCP_SERVER_NAME = "databricks-skill-registry"
EXTERNAL_MCP_SELECTION_PREFIX = "external:"
SQL_MCP_VALUE = "managed:sql"
GENIE_SPACE_SELECTION_PREFIX = "genie-space:"
Expand All @@ -96,14 +99,17 @@
)


def build_mcp_http_entry(url: str) -> dict:
return {
def build_mcp_http_entry(url: str, *, always_load: bool = False) -> dict:
entry: dict[str, Any] = {
"type": "http",
"url": url,
"headers": {
"Authorization": f"Bearer ${{{MCP_AUTH_TOKEN_ENV_VAR}}}",
},
}
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 +1043,9 @@ def apply_mcp_server_changes(
url = server.get("url")
if not isinstance(url, str) or not url:
continue
entry = build_mcp_http_entry(url)
# alwaysLoad (Claude-only) keeps the skills registry's utility tools
# discoverable without an explicit mention; other clients ignore it.
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 @@ -1103,6 +1111,10 @@ def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None:
)


def _skills_entries(servers: list[dict]) -> list[dict]:
return [s for s in servers if s.get("kind") == SKILLS_MCP_KIND]


def _resolve_location_mcp_servers(
workspace: str,
profile: str | None,
Expand All @@ -1113,9 +1125,10 @@ def _resolve_location_mcp_servers(
) -> list[dict]:
"""Build the desired MCP server list for ``--location <cat>.<schema>``.

Strict replacement: the returned list is exactly the mcp-services
discovered at ``location``. Any previously-registered MCP entries outside
that location are removed by ``apply_mcp_server_changes``. Raises ``RuntimeError`` for an invalid
Strict replacement for mcp-services: the returned list is exactly the ones
discovered at ``location`` (any previously-registered mcp-service outside it
is removed by ``apply_mcp_server_changes``), plus any existing skills
connection, preserved untouched. Raises ``RuntimeError`` for an invalid
location (HTTP 404 from the listing API) or any other listing failure.

When ``services`` is given, the discovered set is narrowed to exactly that
Expand Down Expand Up @@ -1174,28 +1187,15 @@ def _resolve_location_mcp_servers(
working_servers.append(original.copy())
else:
working_servers.append(candidate)
return working_servers
return [*working_servers, *_skills_entries(original_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 _setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[str]]:
"""Validate the workspace, resolve configured MCP clients, and prepare auth.

Returns ``(workspace, profile, clients)`` and prints the section header, the
"Configuring for" note, and a warning per configured-but-uninstalled client.
"""
workspace = state.get("workspace")
if not workspace:
raise RuntimeError("Workspace is not configured. Run `ucode configure` first.")
Expand Down Expand Up @@ -1223,14 +1223,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 @@ -1276,20 +1299,24 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None
)

original_mcp_servers: list[dict] = list(state.get("mcp_servers") or [])
original_by_name = _servers_by_name(original_mcp_servers)
# Skills connections are managed by `configure skills`, so keep them out of
# the picker and carry them through untouched.
skills_servers = _skills_entries(original_mcp_servers)
picker_servers = [s for s in original_mcp_servers if s.get("kind") != SKILLS_MCP_KIND]
original_by_name = _servers_by_name(picker_servers)
selections = prompt_for_mcp_server_choices(
available_external_mcp_names,
available_genie_mcp_servers,
available_app_mcp_servers,
original_mcp_servers,
picker_servers,
available_mcp_service_names,
available_vector_search_servers,
available_uc_functions_servers,
)
if selections is None:
return 0

working_mcp_servers: list[dict] = []
working_mcp_servers: list[dict] = list(skills_servers)
working_names: set[str] = set()
add_selections: list[str] = []
for selection in selections:
Expand Down Expand Up @@ -1335,3 +1362,68 @@ 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 _merge_clients(prior: list[str] | None, new: list[str]) -> list[str]:
"""Order-preserving union of a prior client list with newly-configured ones."""
prior = list(prior or [])
return prior + [c for c in new if c not in prior]


def _build_skills_entry(workspace: str, locations: list[str], clients: list[str]) -> dict:
"""Canonical single skills-registry entry. ``skill_locations`` is the source
of truth; the URL is always derived from it, never parsed back."""
return {
"name": SKILLS_MCP_SERVER_NAME,
"kind": SKILLS_MCP_KIND,
"skill_locations": list(locations),
"url": build_skills_mcp_url(workspace, locations),
"auth": f"env:{MCP_AUTH_TOKEN_ENV_VAR}",
"clients": clients,
}


def _resolve_skills_mcp_servers(
workspace: str,
clients: list[str],
locations: list[str],
original_servers: list[dict],
) -> list[dict]:
"""Rebuild the MCP server list around exactly one skills entry.

Drops every prior ``kind=="skills"`` entry and any entry named
``SKILLS_MCP_SERVER_NAME`` (single-connection invariant; also sweeps up a
stray old-named entry via ``apply_mcp_server_changes``), keeps everything
else, and appends one rebuilt entry whose clients merge the prior skills
entry's clients with ``clients``.
"""
prior = next((s for s in original_servers if s.get("kind") == SKILLS_MCP_KIND), None)
merged = _merge_clients((prior or {}).get("clients"), clients)
kept = [
s
for s in original_servers
if s.get("kind") != SKILLS_MCP_KIND and _server_name(s) != SKILLS_MCP_SERVER_NAME
]
return [*kept, _build_skills_entry(workspace, locations, merged)]


def _update_skills_mcp(
state: dict, workspace: str, clients: list[str], locations: list[str]
) -> None:
"""Rebuild the single skills connection for ``locations`` and persist it."""
original = list(state.get("mcp_servers") or [])
working = _resolve_skills_mcp_servers(workspace, clients, locations, original)
changed = apply_mcp_server_changes(original, working, clients)
if changed or original != working:
state["mcp_servers"] = working
save_state(state)
print_success("Saved")


def configure_skills_mcp_command(locations: list[str]) -> int:
"""Set the skills MCP connection's ``skill_locations`` to exactly ``locations``,
replacing any previous set."""
state = load_state()
workspace, _profile, clients = _setup_mcp_clients(state, "Skills MCP")
_update_skills_mcp(state, workspace, clients, locations)
return 0
Loading
Loading