lfx-nightly 0.1.12.dev33__py3-none-any.whl → 0.1.12.dev35__py3-none-any.whl

lfx/base/agents/events.py

@@ -80,7 +80,7 @@ async def handle_on_chain_start(
                 header={"title": "Input", "icon": "MessageSquare"},
             )
             agent_message.content_blocks[0].contents.append(text_content)
-            agent_message = await send_message_method(message=agent_message)
+            agent_message = await send_message_method(message=agent_message, skip_db_update=True)
             start_time = perf_counter()
     return agent_message, start_time
 
@@ -151,7 +151,7 @@ async def handle_on_chain_end(
                 header={"title": "Output", "icon": "MessageSquare"},
             )
             agent_message.content_blocks[0].contents.append(text_content)
-        agent_message = await send_message_method(message=agent_message)
+        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
         start_time = perf_counter()
     return agent_message, start_time
 
@@ -190,7 +190,7 @@ async def handle_on_tool_start(
     tool_blocks_map[tool_key] = tool_content
     agent_message.content_blocks[0].contents.append(tool_content)
 
-    agent_message = await send_message_method(message=agent_message)
+    agent_message = await send_message_method(message=agent_message, skip_db_update=True)
     if agent_message.content_blocks and agent_message.content_blocks[0].contents:
         tool_blocks_map[tool_key] = agent_message.content_blocks[0].contents[-1]
     return agent_message, new_start_time
@@ -210,7 +210,7 @@ async def handle_on_tool_end(
 
     if tool_content and isinstance(tool_content, ToolContent):
         # Call send_message_method first to get the updated message structure
-        agent_message = await send_message_method(message=agent_message)
+        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
         new_start_time = perf_counter()
 
         # Now find and update the tool content in the current message
@@ -258,7 +258,7 @@ async def handle_on_tool_error(
         tool_content.error = event["data"].get("error", "Unknown error")
         tool_content.duration = _calculate_duration(start_time)
         tool_content.header = {"title": f"Error using **{tool_content.name}**", "icon": "Hammer"}
-        agent_message = await send_message_method(message=agent_message)
+        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
         start_time = perf_counter()
     return agent_message, start_time
 
@@ -275,14 +275,14 @@ async def handle_on_chain_stream(
         if output and isinstance(output, str | list):
             agent_message.text = _extract_output_text(output)
         agent_message.properties.state = "complete"
-        agent_message = await send_message_method(message=agent_message)
+        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
         start_time = perf_counter()
     elif isinstance(data_chunk, AIMessageChunk):
         output_text = _extract_output_text(data_chunk.content)
         if output_text and isinstance(agent_message.text, str):
             agent_message.text += output_text
             agent_message.properties.state = "partial"
-            agent_message = await send_message_method(message=agent_message)
+            agent_message = await send_message_method(message=agent_message, skip_db_update=True)
         if not agent_message.text:
             start_time = perf_counter()
     return agent_message, start_time
@@ -346,13 +346,17 @@ async def process_agent_events(
         async for event in agent_executor:
             if event["event"] in TOOL_EVENT_HANDLERS:
                 tool_handler = TOOL_EVENT_HANDLERS[event["event"]]
+                # Use skip_db_update=True during streaming to avoid DB round-trips
                 agent_message, start_time = await tool_handler(
                     event, agent_message, tool_blocks_map, send_message_method, start_time
                 )
             elif event["event"] in CHAIN_EVENT_HANDLERS:
                 chain_handler = CHAIN_EVENT_HANDLERS[event["event"]]
+                # Use skip_db_update=True during streaming to avoid DB round-trips
                 agent_message, start_time = await chain_handler(event, agent_message, send_message_method, start_time)
         agent_message.properties.state = "complete"
+        # Final DB update with the complete message (skip_db_update=False by default)
+        agent_message = await send_message_method(message=agent_message)
     except Exception as e:
         raise ExceptionWithMessageError(agent_message, str(e)) from e
     return await Message.create(**agent_message.model_dump())

lfx/base/mcp/util.py

@@ -28,8 +28,12 @@ HTTP_ERROR_STATUS_CODE = httpx_codes.BAD_REQUEST  # HTTP status code for client
 
 # HTTP status codes used in validation
 HTTP_NOT_FOUND = 404
+HTTP_METHOD_NOT_ALLOWED = 405
+HTTP_NOT_ACCEPTABLE = 406
 HTTP_BAD_REQUEST = 400
 HTTP_INTERNAL_SERVER_ERROR = 500
+HTTP_UNAUTHORIZED = 401
+HTTP_FORBIDDEN = 403
 
 # MCP Session Manager constants
 settings = get_settings_service().settings
@@ -378,8 +382,8 @@ def _validate_node_installation(command: str) -> str:
 
 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'"
+    if mode not in ["Stdio", "Streamable_HTTP", "SSE"]:
+        msg = f"Invalid mode: {mode}. Must be either 'Stdio', 'Streamable_HTTP', or 'SSE'"
         raise ValueError(msg)
 
     if mode == "Stdio" and not command:
@@ -387,8 +391,8 @@ async def _validate_connection_params(mode: str, command: str | None = None, url
         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"
+    if mode in ["Streamable_HTTP", "SSE"] and not url:
+        msg = f"URL is required for {mode} mode"
         raise ValueError(msg)
 
 
@@ -400,6 +404,7 @@ class MCPSessionManager:
     2. Maximum session limits per server to prevent resource exhaustion
     3. Idle timeout for automatic session cleanup
     4. Periodic cleanup of stale sessions
+    5. Transport preference caching to avoid retrying failed transports
     """
 
     def __init__(self):
@@ -410,6 +415,9 @@ class MCPSessionManager:
         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] = {}
+        # Cache which transport works for each server to avoid retrying failed transports
+        # server_key -> "streamable_http" | "sse"
+        self._transport_preference: dict[str, str] = {}
         self._cleanup_task = None
         self._start_cleanup_task()
 
@@ -467,15 +475,16 @@ class MCPSessionManager:
                 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):
+        elif transport_type == "streamable_http" 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)}"
+            return f"streamable_http_{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:
@@ -525,7 +534,7 @@ class MCPSessionManager:
         """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.
+        identity (command + args for stdio, URL for Streamable HTTP) 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)
@@ -578,17 +587,24 @@ class MCPSessionManager:
 
         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)
+            actual_transport = "stdio"
+        elif transport_type == "streamable_http":
+            # Pass the cached transport preference if available
+            preferred_transport = self._transport_preference.get(server_key)
+            session, task, actual_transport = await self._create_streamable_http_session(
+                session_id, connection_params, preferred_transport
+            )
+            # Cache the transport that worked for future connections
+            self._transport_preference[server_key] = actual_transport
         else:
             msg = f"Unknown transport type: {transport_type}"
             raise ValueError(msg)
 
-        # Store session info
+        # Store session info with the actual transport used
         sessions[session_id] = {
             "session": session,
             "task": task,
-            "type": transport_type,
+            "type": actual_transport,
             "last_used": asyncio.get_event_loop().time(),
         }
 
@@ -634,9 +650,9 @@ class MCPSessionManager:
         self._background_tasks.add(task)
         task.add_done_callback(self._background_tasks.discard)
 
-        # Wait for session to be ready
+        # Wait for session to be ready (use longer timeout for remote connections)
         try:
-            session = await asyncio.wait_for(session_future, timeout=10.0)
+            session = await asyncio.wait_for(session_future, timeout=30.0)
         except asyncio.TimeoutError as timeout_err:
             # Clean up the failed task
             if not task.done():
@@ -652,50 +668,136 @@ class MCPSessionManager:
 
         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."""
+    async def _create_streamable_http_session(
+        self, session_id: str, connection_params, preferred_transport: str | None = None
+    ):
+        """Create a new Streamable HTTP session with SSE fallback as a background task to avoid context issues.
+
+        Args:
+            session_id: Unique identifier for this session
+            connection_params: Connection parameters including URL, headers, timeouts
+            preferred_transport: If set to "sse", skip Streamable HTTP and go directly to SSE
+
+        Returns:
+            tuple: (session, task, transport_used) where transport_used is "streamable_http" or "sse"
+        """
         import asyncio
 
         from mcp.client.sse import sse_client
+        from mcp.client.streamable_http import streamablehttp_client
 
         # Create a future to get the session
         session_future: asyncio.Future[ClientSession] = asyncio.Future()
+        # Track which transport succeeded
+        used_transport: list[str] = []
 
         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)
+            streamable_error = None
 
-                        # Keep the session alive until cancelled
-                        import anyio
+            # Skip Streamable HTTP if we know SSE works for this server
+            if preferred_transport != "sse":
+                # Try Streamable HTTP first with a quick timeout
+                try:
+                    await logger.adebug(f"Attempting Streamable HTTP connection for session {session_id}")
+                    # Use a shorter timeout for the initial connection attempt (2 seconds)
+                    async with streamablehttp_client(
+                        url=connection_params["url"],
+                        headers=connection_params["headers"],
+                        timeout=connection_params["timeout_seconds"],
+                    ) as (read, write, _):
+                        session = ClientSession(read, write)
+                        async with session:
+                            # Initialize with a timeout to fail fast
+                            await asyncio.wait_for(session.initialize(), timeout=2.0)
+                            used_transport.append("streamable_http")
+                            await logger.ainfo(f"Session {session_id} connected via Streamable HTTP")
+                            # 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} (Streamable HTTP) is shutting down")
+                except (asyncio.TimeoutError, Exception) as e:  # noqa: BLE001
+                    # If Streamable HTTP fails or times out, try SSE as fallback immediately
+                    streamable_error = e
+                    error_type = "timed out" if isinstance(e, asyncio.TimeoutError) else "failed"
+                    await logger.awarning(
+                        f"Streamable HTTP {error_type} for session {session_id}: {e}. Falling back to SSE..."
+                    )
+            else:
+                await logger.adebug(f"Skipping Streamable HTTP for session {session_id}, using cached SSE preference")
 
-                        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)
+            # Try SSE if Streamable HTTP failed or if SSE is preferred
+            if streamable_error is not None or preferred_transport == "sse":
+                try:
+                    await logger.adebug(f"Attempting SSE connection for session {session_id}")
+                    # Extract SSE read timeout from connection params, default to 30s if not present
+                    sse_read_timeout = connection_params.get("sse_read_timeout_seconds", 30)
+
+                    async with sse_client(
+                        connection_params["url"],
+                        connection_params["headers"],
+                        connection_params["timeout_seconds"],
+                        sse_read_timeout,
+                    ) as (read, write):
+                        session = ClientSession(read, write)
+                        async with session:
+                            await session.initialize()
+                            used_transport.append("sse")
+                            fallback_msg = " (fallback)" if streamable_error else " (preferred)"
+                            await logger.ainfo(f"Session {session_id} connected via SSE{fallback_msg}")
+                            # Signal that session is ready
+                            if not session_future.done():
+                                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} (SSE) is shutting down")
+                except Exception as sse_error:  # noqa: BLE001
+                    # Both transports failed (or just SSE if it was preferred)
+                    if streamable_error:
+                        await logger.aerror(
+                            f"Both Streamable HTTP and SSE failed for session {session_id}. "
+                            f"Streamable HTTP error: {streamable_error}. SSE error: {sse_error}"
+                        )
+                        if not session_future.done():
+                            session_future.set_exception(
+                                ValueError(
+                                    f"Failed to connect via Streamable HTTP ({streamable_error}) or SSE ({sse_error})"
+                                )
+                            )
+                    else:
+                        await logger.aerror(f"SSE connection failed for session {session_id}: {sse_error}")
+                        if not session_future.done():
+                            session_future.set_exception(ValueError(f"Failed to connect via SSE: {sse_error}"))
 
         # 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
+        # Wait for session to be ready (use longer timeout for remote connections)
         try:
-            session = await asyncio.wait_for(session_future, timeout=10.0)
+            session = await asyncio.wait_for(session_future, timeout=30.0)
+            # Log which transport was used
+            if used_transport:
+                transport_used = used_transport[0]
+                await logger.ainfo(f"Session {session_id} successfully established using {transport_used}")
+                return session, task, transport_used
+            # This shouldn't happen, but handle it just in case
+            msg = f"Session {session_id} established but transport not recorded"
+            raise ValueError(msg)
         except asyncio.TimeoutError as timeout_err:
             # Clean up the failed task
             if not task.done():
@@ -705,12 +807,10 @@ class MCPSessionManager:
                 with contextlib.suppress(asyncio.CancelledError):
                     await task
             self._background_tasks.discard(task)
-            msg = f"Timeout waiting for SSE session {session_id} to initialize"
+            msg = f"Timeout waiting for Streamable HTTP/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:
@@ -1056,7 +1156,7 @@ class MCPStdioClient:
         await self.disconnect()
 
 
-class MCPSseClient:
+class MCPStreamableHttpClient:
     def __init__(self, component_cache=None):
         self.session: ClientSession | None = None
         self._connection_params = None
@@ -1080,67 +1180,15 @@ class MCPSseClient:
             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."""
+    async def validate_url(self, url: str | None) -> tuple[bool, str]:
+        """Validate the Streamable HTTP 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:
+        except (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
+        return True, ""
 
     async def _connect_to_server(
         self,
@@ -1149,27 +1197,31 @@ class MCPSseClient:
         timeout_seconds: int = 30,
         sse_read_timeout_seconds: int = 30,
     ) -> list[StructuredTool]:
-        """Connect to MCP server using SSE transport (SDK style)."""
+        """Connect to MCP server using Streamable HTTP transport with SSE fallback (SDK style)."""
         # Validate and sanitize headers early
         validated_headers = _process_headers(headers)
 
         if url is None:
-            msg = "URL is required for SSE mode"
+            msg = "URL is required for StreamableHTTP or 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,
-        }
+        # Only validate URL if we don't have a cached session
+        # This avoids expensive HTTP validation calls when reusing sessions
+        if not self._connected or not self._connection_params:
+            is_valid, error_msg = await self.validate_url(url)
+            if not is_valid:
+                msg = f"Invalid Streamable HTTP or SSE URL ({url}): {error_msg}"
+                raise ValueError(msg)
+            # Store connection parameters for later use in run_tool
+            # Include SSE read timeout for fallback
+            self._connection_params = {
+                "url": url,
+                "headers": validated_headers,
+                "timeout_seconds": timeout_seconds,
+                "sse_read_timeout_seconds": sse_read_timeout_seconds,
+            }
+        elif headers:
+            self._connection_params["headers"] = validated_headers
 
         # If no session context is set, create a default one
         if not self._session_context:
@@ -1177,18 +1229,21 @@ class MCPSseClient:
             import uuid
 
             param_hash = uuid.uuid4().hex[:8]
-            self._session_context = f"default_sse_{param_hash}"
+            self._session_context = f"default_http_{param_hash}"
 
-        # Get or create a persistent session
+        # Get or create a persistent session (will try Streamable HTTP, then SSE fallback)
         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)."""
+    async def connect_to_server(
+        self, url: str, headers: dict[str, str] | None = None, sse_read_timeout_seconds: int = 30
+    ) -> list[StructuredTool]:
+        """Connect to MCP server using Streamable HTTP with SSE fallback transport (SDK style)."""
         return await asyncio.wait_for(
-            self._connect_to_server(url, headers), timeout=get_settings_service().settings.mcp_server_timeout
+            self._connect_to_server(url, headers, sse_read_timeout_seconds=sse_read_timeout_seconds),
+            timeout=get_settings_service().settings.mcp_server_timeout,
         )
 
     def set_session_context(self, context_id: str):
@@ -1204,12 +1259,14 @@ class MCPSseClient:
         # 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")
+        self.session = await session_manager.get_session(
+            self._session_context, self._connection_params, "streamable_http"
+        )
         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
+        # Only relevant for Streamable HTTP or SSE transport
         if not self._connection_params or "url" not in self._connection_params:
             return
 
@@ -1255,7 +1312,7 @@ class MCPSseClient:
             import uuid
 
             param_hash = uuid.uuid4().hex[:8]
-            self._session_context = f"default_sse_{param_hash}"
+            self._session_context = f"default_http_{param_hash}"
 
         max_retries = 2
         last_error_type = None
@@ -1326,7 +1383,7 @@ class MCPSseClient:
                     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}"
+                        cache_key = f"mcp_session_http_{self._session_context}"
                         self._component_cache.delete(cache_key)
                     self._connected = False
                     raise ValueError(msg) from e
@@ -1364,11 +1421,17 @@ class MCPSseClient:
         await self.disconnect()
 
 
+# Backward compatibility: MCPSseClient is now an alias for MCPStreamableHttpClient
+# The new client supports both Streamable HTTP and SSE with automatic fallback
+MCPSseClient = MCPStreamableHttpClient
+
+
 async def update_tools(
     server_name: str,
     server_config: dict,
     mcp_stdio_client: MCPStdioClient | None = None,
-    mcp_sse_client: MCPSseClient | None = None,
+    mcp_streamable_http_client: MCPStreamableHttpClient | None = None,
+    mcp_sse_client: MCPStreamableHttpClient | None = None,  # Backward compatibility
 ) -> tuple[str, list[StructuredTool], dict[str, StructuredTool]]:
     """Fetch server config and update available tools."""
     if server_config is None:
@@ -1377,11 +1440,17 @@ async def update_tools(
         return "", [], {}
     if mcp_stdio_client is None:
         mcp_stdio_client = MCPStdioClient()
-    if mcp_sse_client is None:
-        mcp_sse_client = MCPSseClient()
+
+    # Backward compatibility: accept mcp_sse_client parameter
+    if mcp_streamable_http_client is None:
+        mcp_streamable_http_client = mcp_sse_client if mcp_sse_client is not None else MCPStreamableHttpClient()
 
     # Fetch server config from backend
-    mode = "Stdio" if "command" in server_config else "SSE" if "url" in server_config else ""
+    # Determine mode from config, defaulting to Streamable_HTTP if URL present
+    mode = server_config.get("mode", "")
+    if not mode:
+        mode = "Stdio" if "command" in server_config else "Streamable_HTTP" if "url" in server_config else ""
+
     command = server_config.get("command", "")
     url = server_config.get("url", "")
     tools = []
@@ -1394,7 +1463,7 @@ async def update_tools(
         raise
 
     # Determine connection type and parameters
-    client: MCPStdioClient | MCPSseClient | None = None
+    client: MCPStdioClient | MCPStreamableHttpClient | None = None
     if mode == "Stdio":
         # Stdio connection
         args = server_config.get("args", [])
@@ -1402,10 +1471,10 @@ async def update_tools(
         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
+    elif mode in ["Streamable_HTTP", "SSE"]:
+        # Streamable HTTP connection with SSE fallback
+        tools = await mcp_streamable_http_client.connect_to_server(url, headers=headers)
+        client = mcp_streamable_http_client
     else:
         logger.error(f"Invalid MCP server mode for '{server_name}': {mode}")
         return "", [], {}

lfx/base/models/model.py

@@ -229,7 +229,7 @@ class LCModelComponent(Component):
                             system_message_added = True
                         runnable = prompt | runnable
                     else:
-                        messages.append(input_value.to_lc_message())
+                        messages.append(input_value.to_lc_message(self.name))
             else:
                 messages.append(HumanMessage(content=input_value))
 

lfx/base/tools/flow_tool.py

@@ -6,7 +6,7 @@ from langchain_core.tools import BaseTool, ToolException
 from typing_extensions import override
 
 from lfx.base.flow_processing.utils import build_data_from_result_data, format_flow_output_data
-from lfx.helpers.flow import build_schema_from_inputs, get_arg_names, get_flow_inputs, run_flow
+from lfx.helpers import build_schema_from_inputs, get_arg_names, get_flow_inputs, run_flow
 from lfx.log.logger import logger
 from lfx.utils.async_helpers import run_until_complete
 

lfx/base/tools/run_flow.py

@@ -5,7 +5,7 @@ from lfx.custom.custom_component.component import Component, get_component_toolk
 from lfx.field_typing import Tool
 from lfx.graph.graph.base import Graph
 from lfx.graph.vertex.base import Vertex
-from lfx.helpers.flow import get_flow_inputs
+from lfx.helpers import get_flow_inputs
 from lfx.inputs.inputs import DropdownInput, InputTypes, MessageInput
 from lfx.log.logger import logger
 from lfx.schema.data import Data