diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py new file mode 100644 index 00000000..1d61a765 --- /dev/null +++ b/src/ucode/mcp_connection_login.py @@ -0,0 +1,239 @@ +"""Per-connection login for AI Gateway MCP services, driven by the proxy on a 401. + +A connection-backed AI Gateway MCP service (e.g. ``system.ai.github``) needs a +per-user connection credential before its tools can be used. Until the user has +logged in to the underlying SaaS, AI Gateway answers requests with an HTTP 401 +(RFC 9728 ``WWW-Authenticate``). + +The ``ug mcp-proxy`` bridge (see ``mcp_proxy``) sees that 401 in its httpx auth +flow and, for a connection-backed service, runs the Databricks CLI U2M login +with an RFC 8707 ``resource`` indicator naming the service, then retries the +request. A resource-aware ``/oidc`` drives the connection's own SaaS login +before minting the token, so the credential exists on retry — transparently to +the coding agent, which just sees the connection authenticate and succeed. This +is the behaviour of a generic OAuth MCP bridge (e.g. ``mcp-remote``), done in +ucode with the Databricks CLI so no extra library or per-agent OAuth app is +needed. Requires the CLI ``--resource`` flag (databricks/cli#6621). +""" + +from __future__ import annotations + +import subprocess +import sys +from urllib.parse import quote + +from ucode.databricks import ( + _http_get_json, + _scim_me, + get_databricks_token, + workspace_hostname, +) + +# AI Gateway MCP service endpoints look like +# ``https:///ai-gateway/mcp-services/..``. +AIGW_MCP_SERVICES_SEGMENT = "/ai-gateway/mcp-services/" + +# Login can pop a browser and wait for the user to complete the SaaS login, so +# allow generously more than a token refresh would take. +_LOGIN_TIMEOUT_SECONDS = 300 + +# Connection securable kinds that use a per-user OAuth (U2M) credential — i.e. the +# service needs a connection sign-in. Anything else (PAT/basic/service-managed) +# has no per-user login to drive. Mirrors the webapp's known-OAuth-kinds set. +_OAUTH_U2M_CONNECTION_KINDS = frozenset( + { + "CONNECTION_HTTP_OAUTH_U2M_MAPPING", + "CONNECTION_HTTP_DCR", + "CONNECTION_SLACK_OAUTH_U2M_MAPPING", + } +) + +# Credential-state outcomes for :func:`connection_credential_state`. +CREDENTIAL_PRESENT = "present" # signed in — skip the login +CREDENTIAL_MISSING = "missing" # confirmed no credential yet — run the login +CREDENTIAL_NO_LOGIN = "no_login_needed" # not an OAuth-U2M connection — skip +CREDENTIAL_UNKNOWN = ( + "unknown" # couldn't determine — skip (don't open a browser we're unsure about) +) + + +def connection_from_url(url: str) -> str | None: + """Return the connection FQN of an AI Gateway MCP service URL, or ``None``. + + ``https://ws/ai-gateway/mcp-services/system.ai.github`` -> ``system.ai.github``. + A URL that is not an mcp-services endpoint (or names no service) returns + ``None`` — only connection-backed services get the login-on-401 treatment. + """ + marker = url.find(AIGW_MCP_SERVICES_SEGMENT) + if marker == -1: + return None + tail = url[marker + len(AIGW_MCP_SERVICES_SEGMENT) :] + # Strip any trailing path (``/tools/list``), query, or fragment. + connection = tail.split("/")[0].split("?")[0].split("#")[0] + return connection or None + + +def _strip_connections_prefix(name: str | None) -> str | None: + """UC returns a connection reference as ``connections/``; the REST path + wants the bare name.""" + if not name: + return None + prefix = "connections/" + return name[len(prefix) :] if name.startswith(prefix) else name + + +def connection_credential_state( + resource_url: str, workspace: str, *, profile: str | None = None +) -> str: + """Best-effort per-user credential state for a connection-backed MCP service. + + Gates the connect-time login: the proxy should only run ``databricks auth + login`` (which opens a browser) when the credential is genuinely **missing**, + not on every session for a service the user already signed in to — otherwise + N configured servers would each pop a browser every session. + + Uses the same Unity Catalog REST APIs as ``ug mcp login``: resolve the + service's backing connection, then read the current user's credential + ``provisioning_info.state``. Returns one of ``CREDENTIAL_{PRESENT,MISSING, + NO_LOGIN,UNKNOWN}``. Anything but ``MISSING`` means "don't open a browser": + ``PRESENT``/``NO_LOGIN`` are definitive, and ``UNKNOWN`` (any probe failure) + deliberately fails safe — ``tools/list`` still works, and a tool call can + surface the login later — rather than risk a spurious browser. + """ + full = connection_from_url(resource_url) + if not full: + return CREDENTIAL_NO_LOGIN + try: + token = get_databricks_token(workspace, profile) + except Exception: # noqa: BLE001 - a dead token is reported by _preflight_token, not here + return CREDENTIAL_UNKNOWN + user = (_scim_me(workspace, token) or {}).get("userName") + if not user: + return CREDENTIAL_UNKNOWN + prefix = f"https://{workspace_hostname(workspace)}/api/2.1/unity-catalog" + details, err = _http_get_json(f"{prefix}/mcp-services/{quote(full, safe='')}", token) + if err is not None or not isinstance(details, dict): + return CREDENTIAL_UNKNOWN + source = (details.get("config") or {}).get("source_connection") or {} + if source.get("securable_kind") not in _OAUTH_U2M_CONNECTION_KINDS: + return CREDENTIAL_NO_LOGIN + conn = _strip_connections_prefix(source.get("name")) + service_id = details.get("id") + if not conn or not service_id: + return CREDENTIAL_UNKNOWN + cred_url = ( + f"{prefix}/connections/{quote(conn, safe='')}/user-credentials/" + f"{quote(user, 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 is the authoritative "no credential yet"; any other error is inconclusive. + return CREDENTIAL_MISSING if cred_err.startswith("HTTP 404") else CREDENTIAL_UNKNOWN + if not isinstance(cred, dict): + return CREDENTIAL_UNKNOWN + state = ((cred.get("connection_user_credential") or {}).get("provisioning_info") or {}).get( + "state" + ) + return CREDENTIAL_PRESENT if state == "ACTIVE" else CREDENTIAL_MISSING + + +def _cli_supports_resource_flag(login_binary: str) -> bool: + """Whether `` auth login`` advertises the ``--resource`` flag. + + The connection sign-in needs a Databricks CLI with ``--resource`` + (databricks/cli#6621). An older CLI rejects the flag and the login exits with + a cryptic parse error, so we check ``--help`` up front to give a clear message + instead. Fail-open (assume supported) if ``--help`` can't be run — the real + login attempt will surface any genuine failure.""" + try: + result = subprocess.run( + [login_binary, "auth", "login", "--help"], + check=False, + timeout=20, + capture_output=True, + text=True, + ) + except (OSError, subprocess.TimeoutExpired): + return True + return "--resource" in f"{result.stdout or ''}{result.stderr or ''}" + + +def run_connection_login( + resource_url: str, + workspace: str, + *, + profile: str | None = None, + login_binary: str = "databricks", +) -> tuple[bool, str]: + """Run the CLI U2M login with an RFC 8707 resource indicator for this service. + + ``resource_url`` is the MCP service endpoint (also the proxy's upstream URL); + it is sent as ``--resource`` so a resource-aware ``/oidc`` drives the + connection's SaaS login before issuing the token. Uses the Databricks CLI's + own default client, whose loopback redirect is already registered — no + ``--client-id`` needed. + + The CLI opens the browser to complete the login and prints the authorize URL. + We route its output to **stderr** (never stdout — that is the proxy's MCP + JSON-RPC wire), so a coding agent surfaces it in the server's log and the URL + stays visible when the browser can't open (e.g. a headless remote). Returns + ``(ok, message)``; on failure ``message`` points at that log. + """ + connection = connection_from_url(resource_url) or resource_url + if not _cli_supports_resource_flag(login_binary): + return False, ( + f"the Databricks CLI ('{login_binary}') has no `--resource` flag, so the " + f"'{connection}' connection sign-in can't run. Upgrade the CLI " + "(databricks/cli#6621) and retry." + ) + argv = [ + login_binary, + "auth", + "login", + "--host", + workspace.rstrip("/"), + "--resource", + resource_url, + ] + if profile: + argv += ["--profile", profile] + print( + f"ucode mcp-proxy: '{connection}' needs a one-time connection sign-in. Opening your " + "browser to complete it — if it doesn't open, use the authorization URL printed below.", + file=sys.stderr, + flush=True, + ) + try: + # stdout -> stderr: the CLI's prompts and authorize URL reach the agent's + # MCP log (fd 2) without corrupting this process's stdout (fd 1, the MCP + # JSON-RPC stream). stdin is closed since the flow is browser-driven. + result = subprocess.run( + argv, + check=False, + timeout=_LOGIN_TIMEOUT_SECONDS, + stdin=subprocess.DEVNULL, + stdout=sys.stderr, + stderr=sys.stderr, + ) + except OSError as exc: + return False, f"could not run '{login_binary} auth login': {exc}" + except subprocess.TimeoutExpired: + return False, "connection sign-in timed out waiting for the browser flow to complete" + if result.returncode == 0: + return True, "signed in" + return ( + False, + f"connection sign-in did not complete (CLI exited {result.returncode}; see the log above)", + ) + + +__all__ = [ + "AIGW_MCP_SERVICES_SEGMENT", + "CREDENTIAL_MISSING", + "CREDENTIAL_NO_LOGIN", + "CREDENTIAL_PRESENT", + "CREDENTIAL_UNKNOWN", + "connection_credential_state", + "connection_from_url", + "run_connection_login", +] diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index c0174169..da54c393 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -43,6 +43,12 @@ from mcp.server.stdio import stdio_server from ucode.databricks import ensure_pat_bearer, get_databricks_token +from ucode.mcp_connection_login import ( + CREDENTIAL_MISSING, + connection_credential_state, + connection_from_url, + run_connection_login, +) # Exit code used when the proxy cannot continue. MCP clients surface a non-zero # exit far more usefully than a timeout, so bail out instead of hanging. @@ -117,18 +123,24 @@ def _build_token_auth(workspace: str, profile: str | None): The base class comes from whichever httpx the SDK uses (see ``_httpx``), so the returned auth is accepted by that SDK's ``AsyncClient``. Behaviour is identical across flavours — ``Auth.auth_flow`` has the same generator - contract in httpx and httpx2.""" + contract in httpx and httpx2. + + The bearer is the Databricks *workspace* token, read from the session that + ``serve`` already established via ``databricks auth login --resource`` before + the bridge opened (so a connection-backed service is signed in, credential + and all — see ``serve``). This only *reads* that session's token (refreshing + it as it nears expiry); it never authenticates on its own, so it can't hand + the gateway a token that skips the connection login.""" httpx = _httpx() class _DatabricksTokenAuth(httpx.Auth): def auth_flow(self, request): # get_databricks_token honors the DATABRICKS_BEARER short-circuit and # PAT profiles internally; --use-pat is surfaced via the env ucode set. - # A RuntimeError here means auth is dead (expired refresh token, - # logged-out profile). Raising it from inside auth_flow would tear - # through the transport's task group and stall the process until the - # client times out, so translate it into a terminal ProxyAuthError the - # caller reports cleanly. + # A RuntimeError means the session is dead (expired refresh token, + # logged-out profile). Raising from inside auth_flow would tear through + # the transport's task group and stall the process until the client + # times out, so translate it into a terminal ProxyAuthError. try: token = get_databricks_token(workspace, profile) except RuntimeError as exc: @@ -237,6 +249,30 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool "Set DATABRICKS_BEARER, or reconfigure the profile." ) + # Connection-backed AI Gateway services need a per-user connection credential + # (e.g. a SaaS login) before their tools can be used. Drive that login *here*, + # before the bridge opens — a blocking `databricks auth login --resource + # `, which a resource-aware /oidc routes through the connection's own + # sign-in (/mcp-service-login) before minting the token. The agent blocks on + # "connecting…" while it runs (the browser opens, or the URL is printed to this + # stderr), then the session comes up already authenticated. + # + # But `databricks auth login` always re-runs the browser OAuth — it does NOT + # short-circuit on a cached token — so doing it unconditionally would pop a + # browser for EVERY connection-backed server on EVERY session. So gate it on + # the actual credential state: only log in when the credential is confirmed + # missing. When it's already present (or the state can't be determined) we + # skip the browser entirely — this is what makes N configured servers not + # equal N browser windows per session. PAT profiles have no connection OAuth + # to drive, so they skip it too. + connection = None if use_pat else connection_from_url(url) + if connection is not None: + state = connection_credential_state(url, workspace, profile=profile) + if state == CREDENTIAL_MISSING: + ok, detail = run_connection_login(url, workspace, profile=profile) + if not ok: + _fail_fast(f"connection login for '{connection}' failed: {detail}") + # Pre-flight the token before opening the bridge. Without this, the first # token failure surfaces from inside the transport's task group, where it can # stall the process instead of erroring out. diff --git a/tests/test_mcp_connection_login.py b/tests/test_mcp_connection_login.py new file mode 100644 index 00000000..afe1de91 --- /dev/null +++ b/tests/test_mcp_connection_login.py @@ -0,0 +1,191 @@ +"""Tests for the per-connection MCP login helpers (mcp_connection_login). + +Network-free: URL/connection parsing and the login runner with the subprocess +monkeypatched. +""" + +from __future__ import annotations + +import subprocess + +from ucode import mcp_connection_login as mcl + +WS = "https://ws.staging.cloud.databricks.com" +AIGW_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" + + +class TestConnectionFromUrl: + def test_plain_endpoint(self): + assert mcl.connection_from_url(AIGW_URL) == "system.ai.github" + + def test_with_trailing_path_and_query(self): + assert mcl.connection_from_url(f"{AIGW_URL}/tools/list?x=1") == "system.ai.github" + + def test_non_aigw_url_is_none(self): + assert mcl.connection_from_url(f"{WS}/api/2.0/mcp/functions/system/ai") is None + + def test_missing_service_is_none(self): + assert mcl.connection_from_url(f"{WS}/ai-gateway/mcp-services/") is None + + +class TestConnectionCredentialState: + """The gate that stops the proxy re-running the browser login every session.""" + + _DETAILS = { + "id": "svc-1", + "config": { + "source_connection": { + "name": "connections/github", + "securable_kind": "CONNECTION_HTTP_OAUTH_U2M_MAPPING", + } + }, + } + + def _wire(self, monkeypatch, *, details=_DETAILS, details_err=None, cred=None, cred_err=None): + monkeypatch.setattr(mcl, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcl, "_scim_me", lambda ws, token: {"userName": "u@databricks.com"}) + + def _http(url, token, *a, **k): + if "/mcp-services/" in url: + return details, details_err + if "/user-credentials/" in url: + return cred, cred_err + return None, "unexpected url" + + monkeypatch.setattr(mcl, "_http_get_json", _http) + + def test_present_when_credential_active(self, monkeypatch): + self._wire( + monkeypatch, + cred={"connection_user_credential": {"provisioning_info": {"state": "ACTIVE"}}}, + ) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_PRESENT + + def test_missing_on_404(self, monkeypatch): + self._wire(monkeypatch, cred_err="HTTP 404 Not Found") + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_MISSING + + def test_missing_when_state_not_active(self, monkeypatch): + self._wire( + monkeypatch, + cred={"connection_user_credential": {"provisioning_info": {"state": "PROVISIONING"}}}, + ) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_MISSING + + def test_no_login_for_non_oauth_connection(self, monkeypatch): + details = { + "id": "x", + "config": { + "source_connection": {"name": "connections/c", "securable_kind": "CONNECTION_MYSQL"} + }, + } + self._wire(monkeypatch, details=details) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_NO_LOGIN + + def test_unknown_on_service_lookup_error(self, monkeypatch): + self._wire(monkeypatch, details=None, details_err="HTTP 500") + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_UNKNOWN + + def test_unknown_on_non_404_credential_error(self, monkeypatch): + self._wire(monkeypatch, cred_err="HTTP 403 Forbidden") + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_UNKNOWN + + def test_unknown_when_token_unavailable(self, monkeypatch): + def boom(ws, profile): + raise RuntimeError("no token") + + monkeypatch.setattr(mcl, "get_databricks_token", boom) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_UNKNOWN + + def test_no_login_for_non_aigw_url(self, monkeypatch): + assert ( + mcl.connection_credential_state(f"{WS}/api/2.0/mcp/functions/x", WS) + == mcl.CREDENTIAL_NO_LOGIN + ) + + +class TestRunConnectionLogin: + def _fake_run(self, captured, *, returncode, stderr=""): + def _run(argv, **kwargs): + # The `--resource`-support pre-check runs `auth login --help` first. + if "--help" in argv: + return subprocess.CompletedProcess( + argv, 0, stdout="--resource stringArray", stderr="" + ) + captured.append(argv) + return subprocess.CompletedProcess(argv, returncode, stdout="", stderr=stderr) + + return _run + + def test_success_sends_resource_and_host_without_client_id(self, monkeypatch): + captured: list[list[str]] = [] + monkeypatch.setattr(mcl.subprocess, "run", self._fake_run(captured, returncode=0)) + + ok, message = mcl.run_connection_login(AIGW_URL, WS, profile="p") + + assert ok and message == "signed in" + argv = captured[0] + assert argv[:3] == ["databricks", "auth", "login"] + assert "--resource" in argv and AIGW_URL in argv + assert "--host" in argv and WS in argv + assert "--profile" in argv and "p" in argv + # Uses the CLI's default client (its own registered redirect), so no --client-id. + assert "--client-id" not in argv + + def test_nonzero_exit_reports_failure(self, monkeypatch): + # The CLI's own output streams live to stderr (the agent's MCP log), so on + # failure we return a pointer to that log rather than captured text. + captured: list[list[str]] = [] + monkeypatch.setattr(mcl.subprocess, "run", self._fake_run(captured, returncode=1)) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok + assert "did not complete" in message and "1" in message + + def test_output_is_routed_to_stderr_not_stdout(self, monkeypatch): + # stdout must never be captured to the proxy's stdout (the MCP wire); the + # CLI's URL/prompts go to this process's stderr. + seen: dict = {} + + def _run(argv, **kw): + if "--help" in argv: # the --resource pre-check; not the login call under test + return subprocess.CompletedProcess(argv, 0, stdout="--resource", stderr="") + seen.update(kw) + return subprocess.CompletedProcess(argv, 0) + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, _ = mcl.run_connection_login(AIGW_URL, WS) + assert ok + assert seen.get("stdout") is mcl.sys.stderr + assert seen.get("stderr") is mcl.sys.stderr + assert "capture_output" not in seen + + def test_timeout_is_reported(self, monkeypatch): + def _run(argv, **kwargs): + raise subprocess.TimeoutExpired(argv, 1) + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok and "timed out" in message + + def test_binary_missing_is_reported(self, monkeypatch): + def _run(argv, **kwargs): + raise OSError("not found") + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok and "could not run" in message + + def test_old_cli_without_resource_flag_reports_clearly(self, monkeypatch): + # `auth login --help` lacking `--resource` => an old CLI (no databricks/cli#6621). + # We must report that clearly and never attempt the login (the flag would error). + def _run(argv, **kwargs): + if "--help" in argv: + return subprocess.CompletedProcess( + argv, 0, stdout="usage: login [--host]", stderr="" + ) + raise AssertionError("login must not run when --resource is unsupported") + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok + assert "--resource" in message and "Upgrade" in message diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index 0ded63d9..ec883115 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -10,6 +10,7 @@ import httpx import pytest +from ucode import mcp_connection_login as mcl from ucode import mcp_proxy WS = "https://example.databricks.com" @@ -128,6 +129,9 @@ def boom(ws, profile): list(auth.auth_flow(httpx.Request("POST", URL))) +CONN_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" + + class TestPump: def test_forwards_all_messages_in_order(self): async def scenario() -> list[str]: @@ -392,6 +396,111 @@ def raise_other(func, *args): with pytest.raises(ValueError, match="some transport bug"): mcp_proxy.serve(URL, WS, "p") + def test_connection_backed_url_logs_in_before_the_bridge_when_credential_missing( + self, monkeypatch + ): + # A connection-backed mcp-services URL whose credential is MISSING drives + # `databricks auth login --resource` up front (blocking), then opens the + # bridge — so the session is authenticated before AI Gateway is called. + order: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "connection_credential_state", lambda *a, **k: mcp_proxy.CREDENTIAL_MISSING + ) + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda url, ws, **k: ( + order.append(("login", url, k.get("profile"))) or (True, "signed in") + ), + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: order.append(("bridge",))) + + mcp_proxy.serve(CONN_URL, WS, "p") + + assert order == [("login", CONN_URL, "p"), ("bridge",)] + + def test_present_credential_skips_the_login_and_opens_the_bridge(self, monkeypatch): + # The key scaling fix: an already-signed-in connection must NOT re-run the + # browser login — otherwise N configured servers pop N browsers per session. + logins: list = [] + opened: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "connection_credential_state", lambda *a, **k: mcl.CREDENTIAL_PRESENT + ) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: opened.append("bridge")) + + mcp_proxy.serve(CONN_URL, WS, "p") + + assert logins == [] # no browser + assert opened == ["bridge"] # session still comes up + + def test_unknown_or_no_login_state_skips_the_login(self, monkeypatch): + # UNKNOWN (probe failed) and NO_LOGIN (non-OAuth connection) both fail safe: + # no eager browser — tools/list still works and a tool call can surface it later. + for state in (mcl.CREDENTIAL_UNKNOWN, mcl.CREDENTIAL_NO_LOGIN): + logins: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "connection_credential_state", lambda *a, s=state, **k: s + ) + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, _l=logins, **k: _l.append(1) or (True, ""), + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(CONN_URL, WS, "p") + + assert logins == [], f"state={state} must not open a browser" + + def test_non_connection_url_skips_the_login(self, monkeypatch): + logins: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(URL, WS, "p") # URL is not an mcp-services endpoint + + assert logins == [] + + def test_connection_login_failure_exits_before_the_bridge(self, monkeypatch): + started: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "connection_credential_state", lambda *a, **k: mcp_proxy.CREDENTIAL_MISSING + ) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: started.append("bridge")) + + with pytest.raises(SystemExit) as excinfo: + mcp_proxy.serve(CONN_URL, WS, "p") + + assert excinfo.value.code == mcp_proxy.AUTH_FAILURE_EXIT_CODE + assert started == [] # never opened the bridge + + def test_use_pat_skips_the_connection_login(self, monkeypatch): + logins: list = [] + monkeypatch.setattr(mcp_proxy, "ensure_pat_bearer", lambda profile: True) + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(CONN_URL, WS, "p", use_pat=True) + + assert logins == [] # PAT has no connection OAuth to drive + class TestPreflightToken: def test_passes_through_when_a_token_is_available(self, monkeypatch):