feat: bind auth when opening mcp sessions

src/wf_sources_mcp/client/transport.py

@@ -10,15 +10,29 @@ from mcp.client.session import ClientSession
 from mcp.client.stdio import StdioServerParameters, stdio_client
 from mcp.client.streamable_http import streamable_http_client
 
-from wf_sources_mcp.auth import AuthRecord, mcp_auth_env, mcp_auth_headers
+from wf_api.auth import StoredAuthRecord, auth_record_from_compat
+from wf_sources_mcp.auth import AuthRecord, McpAuthBinder
 from wf_sources_mcp.connections import McpSourceConnection
 from wf_sources_mcp.transports import HttpSourceTransport, StdioSourceTransport
 
 
+def _as_stored_auth(auth: AuthRecord | StoredAuthRecord | None) -> StoredAuthRecord | None:
+    if auth is None or isinstance(auth, StoredAuthRecord):
+        return auth
+    return auth_record_from_compat(
+        id=auth.connection_id,
+        scheme=auth.scheme,
+        payload=auth.payload,
+        metadata={},
+    )
+
+
 @asynccontextmanager
 async def open_mcp_session(
     connection: McpSourceConnection,
-    auth: AuthRecord | None,
+    auth: AuthRecord | StoredAuthRecord | None,
+    *,
+    auth_binder: McpAuthBinder | None = None,
 ) -> AsyncIterator[ClientSession]:
     """Open and initialize an MCP client session for the given connection.
 
@@ -35,11 +49,14 @@ async def open_mcp_session(
     if transport is None:
         raise ValueError(f"connection {connection.id!r} requires metadata.transport")
 
+    binder = auth_binder or McpAuthBinder()
+    stored_auth = _as_stored_auth(auth)
+
     if isinstance(transport, StdioSourceTransport):
-        auth_env = mcp_auth_env(auth)
+        bound = await binder.bind_stdio_auth(stored_auth)
         env = dict(transport.env)
-        if auth_env:
-            env = {**env, **auth_env}
+        if bound.env:
+            env = {**env, **bound.env}
         params = StdioServerParameters(
             command=transport.command,
             args=list(transport.args),
@@ -53,8 +70,11 @@ async def open_mcp_session(
         return
 
     if isinstance(transport, HttpSourceTransport):
-        headers = mcp_auth_headers(auth)
-        http_client = httpx.AsyncClient(headers=headers or None)
+        bound = await binder.bind_http_auth(stored_auth)
+        http_client = httpx.AsyncClient(
+            headers=bound.headers or None,
+            auth=bound.auth,
+        )
         async with http_client:
             async with streamable_http_client(
                 str(transport.url),

tests/wf_sources_mcp/test_client_transport.py

@@ -259,3 +259,63 @@ async def test_http_no_auth_creates_client_with_no_auth_headers(
 
     assert len(captured_clients) == 1
     assert "Authorization" not in captured_clients[0].headers
+
+
+@pytest.mark.asyncio
+async def test_open_mcp_session_uses_binder_for_http_headers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from wf_api.auth import BearerAuth, StoredAuthRecord
+    from wf_sources_mcp.client.transport import open_mcp_session
+
+    captured_clients: list[Any] = []
+
+    class _CapturingClient:
+        def __init__(self, **kwargs: Any) -> None:
+            captured_clients.append(kwargs)
+
+        async def __aenter__(self) -> "_CapturingClient":
+            return self
+
+        async def __aexit__(self, *args: Any) -> None:
+            return None
+
+    import wf_sources_mcp.client.transport as mod
+
+    monkeypatch.setattr(mod.httpx, "AsyncClient", _CapturingClient)
+
+    connection = _http_connection()
+    auth = StoredAuthRecord(
+        id="google.drive.personal",
+        auth=BearerAuth(access_token="token"),
+    )
+
+    async with open_mcp_session(connection, auth):
+        pass
+
+    assert captured_clients[0]["headers"] == {"Authorization": "Bearer token"}
+
+
+@pytest.mark.asyncio
+async def test_open_mcp_session_uses_binder_for_stdio_env() -> None:
+    import wf_sources_mcp.client.transport as mod  # noqa: I001
+    from wf_api.auth import EnvAuth, StoredAuthRecord
+
+    captured_params: list[Any] = []
+
+    @asynccontextmanager
+    async def _capturing_stdio_client(
+        params: Any,
+    ) -> AsyncIterator[tuple[Any, Any]]:
+        captured_params.append(params)
+        yield "read", "write"
+
+    mod.stdio_client = _capturing_stdio_client  # type: ignore[assignment, ty:invalid-assignment]
+
+    connection = _stdio_connection(env={"BASE": "1"})
+    auth = StoredAuthRecord(id="demo.auth", auth=EnvAuth(env={"TOKEN": "abc"}))
+
+    async with open_mcp_session(connection, auth):
+        pass
+
+    assert captured_params[0].env == {"BASE": "1", "TOKEN": "abc"}