From 456a3a67c6cbeec7fdc57279c4edd481ec5af82e Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 17 Sep 2026 17:10:50 +0000 Subject: [PATCH 1/5] Add `ug mcp login` to sign in to connection-backed MCP services `ug mcp login` shows which of the connection-backed AI Gateway MCP services the coding agents are configured to use are already signed in vs. still need a per-user connection sign-in, and runs the sign-in for the ones you pick (interactive picker) or name with `--services` / scope with `--agents`. Reuses #679's building blocks: the developer + workspace-managed server enumeration is extracted from `list_mcp_command` into a shared `configured_mcp_servers_by_name` (behavior-preserving) that both `ug mcp list` and `ug mcp login` call, and the status is rendered with the same rich Table + `status_badge` styling from `ucode.ui`. Per-service status comes from the existing Unity Catalog REST APIs (the ones the `/mcp-service-login` page uses); sign-in is `databricks auth login --resource` (RFC 8707, databricks/cli#6621), so it works for any connection-backed MCP service, not just `system.ai.*`. The credential is per-user and shared across agents, so signing in once unblocks the service for every agent. Co-authored-by: Isaac --- README.md | 28 ++++ src/ucode/cli.py | 44 ++++++ src/ucode/mcp.py | 67 ++++++--- src/ucode/mcp_login.py | 285 +++++++++++++++++++++++++++++++++++++ tests/test_mcp.py | 66 +++++++++ tests/test_mcp_login.py | 305 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 773 insertions(+), 22 deletions(-) create mode 100644 src/ucode/mcp_login.py create mode 100644 tests/test_mcp_login.py diff --git a/README.md b/README.md index b87ad1449..9d666a43c 100644 --- a/README.md +++ b/README.md @@ -127,6 +127,33 @@ profile. V2 AI Gateway servers can be added with typed selectors such as `vector-search:main.docs`, `uc-functions:main.tools`, `external:`, `genie-space:`, or `app:`. +Sign in to connection-backed servers with `ug mcp login`: + +```bash +# Show every configured connection-backed MCP service with its sign-in status, +# and pick which to sign in to. +ug mcp login + +# Sign in to specific services non-interactively (full or short names). +ug mcp login --names system.ai.github,system.ai.slack + +# Scope to specific agents' services. +ug mcp login --agents claude,codex +``` + +Some MCP services (e.g. `system.ai.github`) are backed by a Unity Catalog connection and only vend +their tools once you've completed a one-time per-user sign-in to the underlying SaaS. `ug mcp login` +uses the same configured-server set as `ug mcp list` (including servers delivered through the +agents' OS-managed files), keeping only the connection-backed AI Gateway MCP services, and shows +each one's sign-in status (`signed in` / `needs sign-in`). Sign-in opens your browser to complete +the connection's login (via `databricks auth login`), then mints the credential. The credential is +**per-user and shared across every agent** — signing in once through any agent (or here) unblocks +that MCP service for Claude Code, Cursor, Codex, and the rest. It works for any connection-backed +MCP service, not just `system.ai.*`. + +> Requires a Databricks CLI that supports `--resource` (databricks/cli#6621); `ug mcp login` +> reports a clear message if your CLI is too old. + ## Skills Unity Catalog Skills can be registered as MCP tools or downloaded into local @@ -168,6 +195,7 @@ ug skills remove --location main.default --via mcp | `ug mcp add` | Add MCP servers without removing existing registrations | | `ug mcp remove` | Unregister configured MCP servers | | `ug mcp list` | List configured MCP servers and connection status | +| `ug mcp login` | Sign in to connection-backed MCP services (interactive, or `--names`) | | `ug skills` | Set up the Databricks skills MCP so agents can create and manage skills | | `ug skills list` | List configured skills and how each was configured | | `ug skills add` | Add skill MCP scopes or download skills | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index df0a7e96d..23668c246 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -124,6 +124,7 @@ remove_skills_locations_command, revert_mcp_configs, ) +from ucode.mcp_login import login_mcp_command from ucode.skills_download import ( configure_location_skills_download_command, configure_selected_skills_download_command, @@ -1581,6 +1582,49 @@ def mcp_list( raise typer.Exit(130) from None +@mcp_app.command("login") +def mcp_login( + names: Annotated[ + str | None, + typer.Option( + "--names", + help="Sign in to this comma-separated subset of MCP services non-interactively. " + "Full names like `system.ai.github` or bare short names like `github` both work. " + "Omit --names to show the interactive picker with each service's sign-in status.", + ), + ] = None, + agents: Annotated[ + str | None, + typer.Option( + "--agents", + help="Comma-separated coding agents to scope to (e.g. claude,codex). Without " + "--agents, considers the MCP services configured for every agent.", + ), + ] = None, +) -> None: + """Sign in to the connection-backed MCP services your agents use. + + Shows which configured MCP services are already signed in vs. need a + connection sign-in, and runs the sign-in for the ones you pick (or all named + with --names). Sign-in uses `databricks auth login --resource`, so it + works for any connection-backed MCP service (not just `system.ai.*`). + """ + selected = None if names is None else {s.strip() for s in names.split(",") if s.strip()} + requested_agents = ( + None + if agents is None + else ({a.strip().lower() for a in agents.split(",") if a.strip()} or None) + ) + try: + login_mcp_command(names=selected, agents=requested_agents) + 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 + + @mcp_app.command("web-search") def mcp_web_search_cmd() -> None: """Run the web_search MCP server over stdio. Invoked as a subprocess by Claude Code.""" diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index fd3cfbd97..df87b0670 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2487,6 +2487,47 @@ def _row_status( ) +def configured_mcp_servers_by_name( + state: dict, agents: set[str] | None = None +) -> dict[str, dict[str, Any]]: + """Merge the developer- and workspace-managed MCP servers ug has configured, keyed by + registered name, unioning the agents each is on. Skills connections are excluded (they are + reported/handled separately). ``agents`` drops agents outside that scope, and a server left + with no in-scope agent is omitted. Each value is ``{"server", "clients", "managed"}``. + + Managed servers can be delivered two ways: to fallback state (``managed_mcp_servers``) or, for + Claude/Codex, into the agents' OS-managed files — the latter is the source of truth, so it is + read directly here. Shared by ``ug mcp list`` and ``ug mcp login`` so both (and ``ug status``) + see the same configured-server set regardless of how a managed server was delivered.""" + configured: dict[str, dict[str, Any]] = {} + + def _collect(server: dict, *, managed: bool) -> None: + name = _server_name(server) + if not name or server.get("kind") == SKILLS_MCP_KIND: + return + clients = [ + client for client in _mcp_server_clients(server) if agents is None or client in agents + ] + if not clients: + return + entry = configured.setdefault(name, {"server": server, "clients": [], "managed": managed}) + entry["clients"] = _merge_clients(entry["clients"], clients) + entry["managed"] = entry["managed"] or managed + + for server in state.get("mcp_servers") or []: + _collect(server, managed=False) + for server in state.get("managed_mcp_servers") or []: + _collect(server, managed=True) + # Managed servers delivered through the agents' OS-managed files (Claude/Codex) live in those + # files, not in state, so read them too — otherwise `ug mcp login` would miss them. + for agent, module in (("claude", claude), ("codex", codex)): + if agents is not None and agent not in agents: + continue + for name, url in module.read_managed_mcp_urls().items(): + _collect({"name": name, "url": url, "clients": [agent]}, managed=True) + return configured + + def list_mcp_command(agents: set[str] | None = None) -> int: """`ug mcp list`: show the Databricks MCP servers ug has configured and their live connection status in each coding agent, one row per server. @@ -2518,28 +2559,10 @@ def list_mcp_command(agents: set[str] | None = None) -> int: live = _query_live_statuses(probe_clients) - # Merge developer- and workspace-managed servers by registered name, unioning their agents. - # ``--agents`` drops agents outside the scope, and a server left with no in-scope agent is - # omitted. The skills connection is intentionally excluded — it's reported by the skill commands. - configured: dict[str, dict[str, Any]] = {} - - def _collect(server: dict, *, managed: bool) -> None: - name = _server_name(server) - if not name or server.get("kind") == SKILLS_MCP_KIND: - return - clients = [ - client for client in _mcp_server_clients(server) if agents is None or client in agents - ] - if not clients: - return - entry = configured.setdefault(name, {"server": server, "clients": [], "managed": managed}) - entry["clients"] = _merge_clients(entry["clients"], clients) - entry["managed"] = entry["managed"] or managed - - for server in state.get("mcp_servers") or []: - _collect(server, managed=False) - for server in managed_mcp_servers(state, agents): - _collect(server, managed=True) + # Merge developer- and workspace-managed servers by registered name, unioning their agents + # (shared with `ug mcp login` so both see the same configured-server set, including the servers + # delivered through the agents' OS-managed files). + configured = configured_mcp_servers_by_name(state, agents) if configured: table = Table(box=None, pad_edge=False, header_style="bold") diff --git a/src/ucode/mcp_login.py b/src/ucode/mcp_login.py new file mode 100644 index 000000000..3d479db5d --- /dev/null +++ b/src/ucode/mcp_login.py @@ -0,0 +1,285 @@ +"""`ug mcp login`: sign in to the connection-backed AI Gateway MCP services that +the coding agents are configured to use — a `/mcp`-style status list + login. + +A connection-backed MCP service (e.g. ``system.ai.github``) only vends its tools +once the user holds a per-user connection credential. This command shows, for the +MCP services ucode has registered for the agents, which are already authenticated +and which still need a sign-in, and runs the sign-in for the ones you pick. + +The sign-in is ``databricks auth login --resource `` (RFC 8707), which a +resource-aware ``/oidc`` routes through the connection's own SaaS login +(``/mcp-service-login``) before minting the token — so it works for **any** +connection-backed MCP service, not just ``system.ai.*``. The per-service +credential status comes from the existing Unity Catalog REST APIs (the same ones +the ``/mcp-service-login`` page uses): the mcp-service's backing connection plus +its per-user credential provisioning state. +""" + +from __future__ import annotations + +from urllib.parse import quote + +from rich.table import Table + +from ucode.databricks import ( + _http_get_json, + _scim_me, + get_databricks_token, + workspace_hostname, +) +from ucode.mcp import configured_mcp_servers_by_name +from ucode.mcp_connection_login import connection_from_url, run_connection_login +from ucode.state import load_state +from ucode.ui import ( + console, + muted, + print_heading, + print_kv, + print_note, + print_section, + print_success, + print_warning, + spinner, + status_badge, +) + +# Connection securable kinds that use per-user OAuth (U2M) credentials — i.e. the +# MCP service needs a connection sign-in. Mirrors the webapp's +# `hasGenericAccessTokenFlowKnownKinds`; anything else needs no per-user login. +_OAUTH_U2M_CONNECTION_KINDS = frozenset( + { + "CONNECTION_HTTP_OAUTH_U2M_MAPPING", + "CONNECTION_HTTP_DCR", + "CONNECTION_SLACK_OAUTH_U2M_MAPPING", + } +) + +# Per-service login status. +STATUS_AUTHENTICATED = "authenticated" +STATUS_NEEDS_LOGIN = "needs_login" +STATUS_NO_LOGIN = "no_login_needed" +STATUS_UNKNOWN = "unknown" + + +def _uc_prefix(workspace: str) -> str: + return f"https://{workspace_hostname(workspace)}/api/2.1/unity-catalog" + + +def _strip_connections_prefix(name: str | None) -> str | None: + """UC returns a connection reference as ``connections/``; the REST path + wants the bare name (mirror of the webapp's ``stripConnectionsPrefix``).""" + if not name: + return None + prefix = "connections/" + return name[len(prefix) :] if name.startswith(prefix) else name + + +def mcp_service_login_status(workspace: str, token: str, full_name: str, user_id: str) -> str: + """Per-user login status for one connection-backed MCP service. + + Resolves the service's backing connection (any connection, not just + ``system.ai.*``) and reads the current user's credential provisioning state + from the Unity Catalog REST APIs: + + 1. ``GET /mcp-services/`` → the service id + its ``source_connection`` + (name + securable_kind). A non-OAuth-U2M kind never needs a login. + 2. ``GET /connections//user-credentials/?dependent.mcp_service.id=`` + → ``connection_user_credential.provisioning_info.state``. ``ACTIVE`` means + signed in; the endpoint answers **HTTP 404** ("Credential ... is not found + ... Please login first") when there is no credential yet. ```` is the + numeric workspace user id (the ``connection_user_credential`` key), not the + email — matching the webapp, which keys on ``userId``. + """ + prefix = _uc_prefix(workspace) + details, err = _http_get_json(f"{prefix}/mcp-services/{quote(full_name, safe='')}", token) + if err is not None or not isinstance(details, dict): + return STATUS_UNKNOWN + source = (details.get("config") or {}).get("source_connection") or {} + kind = source.get("securable_kind") + if kind not in _OAUTH_U2M_CONNECTION_KINDS: + return STATUS_NO_LOGIN + conn = _strip_connections_prefix(source.get("name")) + service_id = details.get("id") + if not conn or not service_id: + return STATUS_UNKNOWN + + cred_url = ( + f"{prefix}/connections/{quote(conn, safe='')}/user-credentials/" + f"{quote(str(user_id), safe='')}?dependent.mcp_service.id={quote(str(service_id), safe='')}" + ) + cred, cred_err = _http_get_json(cred_url, token) + if cred_err is not None: + # 404 NOT_FOUND is the authoritative "no credential yet" signal; any other + # error is inconclusive (don't claim authenticated, don't hard-fail). + return STATUS_NEEDS_LOGIN if cred_err.startswith("HTTP 404") else STATUS_UNKNOWN + if not isinstance(cred, dict): + return STATUS_UNKNOWN + state = ((cred.get("connection_user_credential") or {}).get("provisioning_info") or {}).get( + "state" + ) + return STATUS_AUTHENTICATED if state == "ACTIVE" else STATUS_NEEDS_LOGIN + + +def _prompt_login_selection(rows: list[tuple[str, str]]) -> list[str] | None: + """Checklist of connection-backed MCP services with their sign-in status; + the ones that need a login are pre-checked. ``rows`` is ``(full_name, + status)``. Returns the selected service names, or ``None`` if cancelled.""" + import questionary + + label = { + STATUS_AUTHENTICATED: "signed in", + STATUS_NEEDS_LOGIN: "needs sign-in", + STATUS_UNKNOWN: "status unknown", + } + choices = [ + questionary.Choice( + title=f"{full} ({label.get(status, status)})", + value=full, + checked=status != STATUS_AUTHENTICATED, + ) + for full, status in rows + ] + selection = questionary.checkbox( + "Sign in to MCP services (space to toggle, enter to confirm):", choices=choices + ).ask() + return None if selection is None else [str(v) for v in selection] + + +def _login_status_markup(status: str) -> str: + """Colored sign-in status token, using the same `status_badge` styling as the STATUS + column of `ug mcp list`.""" + return { + STATUS_AUTHENTICATED: status_badge("signed in", "ok"), + STATUS_NEEDS_LOGIN: status_badge("needs sign-in", "warn"), + STATUS_NO_LOGIN: muted("no sign-in needed"), + }.get(status, status_badge("status unknown", "warn")) + + +def _connection_mcp_service_entries( + state: dict, agents: set[str] | None +) -> list[tuple[str, str, list[str], bool]]: + """``(full_name, mcp_url, clients, managed)`` for each connection-backed MCP service ug has + configured (developer- or workspace-managed), scoped to ``agents`` when given. + + Uses the shared `configured_mcp_servers_by_name` enumeration so the set matches `ug mcp list`, then + keeps only the AI Gateway mcp-services (the ones that can have a per-user connection login).""" + out: list[tuple[str, str, list[str], bool]] = [] + for entry in configured_mcp_servers_by_name(state, agents).values(): + url = entry["server"].get("url") + full = connection_from_url(url) if isinstance(url, str) else None + if not full: + continue + out.append((full, url, entry["clients"], entry["managed"])) + return out + + +def login_mcp_command(names: set[str] | None = None, agents: set[str] | None = None) -> int: + """`ug mcp login`: show sign-in status for the agents' connection-backed MCP + services and sign in to the selected ones. ``--names`` targets specific + services non-interactively (full ``system.ai.github`` or short ``github``); + ``--agents`` scopes to those agents. Bare, it shows the picker.""" + state = load_state() + workspace = state.get("workspace") + if not workspace: + raise RuntimeError("Workspace is not configured. Run `ug configure` first.") + profile = state.get("profile") + + entries = _connection_mcp_service_entries(state, agents) + if not entries: + scope = "" if agents is None else f" for {', '.join(sorted(agents))}" + print_note(f"No connection-backed MCP services are configured{scope}.") + return 0 + + if names is not None: + entries = [e for e in entries if e[0] in names or e[0].split(".")[-1] in names] + unknown = names - {e[0] for e in entries} - {e[0].split(".")[-1] for e in entries} + if unknown: + print_warning(f"Not configured, skipping: {', '.join(sorted(unknown))}.") + if not entries: + print_note("No matching MCP services to sign in to.") + return 0 + + try: + token = get_databricks_token(workspace, profile) + except Exception as exc: # noqa: BLE001 - surface auth trouble as guidance + raise RuntimeError( + f"Could not get a Databricks token for {workspace}: {exc}. Run `ug configure` first." + ) from exc + # The connection user-credentials API keys on the numeric workspace user id + # (the `connection_user_credential.user_id`), not the email — same as the webapp. + user_id = (_scim_me(workspace, token) or {}).get("id") + if not user_id: + raise RuntimeError("Could not resolve the current Databricks user id.") + + print_section("MCP login") + print_kv("Workspace", workspace) + with spinner("Checking sign-in status..."): + status_by_full = { + full: mcp_service_login_status(workspace, token, full, user_id) + for full, _url, _clients, _managed in entries + } + + url_by_full = {full: url for full, url, _clients, _managed in entries} + # "Pending" = still needs a sign-in OR its status couldn't be determined. UNKNOWN is included + # so an unreachable status API doesn't masquerade as "already signed in" (and so the picker and + # the non-interactive path still act on it), rather than being silently skipped. + pending = [ + full + for full in url_by_full + if status_by_full[full] in (STATUS_NEEDS_LOGIN, STATUS_UNKNOWN) + ] + + # Non-interactive (--names): sign in only to targeted services that still need it — never + # re-run `databricks auth login` (which blocks on a browser) for one that's already signed in. + # Interactive: render the per-service status (same Table style as `ug mcp list`), then a picker. + if names is not None: + targets = [(f, url_by_full[f]) for f in pending] + else: + print_heading("Connection-backed MCP services") + table = Table(box=None, pad_edge=False, header_style="bold") + table.add_column("MCP SERVICE", no_wrap=True) + table.add_column("AGENTS") + table.add_column("SIGN-IN") + for full, _url, clients, managed in entries: + name = full + (" [magenta](managed)[/magenta]" if managed else "") + table.add_row(name, ", ".join(clients), _login_status_markup(status_by_full[full])) + console.print(table) + if not pending: + print_success("All configured MCP services are already signed in.") + return 0 + # Only services that can take a sign-in reach the picker; a NO_LOGIN service (its connection + # needs no per-user login) would otherwise show pre-checked and try to run a pointless login. + selectable = [ + (full, status_by_full[full]) + for full, _u, _c, _m in entries + if status_by_full[full] != STATUS_NO_LOGIN + ] + selected = _prompt_login_selection(selectable) + if not selected: + print_note("Nothing selected.") + return 0 + targets = [(f, url_by_full[f]) for f in selected] + + signed_in = 0 + for full, url in targets: + # `run_connection_login` announces the sign-in on stderr (and prints the + # authorize URL there), so we don't repeat a "Signing in..." note here. + ok, detail = run_connection_login(url, workspace, profile=profile) + if ok: + signed_in += 1 + print_success(f" {full}: {detail}") + else: + print_warning(f" {full}: {detail}") + if signed_in: + print_success(f"Signed in to {signed_in} MCP service(s).") + return 0 + + +__all__ = [ + "login_mcp_command", + "mcp_service_login_status", + "STATUS_AUTHENTICATED", + "STATUS_NEEDS_LOGIN", + "STATUS_NO_LOGIN", + "STATUS_UNKNOWN", +] diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8cef86bab..90d53355f 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -4001,6 +4001,72 @@ def test_authentication_failure_glyph_stays_failed(self): assert mcp.parse_mcp_list_output("claude", out) == {"svc": mcp.LIVE_FAILED} +class TestConfiguredMcpServersByName: + """The shared enumeration used by both `ug mcp list` and `ug mcp login`.""" + + @pytest.fixture(autouse=True) + def _no_managed_files(self, monkeypatch): + # Default: no OS-managed-file servers, so state-only cases are deterministic. + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", dict) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", dict) + + def _state(self): + return { + "mcp_servers": [ + {"name": "system-ai-github", "url": "u1", "clients": ["claude", "codex"]}, + { + "name": mcp.SKILLS_MCP_SERVER_NAME, + "kind": mcp.SKILLS_MCP_KIND, + "clients": ["claude"], + }, + ], + "managed_mcp_servers": [ + {"name": "system-ai-github", "url": "u1", "clients": ["cursor"]}, + {"name": "databricks-genie-abc", "url": "u2", "clients": ["claude"]}, + ], + } + + def test_merges_by_name_and_unions_clients(self): + by_name = mcp.configured_mcp_servers_by_name(self._state()) + assert set(by_name) == {"system-ai-github", "databricks-genie-abc"} + gh = by_name["system-ai-github"] + # Developer + workspace-managed entries unioned; managed flag sticks. + assert gh["clients"] == ["claude", "codex", "cursor"] + assert gh["managed"] is True + + def test_excludes_skills_connection(self): + assert mcp.SKILLS_MCP_SERVER_NAME not in mcp.configured_mcp_servers_by_name(self._state()) + + def test_agents_scope_drops_servers_with_no_in_scope_agent(self): + by_name = mcp.configured_mcp_servers_by_name(self._state(), agents={"cursor"}) + # Only the github service has a cursor client; genie (claude-only) is dropped. + assert set(by_name) == {"system-ai-github"} + assert by_name["system-ai-github"]["clients"] == ["cursor"] + + def test_includes_servers_delivered_via_os_managed_files(self, monkeypatch): + # Managed servers written to the agents' OS-managed files (Claude/Codex) live in those + # files, not state, so `ug mcp login`/`list` must still see them via the managed-file read. + url = f"{WS}/ai-gateway/mcp-services/system.ai.slack" + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + by_name = mcp.configured_mcp_servers_by_name({"mcp_servers": [], "managed_mcp_servers": []}) + assert "system-ai-slack" in by_name + entry = by_name["system-ai-slack"] + # Delivered to both agents' managed files → unioned, flagged managed, URL preserved. + assert entry["clients"] == ["claude", "codex"] + assert entry["managed"] is True + assert entry["server"]["url"] == url + + def test_managed_file_read_respects_agents_scope(self, monkeypatch): + url = f"{WS}/ai-gateway/mcp-services/system.ai.slack" + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + by_name = mcp.configured_mcp_servers_by_name( + {"mcp_servers": [], "managed_mcp_servers": []}, agents={"codex"} + ) + assert by_name["system-ai-slack"]["clients"] == ["codex"] + + class TestListMcpCommand: def _state(self): # A developer-added AI Gateway service on claude+codex, a workspace-managed one, and a diff --git a/tests/test_mcp_login.py b/tests/test_mcp_login.py new file mode 100644 index 000000000..e0f6693d4 --- /dev/null +++ b/tests/test_mcp_login.py @@ -0,0 +1,305 @@ +"""Tests for `ug mcp login` (mcp_login): status classification via the existing +UC REST APIs, the `--resource` login invocation, and command orchestration. + +Network-free: the UC REST calls (`_http_get_json`), the current-user lookup, and +the CLI subprocess are monkeypatched. +""" + +from __future__ import annotations + +import pytest + +from ucode import mcp, mcp_login + +WS = "https://ws.staging.cloud.databricks.com" +FULL = "system.ai.github" +URL = f"{WS}/ai-gateway/mcp-services/{FULL}" +USER = "user@databricks.com" +USER_ID = "1234567890" # numeric workspace user id (the connection-user-credentials key) + +_DETAILS = { + "id": "svc-123", + "config": { + "source_connection": { + "name": "connections/github", + "securable_kind": "CONNECTION_HTTP_OAUTH_U2M_MAPPING", + } + }, +} + + +def _fake_http(details=_DETAILS, details_err=None, cred=None, cred_err=None): + """Return an `_http_get_json` stub that answers the mcp-services details call + and the user-credentials call based on the URL.""" + + def _http(url, token, *args, **kwargs): + if "/mcp-services/" in url: + return details, details_err + if "/user-credentials/" in url: + return cred, cred_err + return None, "unexpected url" + + return _http + + +class TestLoginStatus: + def test_authenticated_when_credential_active(self, monkeypatch): + cred = {"connection_user_credential": {"provisioning_info": {"state": "ACTIVE"}}} + monkeypatch.setattr(mcp_login, "_http_get_json", _fake_http(cred=cred)) + assert ( + mcp_login.mcp_service_login_status(WS, "t", FULL, USER) + == mcp_login.STATUS_AUTHENTICATED + ) + + def test_needs_login_on_404_not_found(self, monkeypatch): + # The user-credentials endpoint answers 404 when there is no credential yet. + monkeypatch.setattr(mcp_login, "_http_get_json", _fake_http(cred_err="HTTP 404 Not Found")) + assert ( + mcp_login.mcp_service_login_status(WS, "t", FULL, USER) == mcp_login.STATUS_NEEDS_LOGIN + ) + + def test_needs_login_when_state_not_active(self, monkeypatch): + cred = {"connection_user_credential": {"provisioning_info": {"state": "PROVISIONING"}}} + monkeypatch.setattr(mcp_login, "_http_get_json", _fake_http(cred=cred)) + assert ( + mcp_login.mcp_service_login_status(WS, "t", FULL, USER) == mcp_login.STATUS_NEEDS_LOGIN + ) + + def test_no_login_needed_for_non_oauth_kind(self, monkeypatch): + details = { + "id": "x", + "config": { + "source_connection": {"name": "connections/c", "securable_kind": "CONNECTION_MYSQL"} + }, + } + monkeypatch.setattr(mcp_login, "_http_get_json", _fake_http(details=details)) + assert mcp_login.mcp_service_login_status(WS, "t", FULL, USER) == mcp_login.STATUS_NO_LOGIN + + def test_unknown_when_details_error(self, monkeypatch): + monkeypatch.setattr( + mcp_login, "_http_get_json", _fake_http(details=None, details_err="HTTP 500") + ) + assert mcp_login.mcp_service_login_status(WS, "t", FULL, USER) == mcp_login.STATUS_UNKNOWN + + def test_unknown_when_credential_error_not_404(self, monkeypatch): + monkeypatch.setattr(mcp_login, "_http_get_json", _fake_http(cred_err="HTTP 403 Forbidden")) + assert mcp_login.mcp_service_login_status(WS, "t", FULL, USER) == mcp_login.STATUS_UNKNOWN + + def test_credential_lookup_keys_on_numeric_user_id_not_email(self, monkeypatch): + # The connection user-credentials API keys on the numeric workspace user id; passing the + # email 404s and mis-reports an already-signed-in service as `needs sign-in`. + seen: dict[str, str] = {} + + def _http(url, token, *args, **kwargs): + if "/mcp-services/" in url: + return _DETAILS, None + seen["cred_url"] = url + return {"connection_user_credential": {"provisioning_info": {"state": "ACTIVE"}}}, None + + monkeypatch.setattr(mcp_login, "_http_get_json", _http) + status = mcp_login.mcp_service_login_status(WS, "t", FULL, USER_ID) + assert status == mcp_login.STATUS_AUTHENTICATED + assert f"/user-credentials/{USER_ID}?" in seen["cred_url"] + assert "user%40" not in seen["cred_url"] # the email was not used as the key + + +class TestLoginCommand: + @pytest.fixture(autouse=True) + def _no_managed_files(self, monkeypatch): + # `configured_mcp_servers_by_name` reads the agents' OS-managed files; stub them empty so + # these state-driven cases don't pick up managed servers from the host. + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", dict) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", dict) + + def _state(self): + return { + "workspace": WS, + "profile": "p", + "mcp_servers": [ + {"name": "system-ai-github", "url": URL, "clients": ["claude", "codex"]}, + { + "name": "databricks-skill-registry", + "url": f"{WS}/ai-gateway/skills/x", + "clients": ["claude"], + }, + ], + } + + def test_services_targets_only_named_and_logs_in(self, monkeypatch): + monkeypatch.setattr(mcp_login, "load_state", lambda: self._state()) + monkeypatch.setattr(mcp_login, "get_databricks_token", lambda ws, p: "tok") + monkeypatch.setattr(mcp_login, "_scim_me", lambda ws, t: {"id": USER_ID, "userName": USER}) + monkeypatch.setattr( + mcp_login, + "mcp_service_login_status", + lambda ws, t, full, u: mcp_login.STATUS_NEEDS_LOGIN, + ) + logged: list[str] = [] + monkeypatch.setattr( + mcp_login, + "run_connection_login", + lambda url, ws, p=None, **k: logged.append(url) or (True, "ok"), + ) + rc = mcp_login.login_mcp_command(names={"github"}) + assert rc == 0 + assert logged == [ + URL + ] # matched by short name; skills entry ignored (not a connection mcp-service) + + def test_agents_scope_excludes_unconfigured_agent(self, monkeypatch): + monkeypatch.setattr(mcp_login, "load_state", lambda: self._state()) + monkeypatch.setattr(mcp_login, "get_databricks_token", lambda ws, p: "tok") + monkeypatch.setattr(mcp_login, "_scim_me", lambda ws, t: {"id": USER_ID, "userName": USER}) + monkeypatch.setattr( + mcp_login, "mcp_service_login_status", lambda *a: mcp_login.STATUS_NEEDS_LOGIN + ) + logged: list[str] = [] + monkeypatch.setattr( + mcp_login, + "run_connection_login", + lambda url, ws, p=None, **k: logged.append(url) or (True, "ok"), + ) + # gemini isn't a client of the github service → nothing to do + rc = mcp_login.login_mcp_command(names={"github"}, agents={"gemini"}) + assert rc == 0 and logged == [] + + def test_no_configured_services_is_noop(self, monkeypatch): + monkeypatch.setattr( + mcp_login, "load_state", lambda: {"workspace": WS, "profile": "p", "mcp_servers": []} + ) + rc = mcp_login.login_mcp_command() + assert rc == 0 + + def test_workspace_managed_service_is_included(self, monkeypatch): + # Uses the shared enumeration, so workspace-managed mcp-services are covered too. + slack_url = f"{WS}/ai-gateway/mcp-services/system.ai.slack" + state = { + "workspace": WS, + "profile": "p", + "mcp_servers": [], + "managed_mcp_servers": [ + {"name": "system-ai-slack", "url": slack_url, "clients": ["claude"]} + ], + } + monkeypatch.setattr(mcp_login, "load_state", lambda: state) + monkeypatch.setattr(mcp_login, "get_databricks_token", lambda ws, p: "tok") + monkeypatch.setattr(mcp_login, "_scim_me", lambda ws, t: {"id": USER_ID, "userName": USER}) + monkeypatch.setattr( + mcp_login, "mcp_service_login_status", lambda *a: mcp_login.STATUS_NEEDS_LOGIN + ) + logged: list[str] = [] + monkeypatch.setattr( + mcp_login, + "run_connection_login", + lambda url, ws, p=None, **k: logged.append(url) or (True, "ok"), + ) + rc = mcp_login.login_mcp_command(names={"slack"}) + assert rc == 0 and logged == [slack_url] + + def test_os_managed_file_service_is_included(self, monkeypatch): + # A managed mcp-service delivered to an agent's OS-managed file (not state) is still seen by + # `ug mcp login` — otherwise you couldn't pre-sign-in to managed servers. + gh_url = f"{WS}/ai-gateway/mcp-services/system.ai.github" + state = {"workspace": WS, "profile": "p", "mcp_servers": [], "managed_mcp_servers": []} + monkeypatch.setattr(mcp_login, "load_state", lambda: state) + monkeypatch.setattr( + mcp.claude, "read_managed_mcp_urls", lambda: {"system-ai-github": gh_url} + ) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", dict) + monkeypatch.setattr(mcp_login, "get_databricks_token", lambda ws, p: "tok") + monkeypatch.setattr(mcp_login, "_scim_me", lambda ws, t: {"id": USER_ID, "userName": USER}) + monkeypatch.setattr( + mcp_login, "mcp_service_login_status", lambda *a: mcp_login.STATUS_NEEDS_LOGIN + ) + logged: list[str] = [] + monkeypatch.setattr( + mcp_login, + "run_connection_login", + lambda url, ws, p=None, **k: logged.append(url) or (True, "ok"), + ) + rc = mcp_login.login_mcp_command(names={"github"}) + assert rc == 0 and logged == [gh_url] + + def _setup(self, monkeypatch, status, *, logged, selection_capture=None): + """Wire load_state/token/user + a fixed status for every service, capturing sign-ins. + `status` is a str (same for all) or a {full_name: status} map.""" + monkeypatch.setattr(mcp_login, "load_state", lambda: self._state()) + monkeypatch.setattr(mcp_login, "get_databricks_token", lambda ws, p: "tok") + monkeypatch.setattr(mcp_login, "_scim_me", lambda ws, t: {"id": USER_ID, "userName": USER}) + resolve = (lambda full: status[full]) if isinstance(status, dict) else (lambda full: status) + monkeypatch.setattr( + mcp_login, "mcp_service_login_status", lambda ws, t, full, u: resolve(full) + ) + monkeypatch.setattr( + mcp_login, + "run_connection_login", + lambda url, ws, p=None, **k: logged.append(url) or (True, "ok"), + ) + if selection_capture is not None: + monkeypatch.setattr( + mcp_login, + "_prompt_login_selection", + lambda rows: selection_capture.extend(rows) or [], + ) + + def test_services_skips_already_signed_in(self, monkeypatch): + # `ug mcp login --names github` must NOT re-run the browser login for an already-signed-in + # service; only NEEDS_LOGIN/UNKNOWN are actionable. + logged: list[str] = [] + self._setup(monkeypatch, mcp_login.STATUS_AUTHENTICATED, logged=logged) + rc = mcp_login.login_mcp_command(names={"github"}) + assert rc == 0 and logged == [] + + def test_services_signs_in_unknown_status(self, monkeypatch): + # An UNKNOWN status (e.g. UC API unreachable) is still actionable, so the non-interactive + # path attempts the sign-in rather than silently skipping it. + logged: list[str] = [] + self._setup(monkeypatch, mcp_login.STATUS_UNKNOWN, logged=logged) + rc = mcp_login.login_mcp_command(names={"github"}) + assert rc == 0 and logged == [URL] + + def test_picker_excludes_no_login_services(self, monkeypatch): + # A NO_LOGIN service (its connection needs no per-user login) must not reach the picker, + # where it would show pre-checked and trigger a pointless sign-in. + state = { + "workspace": WS, + "profile": "p", + "mcp_servers": [ + {"name": "system-ai-github", "url": URL, "clients": ["claude"]}, + { + "name": "system-ai-pg", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.pg", + "clients": ["claude"], + }, + ], + } + monkeypatch.setattr(mcp_login, "load_state", lambda: state) + monkeypatch.setattr(mcp_login, "get_databricks_token", lambda ws, p: "tok") + monkeypatch.setattr(mcp_login, "_scim_me", lambda ws, t: {"id": USER_ID, "userName": USER}) + monkeypatch.setattr( + mcp_login, + "mcp_service_login_status", + lambda ws, t, full, u: mcp_login.STATUS_NO_LOGIN + if full == "system.ai.pg" + else mcp_login.STATUS_NEEDS_LOGIN, + ) + rows: list = [] + monkeypatch.setattr( + mcp_login, "_prompt_login_selection", lambda r: rows.extend(r) or [] + ) + rc = mcp_login.login_mcp_command() + assert rc == 0 + offered = {full for full, _status in rows} + assert offered == {"system.ai.github"} # the NO_LOGIN service is not offered + + def test_unknown_status_is_not_reported_as_all_signed_in(self, monkeypatch): + # If every status is UNKNOWN, the command must still show the picker, not claim everything + # is already signed in and skip it. + captured: list = [] + logged: list[str] = [] + self._setup( + monkeypatch, mcp_login.STATUS_UNKNOWN, logged=logged, selection_capture=captured + ) + rc = mcp_login.login_mcp_command() # interactive (no --names) + assert rc == 0 + assert [full for full, _status in captured] == ["system.ai.github"] # picker was shown From ba1493d0ec491e0ab5d269b1bf022001c5693513 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 22:57:28 +0000 Subject: [PATCH 2/5] docs: simplify the MCP Servers README for non-technical users Rewrite the MCP Servers section in the plain, scannable style of the Skills section: a one-line intro, a single command block with friendly comments (add / list / login / remove), and one plain-language note that you only sign in once. Move the jargon (agent scoping, V2 typed selectors, the mcp-proxy/token mechanics, CLI version requirement) into a collapsible "Advanced options" block so business users aren't hit with it up front. Co-authored-by: Isaac --- README.md | 69 +++++++++++++++++++------------------------------------ 1 file changed, 23 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index 9d666a43c..aa808934c 100644 --- a/README.md +++ b/README.md @@ -96,63 +96,40 @@ Cursor models still run through your Cursor account. ## MCP Servers -Register Databricks MCP servers for configured MCP-capable agents. Cursor Agent -is MCP-only and is included when `cursor-agent` is installed: - -Use `ug mcp add` to add servers without removing existing registrations: +MCP servers let your coding agents use Databricks-governed tools — like GitHub, Slack, +Vector Search, and Genie. `ug` sets them up for all your installed agents at once. ```bash +# Add tools — a whole catalog.schema, or specific ones by name. ug mcp add --location system.ai -ug mcp add --names system.ai.slack,system.ai.github -ug mcp add --agents claude,codex --location system.ai -``` - -Remove configured servers: - -```bash -ug mcp remove -ug mcp remove --agents codex -``` - -List configured servers and their connection status: +ug mcp add --names system.ai.github,system.ai.slack -```bash +# See what's set up, and whether each is signed in. ug mcp list -ug mcp list --agents claude,codex -``` -Every Databricks MCP server is registered as a local stdio server that runs -`ug mcp-proxy`; the proxy refreshes Databricks OAuth tokens from your CLI -profile. V2 AI Gateway servers can be added with typed selectors such as -`vector-search:main.docs`, `uc-functions:main.tools`, `external:`, -`genie-space:`, or `app:`. +# Sign in to a tool that needs it (opens your browser). Plain form shows a picker. +ug mcp login +ug mcp login --names system.ai.github -Sign in to connection-backed servers with `ug mcp login`: +# Remove tools (plain form shows a picker). +ug mcp remove +``` -```bash -# Show every configured connection-backed MCP service with its sign-in status, -# and pick which to sign in to. -ug mcp login +**You only sign in once.** Some tools (like `system.ai.github`) ask you to sign in to the +underlying service the first time. Do it once — through any agent or `ug mcp login` — and the +tool works everywhere: Claude, Cursor, Codex, and the rest. -# Sign in to specific services non-interactively (full or short names). -ug mcp login --names system.ai.github,system.ai.slack +
Advanced options -# Scope to specific agents' services. -ug mcp login --agents claude,codex -``` +- Limit any command to certain agents: add `--agents claude,codex`. +- Add other AI Gateway tools by typed name: `vector-search:main.docs`, + `uc-functions:main.tools`, `external:`, `genie-space:`, `app:`. +- Each tool runs locally through `ug mcp-proxy`, which keeps your Databricks sign-in fresh — + no tokens to copy or manage. +- Sign-in works for any connection-backed tool, not just `system.ai.*`, and needs a recent + Databricks CLI (`ug mcp login` tells you if yours is too old). -Some MCP services (e.g. `system.ai.github`) are backed by a Unity Catalog connection and only vend -their tools once you've completed a one-time per-user sign-in to the underlying SaaS. `ug mcp login` -uses the same configured-server set as `ug mcp list` (including servers delivered through the -agents' OS-managed files), keeping only the connection-backed AI Gateway MCP services, and shows -each one's sign-in status (`signed in` / `needs sign-in`). Sign-in opens your browser to complete -the connection's login (via `databricks auth login`), then mints the credential. The credential is -**per-user and shared across every agent** — signing in once through any agent (or here) unblocks -that MCP service for Claude Code, Cursor, Codex, and the rest. It works for any connection-backed -MCP service, not just `system.ai.*`. - -> Requires a Databricks CLI that supports `--resource` (databricks/cli#6621); `ug mcp login` -> reports a clear message if your CLI is too old. +
## Skills From 49c71db7412ad2fa64d1314816a60fe1c542e3e5 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 23:30:56 +0000 Subject: [PATCH 3/5] docs: add a worked example (add -> list -> login) to MCP Servers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Show non-technical users what the commands actually print — the one-line add confirmation, the ug mcp list status table, and the one-time sign-in — right under the command block. Co-authored-by: Isaac --- README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/README.md b/README.md index aa808934c..7e09b8a56 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,22 @@ ug mcp login --names system.ai.github ug mcp remove ``` +Here's what a typical session looks like — add two tools, check their status, then sign in: + +```text +$ ug mcp add --names system.ai.github,system.ai.slack +✔ Added 2 MCP servers across Claude Code, Codex + +$ ug mcp list + NAME LOCATION AGENTS STATUS + system-ai-github system.ai.github claude, codex needs sign-in + system-ai-slack system.ai.slack claude, codex needs sign-in + +$ ug mcp login --names system.ai.github +Signing in to 'system.ai.github' — opening your browser to finish sign-in… +✔ Signed in to 1 MCP service(s) +``` + **You only sign in once.** Some tools (like `system.ai.github`) ask you to sign in to the underlying service the first time. Do it once — through any agent or `ug mcp login` — and the tool works everywhere: Claude, Cursor, Codex, and the rest. From ea02ed29501e87c7b637cd78d4f85471f5bf1214 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 23:36:42 +0000 Subject: [PATCH 4/5] docs: drop the worked example and internal-detail bullets from MCP Servers Keep the section to the essentials: intro, the command block, the one-time sign-in note, and just the two actionable Advanced options (agent scoping, typed selectors). Removes the sample add/list/login session and the mcp-proxy/CLI-version bullets. Co-authored-by: Isaac --- README.md | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/README.md b/README.md index 7e09b8a56..a72ef1352 100644 --- a/README.md +++ b/README.md @@ -115,22 +115,6 @@ ug mcp login --names system.ai.github ug mcp remove ``` -Here's what a typical session looks like — add two tools, check their status, then sign in: - -```text -$ ug mcp add --names system.ai.github,system.ai.slack -✔ Added 2 MCP servers across Claude Code, Codex - -$ ug mcp list - NAME LOCATION AGENTS STATUS - system-ai-github system.ai.github claude, codex needs sign-in - system-ai-slack system.ai.slack claude, codex needs sign-in - -$ ug mcp login --names system.ai.github -Signing in to 'system.ai.github' — opening your browser to finish sign-in… -✔ Signed in to 1 MCP service(s) -``` - **You only sign in once.** Some tools (like `system.ai.github`) ask you to sign in to the underlying service the first time. Do it once — through any agent or `ug mcp login` — and the tool works everywhere: Claude, Cursor, Codex, and the rest. @@ -140,10 +124,6 @@ tool works everywhere: Claude, Cursor, Codex, and the rest. - Limit any command to certain agents: add `--agents claude,codex`. - Add other AI Gateway tools by typed name: `vector-search:main.docs`, `uc-functions:main.tools`, `external:`, `genie-space:`, `app:`. -- Each tool runs locally through `ug mcp-proxy`, which keeps your Databricks sign-in fresh — - no tokens to copy or manage. -- Sign-in works for any connection-backed tool, not just `system.ai.*`, and needs a recent - Databricks CLI (`ug mcp login` tells you if yours is too old). From d6c47e7c19db1324fb98fb9232c4333f01ca4f2a Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Mon, 28 Sep 2026 22:52:32 +0000 Subject: [PATCH 5/5] style: apply current ruff format to ug mcp login files The formatter version in CI advanced since these files were first written, so ruff-format-check flagged them. Reformat to match; no behavior change. Co-authored-by: Isaac --- src/ucode/mcp_login.py | 4 +--- tests/test_mcp_login.py | 12 ++++++------ 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/src/ucode/mcp_login.py b/src/ucode/mcp_login.py index 3d479db5d..37442ac37 100644 --- a/src/ucode/mcp_login.py +++ b/src/ucode/mcp_login.py @@ -224,9 +224,7 @@ def login_mcp_command(names: set[str] | None = None, agents: set[str] | None = N # so an unreachable status API doesn't masquerade as "already signed in" (and so the picker and # the non-interactive path still act on it), rather than being silently skipped. pending = [ - full - for full in url_by_full - if status_by_full[full] in (STATUS_NEEDS_LOGIN, STATUS_UNKNOWN) + full for full in url_by_full if status_by_full[full] in (STATUS_NEEDS_LOGIN, STATUS_UNKNOWN) ] # Non-interactive (--names): sign in only to targeted services that still need it — never diff --git a/tests/test_mcp_login.py b/tests/test_mcp_login.py index e0f6693d4..7db350b18 100644 --- a/tests/test_mcp_login.py +++ b/tests/test_mcp_login.py @@ -279,14 +279,14 @@ def test_picker_excludes_no_login_services(self, monkeypatch): monkeypatch.setattr( mcp_login, "mcp_service_login_status", - lambda ws, t, full, u: mcp_login.STATUS_NO_LOGIN - if full == "system.ai.pg" - else mcp_login.STATUS_NEEDS_LOGIN, + lambda ws, t, full, u: ( + mcp_login.STATUS_NO_LOGIN + if full == "system.ai.pg" + else mcp_login.STATUS_NEEDS_LOGIN + ), ) rows: list = [] - monkeypatch.setattr( - mcp_login, "_prompt_login_selection", lambda r: rows.extend(r) or [] - ) + monkeypatch.setattr(mcp_login, "_prompt_login_selection", lambda r: rows.extend(r) or []) rc = mcp_login.login_mcp_command() assert rc == 0 offered = {full for full, _status in rows}