Skip to content
Merged
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
20 changes: 19 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,16 @@ OTLP trace export with a refreshing headers helper; no background service is nee
Register Databricks MCP servers for configured MCP-capable agents. Cursor Agent
is MCP-only and is included when `cursor-agent` is installed:

Run `ug mcp add` without selectors to browse a table of services by name, UC location,
and status. ✓ means configured (or managed by the workspace); ✗ means not configured.
Configured and managed services stay visible as read-only rows. `[ ]`/`[x]` indicate
what you select for this operation; `[-]` marks a read-only row.

Highlight a service to see its description, creator, and creation/update dates. Type
to filter names or UC locations, use Space to select, and Enter to save. Page Up/Down
scrolls the detail pane. Narrow terminals shorten cells and show status symbols;
the detail pane retains the full location and status.

Use `ug mcp add` to add servers without removing existing registrations:

```bash
Expand Down Expand Up @@ -226,7 +236,15 @@ this contract does not make shared setup transactional.
## Skills

Unity Catalog Skills can be registered as MCP tools or downloaded into local
agent skill directories.
agent skill directories. Run `ug skills add` without selectors to browse individual
skills in the same table and metadata pane as the MCP browser, plus the bundle name.
✓ Installed requires a recorded download of that exact UC skill with `SKILL.md` still
present at every target agent directory (`--path`, or your home directory). A matching
bundle name from another schema or metastore does not qualify. Unattributed local
bundles and incomplete installations show ✗ Not installed, with an explanation in the
detail pane. Installed skills remain selectable for updates.
Missing metadata is shown explicitly. `Created by` is the UC creator, which may
differ from the author of the bundle.

```bash
# Set up the Databricks skills MCP so your agents can create and manage skills through ug.
Expand Down
25 changes: 22 additions & 3 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -2599,10 +2599,18 @@ def _mcp_service_full_name(service: dict, required_prefix: str) -> str | None:


def list_mcp_services(
workspace: str, token: str, parent: str = "system.ai", *, max_pages: int = 100
workspace: str,
token: str,
parent: str = "system.ai",
*,
max_pages: int = 100,
metadata: dict[str, dict] | None = None,
) -> tuple[list[str], str | None]:
"""List UC MCP services under ``parent`` (a ``<catalog>.<schema>`` ref).

``metadata``, when provided, receives each accepted service's description and audit
fields from the same list response. No per-service requests are made.

Requests the ``BASIC`` view explicitly: only service names are needed to build
the deterministic proxy URLs, and ``BASIC`` omits the source-connection
resolution that makes ``FULL`` costlier (and adds Atlas load) at scale. The
Expand Down Expand Up @@ -2639,6 +2647,11 @@ def list_mcp_services(
full_name = _mcp_service_full_name(service, expected_prefix)
if full_name:
names.append(full_name)
if metadata is not None:
metadata[full_name] = {
key: service.get(key)
for key in ("comment", "created_by", "create_time", "update_time")
}
page_token = data.get("next_page_token") or None
if not page_token or page_token in seen_tokens:
break
Expand Down Expand Up @@ -3304,6 +3317,7 @@ def list_all_mcp_services(
on_progress: Callable[[int, int, int], None] | None = None,
on_services: Callable[[list[str]], None] | None = None,
cancel_event: threading.Event | None = None,
metadata: dict[str, dict] | None = None,
) -> tuple[list[str], str | None]:
"""Return sorted unique MCP-service full names across every `<catalog>.<schema>`
in the workspace. The mcp-services API is one-schema-per-call, so this walks
Expand All @@ -3323,10 +3337,15 @@ def list_all_mcp_services(
names: set[str] = set()

def probe(catalog, schema):
return list_mcp_services(workspace, token, f"{catalog}.{schema}")
details: dict[str, dict] = {}
kwargs = {"metadata": details} if metadata is not None else {}
result = list_mcp_services(workspace, token, f"{catalog}.{schema}", **kwargs)
return result, details

def collect(result, schemas_done, schemas_total):
found, _ = result
(found, _), details = result
if metadata is not None:
metadata.update(details)
new = [n for n in found if n not in names]
names.update(found)
if on_progress is not None:
Expand Down
115 changes: 82 additions & 33 deletions src/ucode/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
scrolling_checkbox,
spinner,
)
from ucode.ui.interactive_picker import CatalogChoice, picker_details

# Workspace-relative path fragments for the V2 AI Gateway MCP endpoints, shared by the URL-shape
# checks (`_is_app_mcp_server`, `_mcp_server_location`) so the set stays in one place.
Expand Down Expand Up @@ -464,7 +465,9 @@ def revert_mcp_configs(state: dict) -> dict[str, bool]:
return results


def discover_mcp_service_names(workspace: str, profile: str | None = None) -> list[str]:
def discover_mcp_service_names(
workspace: str, profile: str | None = None, *, metadata: dict[str, dict] | None = None
) -> list[str]:
"""Curated `system.ai.*` MCP services. Empty list if discovery fails so
callers can fall back to legacy connection discovery without surfacing
every error to the picker.
Expand All @@ -473,7 +476,8 @@ def discover_mcp_service_names(workspace: str, profile: str | None = None) -> li
expired/invalid credential surfaces loudly instead of masquerading as an empty
workspace (AIGTWY-4843)."""
token = get_databricks_token(workspace, profile)
names, reason = list_mcp_services(workspace, token)
kwargs = {"metadata": metadata} if metadata is not None else {}
names, reason = list_mcp_services(workspace, token, **kwargs)
raise_for_invalid_access_token(workspace, reason)
return names

Expand All @@ -484,19 +488,22 @@ def discover_all_mcp_service_names(
on_progress: Callable[[int, int, int], None] | None = None,
on_services: Callable[[list[str]], None] | None = None,
cancel_event: threading.Event | None = None,
metadata: dict[str, dict] | None = None,
) -> list[str]:
"""All MCP services across every `<catalog>.<schema>` in the workspace. This
walks the workspace (see `list_all_mcp_services`) and is the workspace-wide
counterpart to `discover_mcp_service_names`. `on_progress` is forwarded to
the walk for live count reporting, and `on_services` to stream newly-found
service names into the picker as the walk progresses."""
token = get_databricks_token(workspace, profile)
kwargs = {"metadata": metadata} if metadata is not None else {}
names, reason = list_all_mcp_services(
workspace,
token,
on_progress=on_progress,
on_services=on_services,
cancel_event=cancel_event,
**kwargs,
)
raise_for_invalid_access_token(workspace, reason)
return names
Expand Down Expand Up @@ -653,21 +660,47 @@ def _mcp_service_choice(
known_names: set[str],
additive: bool,
managed_names: set[str] | None = None,
) -> questionary.Choice | None:
"""Picker choice for one MCP-service full name (`<catalog>.<schema>.<id>`), or None to hide it.

Shared by the initial `build_mcp_picker_choices` render and the background walk that streams
more services in, so a streamed row is built identically to an up-front one. The picker only
shows rows you can act on, so two kinds are hidden: a managed service (never ug's to add or
remove) and, under `mcp add`, an already-configured one (nothing to add). Replace mode still
lists ug's own as a removable toggle; an unregistered service is an add-choice."""
metadata: dict | None = None,
) -> CatalogChoice:
"""Picker choice for one MCP-service full name (`<catalog>.<schema>.<id>`).

Shared by initial and streamed rows. Configured services stay visible with their
status. In add mode they are read-only; managed services are always read-only.
Replace mode keeps ug's own registrations checked so they can still be removed."""
registered_as = name.replace(".", "-")
display_title = f"MCP: {name}"
if registered_as in (managed_names or set()):
return None
if registered_as in known_names:
return None if additive else _server_choice(registered_as, True, display_title)
return _add_choice(f"{MCP_SERVICE_SELECTION_PREFIX}{name}", display_title)
managed = registered_as in (managed_names or set())
configured = registered_as in known_names or managed
status = "✓ Managed" if managed else "✓ Configured" if configured else "✗ Not configured"
details = metadata or {}
description = (
picker_details(
name,
details.get("comment"),
created_by=details.get("created_by"),
create_time=details.get("create_time"),
update_time=details.get("update_time"),
)
+ f"\nStatus: {status}"
)
location, _, leaf = name.rpartition(".")
return CatalogChoice(
# The title is the filter corpus in table mode. Avoid a common "MCP:" prefix,
# which would make searches such as "mcp" match every service.
title=name,
value=registered_as
if configured
else f"{MCP_ADD_PREFIX}{MCP_SERVICE_SELECTION_PREFIX}{name}",
name=leaf,
location=location,
status=status,
checked=configured and not additive and not managed,
disabled="Managed by workspace"
if managed
else "Already configured"
if configured and additive
else None,
description=description,
)


def build_mcp_picker_choices(
Expand All @@ -680,6 +713,7 @@ def build_mcp_picker_choices(
available_uc_functions_servers: list[dict] | None = None,
additive: bool = False,
managed_names: set[str] | None = None,
service_metadata: dict[str, dict] | None = None,
) -> list[questionary.Choice | questionary.Separator]:
original_by_name = _servers_by_name(original_servers)
managed = managed_names or set()
Expand Down Expand Up @@ -711,7 +745,11 @@ def push(choice: questionary.Choice | None) -> None:
# Picker shows the dotted UC name; state/agents store the dashed form
# (see resolver). The shared helper is also used by the background walk that
# streams more services in, so up-front and streamed rows match exactly.
push(_mcp_service_choice(name, known_names, additive, managed))
push(
_mcp_service_choice(
name, known_names, additive, managed, (service_metadata or {}).get(name)
)
)
displayed_names.add(name.replace(".", "-"))

for name in available_external_names:
Expand Down Expand Up @@ -808,23 +846,24 @@ def prompt_for_mcp_server_choices(
allow_back: bool = False,
additive: bool = False,
managed_names: set[str] | None = None,
service_metadata: dict[str, dict] | None = None,
background_loader: Callable[[Callable[[list[questionary.Choice]], None], threading.Event], None]
| None = None,
) -> list[str] | None | _Back:
"""Show the MCP server picker. Returns the list of selected values, `None`
if cancelled (Ctrl-C), or `_BACK` if `allow_back` and the user pressed Left
to return to the previous wizard step.

``additive`` (``ucode mcp add``) shows already-configured servers as
non-toggleable notes instead of pre-checked, removable boxes.
``additive`` (``ug mcp add``) shows configured services as read-only table rows
instead of pre-checked, removable boxes.

``managed_names`` are servers managed configuration already provides; they show as
non-toggleable in either mode, since there is nothing to add and ug can't remove them here.

``background_loader`` streams more choices in after the picker opens (see
`scrolling_checkbox`) — used to load the workspace-wide MCP-services walk without
blocking on it up front."""
instruction = "(space to toggle, ctrl-a all, enter to save, type to filter)"
instruction = "(space toggle, ^A all, enter save, type filter, pgup/pgdn details)"
if allow_back:
instruction = "(space to toggle, ctrl-a all, ← back, enter to save, type to filter)"
selection = scrolling_checkbox(
Expand All @@ -839,11 +878,14 @@ def prompt_for_mcp_server_choices(
available_uc_functions_servers,
additive=additive,
managed_names=managed_names,
service_metadata=service_metadata,
),
style=picker_style(),
instruction=instruction,
allow_back=allow_back,
background_loader=background_loader,
show_description=True,
table=True,
).ask()
if selection is None:
return None
Expand Down Expand Up @@ -1323,46 +1365,52 @@ def _mcp_services_background_loader(
never blocks the picker from opening. Deduping against already-shown rows (e.g. the fast
`system.ai` list) is handled by the picker's append."""

metadata: dict[str, dict] = {}

def loader(
append: Callable[[list[questionary.Choice]], None], cancel_event: threading.Event
) -> None:
def on_services(new_names: list[str]) -> None:
# A streamed service can be hidden (managed, or already configured under `mcp add`),
# in which case its choice is None — drop those so only real rows reach the picker.
new_choices = [
choice
for name in new_names
if (choice := _mcp_service_choice(name, known_names, additive, managed_names))
is not None
]
if new_choices:
append(new_choices)
append(
[
_mcp_service_choice(
name, known_names, additive, managed_names, metadata.get(name)
)
for name in new_names
]
)

discover_all_mcp_service_names(
workspace, profile, on_services=on_services, cancel_event=cancel_event
workspace,
profile,
on_services=on_services,
cancel_event=cancel_event,
metadata=metadata,
)

return loader


def _discover_selected_mcp_sources(
workspace: str, profile: str | None, sources: set[str]
) -> dict[str, list]:
) -> dict[str, Any]:
"""Discover the picker's sources. The picker searches a single source — MCP services — so
this fetches the fast curated `system.ai` list synchronously (the slow workspace-wide walk
streams in afterward via the picker's background loader). The other keys stay in the returned
dict as empty lists so the picker call is unchanged and can still render/remove
already-registered servers of any type."""
services: list[str] = []
metadata: dict[str, dict] = {}
if MCP_SERVICES_SOURCE in sources:
services = _discover_mcp_source(
"MCP services",
lambda: discover_mcp_service_names(workspace, profile),
lambda: discover_mcp_service_names(workspace, profile, metadata=metadata),
)
return {
"external": [],
"apps": [],
"services": services,
"service_metadata": metadata,
"genie": [],
"vector_search": [],
"uc_functions": [],
Expand Down Expand Up @@ -2060,6 +2108,7 @@ def configure_mcp_command(
additive=append,
managed_names=managed_names,
background_loader=services_loader,
service_metadata=discovered.get("service_metadata"),
)
if selections is None or isinstance(selections, _Back):
# No back-navigation without the source step; `_Back` can't occur, but keep the guard
Expand Down
6 changes: 6 additions & 0 deletions src/ucode/skills_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ class SkillRef:
to preview a skill in the interactive picker. ``metastore_id``, ``skill_id``, and
``uc_update_time`` are UC attribution metadata recorded when the skill is
downloaded (see ``skills_state``); the download flow never reads them.
``created_by`` and ``create_time`` are list metadata for the browser; the UC
creator is not necessarily the author of the skill bundle.
"""

catalog: str
Expand All @@ -61,6 +63,8 @@ class SkillRef:
metastore_id: str | None = None
skill_id: str | None = None
uc_update_time: str | None = None
created_by: str | None = None
create_time: str | None = None

@property
def fqn(self) -> str:
Expand Down Expand Up @@ -124,6 +128,8 @@ def _skill_ref(skill: dict) -> SkillRef | None:
metastore_id=_non_empty_str(skill.get("metastore_id")),
skill_id=_non_empty_str(skill.get("id")),
uc_update_time=_non_empty_str(skill.get("update_time")),
created_by=_non_empty_str(skill.get("created_by")),
create_time=_non_empty_str(skill.get("create_time")),
)


Expand Down
Loading
Loading