From 4d3e1458dab36802b8154f4cf0e5b61a29a8cb5c Mon Sep 17 00:00:00 2001 From: Artur Zdolinski Date: Fri, 14 Mar 2025 13:48:53 +0100 Subject: [PATCH] fix: component MCP Tools (SSE): 'Timeout' (#5638) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 Co-authored-by: Gabriel Luiz Freitas Almeida --- .../base/langflow/components/tools/mcp_sse.py | 33 +++++++++++++------ .../tests/unit/api/v1/test_api_schemas.py | 1 - .../unit/base/tools/test_component_toolkit.py | 1 - .../unit/base/tools/test_toolmodemixin.py | 1 - .../test_structured_output_component.py | 2 +- .../graph/graph/state/test_state_model.py | 1 - .../helpers/test_base_model_from_schema.py | 3 +- src/backend/tests/unit/inputs/test_inputs.py | 1 - src/backend/tests/unit/mock_language_model.py | 3 +- .../unit/serialization/test_serialization.py | 1 - src/backend/tests/unit/test_schema.py | 1 - src/backend/tests/unit/test_template.py | 1 - 12 files changed, 26 insertions(+), 23 deletions(-) diff --git a/src/backend/base/langflow/components/tools/mcp_sse.py b/src/backend/base/langflow/components/tools/mcp_sse.py index 82e6c34e5..af9d857e2 100644 --- a/src/backend/base/langflow/components/tools/mcp_sse.py +++ b/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), ) diff --git a/src/backend/tests/unit/api/v1/test_api_schemas.py b/src/backend/tests/unit/api/v1/test_api_schemas.py index f58bcf45d..2a73afe22 100644 --- a/src/backend/tests/unit/api/v1/test_api_schemas.py +++ b/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 diff --git a/src/backend/tests/unit/base/tools/test_component_toolkit.py b/src/backend/tests/unit/base/tools/test_component_toolkit.py index c9d627973..bf8df7a18 100644 --- a/src/backend/tests/unit/base/tools/test_component_toolkit.py +++ b/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 diff --git a/src/backend/tests/unit/base/tools/test_toolmodemixin.py b/src/backend/tests/unit/base/tools/test_toolmodemixin.py index 059464113..27518360c 100644 --- a/src/backend/tests/unit/base/tools/test_toolmodemixin.py +++ b/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 diff --git a/src/backend/tests/unit/components/helpers/test_structured_output_component.py b/src/backend/tests/unit/components/helpers/test_structured_output_component.py index afe7adf59..8e1d1c32d 100644 --- a/src/backend/tests/unit/components/helpers/test_structured_output_component.py +++ b/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 diff --git a/src/backend/tests/unit/graph/graph/state/test_state_model.py b/src/backend/tests/unit/graph/graph/state/test_state_model.py index 4278637b9..869366b8e 100644 --- a/src/backend/tests/unit/graph/graph/state/test_state_model.py +++ b/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 diff --git a/src/backend/tests/unit/helpers/test_base_model_from_schema.py b/src/backend/tests/unit/helpers/test_base_model_from_schema.py index 3bc2bcd2d..d07a4908e 100644 --- a/src/backend/tests/unit/helpers/test_base_model_from_schema.py +++ b/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: diff --git a/src/backend/tests/unit/inputs/test_inputs.py b/src/backend/tests/unit/inputs/test_inputs.py index eacb2a40b..2a946a2a9 100644 --- a/src/backend/tests/unit/inputs/test_inputs.py +++ b/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 diff --git a/src/backend/tests/unit/mock_language_model.py b/src/backend/tests/unit/mock_language_model.py index 8c06e112b..186d127cf 100644 --- a/src/backend/tests/unit/mock_language_model.py +++ b/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): diff --git a/src/backend/tests/unit/serialization/test_serialization.py b/src/backend/tests/unit/serialization/test_serialization.py index dff548472..e8110906c 100644 --- a/src/backend/tests/unit/serialization/test_serialization.py +++ b/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 diff --git a/src/backend/tests/unit/test_schema.py b/src/backend/tests/unit/test_schema.py index 86aaf4500..97fd930e3 100644 --- a/src/backend/tests/unit/test_schema.py +++ b/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 diff --git a/src/backend/tests/unit/test_template.py b/src/backend/tests/unit/test_template.py index 6a4ad938c..6b2127178 100644 --- a/src/backend/tests/unit/test_template.py +++ b/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