langflow-base-nightly 0.5.1.dev2__py3-none-any.whl → 0.5.1.dev4__py3-none-any.whl

langflow/base/mcp/util.py

@@ -1,1524 +0,0 @@
-import asyncio
-import contextlib
-import inspect
-import os
-import platform
-import re
-import shutil
-import unicodedata
-from collections.abc import Awaitable, Callable
-from typing import Any
-from urllib.parse import urlparse
-from uuid import UUID
-
-import httpx
-from anyio import ClosedResourceError
-from httpx import codes as httpx_codes
-from langchain_core.tools import StructuredTool
-from mcp import ClientSession
-from mcp.shared.exceptions import McpError
-from pydantic import BaseModel, Field, create_model
-from sqlmodel import select
-
-from langflow.logging.logger import logger
-from langflow.services.database.models.flow.model import Flow
-from langflow.services.deps import get_settings_service
-
-HTTP_ERROR_STATUS_CODE = httpx_codes.BAD_REQUEST  # HTTP status code for client errors
-NULLABLE_TYPE_LENGTH = 2  # Number of types in a nullable union (the type itself + null)
-
-# HTTP status codes used in validation
-HTTP_NOT_FOUND = 404
-HTTP_BAD_REQUEST = 400
-HTTP_INTERNAL_SERVER_ERROR = 500
-
-# MCP Session Manager constants
-settings = get_settings_service().settings
-MAX_SESSIONS_PER_SERVER = (
-    settings.mcp_max_sessions_per_server
-)  # Maximum number of sessions per server to prevent resource exhaustion
-SESSION_IDLE_TIMEOUT = settings.mcp_session_idle_timeout  # 5 minutes idle timeout for sessions
-SESSION_CLEANUP_INTERVAL = settings.mcp_session_cleanup_interval  # Cleanup interval in seconds
-# RFC 7230 compliant header name pattern: token = 1*tchar
-# tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*" / "+" / "-" / "." /
-#         "^" / "_" / "`" / "|" / "~" / DIGIT / ALPHA
-HEADER_NAME_PATTERN = re.compile(r"^[!#$%&\'*+\-.0-9A-Z^_`a-z|~]+$")
-
-# Common allowed headers for MCP connections
-ALLOWED_HEADERS = {
-    "authorization",
-    "accept",
-    "accept-encoding",
-    "accept-language",
-    "cache-control",
-    "content-type",
-    "user-agent",
-    "x-api-key",
-    "x-auth-token",
-    "x-custom-header",
-    "x-langflow-session",
-    "x-mcp-client",
-    "x-requested-with",
-}
-
-
-def validate_headers(headers: dict[str, str]) -> dict[str, str]:
-    """Validate and sanitize HTTP headers according to RFC 7230.
-
-    Args:
-        headers: Dictionary of header name-value pairs
-
-    Returns:
-        Dictionary of validated and sanitized headers
-
-    Raises:
-        ValueError: If headers contain invalid names or values
-    """
-    if not headers:
-        return {}
-
-    sanitized_headers = {}
-
-    for name, value in headers.items():
-        if not isinstance(name, str) or not isinstance(value, str):
-            logger.warning(f"Skipping non-string header: {name}={value}")
-            continue
-
-        # Validate header name according to RFC 7230
-        if not HEADER_NAME_PATTERN.match(name):
-            logger.warning(f"Invalid header name '{name}', skipping")
-            continue
-
-        # Normalize header name to lowercase (HTTP headers are case-insensitive)
-        normalized_name = name.lower()
-
-        # Optional: Check against whitelist of allowed headers
-        if normalized_name not in ALLOWED_HEADERS:
-            # For MCP, we'll be permissive and allow non-standard headers
-            # but log a warning for security awareness
-            logger.debug(f"Using non-standard header: {normalized_name}")
-
-        # Check for potential header injection attempts BEFORE sanitizing
-        if "\r" in value or "\n" in value:
-            logger.warning(f"Potential header injection detected in '{name}', skipping")
-            continue
-
-        # Sanitize header value - remove control characters and newlines
-        # RFC 7230: field-value = *( field-content / obs-fold )
-        # We'll remove control characters (0x00-0x1F, 0x7F) except tab (0x09) and space (0x20)
-        sanitized_value = re.sub(r"[\x00-\x08\x0A-\x1F\x7F]", "", value)
-
-        # Remove leading/trailing whitespace
-        sanitized_value = sanitized_value.strip()
-
-        if not sanitized_value:
-            logger.warning(f"Header '{name}' has empty value after sanitization, skipping")
-            continue
-
-        sanitized_headers[normalized_name] = sanitized_value
-
-    return sanitized_headers
-
-
-def sanitize_mcp_name(name: str, max_length: int = 46) -> str:
-    """Sanitize a name for MCP usage by removing emojis, diacritics, and special characters.
-
-    Args:
-        name: The original name to sanitize
-        max_length: Maximum length for the sanitized name
-
-    Returns:
-        A sanitized name containing only letters, numbers, hyphens, and underscores
-    """
-    if not name or not name.strip():
-        return ""
-
-    # Remove emojis using regex pattern
-    emoji_pattern = re.compile(
-        "["
-        "\U0001f600-\U0001f64f"  # emoticons
-        "\U0001f300-\U0001f5ff"  # symbols & pictographs
-        "\U0001f680-\U0001f6ff"  # transport & map symbols
-        "\U0001f1e0-\U0001f1ff"  # flags (iOS)
-        "\U00002500-\U00002bef"  # chinese char
-        "\U00002702-\U000027b0"
-        "\U00002702-\U000027b0"
-        "\U000024c2-\U0001f251"
-        "\U0001f926-\U0001f937"
-        "\U00010000-\U0010ffff"
-        "\u2640-\u2642"
-        "\u2600-\u2b55"
-        "\u200d"
-        "\u23cf"
-        "\u23e9"
-        "\u231a"
-        "\ufe0f"  # dingbats
-        "\u3030"
-        "]+",
-        flags=re.UNICODE,
-    )
-
-    # Remove emojis
-    name = emoji_pattern.sub("", name)
-
-    # Normalize unicode characters to remove diacritics
-    name = unicodedata.normalize("NFD", name)
-    name = "".join(char for char in name if unicodedata.category(char) != "Mn")
-
-    # Replace spaces and special characters with underscores
-    name = re.sub(r"[^\w\s-]", "", name)  # Keep only word chars, spaces, and hyphens
-    name = re.sub(r"[-\s]+", "_", name)  # Replace spaces and hyphens with underscores
-    name = re.sub(r"_+", "_", name)  # Collapse multiple underscores
-
-    # Remove leading/trailing underscores
-    name = name.strip("_")
-
-    # Ensure it starts with a letter or underscore (not a number)
-    if name and name[0].isdigit():
-        name = f"_{name}"
-
-    # Convert to lowercase
-    name = name.lower()
-
-    # Truncate to max length
-    if len(name) > max_length:
-        name = name[:max_length].rstrip("_")
-
-    # If empty after sanitization, provide a default
-    if not name:
-        name = "unnamed"
-
-    return name
-
-
-def create_tool_coroutine(tool_name: str, arg_schema: type[BaseModel], client) -> Callable[..., Awaitable]:
-    async def tool_coroutine(*args, **kwargs):
-        # Get field names from the model (preserving order)
-        field_names = list(arg_schema.model_fields.keys())
-        provided_args = {}
-        # Map positional arguments to their corresponding field names
-        for i, arg in enumerate(args):
-            if i >= len(field_names):
-                msg = "Too many positional arguments provided"
-                raise ValueError(msg)
-            provided_args[field_names[i]] = arg
-        # Merge in keyword arguments
-        provided_args.update(kwargs)
-        # Validate input and fill defaults for missing optional fields
-        try:
-            validated = arg_schema.model_validate(provided_args)
-        except Exception as e:
-            msg = f"Invalid input: {e}"
-            raise ValueError(msg) from e
-
-        try:
-            return await client.run_tool(tool_name, arguments=validated.model_dump())
-        except Exception as e:
-            await logger.aerror(f"Tool '{tool_name}' execution failed: {e}")
-            # Re-raise with more context
-            msg = f"Tool '{tool_name}' execution failed: {e}"
-            raise ValueError(msg) from e
-
-    return tool_coroutine
-
-
-def create_tool_func(tool_name: str, arg_schema: type[BaseModel], client) -> Callable[..., str]:
-    def tool_func(*args, **kwargs):
-        field_names = list(arg_schema.model_fields.keys())
-        provided_args = {}
-        for i, arg in enumerate(args):
-            if i >= len(field_names):
-                msg = "Too many positional arguments provided"
-                raise ValueError(msg)
-            provided_args[field_names[i]] = arg
-        provided_args.update(kwargs)
-        try:
-            validated = arg_schema.model_validate(provided_args)
-        except Exception as e:
-            msg = f"Invalid input: {e}"
-            raise ValueError(msg) from e
-
-        try:
-            loop = asyncio.get_event_loop()
-            return loop.run_until_complete(client.run_tool(tool_name, arguments=validated.model_dump()))
-        except Exception as e:
-            logger.error(f"Tool '{tool_name}' execution failed: {e}")
-            # Re-raise with more context
-            msg = f"Tool '{tool_name}' execution failed: {e}"
-            raise ValueError(msg) from e
-
-    return tool_func
-
-
-def get_unique_name(base_name, max_length, existing_names):
-    name = base_name[:max_length]
-    if name not in existing_names:
-        return name
-    i = 1
-    while True:
-        suffix = f"_{i}"
-        truncated_base = base_name[: max_length - len(suffix)]
-        candidate = f"{truncated_base}{suffix}"
-        if candidate not in existing_names:
-            return candidate
-        i += 1
-
-
-async def get_flow_snake_case(flow_name: str, user_id: str, session, *, is_action: bool | None = None) -> Flow | None:
-    uuid_user_id = UUID(user_id) if isinstance(user_id, str) else user_id
-    stmt = select(Flow).where(Flow.user_id == uuid_user_id).where(Flow.is_component == False)  # noqa: E712
-    flows = (await session.exec(stmt)).all()
-
-    for flow in flows:
-        if is_action and flow.action_name:
-            this_flow_name = sanitize_mcp_name(flow.action_name)
-        else:
-            this_flow_name = sanitize_mcp_name(flow.name)
-
-        if this_flow_name == flow_name:
-            return flow
-    return None
-
-
-def create_input_schema_from_json_schema(schema: dict[str, Any]) -> type[BaseModel]:
-    """Dynamically build a Pydantic model from a JSON schema (with $defs).
-
-    Non-required fields become Optional[...] with default=None.
-    """
-    if schema.get("type") != "object":
-        msg = "Root schema must be type 'object'"
-        raise ValueError(msg)
-
-    defs: dict[str, dict[str, Any]] = schema.get("$defs", {})
-    model_cache: dict[str, type[BaseModel]] = {}
-
-    def resolve_ref(s: dict[str, Any] | None) -> dict[str, Any]:
-        """Follow a $ref chain until you land on a real subschema."""
-        if s is None:
-            return {}
-        while "$ref" in s:
-            ref_name = s["$ref"].split("/")[-1]
-            s = defs.get(ref_name)
-            if s is None:
-                logger.warning(f"Parsing input schema: Definition '{ref_name}' not found")
-                return {"type": "string"}
-        return s
-
-    def parse_type(s: dict[str, Any] | None) -> Any:
-        """Map a JSON Schema subschema to a Python type (possibly nested)."""
-        if s is None:
-            return None
-        s = resolve_ref(s)
-
-        if "anyOf" in s:
-            # Handle common pattern for nullable types (anyOf with string and null)
-            subtypes = [sub.get("type") for sub in s["anyOf"] if isinstance(sub, dict) and "type" in sub]
-
-            # Check if this is a simple nullable type (e.g., str | None)
-            if len(subtypes) == NULLABLE_TYPE_LENGTH and "null" in subtypes:
-                # Get the non-null type
-                non_null_type = next(t for t in subtypes if t != "null")
-                # Map it to Python type
-                if isinstance(non_null_type, str):
-                    return {
-                        "string": str,
-                        "integer": int,
-                        "number": float,
-                        "boolean": bool,
-                        "object": dict,
-                        "array": list,
-                    }.get(non_null_type, Any)
-                return Any
-
-            # For other anyOf cases, use the first non-null type
-            subtypes = [parse_type(sub) for sub in s["anyOf"]]
-            non_null_types = [t for t in subtypes if t is not None and t is not type(None)]
-            if non_null_types:
-                return non_null_types[0]
-            return str
-
-        t = s.get("type", "any")  # Use string "any" as default instead of Any type
-        if t == "array":
-            item_schema = s.get("items", {})
-            schema_type: Any = parse_type(item_schema)
-            return list[schema_type]
-
-        if t == "object":
-            # inline object not in $defs ⇒ anonymous nested model
-            return _build_model(f"AnonModel{len(model_cache)}", s)
-
-        # primitive fallback
-        return {
-            "string": str,
-            "integer": int,
-            "number": float,
-            "boolean": bool,
-            "object": dict,
-            "array": list,
-        }.get(t, Any)
-
-    def _build_model(name: str, subschema: dict[str, Any]) -> type[BaseModel]:
-        """Create (or fetch) a BaseModel subclass for the given object schema."""
-        # If this came via a named $ref, use that name
-        if "$ref" in subschema:
-            refname = subschema["$ref"].split("/")[-1]
-            if refname in model_cache:
-                return model_cache[refname]
-            target = defs.get(refname)
-            if not target:
-                msg = f"Definition '{refname}' not found"
-                raise ValueError(msg)
-            cls = _build_model(refname, target)
-            model_cache[refname] = cls
-            return cls
-
-        # Named anonymous or inline: avoid clashes by name
-        if name in model_cache:
-            return model_cache[name]
-
-        props = subschema.get("properties", {})
-        reqs = set(subschema.get("required", []))
-        fields: dict[str, Any] = {}
-
-        for prop_name, prop_schema in props.items():
-            py_type = parse_type(prop_schema)
-            is_required = prop_name in reqs
-            if not is_required:
-                py_type = py_type | None
-                default = prop_schema.get("default", None)
-            else:
-                default = ...  # required by Pydantic
-
-            fields[prop_name] = (py_type, Field(default, description=prop_schema.get("description")))
-
-        model_cls = create_model(name, **fields)
-        model_cache[name] = model_cls
-        return model_cls
-
-    # build the top - level "InputSchema" from the root properties
-    top_props = schema.get("properties", {})
-    top_reqs = set(schema.get("required", []))
-    top_fields: dict[str, Any] = {}
-
-    for fname, fdef in top_props.items():
-        py_type = parse_type(fdef)
-        if fname not in top_reqs:
-            py_type = py_type | None
-            default = fdef.get("default", None)
-        else:
-            default = ...
-        top_fields[fname] = (py_type, Field(default, description=fdef.get("description")))
-
-    return create_model("InputSchema", **top_fields)
-
-
-def _is_valid_key_value_item(item: Any) -> bool:
-    """Check if an item is a valid key-value dictionary."""
-    return isinstance(item, dict) and "key" in item and "value" in item
-
-
-def _process_headers(headers: Any) -> dict:
-    """Process the headers input into a valid dictionary.
-
-    Args:
-        headers: The headers to process, can be dict, str, or list
-    Returns:
-        Processed and validated dictionary
-    """
-    if headers is None:
-        return {}
-    if isinstance(headers, dict):
-        return validate_headers(headers)
-    if isinstance(headers, list):
-        processed_headers = {}
-        try:
-            for item in headers:
-                if not _is_valid_key_value_item(item):
-                    continue
-                key = item["key"]
-                value = item["value"]
-                processed_headers[key] = value
-        except (KeyError, TypeError, ValueError):
-            return {}  # Return empty dictionary instead of None
-        return validate_headers(processed_headers)
-    return {}
-
-
-def _validate_node_installation(command: str) -> str:
-    """Validate the npx command."""
-    if "npx" in command and not shutil.which("node"):
-        msg = "Node.js is not installed. Please install Node.js to use npx commands."
-        raise ValueError(msg)
-    return command
-
-
-async def _validate_connection_params(mode: str, command: str | None = None, url: str | None = None) -> None:
-    """Validate connection parameters based on mode."""
-    if mode not in ["Stdio", "SSE"]:
-        msg = f"Invalid mode: {mode}. Must be either 'Stdio' or 'SSE'"
-        raise ValueError(msg)
-
-    if mode == "Stdio" and not command:
-        msg = "Command is required for Stdio mode"
-        raise ValueError(msg)
-    if mode == "Stdio" and command:
-        _validate_node_installation(command)
-    if mode == "SSE" and not url:
-        msg = "URL is required for SSE mode"
-        raise ValueError(msg)
-
-
-class MCPSessionManager:
-    """Manages persistent MCP sessions with proper context manager lifecycle.
-
-    Fixed version that addresses the memory leak issue by:
-    1. Session reuse based on server identity rather than unique context IDs
-    2. Maximum session limits per server to prevent resource exhaustion
-    3. Idle timeout for automatic session cleanup
-    4. Periodic cleanup of stale sessions
-    """
-
-    def __init__(self):
-        # Structure: server_key -> {"sessions": {session_id: session_info}, "last_cleanup": timestamp}
-        self.sessions_by_server = {}
-        self._background_tasks = set()  # Keep references to background tasks
-        # Backwards-compatibility maps: which context_id uses which (server_key, session_id)
-        self._context_to_session: dict[str, tuple[str, str]] = {}
-        # Reference count for each active (server_key, session_id)
-        self._session_refcount: dict[tuple[str, str], int] = {}
-        self._cleanup_task = None
-        self._start_cleanup_task()
-
-    def _start_cleanup_task(self):
-        """Start the periodic cleanup task."""
-        if self._cleanup_task is None or self._cleanup_task.done():
-            self._cleanup_task = asyncio.create_task(self._periodic_cleanup())
-            self._background_tasks.add(self._cleanup_task)
-            self._cleanup_task.add_done_callback(self._background_tasks.discard)
-
-    async def _periodic_cleanup(self):
-        """Periodically clean up idle sessions."""
-        while True:
-            try:
-                await asyncio.sleep(SESSION_CLEANUP_INTERVAL)
-                await self._cleanup_idle_sessions()
-            except asyncio.CancelledError:
-                break
-            except (RuntimeError, KeyError, ClosedResourceError, ValueError, asyncio.TimeoutError) as e:
-                # Handle common recoverable errors without stopping the cleanup loop
-                await logger.awarning(f"Error in periodic cleanup: {e}")
-
-    async def _cleanup_idle_sessions(self):
-        """Clean up sessions that have been idle for too long."""
-        current_time = asyncio.get_event_loop().time()
-        servers_to_remove = []
-
-        for server_key, server_data in self.sessions_by_server.items():
-            sessions = server_data.get("sessions", {})
-            sessions_to_remove = []
-
-            for session_id, session_info in sessions.items():
-                if current_time - session_info["last_used"] > SESSION_IDLE_TIMEOUT:
-                    sessions_to_remove.append(session_id)
-
-            # Clean up idle sessions
-            for session_id in sessions_to_remove:
-                await logger.ainfo(f"Cleaning up idle session {session_id} for server {server_key}")
-                await self._cleanup_session_by_id(server_key, session_id)
-
-            # Remove server entry if no sessions left
-            if not sessions:
-                servers_to_remove.append(server_key)
-
-        # Clean up empty server entries
-        for server_key in servers_to_remove:
-            del self.sessions_by_server[server_key]
-
-    def _get_server_key(self, connection_params, transport_type: str) -> str:
-        """Generate a consistent server key based on connection parameters."""
-        if transport_type == "stdio":
-            if hasattr(connection_params, "command"):
-                # Include command, args, and environment for uniqueness
-                command_str = f"{connection_params.command} {' '.join(connection_params.args or [])}"
-                env_str = str(sorted((connection_params.env or {}).items()))
-                key_input = f"{command_str}|{env_str}"
-                return f"stdio_{hash(key_input)}"
-        elif transport_type == "sse" and (isinstance(connection_params, dict) and "url" in connection_params):
-            # Include URL and headers for uniqueness
-            url = connection_params["url"]
-            headers = str(sorted((connection_params.get("headers", {})).items()))
-            key_input = f"{url}|{headers}"
-            return f"sse_{hash(key_input)}"
-
-        # Fallback to a generic key
-        # TODO: add option for streamable HTTP in future.
-        return f"{transport_type}_{hash(str(connection_params))}"
-
-    async def _validate_session_connectivity(self, session) -> bool:
-        """Validate that the session is actually usable by testing a simple operation."""
-        try:
-            # Try to list tools as a connectivity test (this is a lightweight operation)
-            # Use a shorter timeout for the connectivity test to fail fast
-            response = await asyncio.wait_for(session.list_tools(), timeout=3.0)
-        except (asyncio.TimeoutError, ConnectionError, OSError, ValueError) as e:
-            await logger.adebug(f"Session connectivity test failed (standard error): {e}")
-            return False
-        except Exception as e:
-            # Handle MCP-specific errors that might not be in the standard list
-            error_str = str(e)
-            if (
-                "ClosedResourceError" in str(type(e))
-                or "Connection closed" in error_str
-                or "Connection lost" in error_str
-                or "Connection failed" in error_str
-                or "Transport closed" in error_str
-                or "Stream closed" in error_str
-            ):
-                await logger.adebug(f"Session connectivity test failed (MCP connection error): {e}")
-                return False
-            # Re-raise unexpected errors
-            await logger.awarning(f"Unexpected error in connectivity test: {e}")
-            raise
-        else:
-            # Validate that we got a meaningful response
-            if response is None:
-                await logger.adebug("Session connectivity test failed: received None response")
-                return False
-            try:
-                # Check if we can access the tools list (even if empty)
-                tools = getattr(response, "tools", None)
-                if tools is None:
-                    await logger.adebug("Session connectivity test failed: no tools attribute in response")
-                    return False
-            except (AttributeError, TypeError) as e:
-                await logger.adebug(f"Session connectivity test failed while validating response: {e}")
-                return False
-            else:
-                await logger.adebug(f"Session connectivity test passed: found {len(tools)} tools")
-                return True
-
-    async def get_session(self, context_id: str, connection_params, transport_type: str):
-        """Get or create a session with improved reuse strategy.
-
-        The key insight is that we should reuse sessions based on the server
-        identity (command + args for stdio, URL for SSE) rather than the context_id.
-        This prevents creating a new subprocess for each unique context.
-        """
-        server_key = self._get_server_key(connection_params, transport_type)
-
-        # Ensure server entry exists
-        if server_key not in self.sessions_by_server:
-            self.sessions_by_server[server_key] = {"sessions": {}, "last_cleanup": asyncio.get_event_loop().time()}
-
-        server_data = self.sessions_by_server[server_key]
-        sessions = server_data["sessions"]
-
-        # Try to find a healthy existing session
-        for session_id, session_info in sessions.items():
-            session = session_info["session"]
-            task = session_info["task"]
-
-            # Check if session is still alive
-            if not task.done():
-                # Update last used time
-                session_info["last_used"] = asyncio.get_event_loop().time()
-
-                # Quick health check
-                if await self._validate_session_connectivity(session):
-                    await logger.adebug(f"Reusing existing session {session_id} for server {server_key}")
-                    # record mapping & bump ref-count for backwards compatibility
-                    self._context_to_session[context_id] = (server_key, session_id)
-                    self._session_refcount[(server_key, session_id)] = (
-                        self._session_refcount.get((server_key, session_id), 0) + 1
-                    )
-                    return session
-                await logger.ainfo(f"Session {session_id} for server {server_key} failed health check, cleaning up")
-                await self._cleanup_session_by_id(server_key, session_id)
-            else:
-                # Task is done, clean up
-                await logger.ainfo(f"Session {session_id} for server {server_key} task is done, cleaning up")
-                await self._cleanup_session_by_id(server_key, session_id)
-
-        # Check if we've reached the maximum number of sessions for this server
-        if len(sessions) >= MAX_SESSIONS_PER_SERVER:
-            # Remove the oldest session
-            oldest_session_id = min(sessions.keys(), key=lambda x: sessions[x]["last_used"])
-            await logger.ainfo(
-                f"Maximum sessions reached for server {server_key}, removing oldest session {oldest_session_id}"
-            )
-            await self._cleanup_session_by_id(server_key, oldest_session_id)
-
-        # Create new session
-        session_id = f"{server_key}_{len(sessions)}"
-        await logger.ainfo(f"Creating new session {session_id} for server {server_key}")
-
-        if transport_type == "stdio":
-            session, task = await self._create_stdio_session(session_id, connection_params)
-        elif transport_type == "sse":
-            session, task = await self._create_sse_session(session_id, connection_params)
-        else:
-            msg = f"Unknown transport type: {transport_type}"
-            raise ValueError(msg)
-
-        # Store session info
-        sessions[session_id] = {
-            "session": session,
-            "task": task,
-            "type": transport_type,
-            "last_used": asyncio.get_event_loop().time(),
-        }
-
-        # register mapping & initial ref-count for the new session
-        self._context_to_session[context_id] = (server_key, session_id)
-        self._session_refcount[(server_key, session_id)] = 1
-
-        return session
-
-    async def _create_stdio_session(self, session_id: str, connection_params):
-        """Create a new stdio session as a background task to avoid context issues."""
-        import asyncio
-
-        from mcp.client.stdio import stdio_client
-
-        # Create a future to get the session
-        session_future: asyncio.Future[ClientSession] = asyncio.Future()
-
-        async def session_task():
-            """Background task that keeps the session alive."""
-            try:
-                async with stdio_client(connection_params) as (read, write):
-                    session = ClientSession(read, write)
-                    async with session:
-                        await session.initialize()
-                        # Signal that session is ready
-                        session_future.set_result(session)
-
-                        # Keep the session alive until cancelled
-                        import anyio
-
-                        event = anyio.Event()
-                        try:
-                            await event.wait()
-                        except asyncio.CancelledError:
-                            await logger.ainfo(f"Session {session_id} is shutting down")
-            except Exception as e:  # noqa: BLE001
-                if not session_future.done():
-                    session_future.set_exception(e)
-
-        # Start the background task
-        task = asyncio.create_task(session_task())
-        self._background_tasks.add(task)
-        task.add_done_callback(self._background_tasks.discard)
-
-        # Wait for session to be ready
-        try:
-            session = await asyncio.wait_for(session_future, timeout=10.0)
-        except asyncio.TimeoutError as timeout_err:
-            # Clean up the failed task
-            if not task.done():
-                task.cancel()
-                import contextlib
-
-                with contextlib.suppress(asyncio.CancelledError):
-                    await task
-            self._background_tasks.discard(task)
-            msg = f"Timeout waiting for STDIO session {session_id} to initialize"
-            await logger.aerror(msg)
-            raise ValueError(msg) from timeout_err
-
-        return session, task
-
-    async def _create_sse_session(self, session_id: str, connection_params):
-        """Create a new SSE session as a background task to avoid context issues."""
-        import asyncio
-
-        from mcp.client.sse import sse_client
-
-        # Create a future to get the session
-        session_future: asyncio.Future[ClientSession] = asyncio.Future()
-
-        async def session_task():
-            """Background task that keeps the session alive."""
-            try:
-                async with sse_client(
-                    connection_params["url"],
-                    connection_params["headers"],
-                    connection_params["timeout_seconds"],
-                    connection_params["sse_read_timeout_seconds"],
-                ) as (read, write):
-                    session = ClientSession(read, write)
-                    async with session:
-                        await session.initialize()
-                        # Signal that session is ready
-                        session_future.set_result(session)
-
-                        # Keep the session alive until cancelled
-                        import anyio
-
-                        event = anyio.Event()
-                        try:
-                            await event.wait()
-                        except asyncio.CancelledError:
-                            await logger.ainfo(f"Session {session_id} is shutting down")
-            except Exception as e:  # noqa: BLE001
-                if not session_future.done():
-                    session_future.set_exception(e)
-
-        # Start the background task
-        task = asyncio.create_task(session_task())
-        self._background_tasks.add(task)
-        task.add_done_callback(self._background_tasks.discard)
-
-        # Wait for session to be ready
-        try:
-            session = await asyncio.wait_for(session_future, timeout=10.0)
-        except asyncio.TimeoutError as timeout_err:
-            # Clean up the failed task
-            if not task.done():
-                task.cancel()
-                import contextlib
-
-                with contextlib.suppress(asyncio.CancelledError):
-                    await task
-            self._background_tasks.discard(task)
-            msg = f"Timeout waiting for SSE session {session_id} to initialize"
-            await logger.aerror(msg)
-            raise ValueError(msg) from timeout_err
-
-        return session, task
-
-    async def _cleanup_session_by_id(self, server_key: str, session_id: str):
-        """Clean up a specific session by server key and session ID."""
-        if server_key not in self.sessions_by_server:
-            return
-
-        server_data = self.sessions_by_server[server_key]
-        # Handle both old and new session structure
-        if isinstance(server_data, dict) and "sessions" in server_data:
-            sessions = server_data["sessions"]
-        else:
-            # Handle old structure where sessions were stored directly
-            sessions = server_data
-
-        if session_id not in sessions:
-            return
-
-        session_info = sessions[session_id]
-        try:
-            # First try to properly close the session if it exists
-            if "session" in session_info:
-                session = session_info["session"]
-
-                # Try async close first (aclose method)
-                if hasattr(session, "aclose"):
-                    try:
-                        await session.aclose()
-                        await logger.adebug("Successfully closed session %s using aclose()", session_id)
-                    except Exception as e:  # noqa: BLE001
-                        await logger.adebug("Error closing session %s with aclose(): %s", session_id, e)
-
-                # If no aclose, try regular close method
-                elif hasattr(session, "close"):
-                    try:
-                        # Check if close() is awaitable using inspection
-                        if inspect.iscoroutinefunction(session.close):
-                            # It's an async method
-                            await session.close()
-                            await logger.adebug("Successfully closed session %s using async close()", session_id)
-                        else:
-                            # Try calling it and check if result is awaitable
-                            close_result = session.close()
-                            if inspect.isawaitable(close_result):
-                                await close_result
-                                await logger.adebug(
-                                    "Successfully closed session %s using awaitable close()", session_id
-                                )
-                            else:
-                                # It's a synchronous close
-                                await logger.adebug("Successfully closed session %s using sync close()", session_id)
-                    except Exception as e:  # noqa: BLE001
-                        await logger.adebug("Error closing session %s with close(): %s", session_id, e)
-
-            # Cancel the background task which will properly close the session
-            if "task" in session_info:
-                task = session_info["task"]
-                if not task.done():
-                    task.cancel()
-                    try:
-                        await task
-                    except asyncio.CancelledError:
-                        await logger.ainfo(f"Cancelled task for session {session_id}")
-        except Exception as e:  # noqa: BLE001
-            await logger.awarning(f"Error cleaning up session {session_id}: {e}")
-        finally:
-            # Remove from sessions dict
-            del sessions[session_id]
-
-    async def cleanup_all(self):
-        """Clean up all sessions."""
-        # Cancel periodic cleanup task
-        if self._cleanup_task and not self._cleanup_task.done():
-            self._cleanup_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await self._cleanup_task
-
-        # Clean up all sessions
-        for server_key in list(self.sessions_by_server.keys()):
-            server_data = self.sessions_by_server[server_key]
-            # Handle both old and new session structure
-            if isinstance(server_data, dict) and "sessions" in server_data:
-                sessions = server_data["sessions"]
-            else:
-                # Handle old structure where sessions were stored directly
-                sessions = server_data
-
-            for session_id in list(sessions.keys()):
-                await self._cleanup_session_by_id(server_key, session_id)
-
-        # Clear the sessions_by_server structure completely
-        self.sessions_by_server.clear()
-
-        # Clear compatibility maps
-        self._context_to_session.clear()
-        self._session_refcount.clear()
-
-        # Clear all background tasks
-        for task in list(self._background_tasks):
-            if not task.done():
-                task.cancel()
-                with contextlib.suppress(asyncio.CancelledError):
-                    await task
-
-        # Give a bit more time for subprocess transports to clean up
-        # This helps prevent the BaseSubprocessTransport.__del__ warnings
-        await asyncio.sleep(0.5)
-
-    async def _cleanup_session(self, context_id: str):
-        """Backward-compat cleanup by context_id.
-
-        Decrements the ref-count for the session used by *context_id* and only
-        tears the session down when the last context that references it goes
-        away.
-        """
-        mapping = self._context_to_session.get(context_id)
-        if not mapping:
-            await logger.adebug(f"No session mapping found for context_id {context_id}")
-            return
-
-        server_key, session_id = mapping
-        ref_key = (server_key, session_id)
-        remaining = self._session_refcount.get(ref_key, 1) - 1
-
-        if remaining <= 0:
-            await self._cleanup_session_by_id(server_key, session_id)
-            self._session_refcount.pop(ref_key, None)
-        else:
-            self._session_refcount[ref_key] = remaining
-
-        # Remove the mapping for this context
-        self._context_to_session.pop(context_id, None)
-
-
-class MCPStdioClient:
-    def __init__(self, component_cache=None):
-        self.session: ClientSession | None = None
-        self._connection_params = None
-        self._connected = False
-        self._session_context: str | None = None
-        self._component_cache = component_cache
-
-    async def _connect_to_server(self, command_str: str, env: dict[str, str] | None = None) -> list[StructuredTool]:
-        """Connect to MCP server using stdio transport (SDK style)."""
-        from mcp import StdioServerParameters
-
-        command = command_str.split(" ")
-        env_data: dict[str, str] = {"DEBUG": "true", "PATH": os.environ["PATH"], **(env or {})}
-
-        if platform.system() == "Windows":
-            server_params = StdioServerParameters(
-                command="cmd",
-                args=[
-                    "/c",
-                    f"{command[0]} {' '.join(command[1:])} || echo Command failed with exit code %errorlevel% 1>&2",
-                ],
-                env=env_data,
-            )
-        else:
-            server_params = StdioServerParameters(
-                command="bash",
-                args=["-c", f"exec {command_str} || echo 'Command failed with exit code $?' >&2"],
-                env=env_data,
-            )
-
-        # Store connection parameters for later use in run_tool
-        self._connection_params = server_params
-
-        # If no session context is set, create a default one
-        if not self._session_context:
-            # Generate a fallback context based on connection parameters
-            import uuid
-
-            param_hash = uuid.uuid4().hex[:8]
-            self._session_context = f"default_{param_hash}"
-
-        # Get or create a persistent session
-        session = await self._get_or_create_session()
-        response = await session.list_tools()
-        self._connected = True
-        return response.tools
-
-    async def connect_to_server(self, command_str: str, env: dict[str, str] | None = None) -> list[StructuredTool]:
-        """Connect to MCP server using stdio transport (SDK style)."""
-        return await asyncio.wait_for(
-            self._connect_to_server(command_str, env), timeout=get_settings_service().settings.mcp_server_timeout
-        )
-
-    def set_session_context(self, context_id: str):
-        """Set the session context (e.g., flow_id + user_id + session_id)."""
-        self._session_context = context_id
-
-    def _get_session_manager(self) -> MCPSessionManager:
-        """Get or create session manager from component cache."""
-        if not self._component_cache:
-            # Fallback to instance-level session manager if no cache
-            if not hasattr(self, "_session_manager"):
-                self._session_manager = MCPSessionManager()
-            return self._session_manager
-
-        from langflow.services.cache.utils import CacheMiss
-
-        session_manager = self._component_cache.get("mcp_session_manager")
-        if isinstance(session_manager, CacheMiss):
-            session_manager = MCPSessionManager()
-            self._component_cache.set("mcp_session_manager", session_manager)
-        return session_manager
-
-    async def _get_or_create_session(self) -> ClientSession:
-        """Get or create a persistent session for the current context."""
-        if not self._session_context or not self._connection_params:
-            msg = "Session context and connection params must be set"
-            raise ValueError(msg)
-
-        # Use cached session manager to get/create persistent session
-        session_manager = self._get_session_manager()
-        return await session_manager.get_session(self._session_context, self._connection_params, "stdio")
-
-    async def run_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any:
-        """Run a tool with the given arguments using context-specific session.
-
-        Args:
-            tool_name: Name of the tool to run
-            arguments: Dictionary of arguments to pass to the tool
-
-        Returns:
-            The result of the tool execution
-
-        Raises:
-            ValueError: If session is not initialized or tool execution fails
-        """
-        if not self._connected or not self._connection_params:
-            msg = "Session not initialized or disconnected. Call connect_to_server first."
-            raise ValueError(msg)
-
-        # If no session context is set, create a default one
-        if not self._session_context:
-            # Generate a fallback context based on connection parameters
-            import uuid
-
-            param_hash = uuid.uuid4().hex[:8]
-            self._session_context = f"default_{param_hash}"
-
-        max_retries = 2
-        last_error_type = None
-
-        for attempt in range(max_retries):
-            try:
-                await logger.adebug(f"Attempting to run tool '{tool_name}' (attempt {attempt + 1}/{max_retries})")
-                # Get or create persistent session
-                session = await self._get_or_create_session()
-
-                result = await asyncio.wait_for(
-                    session.call_tool(tool_name, arguments=arguments),
-                    timeout=30.0,  # 30 second timeout
-                )
-            except Exception as e:
-                current_error_type = type(e).__name__
-                await logger.awarning(f"Tool '{tool_name}' failed on attempt {attempt + 1}: {current_error_type} - {e}")
-
-                # Import specific MCP error types for detection
-                try:
-                    is_closed_resource_error = isinstance(e, ClosedResourceError)
-                    is_mcp_connection_error = isinstance(e, McpError) and "Connection closed" in str(e)
-                except ImportError:
-                    is_closed_resource_error = "ClosedResourceError" in str(type(e))
-                    is_mcp_connection_error = "Connection closed" in str(e)
-
-                # Detect timeout errors
-                is_timeout_error = isinstance(e, asyncio.TimeoutError | TimeoutError)
-
-                # If we're getting the same error type repeatedly, don't retry
-                if last_error_type == current_error_type and attempt > 0:
-                    await logger.aerror(f"Repeated {current_error_type} error for tool '{tool_name}', not retrying")
-                    break
-
-                last_error_type = current_error_type
-
-                # If it's a connection error (ClosedResourceError or MCP connection closed) and we have retries left
-                if (is_closed_resource_error or is_mcp_connection_error) and attempt < max_retries - 1:
-                    await logger.awarning(
-                        f"MCP session connection issue for tool '{tool_name}', retrying with fresh session..."
-                    )
-                    # Clean up the dead session
-                    if self._session_context:
-                        session_manager = self._get_session_manager()
-                        await session_manager._cleanup_session(self._session_context)
-                    # Add a small delay before retry
-                    await asyncio.sleep(0.5)
-                    continue
-
-                # If it's a timeout error and we have retries left, try once more
-                if is_timeout_error and attempt < max_retries - 1:
-                    await logger.awarning(f"Tool '{tool_name}' timed out, retrying...")
-                    # Don't clean up session for timeouts, might just be a slow response
-                    await asyncio.sleep(1.0)
-                    continue
-
-                # For other errors or no retries left, handle as before
-                if (
-                    isinstance(e, ConnectionError | TimeoutError | OSError | ValueError)
-                    or is_closed_resource_error
-                    or is_mcp_connection_error
-                    or is_timeout_error
-                ):
-                    msg = f"Failed to run tool '{tool_name}' after {attempt + 1} attempts: {e}"
-                    await logger.aerror(msg)
-                    # Clean up failed session from cache
-                    if self._session_context and self._component_cache:
-                        cache_key = f"mcp_session_stdio_{self._session_context}"
-                        self._component_cache.delete(cache_key)
-                    self._connected = False
-                    raise ValueError(msg) from e
-                # Re-raise unexpected errors
-                raise
-            else:
-                await logger.adebug(f"Tool '{tool_name}' completed successfully")
-                return result
-
-        # This should never be reached due to the exception handling above
-        msg = f"Failed to run tool '{tool_name}': Maximum retries exceeded with repeated {last_error_type} errors"
-        await logger.aerror(msg)
-        raise ValueError(msg)
-
-    async def disconnect(self):
-        """Properly close the connection and clean up resources."""
-        # For stdio transport, there is no remote session to terminate explicitly
-        # The session cleanup happens when the background task is cancelled
-
-        # Clean up local session using the session manager
-        if self._session_context:
-            session_manager = self._get_session_manager()
-            await session_manager._cleanup_session(self._session_context)
-
-        # Reset local state
-        self.session = None
-        self._connection_params = None
-        self._connected = False
-        self._session_context = None
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.disconnect()
-
-
-class MCPSseClient:
-    def __init__(self, component_cache=None):
-        self.session: ClientSession | None = None
-        self._connection_params = None
-        self._connected = False
-        self._session_context: str | None = None
-        self._component_cache = component_cache
-
-    def _get_session_manager(self) -> MCPSessionManager:
-        """Get or create session manager from component cache."""
-        if not self._component_cache:
-            # Fallback to instance-level session manager if no cache
-            if not hasattr(self, "_session_manager"):
-                self._session_manager = MCPSessionManager()
-            return self._session_manager
-
-        from langflow.services.cache.utils import CacheMiss
-
-        session_manager = self._component_cache.get("mcp_session_manager")
-        if isinstance(session_manager, CacheMiss):
-            session_manager = MCPSessionManager()
-            self._component_cache.set("mcp_session_manager", session_manager)
-        return session_manager
-
-    async def validate_url(self, url: str | None, headers: dict[str, str] | None = None) -> tuple[bool, str]:
-        """Validate the SSE URL before attempting connection."""
-        try:
-            parsed = urlparse(url)
-            if not parsed.scheme or not parsed.netloc:
-                return False, "Invalid URL format. Must include scheme (http/https) and host."
-
-            async with httpx.AsyncClient() as client:
-                try:
-                    # For SSE endpoints, try a GET request with short timeout
-                    # Many SSE servers don't support HEAD requests and return 404
-                    response = await client.get(
-                        url, timeout=2.0, headers={"Accept": "text/event-stream", **(headers or {})}
-                    )
-
-                    # For SSE, we expect the server to either:
-                    # 1. Start streaming (200)
-                    # 2. Return 404 if HEAD/GET without proper SSE handshake is not supported
-                    # 3. Return other status codes that we should handle gracefully
-
-                    # Don't fail on 404 since many SSE endpoints return this for non-SSE requests
-                    if response.status_code == HTTP_NOT_FOUND:
-                        # This is likely an SSE endpoint that doesn't support regular GET
-                        # Let the actual SSE connection attempt handle this
-                        return True, ""
-
-                    # Fail on client errors except 404, but allow server errors and redirects
-                    if (
-                        HTTP_BAD_REQUEST <= response.status_code < HTTP_INTERNAL_SERVER_ERROR
-                        and response.status_code != HTTP_NOT_FOUND
-                    ):
-                        return False, f"Server returned client error status: {response.status_code}"
-
-                except httpx.TimeoutException:
-                    # Timeout on a short request might indicate the server is trying to stream
-                    # This is actually expected behavior for SSE endpoints
-                    return True, ""
-                except httpx.NetworkError:
-                    return False, "Network error. Could not reach the server."
-                else:
-                    return True, ""
-
-        except (httpx.HTTPError, ValueError, OSError) as e:
-            return False, f"URL validation error: {e!s}"
-
-    async def pre_check_redirect(self, url: str | None, headers: dict[str, str] | None = None) -> str | None:
-        """Check for redirects and return the final URL."""
-        if url is None:
-            return url
-        try:
-            async with httpx.AsyncClient(follow_redirects=False) as client:
-                # Use GET with SSE headers instead of HEAD since many SSE servers don't support HEAD
-                response = await client.get(
-                    url, timeout=2.0, headers={"Accept": "text/event-stream", **(headers or {})}
-                )
-                if response.status_code == httpx.codes.TEMPORARY_REDIRECT:
-                    return response.headers.get("Location", url)
-                # Don't treat 404 as an error here - let the main connection handle it
-        except (httpx.RequestError, httpx.HTTPError) as e:
-            await logger.awarning(f"Error checking redirects: {e}")
-        return url
-
-    async def _connect_to_server(
-        self,
-        url: str | None,
-        headers: dict[str, str] | None = None,
-        timeout_seconds: int = 30,
-        sse_read_timeout_seconds: int = 30,
-    ) -> list[StructuredTool]:
-        """Connect to MCP server using SSE transport (SDK style)."""
-        # Validate and sanitize headers early
-        validated_headers = _process_headers(headers)
-
-        if url is None:
-            msg = "URL is required for SSE mode"
-            raise ValueError(msg)
-        is_valid, error_msg = await self.validate_url(url, validated_headers)
-        if not is_valid:
-            msg = f"Invalid SSE URL ({url}): {error_msg}"
-            raise ValueError(msg)
-
-        url = await self.pre_check_redirect(url, validated_headers)
-
-        # Store connection parameters for later use in run_tool
-        self._connection_params = {
-            "url": url,
-            "headers": validated_headers,
-            "timeout_seconds": timeout_seconds,
-            "sse_read_timeout_seconds": sse_read_timeout_seconds,
-        }
-
-        # If no session context is set, create a default one
-        if not self._session_context:
-            # Generate a fallback context based on connection parameters
-            import uuid
-
-            param_hash = uuid.uuid4().hex[:8]
-            self._session_context = f"default_sse_{param_hash}"
-
-        # Get or create a persistent session
-        session = await self._get_or_create_session()
-        response = await session.list_tools()
-        self._connected = True
-        return response.tools
-
-    async def connect_to_server(self, url: str, headers: dict[str, str] | None = None) -> list[StructuredTool]:
-        """Connect to MCP server using SSE transport (SDK style)."""
-        return await asyncio.wait_for(
-            self._connect_to_server(url, headers), timeout=get_settings_service().settings.mcp_server_timeout
-        )
-
-    def set_session_context(self, context_id: str):
-        """Set the session context (e.g., flow_id + user_id + session_id)."""
-        self._session_context = context_id
-
-    async def _get_or_create_session(self) -> ClientSession:
-        """Get or create a persistent session for the current context."""
-        if not self._session_context or not self._connection_params:
-            msg = "Session context and params must be set"
-            raise ValueError(msg)
-
-        # Use cached session manager to get/create persistent session
-        session_manager = self._get_session_manager()
-        # Cache session so we can access server-assigned session_id later for DELETE
-        self.session = await session_manager.get_session(self._session_context, self._connection_params, "sse")
-        return self.session
-
-    async def _terminate_remote_session(self) -> None:
-        """Attempt to explicitly terminate the remote MCP session via HTTP DELETE (best-effort)."""
-        # Only relevant for SSE transport
-        if not self._connection_params or "url" not in self._connection_params:
-            return
-
-        url: str = self._connection_params["url"]
-
-        # Retrieve session id from the underlying SDK if exposed
-        session_id = None
-        if getattr(self, "session", None) is not None:
-            # Common attributes in MCP python SDK: `session_id` or `id`
-            session_id = getattr(self.session, "session_id", None) or getattr(self.session, "id", None)
-
-        headers: dict[str, str] = dict(self._connection_params.get("headers", {}))
-        if session_id:
-            headers["Mcp-Session-Id"] = str(session_id)
-
-        try:
-            async with httpx.AsyncClient(timeout=5.0) as client:
-                await client.delete(url, headers=headers)
-        except Exception as e:  # noqa: BLE001
-            # DELETE is advisory—log and continue
-            logger.debug(f"Unable to send session DELETE to '{url}': {e}")
-
-    async def run_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any:
-        """Run a tool with the given arguments using context-specific session.
-
-        Args:
-            tool_name: Name of the tool to run
-            arguments: Dictionary of arguments to pass to the tool
-
-        Returns:
-            The result of the tool execution
-
-        Raises:
-            ValueError: If session is not initialized or tool execution fails
-        """
-        if not self._connected or not self._connection_params:
-            msg = "Session not initialized or disconnected. Call connect_to_server first."
-            raise ValueError(msg)
-
-        # If no session context is set, create a default one
-        if not self._session_context:
-            # Generate a fallback context based on connection parameters
-            import uuid
-
-            param_hash = uuid.uuid4().hex[:8]
-            self._session_context = f"default_sse_{param_hash}"
-
-        max_retries = 2
-        last_error_type = None
-
-        for attempt in range(max_retries):
-            try:
-                await logger.adebug(f"Attempting to run tool '{tool_name}' (attempt {attempt + 1}/{max_retries})")
-                # Get or create persistent session
-                session = await self._get_or_create_session()
-
-                # Add timeout to prevent hanging
-                import asyncio
-
-                result = await asyncio.wait_for(
-                    session.call_tool(tool_name, arguments=arguments),
-                    timeout=30.0,  # 30 second timeout
-                )
-            except Exception as e:
-                current_error_type = type(e).__name__
-                await logger.awarning(f"Tool '{tool_name}' failed on attempt {attempt + 1}: {current_error_type} - {e}")
-
-                # Import specific MCP error types for detection
-                try:
-                    from anyio import ClosedResourceError
-                    from mcp.shared.exceptions import McpError
-
-                    is_closed_resource_error = isinstance(e, ClosedResourceError)
-                    is_mcp_connection_error = isinstance(e, McpError) and "Connection closed" in str(e)
-                except ImportError:
-                    is_closed_resource_error = "ClosedResourceError" in str(type(e))
-                    is_mcp_connection_error = "Connection closed" in str(e)
-
-                # Detect timeout errors
-                is_timeout_error = isinstance(e, asyncio.TimeoutError | TimeoutError)
-
-                # If we're getting the same error type repeatedly, don't retry
-                if last_error_type == current_error_type and attempt > 0:
-                    await logger.aerror(f"Repeated {current_error_type} error for tool '{tool_name}', not retrying")
-                    break
-
-                last_error_type = current_error_type
-
-                # If it's a connection error (ClosedResourceError or MCP connection closed) and we have retries left
-                if (is_closed_resource_error or is_mcp_connection_error) and attempt < max_retries - 1:
-                    await logger.awarning(
-                        f"MCP session connection issue for tool '{tool_name}', retrying with fresh session..."
-                    )
-                    # Clean up the dead session
-                    if self._session_context:
-                        session_manager = self._get_session_manager()
-                        await session_manager._cleanup_session(self._session_context)
-                    # Add a small delay before retry
-                    await asyncio.sleep(0.5)
-                    continue
-
-                # If it's a timeout error and we have retries left, try once more
-                if is_timeout_error and attempt < max_retries - 1:
-                    await logger.awarning(f"Tool '{tool_name}' timed out, retrying...")
-                    # Don't clean up session for timeouts, might just be a slow response
-                    await asyncio.sleep(1.0)
-                    continue
-
-                # For other errors or no retries left, handle as before
-                if (
-                    isinstance(e, ConnectionError | TimeoutError | OSError | ValueError)
-                    or is_closed_resource_error
-                    or is_mcp_connection_error
-                    or is_timeout_error
-                ):
-                    msg = f"Failed to run tool '{tool_name}' after {attempt + 1} attempts: {e}"
-                    await logger.aerror(msg)
-                    # Clean up failed session from cache
-                    if self._session_context and self._component_cache:
-                        cache_key = f"mcp_session_sse_{self._session_context}"
-                        self._component_cache.delete(cache_key)
-                    self._connected = False
-                    raise ValueError(msg) from e
-                # Re-raise unexpected errors
-                raise
-            else:
-                await logger.adebug(f"Tool '{tool_name}' completed successfully")
-                return result
-
-        # This should never be reached due to the exception handling above
-        msg = f"Failed to run tool '{tool_name}': Maximum retries exceeded with repeated {last_error_type} errors"
-        await logger.aerror(msg)
-        raise ValueError(msg)
-
-    async def disconnect(self):
-        """Properly close the connection and clean up resources."""
-        # Attempt best-effort remote session termination first
-        await self._terminate_remote_session()
-
-        # Clean up local session using the session manager
-        if self._session_context:
-            session_manager = self._get_session_manager()
-            await session_manager._cleanup_session(self._session_context)
-
-        # Reset local state
-        self.session = None
-        self._connection_params = None
-        self._connected = False
-        self._session_context = None
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.disconnect()
-
-
-async def update_tools(
-    server_name: str,
-    server_config: dict,
-    mcp_stdio_client: MCPStdioClient | None = None,
-    mcp_sse_client: MCPSseClient | None = None,
-) -> tuple[str, list[StructuredTool], dict[str, StructuredTool]]:
-    """Fetch server config and update available tools."""
-    if server_config is None:
-        server_config = {}
-    if not server_name:
-        return "", [], {}
-    if mcp_stdio_client is None:
-        mcp_stdio_client = MCPStdioClient()
-    if mcp_sse_client is None:
-        mcp_sse_client = MCPSseClient()
-
-    # Fetch server config from backend
-    mode = "Stdio" if "command" in server_config else "SSE" if "url" in server_config else ""
-    command = server_config.get("command", "")
-    url = server_config.get("url", "")
-    tools = []
-    headers = _process_headers(server_config.get("headers", {}))
-
-    try:
-        await _validate_connection_params(mode, command, url)
-    except ValueError as e:
-        logger.error(f"Invalid MCP server configuration for '{server_name}': {e}")
-        raise
-
-    # Determine connection type and parameters
-    client: MCPStdioClient | MCPSseClient | None = None
-    if mode == "Stdio":
-        # Stdio connection
-        args = server_config.get("args", [])
-        env = server_config.get("env", {})
-        full_command = " ".join([command, *args])
-        tools = await mcp_stdio_client.connect_to_server(full_command, env)
-        client = mcp_stdio_client
-    elif mode == "SSE":
-        # SSE connection
-        tools = await mcp_sse_client.connect_to_server(url, headers=headers)
-        client = mcp_sse_client
-    else:
-        logger.error(f"Invalid MCP server mode for '{server_name}': {mode}")
-        return "", [], {}
-
-    if not tools or not client or not client._connected:
-        logger.warning(f"No tools available from MCP server '{server_name}' or connection failed")
-        return "", [], {}
-
-    tool_list = []
-    tool_cache: dict[str, StructuredTool] = {}
-    for tool in tools:
-        if not tool or not hasattr(tool, "name"):
-            continue
-        try:
-            args_schema = create_input_schema_from_json_schema(tool.inputSchema)
-            if not args_schema:
-                logger.warning(f"Could not create schema for tool '{tool.name}' from server '{server_name}'")
-                continue
-
-            tool_obj = StructuredTool(
-                name=tool.name,
-                description=tool.description or "",
-                args_schema=args_schema,
-                func=create_tool_func(tool.name, args_schema, client),
-                coroutine=create_tool_coroutine(tool.name, args_schema, client),
-                tags=[tool.name],
-                metadata={"server_name": server_name},
-            )
-            tool_list.append(tool_obj)
-            tool_cache[tool.name] = tool_obj
-        except (ConnectionError, TimeoutError, OSError, ValueError) as e:
-            logger.error(f"Failed to create tool '{tool.name}' from server '{server_name}': {e}")
-            msg = f"Failed to create tool '{tool.name}' from server '{server_name}': {e}"
-            raise ValueError(msg) from e
-
-    logger.info(f"Successfully loaded {len(tool_list)} tools from MCP server '{server_name}'")
-    return mode, tool_list, tool_cache