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

[sunishsheth2009]

Dependency merged: databricks auth login --resource shipped in databricks/cli#6763, so the connection sign-in works with a current Databricks CLI. Rebased on main.

#654 (ug mcp login) is stacked on this and reuses the mcp_connection_login module introduced here (connection_from_url, run_connection_login). Merge this first.

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 — login is driven lazily, on a 401 (not eagerly at startup)

A connection-backed mcp-services endpoint needs a per-user connection credential before its tools can be used; until then AI Gateway answers with an RFC 9728 401. The proxy handles that inside its httpx Auth hook:

1. The bridge comes up immediately — tools/list needs no credential, so nothing blocks startup.
2. On the first request that gets a 401 from a connection-backed service, the auth flow runs a one-time databricks auth login --resource <mcp-url> and retries. 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.
3. The blocking login runs off the event loop (anyio.to_thread.run_sync in async_auth_flow), so the bridge's other pumps aren't starved while the user completes the sign-in. It fires at most once per session; an already-signed-in service never prompts; PAT profiles have no connection OAuth to drive and skip it entirely.

This is deliberately lazy: a browser opens only for a service whose tool you actually call, one at a time — not N browsers at startup for every configured server. The proxy never writes to stdout (the MCP JSON-RPC wire); the CLI's prompts and authorize URL go to stderr, so a headless/remote box still shows the URL to open manually.

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, and an eager connect-time login that ran a blocking databricks auth login before opening the bridge. The eager form popped a browser at every launch and, with N configured connection-backed servers, opened N browsers at once and stalled the agent's MCP startup until each timed out. Driving the login lazily on a real 401, off the event loop, once per session, is the minimal generic realization and matches how mcp-remote behaves.

Dependencies

• CLI --resource: Allow OAuth U2M logins to send RFC 8707 resource indicators cli#6763 (merged). An older CLI without the flag fails fast with a clear "upgrade your CLI" message.
• /oidc resource-indicator handling (login §2) + /mcp-service-login return_to (§3) — the server side of this flow. AI Gateway §1 (the 401) is already on staging.

Scope

• Connection-backed mcp-services endpoints get the lazy login-on-401. PAT profiles skip it (no connection OAuth to drive).
• Plain Databricks MCP services keep the existing transparent per-request 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; old-CLI clear error). test_mcp_proxy — the lazy auth flow (no login when the response isn't a connection 401; login-then-retry on a 401; login fires at most once; non-connection URLs and --use-pat never log in). uv run pytest tests/test_mcp*.py green; ruff + ty clean.

This pull request and its description were written by Isaac.

Test evidence

Local run on c28a77b (the pushed head, rebased on latest main): lint, type-check, and the full proxy + connection-login suites all green. CI on this PR is also green (Unit tests, Agent launch tests for all six agents, Gateway API, Integration/Installation).

### ruff check       → All checks passed!
### ty typecheck     → All checks passed!

### pytest tests/test_mcp_connection_login.py tests/test_mcp_proxy.py
============================= test session starts ==============================
platform linux -- Python 3.12.13, pytest-9.1.1
collected 46 items

tests/test_mcp_connection_login.py ..........                            [ 21%]
tests/test_mcp_proxy.py ....................................             [100%]
============================== 46 passed in 1.31s ==============================

All 46 tests (click to expand)

TestConnectionFromUrl:
  test_plain_endpoint
  test_with_trailing_path_and_query
  test_non_aigw_url_is_none
  test_missing_service_is_none
TestRunConnectionLogin:
  test_success_sends_resource_and_host_without_client_id
  test_nonzero_exit_reports_failure
  test_output_is_routed_to_stderr_not_stdout
  test_timeout_is_reported
  test_binary_missing_is_reported
  test_old_cli_without_resource_flag_reports_clearly
test_mcp_proxy (SDK/httpx wiring):
  test_httpx_is_a_direct_runtime_dependency
  test_mcp_dependency_is_uncapped
  test_selected_httpx_matches_the_installed_mcp_sdk
  test_proxy_imports_the_streamable_http_client_shared_by_both_majors
TestDatabricksTokenAuth (the lazy login-on-401 core):
  test_injects_bearer_from_minted_token
  test_auth_is_an_instance_of_the_selected_httpx_auth
  test_calls_get_token_with_workspace_and_profile
  test_mints_a_fresh_token_per_request
  test_auth_flow_yields_the_same_request
  test_dead_auth_becomes_a_terminal_proxy_auth_error
  test_on_401_from_connection_service_runs_login_then_retries
  test_non_connection_401_is_not_retried
  test_success_response_never_logs_in
  test_login_runs_at_most_once_per_session
  test_use_pat_never_drives_connection_login
  test_login_failure_becomes_a_proxy_auth_error
  test_async_auth_flow_drives_login_on_401
TestPump / TestServe / TestPreflightToken:
  test_forwards_all_messages_in_order
  test_closes_destination_when_source_exhausts
  test_client_errors_are_forwarded
  test_upstream_errors_are_raised
  test_upstream_eof_is_an_error
  test_run_uses_mcp_http_defaults
  test_runs_the_bridge_with_parsed_args
  test_defaults_profile_none
  test_use_pat_exports_the_bearer_before_serving
  test_use_pat_without_a_resolvable_pat_exits_before_serving
  test_oauth_path_never_touches_pat
  test_preflights_auth_before_opening_the_bridge
  test_dead_auth_exits_fast_without_starting_the_bridge
  test_auth_expiring_mid_session_exits_with_the_actionable_message
  test_transport_failure_exits_with_a_one_line_message
  test_non_auth_failures_still_propagate
  test_passes_through_when_a_token_is_available
  test_surfaces_the_cli_error_message
  test_checks_the_same_workspace_and_profile_the_bridge_will_use

[xsh310]

Might be helpful to add a video recording of this