fix: component MCP Tools (SSE): 'Timeout'

src/backend/base/langflow/components/tools/mcp_sse.py

@@ -1,4 +1,5 @@
 # from langflow.field_typing import Data
+import asyncio
 from contextlib import AsyncExitStack
 
 import httpx
@@ -10,7 +11,6 @@
 from langflow.custom import Component
 from langflow.field_typing import Tool
 from langflow.io import MessageTextInput, Output
-from langflow.utils.async_helpers import timeout_context
 
 # Define constant for status code
 HTTP_TEMPORARY_REDIRECT = 307
@@ -39,18 +39,22 @@ async def connect_to_server(
             headers = {}
         url = await self.pre_check_redirect(url)
 
-        async with timeout_context(timeout_seconds):
-            sse_transport = await self.exit_stack.enter_async_context(
-                sse_client(url, headers, timeout_seconds, sse_read_timeout_seconds)
-            )
-            self.sse, self.write = sse_transport
-            self.session = await self.exit_stack.enter_async_context(ClientSession(self.sse, self.write))
+        try:
+            async with asyncio.timeout(timeout_seconds):
+                sse_transport = await self.exit_stack.enter_async_context(
+                    sse_client(url, headers, timeout_seconds, sse_read_timeout_seconds)
+                )
+                self.sse, self.write = sse_transport
+                self.session = await self.exit_stack.enter_async_context(ClientSession(self.sse, self.write))
 
-            await self.session.initialize()
+                await self.session.initialize()
 
-            # List available tools
-            response = await self.session.list_tools()
-            return response.tools
+                # List available tools
+                response = await self.session.list_tools()
+                return response.tools
+        except asyncio.TimeoutError as err:
+            error_message = f"Connection to {url} timed out after {timeout_seconds} seconds"
+            raise TimeoutError(error_message) from err
 
 
 class MCPSse(Component):
@@ -89,6 +93,7 @@ async def build_output(self) -> list[Tool]:
                 Tool(
                     name=tool.name,  # maybe format this
                     description=tool.description,
+                    args_schema=args_schema,
                     coroutine=create_tool_coroutine(tool.name, args_schema, self.client.session),
                     func=create_tool_func(tool.name, self.client.session),
                 )