feat: Enhance async handling in simple_agent and related tests

- Updated `get_graph` function in `simple_agent.py` to utilize async component initialization for improved responsiveness.
- Modified test cases in `test_simple_agent_in_lfx_run.py` to validate the async behavior of `get_graph`.
- Refactored various test functions across multiple files to support async execution, ensuring compatibility with the new async workflows.
- Improved documentation for async functions to clarify their purpose and usage.

Author: ogabrielluiz

src/backend/tests/data/simple_agent.py

@@ -42,7 +42,7 @@ async def get_graph() -> Graph:
     chat_input = cp.ChatInput()
     agent = cp.AgentComponent()
 
-    # For this example, we'll use a calculator tool that can be built synchronously
+    # Use URLComponent for web search capabilities
     url_component = cp.URLComponent()
     tools = await url_component.to_toolkit()
 

src/backend/tests/unit/test_simple_agent_in_lfx_run.py

@@ -22,14 +22,15 @@ def simple_agent_script_content(self):
         return '''"""A simple agent flow example for Langflow.
 
 This script demonstrates how to set up a conversational agent using Langflow's
-Agent component with web search capabilities.
+Agent component with proper async handling.
 
 Features:
 - Uses the new flattened component access (cp.AgentComponent instead of deep imports)
 - Configures logging to 'langflow.log' at INFO level
 - Creates an agent with OpenAI GPT model
-- Provides web search tools via URLComponent
 - Connects ChatInput → Agent → ChatOutput
+- Uses async get_graph() function for proper async handling
+- Demonstrates the new async script loading pattern
 
 Usage:
     uv run lfx run simple_agent.py "How are you?"
@@ -43,27 +44,42 @@ def simple_agent_script_content(self):
 from lfx.graph import Graph
 from lfx.log.logger import LogConfig
 
-log_config = LogConfig(
-    log_level="INFO",
-    log_file=Path("langflow.log"),
-)
-
-# Showcase the new flattened component access - no need for deep imports!
-chat_input = cp.ChatInput()
-agent = cp.AgentComponent()
-url_component = cp.URLComponent()
-tools = url_component.to_toolkit()
-
-agent.set(
-    model_name="gpt-4o-mini",
-    agent_llm="OpenAI",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    input_value=chat_input.message_response,
-    tools=tools,
-)
-chat_output = cp.ChatOutput().set(input_value=agent.message_response)
-
-graph = Graph(chat_input, chat_output, log_config=log_config)
+
+async def get_graph() -> Graph:
+    """Create and return the graph with async component initialization.
+
+    This function properly handles async component initialization without
+    blocking the module loading process. The script loader will detect this
+    async function and handle it appropriately using run_until_complete.
+
+    Returns:
+        Graph: The configured graph with ChatInput → Agent → ChatOutput flow
+    """
+    log_config = LogConfig(
+        log_level="INFO",
+        log_file=Path("langflow.log"),
+    )
+
+    # Showcase the new flattened component access - no need for deep imports!
+    chat_input = cp.ChatInput()
+    agent = cp.AgentComponent()
+
+    # Use URLComponent for web search capabilities
+    url_component = cp.URLComponent()
+
+    # Properly handle async component initialization
+    tools = await url_component.to_toolkit()
+
+    agent.set(
+        model_name="gpt-4o-mini",
+        agent_llm="OpenAI",
+        api_key=os.getenv("OPENAI_API_KEY"),
+        input_value=chat_input.message_response,
+        tools=tools,
+    )
+    chat_output = cp.ChatOutput().set(input_value=agent.message_response)
+
+    return Graph(chat_input, chat_output, log_config=log_config)
 '''
 
     @pytest.fixture
@@ -96,10 +112,11 @@ def test_agent_script_structure_and_syntax(self, simple_agent_script_content):
         assert "cp.AgentComponent()" in simple_agent_script_content
         assert "cp.URLComponent()" in simple_agent_script_content
         assert "cp.ChatOutput()" in simple_agent_script_content
-        assert "url_component.to_toolkit()" in simple_agent_script_content
+        assert "async def get_graph()" in simple_agent_script_content
+        assert "await url_component.to_toolkit()" in simple_agent_script_content
         assert 'model_name="gpt-4o-mini"' in simple_agent_script_content
         assert 'agent_llm="OpenAI"' in simple_agent_script_content
-        assert "Graph(chat_input, chat_output" in simple_agent_script_content
+        assert "return Graph(chat_input, chat_output" in simple_agent_script_content
 
     def test_agent_script_file_validation(self, simple_agent_script_file):
         """Test that the agent script file exists and has valid content."""
@@ -111,7 +128,8 @@ def test_agent_script_file_validation(self, simple_agent_script_file):
         content = simple_agent_script_file.read_text()
         assert "from lfx import components as cp" in content
         assert "cp.AgentComponent()" in content
-        assert "Graph(chat_input, chat_output" in content
+        assert "async def get_graph()" in content
+        assert "return Graph(chat_input, chat_output" in content
 
     def test_agent_script_supports_formats(self, simple_agent_script_file):
         """Test that the script supports logging configuration."""

src/lfx/tests/unit/cli/test_common.py

@@ -168,7 +168,8 @@ def test_flow_id_from_path_different_paths(self):
 class TestLoadGraph:
     """Test graph loading functionality."""
 
-    def test_load_graph_from_path_success(self):
+    @pytest.mark.asyncio
+    async def test_load_graph_from_path_success(self):
         """Test successful graph loading from JSON."""
         mock_graph = MagicMock()
         mock_graph.nodes = [1, 2, 3]
@@ -177,21 +178,22 @@ def test_load_graph_from_path_success(self):
             verbose_print = Mock()
             path = Path("/test/flow.json")
 
-            result = load_graph_from_path(path, ".json", verbose_print, verbose=True)
+            result = await load_graph_from_path(path, ".json", verbose_print, verbose=True)
 
             assert result == mock_graph
             mock_load_flow.assert_called_once_with(path, disable_logs=False)
             verbose_print.assert_any_call(f"Analyzing JSON flow: {path}")
             verbose_print.assert_any_call("Loading JSON flow...")
 
-    def test_load_graph_from_path_failure(self):
+    @pytest.mark.asyncio
+    async def test_load_graph_from_path_failure(self):
         """Test graph loading failure."""
         with patch("lfx.cli.common.load_flow_from_json", side_effect=Exception("Load error")) as mock_load_flow:
             verbose_print = Mock()
             path = Path("/test/flow.json")
 
             with pytest.raises(typer.Exit) as exc_info:
-                load_graph_from_path(path, ".json", verbose_print, verbose=False)
+                await load_graph_from_path(path, ".json", verbose_print, verbose=False)
 
             assert exc_info.value.exit_code == 1
             mock_load_flow.assert_called_once_with(path, disable_logs=True)

src/lfx/tests/unit/cli/test_script_loader.py

@@ -170,10 +170,10 @@ def test_validate_graph_instance_missing_chat_output(self):
 class TestLoadGraphFromScript:
     """Test loading graph from script functionality."""
 
-    def test_load_graph_from_script_success(self, simple_chat_py):
+    async def test_load_graph_from_script_success(self, simple_chat_py):
         """Test successful graph loading from script with real Graph object."""
         # Use the existing test data file
-        graph = load_graph_from_script(simple_chat_py)
+        graph = await load_graph_from_script(simple_chat_py)
 
         # Verify it's a real Graph instance
         from lfx.graph import Graph
@@ -185,24 +185,78 @@ def test_load_graph_from_script_success(self, simple_chat_py):
         assert "ChatInput" in component_names
         assert "ChatOutput" in component_names
 
-    def test_load_graph_from_script_no_graph_variable(self):
+    async def test_load_graph_from_script_no_graph_variable(self):
         """Test error when script has no graph variable."""
         with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
             f.write("other_var = 123\n")
             script_path = Path(f.name)
 
         try:
-            with pytest.raises(RuntimeError, match="No 'graph' variable found"):
-                load_graph_from_script(script_path)
+            with pytest.raises(RuntimeError, match="No 'graph' variable or 'get_graph\\(\\)' function found"):
+                await load_graph_from_script(script_path)
         finally:
             script_path.unlink()
 
-    def test_load_graph_from_script_import_error(self):
+    async def test_load_graph_from_script_import_error(self):
         """Test error handling for import errors."""
         script_path = Path("/non/existent/script.py")
 
         with pytest.raises(RuntimeError, match="Error executing script"):
-            load_graph_from_script(script_path)
+            await load_graph_from_script(script_path)
+
+    async def test_load_graph_from_script_with_async_get_graph(self):
+        """Test loading graph from script with async get_graph function."""
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write("from lfx.components.input_output import ChatInput, ChatOutput\n")
+            f.write("from lfx.graph import Graph\n")
+            f.write("\n")
+            f.write("async def get_graph():\n")
+            f.write("    chat_input = ChatInput()\n")
+            f.write("    chat_output = ChatOutput().set(input_value=chat_input.message_response)\n")
+            f.write("    return Graph(chat_input, chat_output)\n")
+            script_path = Path(f.name)
+
+        try:
+            graph = await load_graph_from_script(script_path)
+
+            # Verify it's a real Graph instance
+            from lfx.graph import Graph
+
+            assert isinstance(graph, Graph)
+
+            # Verify it has the expected components
+            component_names = {v.custom_component.__class__.__name__ for v in graph.vertices}
+            assert "ChatInput" in component_names
+            assert "ChatOutput" in component_names
+        finally:
+            script_path.unlink()
+
+    async def test_load_graph_from_script_with_sync_get_graph(self):
+        """Test loading graph from script with sync get_graph function."""
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write("from lfx.components.input_output import ChatInput, ChatOutput\n")
+            f.write("from lfx.graph import Graph\n")
+            f.write("\n")
+            f.write("def get_graph():\n")
+            f.write("    chat_input = ChatInput()\n")
+            f.write("    chat_output = ChatOutput().set(input_value=chat_input.message_response)\n")
+            f.write("    return Graph(chat_input, chat_output)\n")
+            script_path = Path(f.name)
+
+        try:
+            graph = await load_graph_from_script(script_path)
+
+            # Verify it's a real Graph instance
+            from lfx.graph import Graph
+
+            assert isinstance(graph, Graph)
+
+            # Verify it has the expected components
+            component_names = {v.custom_component.__class__.__name__ for v in graph.vertices}
+            assert "ChatInput" in component_names
+            assert "ChatOutput" in component_names
+        finally:
+            script_path.unlink()
 
 
 class TestResultExtraction:
@@ -505,10 +559,10 @@ def test_find_graph_variable_file_not_found(self):
 class TestIntegrationWithRealFlows:
     """Integration tests using real flows and minimal mocking."""
 
-    def test_load_and_validate_real_script(self, simple_chat_py):
+    async def test_load_and_validate_real_script(self, simple_chat_py):
         """Test loading and validating a real script file."""
         # Load the real graph from the script
-        graph = load_graph_from_script(simple_chat_py)
+        graph = await load_graph_from_script(simple_chat_py)
 
         # Verify it's a real Graph
         from lfx.graph import Graph
@@ -523,7 +577,7 @@ def test_load_and_validate_real_script(self, simple_chat_py):
     async def test_execute_real_flow_with_results(self, simple_chat_py):
         """Test executing a real flow and extracting results."""
         # Load the real graph
-        graph = load_graph_from_script(simple_chat_py)
+        graph = await load_graph_from_script(simple_chat_py)
 
         # Execute the graph with real input
         from lfx.graph.schema import RunOutputs

src/lfx/tests/unit/cli/test_serve_components.py

@@ -200,7 +200,8 @@ def verbose_print(msg):
             validate_script_path("/nonexistent/path.json", verbose_print)
 
     @patch("lfx.cli.common.load_flow_from_json")
-    def test_load_graph_from_path_success(self, mock_load_flow):
+    @pytest.mark.asyncio
+    async def test_load_graph_from_path_success(self, mock_load_flow):
         """Test successful graph loading."""
         mock_graph = Mock()
         mock_load_flow.return_value = mock_graph
@@ -212,12 +213,13 @@ def test_load_graph_from_path_success(self, mock_load_flow):
             def verbose_print(msg):
                 pass  # Real function
 
-            graph = load_graph_from_path(Path(tmp.name), ".json", verbose_print, verbose=True)
+            graph = await load_graph_from_path(Path(tmp.name), ".json", verbose_print, verbose=True)
             assert graph == mock_graph
             mock_load_flow.assert_called_once_with(Path(tmp.name), disable_logs=False)
 
     @patch("lfx.cli.common.load_flow_from_json")
-    def test_load_graph_from_path_error(self, mock_load_flow):
+    @pytest.mark.asyncio
+    async def test_load_graph_from_path_error(self, mock_load_flow):
         """Test graph loading with error."""
         mock_load_flow.side_effect = Exception("Parse error")
 
@@ -229,7 +231,7 @@ def verbose_print(msg):
                 pass  # Real function
 
             with pytest.raises(typer.Exit):
-                load_graph_from_path(Path(tmp.name), ".json", verbose_print, verbose=False)
+                await load_graph_from_path(Path(tmp.name), ".json", verbose_print, verbose=False)
             mock_load_flow.assert_called_once_with(Path(tmp.name), disable_logs=True)
 
 
@@ -429,7 +431,8 @@ class TestIntegration:
     """Integration tests combining multiple components."""
 
     @patch("lfx.cli.common.load_flow_from_json")
-    def test_full_app_integration(self, mock_load_flow):
+    @pytest.mark.asyncio
+    async def test_full_app_integration(self, mock_load_flow):
         """Test full app integration with realistic data."""
         # Setup real graph
         real_graph = create_real_graph()
@@ -447,7 +450,7 @@ def verbose_print(msg):
                 pass  # Real function
 
             mock_verbose_print = verbose_print
-            loaded_graph = load_graph_from_path(flow_path, ".json", mock_verbose_print)
+            loaded_graph = await load_graph_from_path(flow_path, ".json", mock_verbose_print)
             assert loaded_graph == real_graph
 
             # Test flow ID generation