mcp-proxy: support both mcp 1.x and 2.x (drop the <2 cap) (#307)

pyproject.toml

@@ -14,14 +14,15 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
   "databricks-sql-connector>=3.6.0",
+  # httpx is the HTTP stack for mcp 1.x. mcp 2.x uses httpx2 instead (pulled in
+  # transitively by mcp itself); `ucode.mcp_proxy` selects whichever the
+  # installed SDK uses at runtime, so no direct httpx2 dependency is declared.
   "httpx>=0.27.1",
   # `ucode mcp-proxy` bridges a client's stdio MCP transport to a Databricks
   # streamable-HTTP MCP endpoint, injecting a freshly-minted OAuth bearer per
-  # request. Uses the official MCP SDK's stdio server + streamable-HTTP client.
-  # Capped below 2.0: mcp 2.x swaps httpx for httpx2 and renames
-  # `streamablehttp_client` to `streamable_http_client`, so `ucode.mcp_proxy`
-  # fails to import against it. Lift once the proxy is ported.
-  "mcp>=1.28.0,<2",
+  # request, via the SDK's stdio server + streamable-HTTP client. Works against
+  # both mcp 1.x (httpx) and mcp 2.x (httpx2) — see mcp_proxy for the shared path.
+  "mcp>=1.28.0",
   "questionary>=2.0.0",
   "tomlkit>=0.13.0",
   "typer>=0.12.0",

src/ucode/cli.py

@@ -1080,15 +1080,21 @@ def mcp_proxy_cmd(
         str | None, typer.Option("--profile", help="Databricks CLI profile.")
     ] = None,
     use_pat: Annotated[
-        bool, typer.Option("--use-pat", help="Use the profile's static PAT instead of OAuth.")
+        bool,
+        typer.Option(
+            "--use-pat",
+            help="Authenticate with the profile's static personal access token (from "
+            "~/.databrickscfg) instead of OAuth. Set automatically for workspaces configured "
+            "with `ucode configure --profiles <name> --use-pat`.",
+        ),
     ] = False,
 ) -> None:
     """Bridge a coding agent's stdio MCP transport to a Databricks MCP endpoint.
 
     Each configured client spawns this as a local stdio MCP server (see
     `ucode configure mcp`); it forwards messages to ``--url`` and injects a
-    freshly-minted OAuth bearer on every upstream request, so the token never
-    expires mid-session. Not meant for interactive use — the agent manages this
+    freshly-minted token on every upstream request, so it never expires
+    mid-session. Not meant for interactive use — the agent manages this
     process's lifecycle."""
     from ucode.mcp_proxy import serve
 

src/ucode/mcp_proxy.py

@@ -7,13 +7,21 @@
 — ucode owns no long-lived process and no background refresh thread. The proxy
 speaks stdio to the agent and streamable-HTTP to Databricks, and mints a fresh
 token from the Databricks CLI profile on **every** upstream HTTP request via an
-``httpx.Auth`` hook, so the bearer never goes stale mid-session.
+httpx ``Auth`` hook, so the bearer never goes stale mid-session.
 
 This replaces the previous per-client header auth (static ``Bearer
 ${OAUTH_TOKEN}``, Claude ``headersHelper``, Cursor literal-token rewrites): one
 uniform mechanism, token refresh in a single place, and the proxy is an
 invisible implementation detail baked into each client's config.
 
+MCP SDK compatibility (1.x and 2.x). The proxy uses the SDK's *2.x-native* call
+shape — ``streamable_http_client(url, http_client=<AsyncClient>)`` — which both
+mcp 1.28+ and mcp 2.x export (1.x's older ``streamablehttp_client`` is a thin
+deprecated shim over it). The only thing that differs across the major versions
+is the HTTP library: mcp 1.x builds on ``httpx``, mcp 2.x on ``httpx2``. We
+resolve whichever the installed SDK uses (see ``_httpx``) and build the client
+and auth from that, so a single code path works against both — no version cap.
+
 Auth failures are terminal and are reported *fast*. When the Databricks CLI
 can't mint a token (expired refresh token, logged-out profile), the proxy prints
 the CLI's own message to stderr and exits ``AUTH_FAILURE_EXIT_CODE`` rather than
@@ -26,21 +34,40 @@
 from __future__ import annotations
 
 import sys
+from types import ModuleType
 
 import anyio
-import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.streamable_http import streamable_http_client
 from mcp.server.stdio import stdio_server
 
-from ucode.databricks import get_databricks_token
+from ucode.databricks import ensure_pat_bearer, get_databricks_token
 
 # Exit code used when the proxy cannot authenticate. MCP clients surface a
 # non-zero exit far more usefully than a startup timeout, so bail out with this
 # instead of letting the process hang until the client's timeout fires.
 AUTH_FAILURE_EXIT_CODE = 2
 
 
+def _httpx() -> ModuleType:
+    """Return the httpx module the installed MCP SDK is built on.
+
+    mcp 2.x moved from ``httpx`` to ``httpx2`` and passes an ``AsyncClient`` of
+    that flavor into ``streamable_http_client``. We must construct our client and
+    ``Auth`` from the *same* module the SDK uses, or the transport rejects it.
+    Prefer ``httpx2`` (mcp 2.x) and fall back to ``httpx`` (mcp 1.x)."""
+    # Imported dynamically: httpx2 ships only with mcp 2.x, so a static `import
+    # httpx2` is unresolvable under the mcp 1.x that's pinned for type-checking
+    # and CI. importlib keeps the type checker out of it while the runtime picks
+    # the right flavor.
+    from importlib import import_module
+
+    try:
+        return import_module("httpx2")
+    except ImportError:
+        return import_module("httpx")
+
+
 class ProxyAuthError(RuntimeError):
     """The proxy could not mint a Databricks token, so it cannot serve requests.
 
@@ -58,31 +85,32 @@ def _fail_fast(message: str) -> None:
     raise SystemExit(AUTH_FAILURE_EXIT_CODE)
 
 
-class _DatabricksTokenAuth(httpx.Auth):
-    """Injects a fresh Databricks OAuth bearer on every request.
+def _build_token_auth(workspace: str, profile: str | None):
+    """Build an httpx ``Auth`` that injects a fresh bearer on every request.
 
-    ``get_databricks_token`` returns a cached token and refreshes it when
-    expired, so calling it per request keeps auth current without the proxy
-    tracking token lifetimes itself."""
+    The base class comes from whichever httpx the SDK uses (see ``_httpx``), so
+    the returned auth is accepted by that SDK's ``AsyncClient``. Behaviour is
+    identical across flavours — ``Auth.auth_flow`` has the same generator
+    contract in httpx and httpx2."""
+    httpx = _httpx()
 
-    def __init__(self, workspace: str, profile: str | None, *, use_pat: bool) -> None:
-        self._workspace = workspace
-        self._profile = profile
-        self._use_pat = use_pat
+    class _DatabricksTokenAuth(httpx.Auth):
+        def auth_flow(self, request):
+            # get_databricks_token honors the DATABRICKS_BEARER short-circuit and
+            # PAT profiles internally; --use-pat is surfaced via the env ucode set.
+            # A RuntimeError here means auth is dead (expired refresh token,
+            # logged-out profile). Raising it from inside auth_flow would tear
+            # through the transport's task group and stall the process until the
+            # client times out, so translate it into a terminal ProxyAuthError the
+            # caller reports cleanly.
+            try:
+                token = get_databricks_token(workspace, profile)
+            except RuntimeError as exc:
+                raise ProxyAuthError(str(exc)) from exc
+            request.headers["Authorization"] = f"Bearer {token}"
+            yield request
 
-    def auth_flow(self, request: httpx.Request):
-        # get_databricks_token honors the DATABRICKS_BEARER short-circuit and PAT
-        # profiles internally; --use-pat is surfaced via the env ucode already set.
-        # A RuntimeError here means auth is dead (expired refresh token, logged-out
-        # profile). Raising it from inside httpx's auth_flow would tear through the
-        # transport's task group and stall the process until the client times out,
-        # so translate it into a terminal ProxyAuthError the caller reports cleanly.
-        try:
-            token = get_databricks_token(self._workspace, self._profile)
-        except RuntimeError as exc:
-            raise ProxyAuthError(str(exc)) from exc
-        request.headers["Authorization"] = f"Bearer {token}"
-        yield request
+    return _DatabricksTokenAuth()
 
 
 async def _pump(
@@ -98,14 +126,23 @@ async def _pump(
             await dest.send(message)
 
 
-async def _run(url: str, workspace: str, profile: str | None, use_pat: bool) -> None:
-    auth = _DatabricksTokenAuth(workspace, profile, use_pat=use_pat)
-    async with streamablehttp_client(url, auth=auth) as (http_read, http_write, _get_session_id):
-        async with stdio_server() as (stdio_read, stdio_write):
-            # Bidirectional bridge: client stdin -> Databricks, Databricks -> client stdout.
-            async with anyio.create_task_group() as tg:
-                tg.start_soon(_pump, stdio_read, http_write)
-                tg.start_soon(_pump, http_read, stdio_write)
+async def _run(url: str, workspace: str, profile: str | None) -> None:
+    httpx = _httpx()
+    auth = _build_token_auth(workspace, profile)
+    # 2.x-native shape: hand the transport a pre-built AsyncClient carrying our
+    # per-request auth. Works on mcp 1.28+ and 2.x; `streamable_http_client`
+    # yields a (read, write) pair in both.
+    async with httpx.AsyncClient(auth=auth) as http_client:
+        async with streamable_http_client(url, http_client=http_client) as streams:
+            # mcp 1.x yields (read, write, get_session_id); mcp 2.x drops the
+            # trailing callback and yields (read, write). Take the first two
+            # positionally so both arities work — we don't use get_session_id.
+            http_read, http_write = streams[0], streams[1]
+            async with stdio_server() as (stdio_read, stdio_write):
+                # Bidirectional bridge: client stdin -> Databricks, Databricks -> client stdout.
+                async with anyio.create_task_group() as tg:
+                    tg.start_soon(_pump, stdio_read, http_write)
+                    tg.start_soon(_pump, http_read, stdio_write)
 
 
 def _preflight_token(workspace: str, profile: str | None) -> None:
@@ -138,7 +175,19 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool
 
     Authentication is checked up front: a dead profile is a terminal condition,
     and failing here (fast, with the CLI's own message) is far better than
-    letting the client wait out its MCP startup timeout with no explanation."""
+    letting the client wait out its MCP startup timeout with no explanation.
+
+    ``use_pat`` selects static personal-access-token auth: ``databricks auth
+    token`` only reads OAuth caches, so a PAT profile's token must be exported as
+    ``DATABRICKS_BEARER`` first (``ensure_pat_bearer``) — then every per-request
+    mint takes that short-circuit. OAuth needs no such step."""
+    if use_pat and not ensure_pat_bearer(profile):
+        _fail_fast(
+            "--use-pat is set but no personal access token was found for profile "
+            f"'{profile or '<none>'}' in ~/.databrickscfg (expected auth_type = pat). "
+            "Set DATABRICKS_BEARER, or reconfigure the profile."
+        )
+
     # Pre-flight the token before opening the bridge. Without this, the first
     # token failure surfaces from inside the transport's task group, where it can
     # stall the process instead of erroring out.
@@ -148,7 +197,7 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool
         _fail_fast(str(exc))
 
     try:
-        anyio.run(_run, url, workspace, profile, use_pat)
+        anyio.run(_run, url, workspace, profile)
     except BaseException as exc:  # noqa: BLE001 - re-raised unless it's an auth failure
         # The token can still expire mid-session; report that the same way
         # rather than letting the ExceptionGroup surface as a hang or traceback.

tests/test_cli.py

@@ -3038,3 +3038,42 @@ def test_not_checked_when_the_env_var_is_off(self, monkeypatch, env_value):
         result, calls, _ = self._launch(monkeypatch, managed=None)
         assert result.exit_code == 0, result.output
         assert calls == []
+
+
+class TestMcpProxyCmdForwardsUsePat:
+    """`ucode mcp-proxy` forwards the PAT choice to `serve`, which owns the
+    actual PAT resolution. Behavior of that resolution lives in test_mcp_proxy."""
+
+    def _invoke(self, monkeypatch, *, flag, state):
+        captured: dict = {}
+        monkeypatch.setattr("ucode.cli.load_state", lambda: state)
+        monkeypatch.setattr(
+            "ucode.mcp_proxy.serve",
+            lambda *a, **kw: captured.update(args=a, kwargs=kw),
+        )
+        args = ["mcp-proxy", "--url", "https://x/mcp", "--host", "https://x"]
+        if flag:
+            args.append("--use-pat")
+        result = runner.invoke(app, args)
+        return result, captured
+
+    def test_flag_forwards_use_pat_true(self, monkeypatch):
+        result, captured = self._invoke(
+            monkeypatch, flag=True, state={"workspace": "https://x", "profile": "p"}
+        )
+        assert result.exit_code == 0, result.output
+        assert captured["kwargs"]["use_pat"] is True
+
+    def test_saved_use_pat_state_forwards_true(self, monkeypatch):
+        # A workspace configured with --use-pat persists use_pat=True; the proxy
+        # honors it without the flag being repeated.
+        result, captured = self._invoke(
+            monkeypatch, flag=False, state={"workspace": "https://x", "use_pat": True}
+        )
+        assert result.exit_code == 0, result.output
+        assert captured["kwargs"]["use_pat"] is True
+
+    def test_no_flag_and_no_state_forwards_false(self, monkeypatch):
+        result, captured = self._invoke(monkeypatch, flag=False, state={"workspace": "https://x"})
+        assert result.exit_code == 0, result.output
+        assert captured["kwargs"]["use_pat"] is False