Langchain MCP Adapters Fix (#245)

Author: aravind-segu

databricks_mcp/src/databricks_mcp/oauth_provider.py

@@ -46,6 +46,7 @@ class DatabricksOAuthClientProvider(OAuthClientProvider):
     """
 
     def __init__(self, workspace_client: WorkspaceClient):
+        self.workspace_client = workspace_client
         self.databricks_token_storage = DatabricksTokenStorage(workspace_client)
 
         super().__init__(

integrations/langchain/src/databricks_langchain/multi_server_mcp_client.py

@@ -1,11 +1,39 @@
 from typing import Any, Callable, List, Union
 
+import httpx
 from databricks.sdk import WorkspaceClient
 from databricks_mcp.oauth_provider import DatabricksOAuthClientProvider
 from langchain_mcp_adapters.client import MultiServerMCPClient
+from langchain_mcp_adapters.sessions import McpHttpClientFactory
 from pydantic import BaseModel, ConfigDict, Field
 
 
+class DatabricksMcpHttpClientFactory(McpHttpClientFactory):
+    def __call__(
+        self,
+        headers: dict[str, str] | None = None,
+        timeout: httpx.Timeout | None = None,
+        auth: httpx.Auth | None = None,
+    ) -> httpx.AsyncClient:
+        if isinstance(auth, DatabricksOAuthClientProvider) and auth.workspace_client is not None:
+            # Currently DatabricksOAuthClientProvider does not do a full U2M.
+            # Therefore a new fresh token is only retrieved on the first initialization.
+            #  As this factory is called for each request, we are reinitailizing the
+            # DatabricksOAuthClientProvider with the original workspace client to get a new token
+
+            return httpx.AsyncClient(
+                headers=headers,
+                timeout=timeout,
+                auth=DatabricksOAuthClientProvider(auth.workspace_client),
+            )
+        else:
+            return httpx.AsyncClient(
+                headers=headers,
+                timeout=timeout,
+                auth=auth,
+            )
+
+
 class MCPServer(BaseModel):
     """
     Base configuration for an MCP server connection using streamable HTTP transport.
@@ -234,6 +262,7 @@ def to_connection_dict(self) -> dict[str, Any]:
 
         # Add Databricks auth provider
         data["auth"] = self._auth_provider
+        data["httpx_client_factory"] = DatabricksMcpHttpClientFactory()
 
         return data
 

integrations/langchain/tests/unit_tests/test_multi_server_mcp_client.py

@@ -8,6 +8,7 @@
 from databricks.sdk import WorkspaceClient
 
 from databricks_langchain.multi_server_mcp_client import (
+    DatabricksMcpHttpClientFactory,
     DatabricksMCPServer,
     DatabricksMultiServerMCPClient,
     MCPServer,
@@ -144,6 +145,29 @@ def test_databricks_server_accepts_extra_params(self):
             assert connection_dict["timeout"] == 45.0
             assert connection_dict["headers"] == {"X-Custom": "header"}
 
+    def test_databricks_server_includes_http_factory(self):
+        """Test that DatabricksMCPServer includes the custom HTTP client factory."""
+        mock_workspace_client = create_autospec(WorkspaceClient, instance=True)
+
+        with patch(
+            "databricks_langchain.multi_server_mcp_client.DatabricksOAuthClientProvider"
+        ) as mock_auth:
+            mock_auth_instance = MagicMock()
+            mock_auth.return_value = mock_auth_instance
+
+            server = DatabricksMCPServer(
+                name="databricks",
+                url="https://databricks.com/mcp",
+                workspace_client=mock_workspace_client,
+            )
+
+            connection_dict = server.to_connection_dict()
+
+            assert "httpx_client_factory" in connection_dict
+            assert isinstance(
+                connection_dict["httpx_client_factory"], DatabricksMcpHttpClientFactory
+            )
+
 
 class TestDatabricksMultiServerMCPClient:
     """Tests for the DatabricksMultiServerMCPClient class."""