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

* Update mcp_sse.py

Uses Python's built-in asyncio.timeout() context manager
Properly handles timeout exceptions
Maintains the same functionality but with correct async context management

* [autofix.ci] apply automated fixes

* add asyncio

+clean up comment

* [autofix.ci] apply automated fixes

* missing arg_schema in Tool

Missing args_schema inside cause that tools are generated without input schema and are not able to be properly executed as agent know tool, but dost know what input field tool have.

Same problem looks to be in MCP STDIO.

* fix Ruff Check

Line 56:

Error: B904 Within an `except` clause, raise exceptions with `raise ... from err` or `raise ... from None` to distinguish them from errors in exception handling
TRY003 Avoid specifying long messages outside the exception class
EM102 Exception must not use an f-string literal, assign to variable first

* [autofix.ci] apply automated fixes

* remove asyncio.timeout

Remove asyncio.timeout() (not valid for Py3.10) and replace it by asyncio.wait_for()

* [autofix.ci] apply automated fixes

* Ruff (TRY300)

Move return response.tools inside an else block. This makes it clearer that tools are returned only if the connection is successful, and not if a TimeoutError occurs.

* fix: add session initialization check in MCPSseClient

Added a check to ensure the session is initialized before attempting to list tools, raising a ValueError with a descriptive message if the session is None. This improves error handling and robustness of the MCPSseClient class.

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Sebastián Estévez <estevezsebastian@gmail.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

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.components.tools.mcp_stdio import create_input_schema_from_json_sc
 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
@@ -38,20 +38,32 @@ class MCPSseClient:
         if headers is None:
             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)
+        try:
+            await asyncio.wait_for(
+                self._connect_with_timeout(url, headers, timeout_seconds, sse_read_timeout_seconds),
+                timeout=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()
-
             # List available tools
+            if self.session is None:
+                msg = "Session not initialized"
+                raise ValueError(msg)
             response = await self.session.list_tools()
+        except asyncio.TimeoutError as err:
+            error_message = f"Connection to {url} timed out after {timeout_seconds} seconds"
+            raise TimeoutError(error_message) from err
+        else:  # Only executed if no TimeoutError
             return response.tools
 
+    async def _connect_with_timeout(
+        self, url: str, headers: dict[str, str] | None, timeout_seconds: int, sse_read_timeout_seconds: int
+    ):
+        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()
+
 
 class MCPSse(Component):
     client = MCPSseClient()
@@ -89,6 +101,7 @@ class MCPSse(Component):
                 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),
                 )

src/backend/tests/unit/api/v1/test_api_schemas.py

@@ -6,7 +6,6 @@ from langflow.api.v1.schemas import ResultDataResponse, VertexBuildResponse
 from langflow.schema.schema import OutputValue
 from langflow.serialization import serialize
 from langflow.services.tracing.schema import Log
-
 from pydantic import BaseModel
 
 # Use a smaller test size for hypothesis

src/backend/tests/unit/base/tools/test_component_toolkit.py

@@ -8,7 +8,6 @@ from langflow.components.outputs.chat import ChatOutput
 from langflow.components.tools.calculator import CalculatorToolComponent
 from langflow.graph import Graph
 from langflow.schema.data import Data
-
 from pydantic import BaseModel
 
 

src/backend/tests/unit/base/tools/test_toolmodemixin.py

@@ -21,7 +21,6 @@ from langflow.io import (
     TableInput,
 )
 from langflow.schema import Data
-
 from pydantic import BaseModel
 
 

src/backend/tests/unit/components/helpers/test_structured_output_component.py

@@ -5,8 +5,8 @@ import pytest
 from langflow.components.helpers.structured_output import StructuredOutputComponent
 from langflow.helpers.base_model import build_model_from_schema
 from langflow.inputs.inputs import TableInput
-
 from pydantic import BaseModel
+
 from tests.base import ComponentTestBaseWithoutClient
 from tests.unit.mock_language_model import MockLanguageModel
 

src/backend/tests/unit/graph/graph/state/test_state_model.py

@@ -5,7 +5,6 @@ from langflow.graph import Graph
 from langflow.graph.graph.constants import Finish
 from langflow.graph.state.model import create_state_model
 from langflow.template.field.base import UNDEFINED
-
 from pydantic import Field
 
 

src/backend/tests/unit/helpers/test_base_model_from_schema.py

@@ -4,9 +4,8 @@ from typing import Any
 
 import pytest
 from langflow.helpers.base_model import build_model_from_schema
-from pydantic_core import PydanticUndefined
-
 from pydantic import BaseModel
+from pydantic_core import PydanticUndefined
 
 
 class TestBuildModelFromSchema:

src/backend/tests/unit/inputs/test_inputs.py

@@ -23,7 +23,6 @@ from langflow.inputs.inputs import (
 )
 from langflow.inputs.utils import instantiate_input
 from langflow.schema.message import Message
-
 from pydantic import ValidationError
 
 

src/backend/tests/unit/mock_language_model.py

@@ -1,9 +1,8 @@
 from unittest.mock import MagicMock
 
 from langchain_core.language_models import BaseLanguageModel
-from typing_extensions import override
-
 from pydantic import BaseModel, Field
+from typing_extensions import override
 
 
 class MockLanguageModel(BaseLanguageModel, BaseModel):

src/backend/tests/unit/serialization/test_serialization.py

@@ -9,7 +9,6 @@ from hypothesis import strategies as st
 from langchain_core.documents import Document
 from langflow.serialization.constants import MAX_ITEMS_LENGTH, MAX_TEXT_LENGTH
 from langflow.serialization.serialization import serialize, serialize_or_str
-
 from pydantic import BaseModel as PydanticBaseModel
 from pydantic.v1 import BaseModel as PydanticV1BaseModel
 

src/backend/tests/unit/test_schema.py

@@ -7,7 +7,6 @@ from langflow.schema.data import Data
 from langflow.template import Input, Output
 from langflow.template.field.base import UNDEFINED
 from langflow.type_extraction.type_extraction import post_process_type
-
 from pydantic import ValidationError
 
 

src/backend/tests/unit/test_template.py

@@ -2,7 +2,6 @@ import importlib
 
 import pytest
 from langflow.utils.util import build_template_from_function, get_base_classes, get_default_factory
-
 from pydantic import BaseModel