fix(api): wait for workflow worker cleanup before task completion (#39614)

api/core/app/apps/advanced_chat/app_generator.py

@@ -616,23 +616,34 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
             message_snapshot = MessageSnapshot.from_message(message)
             session.close()
 
-            # return response or stream generator
-            response = self._handle_advanced_chat_response(
-                application_generate_entity=application_generate_entity,
-                workflow=workflow_snapshot,
-                queue_manager=queue_manager,
-                conversation=conversation_snapshot,
-                message=message_snapshot,
-                user=user,
-                stream=stream,
-                draft_var_saver_factory=self._get_draft_var_saver_factory(
-                    invoke_from,
-                    account=user,
-                    tenant_id=application_generate_entity.app_config.tenant_id,
-                ),
-            )
+            try:
+                response = self._handle_advanced_chat_response(
+                    application_generate_entity=application_generate_entity,
+                    workflow=workflow_snapshot,
+                    queue_manager=queue_manager,
+                    conversation=conversation_snapshot,
+                    message=message_snapshot,
+                    user=user,
+                    stream=stream,
+                    draft_var_saver_factory=self._get_draft_var_saver_factory(
+                        invoke_from,
+                        account=user,
+                        tenant_id=application_generate_entity.app_config.tenant_id,
+                    ),
+                )
+                converted_response = AdvancedChatAppGenerateResponseConverter.convert(
+                    response=response,
+                    invoke_from=invoke_from,
+                )
+            except BaseException:
+                self._join_worker_thread(worker_thread)
+                raise
 
-            return AdvancedChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
+            if isinstance(converted_response, Generator):
+                return self._wrap_stream_with_worker_thread_join(converted_response, worker_thread)
+
+            self._join_worker_thread(worker_thread)
+            return converted_response
 
     def _generate_worker(
         self,

api/core/app/apps/base_app_generator.py

@@ -1,3 +1,5 @@
+import logging
+import threading
 from collections.abc import Generator, Mapping, Sequence
 from contextlib import AbstractContextManager, nullcontext
 from typing import TYPE_CHECKING, Any, Union, final
@@ -23,6 +25,10 @@ from services.workflow_draft_variable_service import DraftVariableSaver as Draft
 if TYPE_CHECKING:
     from graphon.variables.input_entities import VariableEntity
 
+logger = logging.getLogger(__name__)
+
+_WORKER_THREAD_JOIN_TIMEOUT_SECONDS = 300
+
 
 @final
 class _DebuggerDraftVariableSaver:
@@ -64,6 +70,29 @@ class _DebuggerDraftVariableSaver:
 class BaseAppGenerator:
     _file_access_controller: DatabaseFileAccessController = DatabaseFileAccessController()
 
+    @staticmethod
+    def _join_worker_thread(worker_thread: threading.Thread) -> None:
+        # Bound the wait so a leaked app worker cannot occupy an execution slot indefinitely.
+        worker_thread.join(timeout=_WORKER_THREAD_JOIN_TIMEOUT_SECONDS)
+        if worker_thread.is_alive():
+            logger.warning(
+                "Possible app worker thread leak: thread_name=%s timeout_seconds=%s; "
+                "continuing without waiting further to avoid occupying an execution slot indefinitely",
+                worker_thread.name,
+                _WORKER_THREAD_JOIN_TIMEOUT_SECONDS,
+            )
+
+    @staticmethod
+    def _wrap_stream_with_worker_thread_join[ResponseT](
+        response_stream: Generator[ResponseT, None, None],
+        worker_thread: threading.Thread,
+    ) -> Generator[ResponseT, None, None]:
+        """Keep the producer owned by the response stream until both finish."""
+        try:
+            yield from response_stream
+        finally:
+            BaseAppGenerator._join_worker_thread(worker_thread)
+
     @staticmethod
     def _bind_file_access_scope(
         *,

api/core/app/apps/pipeline/pipeline_generator.py

@@ -351,17 +351,28 @@ class PipelineGenerator(BaseAppGenerator):
                 user,
                 tenant_id=pipeline.tenant_id,
             )
-            # return response or stream generator
-            response = self._handle_response(
-                application_generate_entity=application_generate_entity,
-                workflow=workflow,
-                queue_manager=queue_manager,
-                user=user,
-                stream=streaming,
-                draft_var_saver_factory=draft_var_saver_factory,
-            )
+            try:
+                response = self._handle_response(
+                    application_generate_entity=application_generate_entity,
+                    workflow=workflow,
+                    queue_manager=queue_manager,
+                    user=user,
+                    stream=streaming,
+                    draft_var_saver_factory=draft_var_saver_factory,
+                )
+                converted_response = WorkflowAppGenerateResponseConverter.convert(
+                    response=response,
+                    invoke_from=invoke_from,
+                )
+            except BaseException:
+                self._join_worker_thread(worker_thread)
+                raise
 
-            return WorkflowAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
+            if isinstance(converted_response, Generator):
+                return self._wrap_stream_with_worker_thread_join(converted_response, worker_thread)
+
+            self._join_worker_thread(worker_thread)
+            return converted_response
 
     def single_iteration_generate(
         self,

api/core/app/apps/workflow/app_generator.py

@@ -405,17 +405,28 @@ class WorkflowAppGenerator(BaseAppGenerator):
                 tenant_id=app_model.tenant_id,
             )
 
-            # return response or stream generator
-            response = self._handle_response(
-                application_generate_entity=application_generate_entity,
-                workflow=workflow,
-                queue_manager=queue_manager,
-                user=user,
-                draft_var_saver_factory=draft_var_saver_factory,
-                stream=streaming,
-            )
+            try:
+                response = self._handle_response(
+                    application_generate_entity=application_generate_entity,
+                    workflow=workflow,
+                    queue_manager=queue_manager,
+                    user=user,
+                    draft_var_saver_factory=draft_var_saver_factory,
+                    stream=streaming,
+                )
+                converted_response = WorkflowAppGenerateResponseConverter.convert(
+                    response=response,
+                    invoke_from=invoke_from,
+                )
+            except BaseException:
+                self._join_worker_thread(worker_thread)
+                raise
 
-            return WorkflowAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
+            if isinstance(converted_response, Generator):
+                return self._wrap_stream_with_worker_thread_join(converted_response, worker_thread)
+
+            self._join_worker_thread(worker_thread)
+            return converted_response
 
     def single_iteration_generate(
         self,

api/tasks/app_generate/workflow_execute_task.py

@@ -457,7 +457,7 @@ def _publish_streaming_response(
 @shared_task(queue=WORKFLOW_BASED_APP_EXECUTION_QUEUE)
 def workflow_based_app_execution_task(
     payload: str,
-) -> Generator[Mapping[str, Any] | str, None, None] | Mapping[str, Any] | None:
+) -> Mapping[str, Any] | None:
     exec_params = AppExecutionParams.model_validate_json(payload)
 
     logger.info("workflow_based_app_execution_task run with params: %s", exec_params)

api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py

@@ -442,6 +442,13 @@ class TestAdvancedChatAppGeneratorInternals:
             def start(self):
                 thread_data["started"] = True
 
+            def join(self, timeout):
+                thread_data["joined"] = True
+                thread_data["join_timeout"] = timeout
+
+            def is_alive(self):
+                return False
+
         monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.threading.Thread", _Thread)
         monkeypatch.setattr(
             "core.app.apps.advanced_chat.app_generator.db", SimpleNamespace(engine=object(), session=db_session)
@@ -475,6 +482,8 @@ class TestAdvancedChatAppGeneratorInternals:
 
         assert response["response"] == {"raw": True}
         assert thread_data["started"] is True
+        assert thread_data["joined"] is True
+        assert thread_data["join_timeout"] == 300
         assert "pause-layer" in thread_data["kwargs"]["graph_engine_layers"]
         assert generator._dialogue_count == 3
         assert init_records.call_args.kwargs["session"] is db_session
@@ -542,6 +551,13 @@ class TestAdvancedChatAppGeneratorInternals:
             def start(self):
                 thread_data["started"] = True
 
+            def join(self, timeout):
+                thread_data["joined"] = True
+                thread_data["join_timeout"] = timeout
+
+            def is_alive(self):
+                return False
+
         monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.threading.Thread", _Thread)
         monkeypatch.setattr(
             "core.app.apps.advanced_chat.app_generator.db", SimpleNamespace(engine=object(), session=db_session)
@@ -574,6 +590,8 @@ class TestAdvancedChatAppGeneratorInternals:
         init_records.assert_not_called()
         get_thread_messages_length.assert_called_once_with(conversation.id, session=db_session)
         assert thread_data["started"] is True
+        assert thread_data["joined"] is True
+        assert thread_data["join_timeout"] == 300
         db_session.commit.assert_not_called()
         db_session.refresh.assert_not_called()
         db_session.close.assert_called_once()

api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py

@@ -435,6 +435,7 @@ def test_generate_success_returns_converted(generator, mocker: MockerFixture):
     mocker.patch.object(module, "PipelineQueueManager", return_value=queue_manager)
 
     worker_thread = MagicMock()
+    worker_thread.is_alive.return_value = False
     mocker.patch.object(module.threading, "Thread", return_value=worker_thread)
 
     mocker.patch.object(generator, "_get_draft_var_saver_factory", return_value=MagicMock())
@@ -461,6 +462,7 @@ def test_generate_success_returns_converted(generator, mocker: MockerFixture):
     )
 
     assert result == "converted"
+    worker_thread.join.assert_called_once_with(timeout=300)
 
 
 def test_single_iteration_generate_validates_inputs(generator, mocker: MockerFixture):

api/tests/unit_tests/core/app/apps/test_base_app_generator.py

@@ -1,3 +1,6 @@
+import logging
+from unittest.mock import Mock
+
 import pytest
 
 from core.app.apps.base_app_generator import BaseAppGenerator
@@ -369,6 +372,58 @@ def test_validate_inputs_optional_file_with_empty_string_ignores_default():
 
 
 class TestBaseAppGeneratorExtras:
+    def test_wrap_stream_joins_worker_after_stream_exhaustion(self):
+        base_app_generator = BaseAppGenerator()
+        worker_thread = Mock()
+        worker_thread.is_alive.return_value = False
+
+        def response_stream():
+            yield {"event": "workflow_finished"}
+
+        managed_stream = base_app_generator._wrap_stream_with_worker_thread_join(
+            response_stream(),
+            worker_thread,
+        )
+
+        assert next(managed_stream) == {"event": "workflow_finished"}
+        worker_thread.join.assert_not_called()
+
+        with pytest.raises(StopIteration):
+            next(managed_stream)
+
+        worker_thread.join.assert_called_once_with(timeout=300)
+
+    def test_wrap_stream_joins_worker_when_stream_closes(self):
+        base_app_generator = BaseAppGenerator()
+        worker_thread = Mock()
+        worker_thread.is_alive.return_value = False
+
+        def response_stream():
+            yield {"event": "workflow_started"}
+            yield {"event": "workflow_finished"}
+
+        managed_stream = base_app_generator._wrap_stream_with_worker_thread_join(
+            response_stream(),
+            worker_thread,
+        )
+
+        assert next(managed_stream) == {"event": "workflow_started"}
+        managed_stream.close()
+
+        worker_thread.join.assert_called_once_with(timeout=300)
+
+    def test_join_worker_thread_warns_when_thread_remains_alive(self, caplog: pytest.LogCaptureFixture):
+        worker_thread = Mock()
+        worker_thread.name = "leaked-app-worker"
+        worker_thread.is_alive.return_value = True
+
+        with caplog.at_level(logging.WARNING, logger="core.app.apps.base_app_generator"):
+            BaseAppGenerator._join_worker_thread(worker_thread)
+
+        worker_thread.join.assert_called_once_with(timeout=300)
+        assert "Possible app worker thread leak" in caplog.text
+        assert "leaked-app-worker" in caplog.text
+
     def test_prepare_user_inputs_converts_files_and_lists(self, monkeypatch: pytest.MonkeyPatch):
         base_app_generator = BaseAppGenerator()
 

api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py

@@ -211,6 +211,13 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture):
         def start(self):
             return None
 
+        def join(self, timeout):
+            worker_kwargs["joined"] = True
+            worker_kwargs["join_timeout"] = timeout
+
+        def is_alive(self):
+            return False
+
     mocker.patch("core.app.apps.workflow.app_generator.threading.Thread", DummyThread)
 
     app_model = SimpleNamespace(mode="workflow", tenant_id="tenant")
@@ -244,6 +251,8 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture):
     assert result == "converted"
     assert worker_kwargs["kwargs"]["graph_engine_layers"] == ("base-layer", pause_layer)
     assert worker_kwargs["kwargs"]["graph_runtime_state"] is graph_runtime_state
+    assert worker_kwargs["joined"] is True
+    assert worker_kwargs["join_timeout"] == 300
     assert draft_saver_factory.call_args.kwargs["tenant_id"] == app_model.tenant_id
 
 
@@ -286,6 +295,8 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
 
     mocker.patch("core.app.apps.workflow.app_generator.WorkflowAppRunner", side_effect=runner_ctor)
 
+    worker_lifecycle: dict[str, bool] = {}
+
     class ImmediateThread:
         def __init__(self, target, kwargs):
             target(**kwargs)
@@ -293,6 +304,13 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
         def start(self):
             return None
 
+        def join(self, timeout):
+            worker_lifecycle["joined"] = True
+            worker_lifecycle["join_timeout"] = timeout
+
+        def is_alive(self):
+            return False
+
     mocker.patch("core.app.apps.workflow.app_generator.threading.Thread", ImmediateThread)
 
     mocker.patch(
@@ -331,5 +349,7 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
     )
 
     assert result == "raw-response"
+    assert worker_lifecycle["joined"] is True
+    assert worker_lifecycle["join_timeout"] == 300
     runner_instance.run.assert_called_once()
     queue_manager.graph_runtime_state = runtime_state

api/tests/unit_tests/core/app/apps/workflow/test_active_workflow_tasks.py

@@ -1,5 +1,9 @@
+import threading
+from collections.abc import Generator
+
 import pytest
 
+from core.app.apps.base_app_generator import BaseAppGenerator
 from core.app.apps.workflow.active_workflow_tasks import (
     active_workflow_task,
     get_active_workflow_task_count,
@@ -28,3 +32,51 @@ def test_active_workflow_task_rejects_duplicate_task_id() -> None:
         with pytest.raises(ValueError, match="already active"):
             with active_workflow_task("task-a"):
                 pass
+
+
+def test_managed_stream_waits_for_active_worker_cleanup() -> None:
+    worker_started = threading.Event()
+    release_worker = threading.Event()
+    stream_exhausted = threading.Event()
+    consumer_finished = threading.Event()
+    consumer_errors: list[BaseException] = []
+
+    def run_worker() -> None:
+        with active_workflow_task("task-a"):
+            worker_started.set()
+            release_worker.wait()
+
+    def response_stream() -> Generator[dict[str, str], None, None]:
+        yield {"event": "workflow_finished"}
+        stream_exhausted.set()
+
+    worker_thread = threading.Thread(target=run_worker)
+    worker_thread.start()
+    assert worker_started.wait(timeout=2)
+
+    managed_stream = BaseAppGenerator._wrap_stream_with_worker_thread_join(response_stream(), worker_thread)
+    assert next(managed_stream) == {"event": "workflow_finished"}
+
+    def finish_stream() -> None:
+        try:
+            list(managed_stream)
+        except BaseException as exc:
+            consumer_errors.append(exc)
+        finally:
+            consumer_finished.set()
+
+    consumer_thread = threading.Thread(target=finish_stream)
+    consumer_thread.start()
+    try:
+        assert stream_exhausted.wait(timeout=2)
+        assert not consumer_finished.is_set()
+        assert get_active_workflow_task_count() == 1
+    finally:
+        release_worker.set()
+        consumer_thread.join(timeout=2)
+        worker_thread.join(timeout=2)
+
+    assert not consumer_thread.is_alive()
+    assert not worker_thread.is_alive()
+    assert consumer_errors == []
+    assert get_active_workflow_task_count() == 0

api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py

@@ -15,6 +15,70 @@ from models.model import AppMode
 
 
 class TestWorkflowAppGeneratorValidation:
+    def test_generate_stream_joins_worker_after_response_exhaustion(self, monkeypatch: pytest.MonkeyPatch):
+        generator = WorkflowAppGenerator()
+        worker_thread = Mock()
+        worker_thread.is_alive.return_value = False
+        app_config = WorkflowUIBasedAppConfig(
+            tenant_id="tenant",
+            app_id="app",
+            app_mode=AppMode.WORKFLOW,
+            additional_features=AppAdditionalFeatures(),
+            variables=[],
+            workflow_id="workflow-id",
+        )
+        application_generate_entity = WorkflowAppGenerateEntity.model_construct(
+            task_id="task",
+            app_config=app_config,
+            inputs={},
+            files=[],
+            user_id="user",
+            stream=True,
+            invoke_from=InvokeFrom.WEB_APP,
+            extras={},
+        )
+
+        def response_stream():
+            yield {"event": "workflow_finished"}
+
+        monkeypatch.setattr(generator, "_bind_file_access_scope", lambda **kwargs: contextlib.nullcontext())
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.WorkflowAppQueueManager",
+            lambda **kwargs: SimpleNamespace(**kwargs),
+        )
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.current_app",
+            SimpleNamespace(_get_current_object=lambda: SimpleNamespace(name="flask")),
+        )
+        monkeypatch.setattr("core.app.apps.workflow.app_generator.contextvars.copy_context", lambda: "ctx")
+        monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", lambda **kwargs: worker_thread)
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.db",
+            SimpleNamespace(session=SimpleNamespace(close=Mock())),
+        )
+        monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory")
+        monkeypatch.setattr(generator, "_handle_response", lambda **kwargs: response_stream())
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert",
+            lambda response, invoke_from: response,
+        )
+
+        managed_stream = generator._generate(
+            app_model=SimpleNamespace(mode=AppMode.WORKFLOW, tenant_id="tenant"),
+            workflow=SimpleNamespace(id="workflow-id"),
+            user=SimpleNamespace(id="user"),
+            application_generate_entity=application_generate_entity,
+            invoke_from=InvokeFrom.WEB_APP,
+            workflow_execution_repository=SimpleNamespace(),
+            workflow_node_execution_repository=SimpleNamespace(),
+            streaming=True,
+        )
+
+        worker_thread.start.assert_called_once_with()
+        worker_thread.join.assert_not_called()
+        assert list(managed_stream) == [{"event": "workflow_finished"}]
+        worker_thread.join.assert_called_once_with(timeout=300)
+
     def test_ensure_snippet_start_node_returns_original_for_non_snippet_workflow(self):
         workflow = SimpleNamespace(kind_or_standard="workflow")
         session = SimpleNamespace(scalar=Mock())