diff --git a/README.md b/README.md index 7234ecaf..c5769e22 100644 --- a/README.md +++ b/README.md @@ -274,6 +274,35 @@ summarized as a per-agent count. Skills connections are managed separately (via `ug configure skills`) and aren't listed here. It reads local state plus each installed agent's `mcp list`, so it needs no Databricks login. +#### Sign in to connection-backed servers + +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. Use +`ug mcp login` to see which of your configured MCP services are signed in vs. still need a +sign-in, and to complete the sign-in: + +```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 --services system.ai.github,system.ai.slack + +# Scope to specific agents' services. +ug mcp login --agents claude,codex +``` + +It uses the same configured-server set as `ug mcp list`, 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 (optional) Configure Unity Catalog Skills for your coding tools with `ug configure skills`: @@ -438,6 +467,8 @@ The output looks like: | `ug mcp remove --agents codex` | Unregister selected servers from specific agents only | | `ug mcp list` | List configured MCP servers and their live per-agent connection status | | `ug mcp list --agents claude` | Show the connection-status report for specific agents only | +| `ug mcp login` | Show configured connection-backed MCP services' sign-in status and sign in to the ones you pick | +| `ug mcp login --services system.ai.github` | Sign in to specific connection-backed MCP service(s) non-interactively | | `ug configure skills` | Register the skills MCP connection (utility tools only); no skills download | | `ug configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | | `ug configure skills --skill main.default.my-skill` | Download named skills by fully-qualified name (comma-separated; may span schemas) | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3a5d8c40..b0839fe0 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -114,6 +114,7 @@ revert_mcp_configs, skill_locations_for_client, ) +from ucode.mcp_login import login_mcp_command from ucode.skills_download import ( configure_location_skills_download_command, configure_selected_skills_download_command, @@ -1232,6 +1233,49 @@ def mcp_list( raise typer.Exit(130) from None +@mcp_app.command("login") +def mcp_login( + services: Annotated[ + str | None, + typer.Option( + "--services", + 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 --services 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 --services). Sign-in uses `databricks auth login --resource`, so it + works for any connection-backed MCP service (not just `system.ai.*`). + """ + selected = None if services is None else {s.strip() for s in services.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(services=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 5444be3c..fd20ce01 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1977,6 +1977,37 @@ 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"}``. + + Shared by ``ug mcp list`` and ``ug mcp login`` so both see the same configured-server set.""" + 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) + 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. @@ -2008,28 +2039,9 @@ 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 state.get("managed_mcp_servers") or []: - _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). + 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 00000000..0c290de7 --- /dev/null +++ b/src/ucode/mcp_login.py @@ -0,0 +1,315 @@ +"""`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 + +import subprocess +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 ( + AIGW_MCP_SERVICES_PATH, + configured_mcp_servers_by_name, +) +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 mcp_service_full_name_from_url(url: str) -> str | None: + """Extract the ``..`` name from an AI Gateway + mcp-services URL, or ``None`` if it isn't one.""" + if not isinstance(url, str) or AIGW_MCP_SERVICES_PATH not in url: + return None + tail = url.split(AIGW_MCP_SERVICES_PATH, 1)[1].strip("/") + # The service name is the first path segment after the mcp-services prefix. + name = tail.split("/", 1)[0].split("?", 1)[0] + return name or None + + +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_identity: 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. + """ + 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(user_identity, 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 run_connection_login( + mcp_url: str, workspace: str, profile: str | None = None, *, login_binary: str = "databricks" +) -> tuple[bool, str]: + """Run the Databricks CLI U2M login with an RFC 8707 ``--resource`` indicator + for one MCP service. A resource-aware ``/oidc`` routes it through the + connection's own SaaS sign-in before minting the token. The CLI opens the + browser and prints the authorization URL; this blocks until it completes. + + Returns ``(ok, detail)``. Requires a Databricks CLI with ``--resource`` + (databricks/cli#6621); an older CLI rejects the flag — surfaced as a clear + failure rather than a cryptic one. + """ + cmd = [login_binary, "auth", "login", "--host", workspace, "--resource", mcp_url] + if profile: + cmd += ["--profile", profile] + try: + result = subprocess.run(cmd, check=False, text=True, capture_output=True, timeout=600) + except FileNotFoundError: + return False, f"'{login_binary}' not found on PATH" + except subprocess.TimeoutExpired: + return False, "login timed out" + if result.returncode == 0: + return True, "signed in" + stderr = (result.stderr or "") + (result.stdout or "") + if "--resource" in stderr and ("unknown flag" in stderr or "unknown shorthand" in stderr): + return False, ( + "your Databricks CLI does not support `--resource` (needs databricks/cli#6621). " + "Upgrade the CLI and retry." + ) + return False, f"`{login_binary} auth login` failed (exit {result.returncode})" + + +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 = mcp_service_full_name_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(services: 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. ``--services`` 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 services is not None: + entries = [e for e in entries if e[0] in services or e[0].split(".")[-1] in services] + unknown = services - {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 + user_identity = (_scim_me(workspace, token) or {}).get("userName") + if not user_identity: + raise RuntimeError("Could not resolve the current Databricks user identity.") + + 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_identity) + for full, _url, _clients, _managed in entries + } + + url_by_full = {full: url for full, url, _clients, _managed in entries} + login_needing = [full for full in url_by_full if status_by_full[full] == STATUS_NEEDS_LOGIN] + + # Non-interactive (--services): sign in to every targeted service that needs it. + # Interactive: render the per-service status (same Table style as `ug mcp list`), then a picker. + if services is not None: + targets = [(f, url_by_full[f]) for f in url_by_full if status_by_full[f] != STATUS_NO_LOGIN] + 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 login_needing: + print_success("All configured MCP services are already signed in.") + return 0 + selected = _prompt_login_selection( + [(full, status_by_full[full]) for full, _u, _c, _m in entries] + ) + 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: + print_note(f"Signing in to {full}...") + ok, detail = run_connection_login(url, workspace, 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", + "run_connection_login", + "mcp_service_full_name_from_url", + "STATUS_AUTHENTICATED", + "STATUS_NEEDS_LOGIN", + "STATUS_NO_LOGIN", + "STATUS_UNKNOWN", +] diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 43b49c0c..ac8b9f79 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -3213,6 +3213,43 @@ def test_keyword_fallback_used_when_no_glyph(self): } +class TestConfiguredMcpServersByName: + """The shared enumeration used by both `ug mcp list` and `ug mcp login`.""" + + 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"] + + 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 00000000..ce3aa6cd --- /dev/null +++ b/tests/test_mcp_login.py @@ -0,0 +1,226 @@ +"""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 + +from unittest.mock import MagicMock + +from ucode import 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" + +_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 TestFullNameFromUrl: + def test_extracts_service_name(self): + assert mcp_login.mcp_service_full_name_from_url(URL) == FULL + + def test_ignores_query_and_trailing(self): + assert mcp_login.mcp_service_full_name_from_url(f"{URL}/?x=1") == FULL + + def test_non_mcp_service_url_is_none(self): + assert mcp_login.mcp_service_full_name_from_url(f"{WS}/api/2.0/mcp/external/foo") is None + + +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 + + +class TestRunConnectionLogin: + def test_success(self, monkeypatch): + monkeypatch.setattr( + mcp_login.subprocess, + "run", + lambda *a, **k: MagicMock(returncode=0, stderr="", stdout=""), + ) + ok, _ = mcp_login.run_connection_login(URL, WS, "p") + assert ok is True + + def test_old_cli_without_resource_flag_gives_clear_error(self, monkeypatch): + monkeypatch.setattr( + mcp_login.subprocess, + "run", + lambda *a, **k: MagicMock(returncode=1, stderr="unknown flag: --resource", stdout=""), + ) + ok, detail = mcp_login.run_connection_login(URL, WS, "p") + assert ok is False + assert "does not support `--resource`" in detail + + def test_missing_binary(self, monkeypatch): + def boom(*a, **k): + raise FileNotFoundError() + + monkeypatch.setattr(mcp_login.subprocess, "run", boom) + ok, detail = mcp_login.run_connection_login(URL, WS, "p") + assert ok is False and "not found" in detail + + def test_passes_host_and_resource(self, monkeypatch): + seen = {} + + def capture(cmd, *a, **k): + seen["cmd"] = cmd + return MagicMock(returncode=0, stderr="", stdout="") + + monkeypatch.setattr(mcp_login.subprocess, "run", capture) + mcp_login.run_connection_login(URL, WS, "prof") + assert seen["cmd"][:3] == ["databricks", "auth", "login"] + assert "--host" in seen["cmd"] and WS in seen["cmd"] + assert "--resource" in seen["cmd"] and URL in seen["cmd"] + assert seen["cmd"][-2:] == ["--profile", "prof"] + + +class TestLoginCommand: + 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: {"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(services={"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: {"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(services={"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: {"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(services={"slack"}) + assert rc == 0 and logged == [slack_url]