From ad880f5b0a753c16be7f6951f8ff4c83e4924c55 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 6 Oct 2026 22:49:25 +0000 Subject: [PATCH 1/4] mcp: extract shared configured_mcp_servers_by_name enumerator Factor the configured-MCP-server enumeration (merge developer- + workspace-managed servers by registered name, union each server's agents, and read the Claude/Codex OS-managed files) out of `list_mcp_command` into a shared `configured_mcp_servers_by_name`, so `ug status` can reuse it instead of carrying a parallel copy that can drift. No behavior change for `ug mcp list`. Previously this refactor was bundled with the `ug mcp login` command; carved out here so it (and the `ug status` cleanup) can land independently of that command. Co-authored-by: Isaac --- src/ucode/mcp.py | 67 +++++++++++++++++++++++++++++++---------------- tests/test_mcp.py | 66 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 22 deletions(-) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 9fc1616c8..2001a5172 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2493,6 +2493,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. @@ -2524,28 +2565,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/tests/test_mcp.py b/tests/test_mcp.py index 8cef86bab..90d53355f 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -4001,6 +4001,72 @@ def test_authentication_failure_glyph_stays_failed(self): assert mcp.parse_mcp_list_output("claude", out) == {"svc": mcp.LIVE_FAILED} +class TestConfiguredMcpServersByName: + """The shared enumeration used by both `ug mcp list` and `ug mcp login`.""" + + @pytest.fixture(autouse=True) + def _no_managed_files(self, monkeypatch): + # Default: no OS-managed-file servers, so state-only cases are deterministic. + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", dict) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", dict) + + def _state(self): + return { + "mcp_servers": [ + {"name": "system-ai-github", "url": "u1", "clients": ["claude", "codex"]}, + { + "name": mcp.SKILLS_MCP_SERVER_NAME, + "kind": mcp.SKILLS_MCP_KIND, + "clients": ["claude"], + }, + ], + "managed_mcp_servers": [ + {"name": "system-ai-github", "url": "u1", "clients": ["cursor"]}, + {"name": "databricks-genie-abc", "url": "u2", "clients": ["claude"]}, + ], + } + + def test_merges_by_name_and_unions_clients(self): + by_name = mcp.configured_mcp_servers_by_name(self._state()) + assert set(by_name) == {"system-ai-github", "databricks-genie-abc"} + gh = by_name["system-ai-github"] + # Developer + workspace-managed entries unioned; managed flag sticks. + assert gh["clients"] == ["claude", "codex", "cursor"] + assert gh["managed"] is True + + def test_excludes_skills_connection(self): + assert mcp.SKILLS_MCP_SERVER_NAME not in mcp.configured_mcp_servers_by_name(self._state()) + + def test_agents_scope_drops_servers_with_no_in_scope_agent(self): + by_name = mcp.configured_mcp_servers_by_name(self._state(), agents={"cursor"}) + # Only the github service has a cursor client; genie (claude-only) is dropped. + assert set(by_name) == {"system-ai-github"} + assert by_name["system-ai-github"]["clients"] == ["cursor"] + + def test_includes_servers_delivered_via_os_managed_files(self, monkeypatch): + # Managed servers written to the agents' OS-managed files (Claude/Codex) live in those + # files, not state, so `ug mcp login`/`list` must still see them via the managed-file read. + url = f"{WS}/ai-gateway/mcp-services/system.ai.slack" + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + by_name = mcp.configured_mcp_servers_by_name({"mcp_servers": [], "managed_mcp_servers": []}) + assert "system-ai-slack" in by_name + entry = by_name["system-ai-slack"] + # Delivered to both agents' managed files → unioned, flagged managed, URL preserved. + assert entry["clients"] == ["claude", "codex"] + assert entry["managed"] is True + assert entry["server"]["url"] == url + + def test_managed_file_read_respects_agents_scope(self, monkeypatch): + url = f"{WS}/ai-gateway/mcp-services/system.ai.slack" + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", lambda: {"system-ai-slack": url}) + by_name = mcp.configured_mcp_servers_by_name( + {"mcp_servers": [], "managed_mcp_servers": []}, agents={"codex"} + ) + assert by_name["system-ai-slack"]["clients"] == ["codex"] + + class TestListMcpCommand: def _state(self): # A developer-added AI Gateway service on claude+codex, a workspace-managed one, and a From 82eec3a80d421a317be6ee7ea871425d4ea53178 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 18:42:10 +0000 Subject: [PATCH 2/4] ug status: count MCP servers via the shared enumerator (no re-implementation) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ug status` re-implemented the configured-MCP-server enumeration inline: it merged mcp_servers + managed_mcp_servers, deduped by name, and separately read the Claude/Codex OS-managed files — a parallel copy of `configured_mcp_servers_by_name`, the function `ug mcp list` and `ug mcp login` already use. The two could silently drift (the code even carried a comment asking a reader to keep them in agreement). Route the per-agent MCP count through `configured_mcp_servers_by_name` so all three commands read one enumerator. Net -9 lines and the count now matches `ug mcp list` by construction, including servers delivered through an agent's OS-managed file. No behavior change; existing TestStatus coverage (incl. managed-server + dedupe) stays green. Co-authored-by: Isaac --- src/ucode/cli.py | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 177805e03..e5813290a 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` / `ug mcp login` use, keyed by name with each server's + # agents — so the per-agent counts below can't drift from what those commands report. + 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): From 562659aa2accbf5c5e0e5f790c0f121619ad0989 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 22:21:33 +0000 Subject: [PATCH 3/4] mcp-proxy: async-only auth (drop the dead sync auth_flow) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The proxy only ever drives the httpx Auth from an AsyncClient, so the sync auth_flow was never exercised at runtime — it duplicated the async flow and existed only for the tests to drive synchronously. Remove it and override sync_auth_flow to raise a clear ProxyAuthError, so a future sync client fails loudly instead of silently falling back to httpx's no-op default flow (which would skip the bearer). Tests now drive async_auth_flow through one _drive helper (anyio.run), and a new test_sync_auth_flow_is_rejected covers the guard. Co-authored-by: Isaac --- src/ucode/mcp_proxy.py | 16 ++++----- tests/test_mcp_proxy.py | 79 ++++++++++++++++++++++------------------- 2 files changed, 49 insertions(+), 46 deletions(-) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index e14c91edf..15e6eb47a 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -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_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" From 3af1661db9e1a689de9bb678edddb1c6b8eaf489 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 6 Oct 2026 23:16:21 +0000 Subject: [PATCH 4/4] Address review: reuse managed_mcp_servers helper; drop stale ug mcp login refs - configured_mcp_servers_by_name now delegates managed-server enumeration to the existing managed_mcp_servers(state, agents) helper instead of hand-rolling the claude/codex OS-managed-file read. Restores the isinstance guard and uses _MANAGED_FILE_AGENTS, so ug status can't drift from `ug mcp list` / the add picker. - Drop references to the (closed, nonexistent) `ug mcp login` command in docstrings and comments; the real callers are `ug mcp list` and `ug status`. - _build_token_auth docstring: auth is async-only now (the sync auth_flow raises). Co-authored-by: Isaac --- src/ucode/cli.py | 4 ++-- src/ucode/mcp.py | 21 ++++++++------------- src/ucode/mcp_proxy.py | 6 +++--- tests/test_mcp.py | 4 ++-- 4 files changed, 15 insertions(+), 20 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index e5813290a..cfbe67eea 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1168,8 +1168,8 @@ def status() -> int: state = load_state() workspace = state.get("workspace") managed_configs = state.get("managed_configs") or {} - # The one enumerator `ug mcp list` / `ug mcp login` use, keyed by name with each server's - # agents — so the per-agent counts below can't drift from what those commands report. + # 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) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 2001a5172..fd4eaab09 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2501,10 +2501,10 @@ def configured_mcp_servers_by_name( 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.""" + 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: @@ -2522,15 +2522,10 @@ def _collect(server: dict, *, managed: bool) -> None: for server in state.get("mcp_servers") or []: _collect(server, managed=False) - for server in state.get("managed_mcp_servers") or []: + # 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) - # 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 @@ -2566,7 +2561,7 @@ 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 - # (shared with `ug mcp login` so both see the same configured-server set, including the servers + # (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) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 15e6eb47a..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) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 90d53355f..85b5ddcb8 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -4002,7 +4002,7 @@ def test_authentication_failure_glyph_stays_failed(self): class TestConfiguredMcpServersByName: - """The shared enumeration used by both `ug mcp list` and `ug mcp login`.""" + """The shared enumeration used by both `ug mcp list` and `ug status`.""" @pytest.fixture(autouse=True) def _no_managed_files(self, monkeypatch): @@ -4045,7 +4045,7 @@ def test_agents_scope_drops_servers_with_no_in_scope_agent(self): 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. + # 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})