diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 177805e03..cfbe67eea 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -107,7 +107,6 @@ ) from ucode.mcp import ( MCP_CLIENTS, - SKILLS_MCP_KIND, McpServiceListingRateLimited, add_mcp_command, add_skills_command, @@ -116,8 +115,8 @@ configure_mcp_command, configure_skills_mcp_picker_command, configured_mcp_clients, + configured_mcp_servers_by_name, list_mcp_command, - managed_mcp_server_names, purge_cross_workspace_mcp_residue, reconcile_managed_mcp_servers, remove_mcp_command, @@ -1169,8 +1168,9 @@ def status() -> int: state = load_state() workspace = state.get("workspace") managed_configs = state.get("managed_configs") or {} - # Both developer- and workspace-managed servers, so the count agrees with `ug mcp list`. - mcp_servers = (state.get("mcp_servers") or []) + (state.get("managed_mcp_servers") or []) + # The one enumerator `ug mcp list` uses, keyed by name with each server's agents — so the + # per-agent counts below can't drift from what `ug mcp list` reports. + configured_mcp = configured_mcp_servers_by_name(state) cached_managed = load_managed_state(workspace) if workspace else None managed, managed_freshness = _live_status_managed_state(state, cached_managed) configured_tools = ( @@ -1242,18 +1242,11 @@ def status() -> int: ) if tool in MCP_CLIENTS: # High-level overview: just a count per agent. `ug mcp list` (see the note below) shows - # the per-server detail and live connection status, so status stays scannable. Dedupe by - # name so a server present in both mcp_servers and managed_mcp_servers isn't double-counted. - mcp_names = { - server.get("name") - for server in mcp_servers - if tool in (server.get("clients") or []) - and server.get("name") - and server.get("kind") != SKILLS_MCP_KIND - } - # Managed servers ug delivers through an OS-managed file live in that file, not state. - mcp_names |= managed_mcp_server_names(state, {tool}) - rows.append(("MCP servers", str(len(mcp_names)))) + # the per-server detail and live connection status, so status stays scannable. The shared + # enumerator already dedupes by name and folds in servers delivered through an agent's + # OS-managed file (Claude/Codex), so this count matches `ug mcp list` by construction. + mcp_count = sum(1 for entry in configured_mcp.values() if tool in entry["clients"]) + rows.append(("MCP servers", str(mcp_count))) rows.append(("Skills", str(skill_counts_by_agent.get(tool, 0)))) base_url = state.get("base_urls", {}).get(tool) if isinstance(base_url, dict): diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 9fc1616c8..fd4eaab09 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2493,6 +2493,42 @@ 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 (the ``managed_mcp_servers`` fallback state and each agent's OS-managed file) + come from the shared ``managed_mcp_servers`` helper, so this stays in agreement with ``ug mcp + list`` and the ``ug mcp add`` picker. Used by ``ug mcp list`` and ``ug status`` so both 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) + # Fallback ``managed_mcp_servers`` state + each agent's OS-managed file, via the shared helper + # (keeps the isinstance guard and _MANAGED_FILE_AGENTS set, so this can't drift from `ug mcp list`). + for server in managed_mcp_servers(state, agents): + _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. @@ -2524,28 +2560,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 status` 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_proxy.py b/src/ucode/mcp_proxy.py index e14c91edf..e2cf2fc22 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -132,9 +132,9 @@ def _build_token_auth(url: str, workspace: str, profile: str | None, *, use_pat: an already-signed-in service never prompts. The login runs at most once per session; PAT profiles have no connection OAuth to drive, so they never do it. - The base class comes from whichever httpx the SDK uses (see ``_httpx``); the - sync and async flavours share the same generator contract. The async client - uses ``async_auth_flow``, so the blocking login runs off the event loop.""" + The base class comes from whichever httpx the SDK uses (see ``_httpx``). Auth + is async-only: the sync ``auth_flow`` raises, and the async client uses + ``async_auth_flow``, so the blocking login runs off the event loop.""" httpx = _httpx() # PAT auth has no interactive OAuth to drive, so never treat it as connection-backed. connection = None if use_pat else connection_from_url(url) @@ -182,15 +182,13 @@ def _connection_login_or_fail() -> bool: return True class _DatabricksTokenAuth(httpx.Auth): - def auth_flow(self, request): - _mint(request) - response = yield request - if not _needs_connection_login(response): - return - if not _connection_login_or_fail(): - return - _mint(request) - yield request + def sync_auth_flow(self, request): + # The proxy always drives this Auth from an httpx AsyncClient (see `_run`), so only + # `async_auth_flow` is implemented. Fail loudly on the sync entry point so a future sync + # client can't silently fall back to httpx's no-op default flow (which skips the bearer) + # instead of injecting the token. + raise ProxyAuthError("mcp-proxy auth is async-only; use it with an httpx AsyncClient") + yield request # unreachable — present only so httpx treats this as a generator async def async_auth_flow(self, request): _mint(request) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8cef86bab..85b5ddcb8 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 status`.""" + + @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 status`/`ug mcp 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_proxy.py b/tests/test_mcp_proxy.py index 20e11dd4b..cbfc857df 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -64,8 +64,7 @@ def test_injects_bearer_from_minted_token(self, monkeypatch): auth = mcp_proxy._build_token_auth(URL, WS, "uc-dogfood") request = httpx.Request("POST", URL) - # auth_flow is a generator that yields the (mutated) request. - list(auth.auth_flow(request)) + self._drive(auth, request) assert request.headers["Authorization"] == "Bearer tok-123" @@ -86,7 +85,7 @@ def test_calls_get_token_with_workspace_and_profile(self, monkeypatch): ) auth = mcp_proxy._build_token_auth(URL, WS, "myprofile") - list(auth.auth_flow(httpx.Request("POST", URL))) + self._drive(auth, httpx.Request("POST", URL)) assert calls == [(WS, "myprofile")] @@ -99,8 +98,8 @@ def test_mints_a_fresh_token_per_request(self, monkeypatch): r1 = httpx.Request("POST", URL) r2 = httpx.Request("POST", URL) - list(auth.auth_flow(r1)) - list(auth.auth_flow(r2)) + self._drive(auth, r1) + self._drive(auth, r2) assert r1.headers["Authorization"] == "Bearer first" assert r2.headers["Authorization"] == "Bearer second" @@ -110,7 +109,7 @@ def test_auth_flow_yields_the_same_request(self, monkeypatch): auth = mcp_proxy._build_token_auth(URL, WS, None) request = httpx.Request("POST", URL) - yielded = list(auth.auth_flow(request)) + yielded = self._drive(auth, request) assert yielded == [request] @@ -125,21 +124,30 @@ def boom(ws, profile): auth = mcp_proxy._build_token_auth(URL, WS, "p") with pytest.raises(mcp_proxy.ProxyAuthError, match="databricks auth login"): - list(auth.auth_flow(httpx.Request("POST", URL))) + self._drive(auth, httpx.Request("POST", URL)) # --- lazy on-401 connection login --------------------------------------- @staticmethod - def _run_flow(auth, request, response): - """Drive the sync auth_flow, feeding `response` after the first yield. - Returns the yielded requests (1 = no retry, 2 = logged in and retried).""" - gen = auth.auth_flow(request) - yielded = [next(gen)] - try: - yielded.append(gen.send(response)) - except StopIteration: - pass - return yielded + def _drive(auth, request, response=None): + """Drive `async_auth_flow` to completion, feeding `response` after the first yield. + Returns the yielded requests (1 = no retry, 2 = logged in and retried). The proxy only ever + uses the async client, so this exercises the real path (the blocking login runs off the + event loop via anyio.to_thread).""" + + async def run(): + gen = auth.async_auth_flow(request) + yielded: list = [] + try: + yielded.append(await gen.__anext__()) + yielded.append(await gen.asend(response)) + except StopAsyncIteration: + pass + finally: + await gen.aclose() + return yielded + + return anyio.run(run) def test_on_401_from_connection_service_runs_login_then_retries(self, monkeypatch): monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") @@ -151,7 +159,7 @@ def test_on_401_from_connection_service_runs_login_then_retries(self, monkeypatc lambda url, ws, **k: logins.append((url, k.get("profile"))) or (True, "signed in"), ) auth = mcp_proxy._build_token_auth(CONN_URL, WS, "p") - yielded = self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + yielded = self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) assert logins == [(CONN_URL, "p")] # login driven once, only on the 401 assert len(yielded) == 2 # retried after signing in assert yielded[1].headers["Authorization"] == "Bearer tok" @@ -163,7 +171,7 @@ def test_non_connection_401_is_not_retried(self, monkeypatch): mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") ) auth = mcp_proxy._build_token_auth(URL, WS, "p") # URL is not an mcp-services endpoint - yielded = self._run_flow(auth, httpx.Request("POST", URL), httpx.Response(401)) + yielded = self._drive(auth, httpx.Request("POST", URL), httpx.Response(401)) assert logins == [] and len(yielded) == 1 def test_success_response_never_logs_in(self, monkeypatch): @@ -173,7 +181,7 @@ def test_success_response_never_logs_in(self, monkeypatch): mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") ) auth = mcp_proxy._build_token_auth(CONN_URL, WS, "p") - yielded = self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(200)) + yielded = self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(200)) assert logins == [] and len(yielded) == 1 # tools/list etc. never trigger a browser def test_no_login_service_does_not_drive_login_on_401(self, monkeypatch): @@ -187,7 +195,7 @@ def test_no_login_service_does_not_drive_login_on_401(self, monkeypatch): ) ws_url = f"{WS}/ai-gateway/mcp-services/system.ai.web_search" auth = mcp_proxy._build_token_auth(ws_url, WS, "p") - yielded = self._run_flow(auth, httpx.Request("POST", ws_url), httpx.Response(401)) + yielded = self._drive(auth, httpx.Request("POST", ws_url), httpx.Response(401)) assert logins == [] and len(yielded) == 1 # token-only; no resource login driven def test_login_runs_at_most_once_per_session(self, monkeypatch): @@ -199,8 +207,8 @@ def test_login_runs_at_most_once_per_session(self, monkeypatch): ) auth = mcp_proxy._build_token_auth(CONN_URL, WS, "p") # Two 401s in the same session — the browser login must fire only once. - self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) - self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) assert logins == [1] def test_use_pat_never_drives_connection_login(self, monkeypatch): @@ -210,7 +218,7 @@ def test_use_pat_never_drives_connection_login(self, monkeypatch): mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") ) auth = mcp_proxy._build_token_auth(CONN_URL, WS, "p", use_pat=True) - yielded = self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + yielded = self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) assert logins == [] and len(yielded) == 1 # PAT has no connection OAuth to drive def test_login_failure_becomes_a_proxy_auth_error(self, monkeypatch): @@ -220,10 +228,8 @@ def test_login_failure_becomes_a_proxy_auth_error(self, monkeypatch): mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") ) auth = mcp_proxy._build_token_auth(CONN_URL, WS, "p") - gen = auth.auth_flow(httpx.Request("POST", CONN_URL)) - next(gen) with pytest.raises(mcp_proxy.ProxyAuthError, match="user cancelled"): - gen.send(httpx.Response(401)) + self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) def test_async_auth_flow_drives_login_on_401(self, monkeypatch): # The proxy uses the async client, so async_auth_flow is the real path; the @@ -235,18 +241,17 @@ def test_async_auth_flow_drives_login_on_401(self, monkeypatch): mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") ) auth = mcp_proxy._build_token_auth(CONN_URL, WS, "p") + self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + assert logins == [1] - async def scenario(): - gen = auth.async_auth_flow(httpx.Request("POST", CONN_URL)) - await gen.__anext__() - try: - await gen.asend(httpx.Response(401)) - except StopAsyncIteration: - pass - await gen.aclose() - return logins + def test_sync_auth_flow_is_rejected(self, monkeypatch): + # The proxy is async-only; a sync client must fail loudly rather than silently skip the + # bearer via httpx's no-op default flow. + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "t") + auth = mcp_proxy._build_token_auth(URL, WS, None) - assert anyio.run(scenario) == [1] + with pytest.raises(mcp_proxy.ProxyAuthError, match="async-only"): + list(auth.sync_auth_flow(httpx.Request("POST", URL))) CONN_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github"