MCP follow-ups: status count via shared enumerator + async-only proxy auth - #810
Open
sunishsheth2009 wants to merge 13 commits into
Open
sunishsheth2009 wants to merge 13 commits into
sunishsheth2009 wants to merge 13 commits into
Conversation
sunishsheth2009
force-pushed
the
mcp-status-single-enumerator
branch
from
September 23, 2026 20:47
70b6f0f to
911bc79
Compare
Give every coding agent a working login flow for connection-backed AI Gateway mcp-services endpoints, done in the stdio proxy every agent already spawns — no new library (cf. mcp-remote) and no per-agent OAuth app. When AI Gateway has no per-user connection credential it answers with HTTP 401 (RFC 9728). The proxy's httpx auth hook already sees every response, so on a 401 for a connection-backed URL it runs the Databricks CLI U2M login with an RFC 8707 resource indicator (`databricks auth login --resource <mcp-url>`, using the CLI's own registered redirect — no --client-id), then retries with a fresh token. A resource-aware /oidc drives the connection's SaaS login before minting the token, so the retry succeeds — transparently to the agent, which just sees the request authenticate rather than a failed tools/list. A later credential revoke re-triggers the login on the next 401. New module mcp_connection_login holds connection_from_url + run_connection_login; mcp_proxy._build_token_auth gains the login-on-401 retry. Unit-tested. Depends on the CLI --resource flag (databricks/cli#6621) and /oidc resource handling (login). Co-authored-by: Isaac <no-reply@databricks.com>
Fixes the two defects that made the proxy hang "connecting…" on a 401 instead of behaving like a generic OAuth MCP bridge (mcp-remote): 1. Non-blocking: the browser login ran via a synchronous subprocess inside the sync httpx auth_flow, which the async client executes on the event-loop thread — freezing the transport (stdio pumps included) for the whole login. Add async_auth_flow that offloads run_connection_login to a worker thread (anyio.to_thread.run_sync), so the loop stays responsive and cancellable while the user completes the browser flow. sync auth_flow kept for parity; both share the decision + login logic. 2. Visible URL: run_connection_login captured the CLI's output, hiding the authorize URL. Route the CLI's stdout+stderr to the proxy's stderr (fd 2, the agent's MCP log) — never fd 1 (the JSON-RPC wire) — and let the CLI open the browser, so the login is discoverable exactly like mcp-remote's. Co-authored-by: Isaac <no-reply@databricks.com>
Makes the generic proxy behave like a generic OAuth MCP bridge (mcp-remote): the connection login happens while the agent shows "connecting…", and the browser opens on its own — instead of racing the agent's tools/list timeout or burying the URL. - Connect-time login: before opening the bridge, serve() probes the connection (_connection_login_required: a lightweight initialize + tools/list to AI Gateway); on a 401 it drives run_connection_login *then*, so the agent's session comes up already authenticated. The on-401 retry in the auth hook stays as a mid-session fallback (credential revoked while connected). PAT profiles skip it (no connection OAuth). - Browser auto-open: the login inherits the environment (incl. $BROWSER), so databricks-cli opens the browser on the user's machine; the authorize URL is the printed fallback. Co-authored-by: Isaac <no-reply@databricks.com>
… the probe Simplify to what a generic OAuth MCP bridge (mcp-remote) does: authenticate at connect, then serve. Before opening the bridge, serve() runs a blocking `databricks auth login --resource <mcp-url>` for connection-backed services — the databricks-cli equivalent of mcp-remote's in-process OAuth, where --resource also routes /oidc through the connection sign-in (/mcp-service-login). The agent blocks on "connecting…" while it runs (browser opens via $BROWSER, or the URL is printed), then the session comes up authenticated, so AI Gateway is never asked to elicit a login. Idempotent: once signed in it returns immediately. Removes the redundant startup probe (Claude already fires initialize+tools/list; the proxy shouldn't duplicate that) and the on-401 retry inside the auth hook (the connect-time login makes it unnecessary). _build_token_auth is back to a plain per-request bearer read of the session that login established. Co-authored-by: Isaac <no-reply@databricks.com>
The connect-time connection login runs `databricks auth login --resource`; an older CLI without databricks/cli#6621 rejects the flag and the login exits with a cryptic parse error (seen as a bare "connection closed" in the agent). Add a one-time `auth login --help` pre-check and, when `--resource` is absent, return a clear "upgrade your Databricks CLI" message instead of attempting the doomed login. Fail-open if --help can't run. Co-authored-by: Isaac <no-reply@databricks.com>
…lly missing The connect-time login ran `databricks auth login --resource` unconditionally on every proxy startup for every connection-backed server. But `databricks auth login` always re-runs the browser OAuth (it does not short-circuit on a cached token), so N configured servers meant N browser windows on every session — and concurrent `auth login` runs also race on the shared token cache. Gate the login on the actual per-user credential state (new `connection_credential_state`, using the same Unity Catalog REST APIs as `ug mcp login`): only run the login when the credential is confirmed MISSING. When it is already PRESENT — or the state can't be determined (UNKNOWN) or the connection isn't OAuth-U2M (NO_LOGIN) — skip the browser entirely. `tools/list` still works without the credential, so UNKNOWN fails safe rather than risking a spurious browser. This makes repeat sessions open zero browsers and first sessions prompt only for the servers genuinely needing a sign-in. Co-authored-by: Isaac <no-reply@databricks.com>
… startup The connect-time login blocked the bridge from opening until it finished. With N configured connection-backed servers all missing their credential, an agent spawns N proxies at once, each launches a blocking `databricks auth login`, and they stall on browser flows / racing loopback ports until every one hits the agent's MCP startup timeout — all N servers fail to start (observed with 5 on Codex, all timing out at 120s). Move the login into the httpx auth flow, triggered by a 401: the bridge opens immediately (tools/list needs no credential and returns the full list), and `databricks auth login --resource` runs only when a request actually gets a 401 — i.e. when a tool you invoke needs the per-user connection credential. So a browser opens only for a service you use, one at a time; an already-signed-in service never prompts; and startup never blocks. The login runs at most once per session, off the event loop (async_auth_flow + anyio.to_thread) so the bridge's pumps aren't starved. PAT profiles have no connection OAuth, so they skip it. This drops the eager credential-state gate added in the previous commit (the lazy 401 signal supersedes it) along with its UC-REST probe. Co-authored-by: Isaac <no-reply@databricks.com>
…gin` Drop the proxy-specific `ucode mcp-proxy:` prefix from the connect-time sign-in message so `run_connection_login` / `connection_from_url` can be reused verbatim by `ug mcp login` (PR databricks#654), which stacks on this one. Behavior is unchanged for the proxy. Co-authored-by: Isaac <no-reply@databricks.com>
… parse Consolidate the AI Gateway mcp-services URL knowledge that had drifted into three copies: the path-segment literal lived in databricks.py (URL building), mcp.py (AIGW_MCP_SERVICES_PATH, URL-shape classification), and mcp_connection_login.py (AIGW_MCP_SERVICES_SEGMENT, connection-backed detection). Define AIGW_MCP_SERVICES_SEGMENT once in databricks.py (where the URL is built) and import it into mcp.py and mcp_connection_login.py. Also reuse connection_from_url in mcp._mcp_server_location instead of a parallel url.split(), so the URL -> FQN parse lives in one place too. No behavior change (the parser additionally strips a trailing slash/query, which is strictly more correct for the label). Co-authored-by: Isaac <no-reply@databricks.com>
… nit) Clearer name for the closure that drives the per-connection sign-in on a 401. Co-authored-by: Isaac <no-reply@databricks.com>
`ug mcp login` shows which of the connection-backed AI Gateway MCP services the coding agents are configured to use are already signed in vs. still need a per-user connection sign-in, and runs the sign-in for the ones you pick (interactive picker) or name with `--services` / scope with `--agents`. Reuses databricks#679's building blocks: the developer + workspace-managed server enumeration is extracted from `list_mcp_command` into a shared `configured_mcp_servers_by_name` (behavior-preserving) that both `ug mcp list` and `ug mcp login` call, and the status is rendered with the same rich Table + `status_badge` styling from `ucode.ui`. Per-service status comes from the existing Unity Catalog REST APIs (the ones the `/mcp-service-login` page uses); sign-in is `databricks auth login --resource` (RFC 8707, databricks/cli#6621), so it works for any connection-backed MCP service, not just `system.ai.*`. The credential is per-user and shared across agents, so signing in once unblocks the service for every agent. Co-authored-by: Isaac <no-reply@databricks.com>
…ntation) `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 <no-reply@databricks.com>
sunishsheth2009
force-pushed
the
mcp-status-single-enumerator
branch
from
September 23, 2026 20:53
911bc79 to
c0529c2
Compare
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 <no-reply@databricks.com>
sunishsheth2009
changed the title
sunishsheth2009
marked this pull request as ready for review
September 23, 2026 22:22
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-ups split out of #557/#654 to keep those focused (per reviewer request). No behavior change; both are DRY/dead-code cleanups.
1.
ug status: count MCP servers via the shared enumeratorug statusre-implemented the "which MCP servers are configured" enumeration inline (mergedmcp_servers + managed_mcp_servers, deduped by name, separately read the Claude/Codex OS-managed files) — a parallel copy ofconfigured_mcp_servers_by_name, the enumeratorug mcp listandug mcp loginalready use. Routed the per-agent count through it, so all three commands read one source of truth and the count can't drift. Net −9 lines.revert_mcp_configswas intentionally left alone: its narrowermcp_servers + managed_mcp_serversmerge is the user-scope set it removes; OS-managed files are reverted separately byclaude.revert_managed_settings()/codex.revert_managed_config().2.
ug mcp-proxy: async-only auth (drop the dead syncauth_flow)The proxy only ever drives its httpx
Authfrom anAsyncClient, so the syncauth_flownever ran at runtime — it duplicated the async flow and existed only so the tests could drive it synchronously. Removed it;sync_auth_flownow raises a clearProxyAuthErrorso a future sync client fails loudly instead of silently falling back to httpx's no-op default (which would skip the bearer). Tests driveasync_auth_flowthrough one_drivehelper (anyio.run), andtest_sync_auth_flow_is_rejectedcovers the guard. (Resolves the "do we still need the sync one?" thread on #557.)Testing
ruff+tyclean.test_mcp_proxy.py(36) green after the async-driver rewrite;TestStatus/TestStatusLiveModels(8) green for the status change. No production behavior change in either.This pull request and its description were written by Isaac.