diff --git a/README.md b/README.md index b87ad144..a72ef135 100644 --- a/README.md +++ b/README.md @@ -96,36 +96,36 @@ 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 -``` +ug mcp add --names system.ai.github,system.ai.slack -Remove configured servers: +# See what's set up, and whether each is signed in. +ug mcp list -```bash +# 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 + +# Remove tools (plain form shows a picker). ug mcp remove -ug mcp remove --agents codex ``` -List configured servers and their connection status: +**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. -```bash -ug mcp list -ug mcp list --agents claude,codex -``` +
Advanced options + +- 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:`. -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:`. +
## Skills @@ -168,6 +168,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 df0a7e96..23668c24 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 fd3cfbd9..df87b067 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 00000000..37442ac3 --- /dev/null +++ b/src/ucode/mcp_login.py @@ -0,0 +1,283 @@ +"""`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 8cef86ba..90d53355 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 00000000..7db350b1 --- /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