Fix dispatcher idle hang and add pytest timeouts (#26998)

api/core/workflow/graph_engine/orchestration/dispatcher.py

@@ -99,6 +99,8 @@ class Dispatcher:
                         self._execution_coordinator.check_commands()
                     self._event_queue.task_done()
                 except queue.Empty:
+                    # Process commands even when no new events arrive so abort requests are not missed
+                    self._execution_coordinator.check_commands()
                     # Check if execution is complete
                     if self._execution_coordinator.is_execution_complete():
                         break

api/pyproject.toml

@@ -166,6 +166,7 @@ dev = [
     "mypy~=1.17.1",
     # "locust>=2.40.4",  # Temporarily removed due to compatibility issues. Uncomment when resolved.
     "sseclient-py>=1.8.0",
+    "pytest-timeout>=2.4.0",
 ]
 
 ############################################################

api/tests/unit_tests/core/workflow/graph_engine/test_dispatcher.py

@@ -95,10 +95,10 @@ def _make_succeeded_event() -> NodeRunSucceededEvent:
     )
 
 
-def test_dispatcher_checks_commands_after_node_completion() -> None:
+def test_dispatcher_checks_commands_during_idle_and_on_completion() -> None:
-    """Dispatcher should only check commands after node completion events."""
+    """Dispatcher polls commands when idle and re-checks after completion events."""
     started_checks = _run_dispatcher_for_event(_make_started_event())
     succeeded_checks = _run_dispatcher_for_event(_make_succeeded_event())
 
-    assert started_checks == 0
+    assert started_checks == 1
-    assert succeeded_checks == 1
+    assert succeeded_checks == 2

api/tests/unit_tests/services/test_metadata_bug_complete.py

@@ -41,7 +41,10 @@ class TestMetadataBugCompleteValidation:
         mock_user.current_tenant_id = "tenant-123"
         mock_user.id = "user-456"
 
-        with patch("services.metadata_service.current_user", mock_user):
+        with patch(
+            "services.metadata_service.current_account_with_tenant",
+            return_value=(mock_user, mock_user.current_tenant_id),
+        ):
             # Should crash with TypeError
             with pytest.raises(TypeError, match="object of type 'NoneType' has no len"):
                 MetadataService.create_metadata("dataset-123", mock_metadata_args)
@@ -51,7 +54,10 @@ class TestMetadataBugCompleteValidation:
         mock_user.current_tenant_id = "tenant-123"
         mock_user.id = "user-456"
 
-        with patch("services.metadata_service.current_user", mock_user):
+        with patch(
+            "services.metadata_service.current_account_with_tenant",
+            return_value=(mock_user, mock_user.current_tenant_id),
+        ):
             with pytest.raises(TypeError, match="object of type 'NoneType' has no len"):
                 MetadataService.update_metadata_name("dataset-123", "metadata-456", None)
 

api/tests/unit_tests/services/test_metadata_nullable_bug.py

@@ -29,7 +29,10 @@ class TestMetadataNullableBug:
         mock_user.current_tenant_id = "tenant-123"
         mock_user.id = "user-456"
 
-        with patch("services.metadata_service.current_user", mock_user):
+        with patch(
+            "services.metadata_service.current_account_with_tenant",
+            return_value=(mock_user, mock_user.current_tenant_id),
+        ):
             # This should crash with TypeError when calling len(None)
             with pytest.raises(TypeError, match="object of type 'NoneType' has no len"):
                 MetadataService.create_metadata("dataset-123", mock_metadata_args)
@@ -40,7 +43,10 @@ class TestMetadataNullableBug:
         mock_user.current_tenant_id = "tenant-123"
         mock_user.id = "user-456"
 
-        with patch("services.metadata_service.current_user", mock_user):
+        with patch(
+            "services.metadata_service.current_account_with_tenant",
+            return_value=(mock_user, mock_user.current_tenant_id),
+        ):
             # This should crash with TypeError when calling len(None)
             with pytest.raises(TypeError, match="object of type 'NoneType' has no len"):
                 MetadataService.update_metadata_name("dataset-123", "metadata-456", None)
@@ -88,7 +94,10 @@ class TestMetadataNullableBug:
         mock_user.current_tenant_id = "tenant-123"
         mock_user.id = "user-456"
 
-        with patch("services.metadata_service.current_user", mock_user):
+        with patch(
+            "services.metadata_service.current_account_with_tenant",
+            return_value=(mock_user, mock_user.current_tenant_id),
+        ):
             # Step 4: Service layer crashes on len(None)
             with pytest.raises(TypeError, match="object of type 'NoneType' has no len"):
                 MetadataService.create_metadata("dataset-123", mock_metadata_args)

api/uv.lock

@@ -1394,6 +1394,7 @@ dev = [
     { name = "pytest-cov" },
     { name = "pytest-env" },
     { name = "pytest-mock" },
+    { name = "pytest-timeout" },
     { name = "ruff" },
     { name = "scipy-stubs" },
     { name = "sseclient-py" },
@@ -1583,6 +1584,7 @@ dev = [
     { name = "pytest-cov", specifier = "~=4.1.0" },
     { name = "pytest-env", specifier = "~=1.1.3" },
     { name = "pytest-mock", specifier = "~=3.14.0" },
+    { name = "pytest-timeout", specifier = ">=2.4.0" },
     { name = "ruff", specifier = "~=0.14.0" },
     { name = "scipy-stubs", specifier = ">=1.15.3.0" },
     { name = "sseclient-py", specifier = ">=1.8.0" },
@@ -4979,6 +4981,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b2/05/77b60e520511c53d1c1ca75f1930c7dd8e971d0c4379b7f4b3f9644685ba/pytest_mock-3.14.1-py3-none-any.whl", hash = "sha256:178aefcd11307d874b4cd3100344e7e2d888d9791a6a1d9bfe90fbc1b74fd1d0", size = 9923, upload-time = "2025-05-26T13:58:43.487Z" },
 ]
 
+[[package]]
+name = "pytest-timeout"
+version = "2.4.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ac/82/4c9ecabab13363e72d880f2fb504c5f750433b2b6f16e99f4ec21ada284c/pytest_timeout-2.4.0.tar.gz", hash = "sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a", size = 17973, upload-time = "2025-05-05T19:44:34.99Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fa/b6/3127540ecdf1464a00e5a01ee60a1b09175f6913f0644ac748494d9c4b21/pytest_timeout-2.4.0-py3-none-any.whl", hash = "sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2", size = 14382, upload-time = "2025-05-05T19:44:33.502Z" },
+]
+
 [[package]]
 name = "python-calamine"
 version = "0.5.3"

dev/pytest/pytest_artifacts.sh

@@ -4,4 +4,6 @@ set -x
 SCRIPT_DIR="$(dirname "$(realpath "$0")")"
 cd "$SCRIPT_DIR/../.."
 
-pytest api/tests/artifact_tests/
+PYTEST_TIMEOUT="${PYTEST_TIMEOUT:-120}"
+
+pytest --timeout "${PYTEST_TIMEOUT}" api/tests/artifact_tests/

dev/pytest/pytest_model_runtime.sh

@@ -4,7 +4,9 @@ set -x
 SCRIPT_DIR="$(dirname "$(realpath "$0")")"
 cd "$SCRIPT_DIR/../.."
 
-pytest api/tests/integration_tests/model_runtime/anthropic \
+PYTEST_TIMEOUT="${PYTEST_TIMEOUT:-180}"
+
+pytest --timeout "${PYTEST_TIMEOUT}" api/tests/integration_tests/model_runtime/anthropic \
   api/tests/integration_tests/model_runtime/azure_openai \
   api/tests/integration_tests/model_runtime/openai api/tests/integration_tests/model_runtime/chatglm \
   api/tests/integration_tests/model_runtime/google api/tests/integration_tests/model_runtime/xinference \

dev/pytest/pytest_testcontainers.sh

@@ -4,4 +4,6 @@ set -x
 SCRIPT_DIR="$(dirname "$(realpath "$0")")"
 cd "$SCRIPT_DIR/../.."
 
-pytest api/tests/test_containers_integration_tests
+PYTEST_TIMEOUT="${PYTEST_TIMEOUT:-120}"
+
+pytest --timeout "${PYTEST_TIMEOUT}" api/tests/test_containers_integration_tests

dev/pytest/pytest_tools.sh

@@ -4,4 +4,6 @@ set -x
 SCRIPT_DIR="$(dirname "$(realpath "$0")")"
 cd "$SCRIPT_DIR/../.."
 
-pytest api/tests/integration_tests/tools
+PYTEST_TIMEOUT="${PYTEST_TIMEOUT:-120}"
+
+pytest --timeout "${PYTEST_TIMEOUT}" api/tests/integration_tests/tools

dev/pytest/pytest_unit_tests.sh

@@ -4,5 +4,7 @@ set -x
 SCRIPT_DIR="$(dirname "$(realpath "$0")")"
 cd "$SCRIPT_DIR/../.."
 
+PYTEST_TIMEOUT="${PYTEST_TIMEOUT:-20}"
+
 # libs
-pytest api/tests/unit_tests
+pytest --timeout "${PYTEST_TIMEOUT}" api/tests/unit_tests

dev/pytest/pytest_vdb.sh

@@ -4,7 +4,9 @@ set -x
 SCRIPT_DIR="$(dirname "$(realpath "$0")")"
 cd "$SCRIPT_DIR/../.."
 
-pytest api/tests/integration_tests/vdb/chroma \
+PYTEST_TIMEOUT="${PYTEST_TIMEOUT:-180}"
+
+pytest --timeout "${PYTEST_TIMEOUT}" api/tests/integration_tests/vdb/chroma \
   api/tests/integration_tests/vdb/milvus \
   api/tests/integration_tests/vdb/pgvecto_rs \
   api/tests/integration_tests/vdb/pgvector \

dev/pytest/pytest_workflow.sh

@@ -4,4 +4,6 @@ set -x
 SCRIPT_DIR="$(dirname "$(realpath "$0")")"
 cd "$SCRIPT_DIR/../.."
 
-pytest api/tests/integration_tests/workflow
+PYTEST_TIMEOUT="${PYTEST_TIMEOUT:-120}"
+
+pytest --timeout "${PYTEST_TIMEOUT}" api/tests/integration_tests/workflow