refac

Author: tjbck

backend/open_webui/utils/mcp/client.py

@@ -9,29 +9,42 @@
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
 import httpx
-from open_webui.env import AIOHTTP_CLIENT_SESSION_TOOL_SERVER_SSL
+from open_webui.env import AIOHTTP_CLIENT_SESSION_TOOL_SERVER_SSL, AIOHTTP_CLIENT_TIMEOUT_TOOL_SERVER
 
 
-def create_insecure_httpx_client(headers=None, timeout=None, auth=None):
-    """Create an httpx AsyncClient with SSL verification disabled.
+def _build_httpx_client(headers=None, timeout=None, auth=None, verify=True):
+    """Create an httpx AsyncClient for MCP transport.
+
+    Falls back to AIOHTTP_CLIENT_TIMEOUT_TOOL_SERVER when the caller
+    (i.e. the MCP SDK) does not supply an explicit timeout.
 
-    Note: verify=False must be passed at construction time because httpx
+    Note: verify must be passed at construction time because httpx
     configures the SSL context during __init__. Setting client.verify = False
     after construction does not affect the underlying transport's SSL context.
     """
     kwargs = {
         'follow_redirects': True,
-        'verify': False,
+        'verify': verify,
     }
     if timeout is not None:
         kwargs['timeout'] = timeout
+    elif AIOHTTP_CLIENT_TIMEOUT_TOOL_SERVER is not None:
+        kwargs['timeout'] = float(AIOHTTP_CLIENT_TIMEOUT_TOOL_SERVER)
     if headers is not None:
         kwargs['headers'] = headers
     if auth is not None:
         kwargs['auth'] = auth
     return httpx.AsyncClient(**kwargs)
 
 
+def create_httpx_client(headers=None, timeout=None, auth=None):
+    return _build_httpx_client(headers=headers, timeout=timeout, auth=auth, verify=True)
+
+
+def create_insecure_httpx_client(headers=None, timeout=None, auth=None):
+    return _build_httpx_client(headers=headers, timeout=timeout, auth=auth, verify=False)
+
+
 class MCPClient:
     def __init__(self):
         self.session: Optional[ClientSession] = None
@@ -40,14 +53,11 @@ def __init__(self):
     async def connect(self, url: str, headers: Optional[dict] = None):
         async with AsyncExitStack() as exit_stack:
             try:
-                if AIOHTTP_CLIENT_SESSION_TOOL_SERVER_SSL:
-                    self._streams_context = streamablehttp_client(url, headers=headers)
-                else:
-                    self._streams_context = streamablehttp_client(
-                        url,
-                        headers=headers,
-                        httpx_client_factory=create_insecure_httpx_client,
-                    )
+                self._streams_context = streamablehttp_client(
+                    url,
+                    headers=headers,
+                    httpx_client_factory=create_httpx_client if AIOHTTP_CLIENT_SESSION_TOOL_SERVER_SSL else create_insecure_httpx_client,
+                )
 
                 transport = await exit_stack.enter_async_context(self._streams_context)
                 read_stream, write_stream, _ = transport