diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index 3c677804028..912641b5e8d 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/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, diff --git a/api/core/app/apps/base_app_generator.py b/api/core/app/apps/base_app_generator.py index 2762f99301d..71831bbb6a7 100644 --- a/api/core/app/apps/base_app_generator.py +++ b/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( *, diff --git a/api/core/app/apps/pipeline/pipeline_generator.py b/api/core/app/apps/pipeline/pipeline_generator.py index 3eb93e7c08a..9e97ce836ca 100644 --- a/api/core/app/apps/pipeline/pipeline_generator.py +++ b/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, diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index fb5393d7730..db1c49cd6d0 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/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, diff --git a/api/tasks/app_generate/workflow_execute_task.py b/api/tasks/app_generate/workflow_execute_task.py index d76066a8aa7..9bc09bac781 100644 --- a/api/tasks/app_generate/workflow_execute_task.py +++ b/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) diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py index f347a5fae7e..004906a5de3 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py +++ b/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() diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py index f2b8179160b..9000cc94ed4 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py +++ b/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): diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_generator.py b/api/tests/unit_tests/core/app/apps/test_base_app_generator.py index 8e7468bb0b8..fe07e420198 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_generator.py +++ b/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() diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py index 8f5cb2b8115..8fc8959cd23 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py +++ b/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 diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_active_workflow_tasks.py b/api/tests/unit_tests/core/app/apps/workflow/test_active_workflow_tasks.py index c50b16533ff..769dd0c092d 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_active_workflow_tasks.py +++ b/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 diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py index 3509c349aef..279b886cdf0 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py +++ b/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())