Skip to content

mcp-proxy: drive per-connection login on a 401 (generic, all agents) - #557

Open
sunishsheth2009 wants to merge 4 commits into
databricks:mainfrom
sunishsheth2009:mcp-connection-login-proxy
Open

mcp-proxy: drive per-connection login on a 401 (generic, all agents)#557
sunishsheth2009 wants to merge 4 commits into
databricks:mainfrom
sunishsheth2009:mcp-connection-login-proxy

Conversation

@sunishsheth2009

@sunishsheth2009 sunishsheth2009 commented Sep 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

What

Give every coding agent (codex, cursor, gemini, claude, …) a working login flow for connection-backed AI Gateway mcp-services endpoints (e.g. system.ai.github) — implemented in the ug mcp-proxy stdio bridge every agent already spawns. No new library (this is what a generic OAuth MCP bridge like mcp-remote does, done in ucode with the Databricks CLI), and no per-agent OAuth app.

How

A connection-backed mcp-services endpoint needs a per-user connection credential before its tools can be used. The proxy drives that login at connect time, before it opens the stdio↔HTTP bridge:

  1. serve() detects a connection-backed URL and runs a blocking databricks auth login --resource <mcp-url>. The RFC 8707 resource indicator makes a resource-aware /oidc route the browser through the connection's own SaaS login (/mcp-service-login) before minting the token. It uses the CLI's own default client + registered loopback redirect — no --client-id, no auth-side redirect change.
  2. Once the login completes, the credential exists server-side, so the bridge opens already authenticated and AI Gateway is only ever called with a valid credential (it never has to elicit a login). The login is idempotent — once signed in it returns immediately.

This is the mcp-remote pattern: the CLI opens the browser natively (macOS/Windows/Linux-desktop) and, when it can't (a headless/remote box), prints the authorize URL to stderr so the user can open it manually. The proxy never writes to stdout (the MCP JSON-RPC wire).

Why this shape (supersedes the earlier bridges in this PR's history)

Earlier iterations tried a proxy that faked a sign_in tool / rewrote tools/list, a login-on-401 httpx auth-hook retry, and a launch-time check — each either more complex or subtly wrong (the auth hook ran a blocking subprocess inside the event loop; launch-time fired a login per connection at every launch). Driving the login once at connect, and letting the agent block on "connecting…" while it runs, is the minimal generic realization and matches how mcp-remote behaves.

Dependencies

Scope

  • Connection-backed mcp-services endpoints get the connect-time login. PAT profiles skip it (no connection OAuth to drive).
  • Plain Databricks MCP services keep the existing transparent token-injection (no extra login).

Tests

test_mcp_connection_login — connection-FQN parsing, and the CLI login runner with the subprocess mocked (sends --resource/--host and no --client-id; routes output to stderr; reports nonzero-exit, timeout, and missing-binary). test_mcp_proxy — connect-time login ordering (login → bridge), non-connection URLs skip it, a login failure exits before the bridge, and --use-pat skips it. uv run pytest tests/test_mcp*.py green; ruff + ty clean.

This pull request and its description were written by Isaac.

This was referenced Sep 10, 2026
Closed
Closed
@sunishsheth2009 Isaac
mcp-proxy: drive per-connection login on a 401 (generic, all agents)
79ee16c 
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>
@sunishsheth2009
sunishsheth2009 force-pushed the mcp-connection-login-proxy branch from 42235ff to 79ee16c Compare September 10, 2026 23:44
@sunishsheth2009 sunishsheth2009 changed the title mcp-proxy: generic per-connection sign-in for AI Gateway MCP services mcp-proxy: drive per-connection login on a 401 (generic, all agents) Sep 10, 2026
sunishsheth2009 and others added 3 commits September 11, 2026 01:43
@sunishsheth2009 Isaac
mcp-proxy: make the connection login non-blocking and surface its URL
0cff035 
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>
@sunishsheth2009 Isaac
mcp-proxy: drive the connection login at connect-time (mcp-remote style)
81cf9d7 
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>
@sunishsheth2009 Isaac
mcp-proxy: connection login at connect (the mcp-remote pattern), drop…
35c1a1d 
… 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>
@sunishsheth2009
sunishsheth2009 force-pushed the mcp-connection-login-proxy branch 2 times, most recently from 9dec3fa to 35c1a1d Compare September 11, 2026 04:01
@sunishsheth2009
sunishsheth2009 marked this pull request as ready for review September 11, 2026 04:05
@sunishsheth2009 sunishsheth2009 mentioned this pull request Sep 11, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@sunishsheth2009